import os
import sys
from pymatgen.core.structure import Structure
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from tabulate import tabulate

infile = None
prec = 1.0e-3

argv = sys.argv
narg = len(argv)
if narg >= 2:
    infile = argv[1]
if narg >= 3:
    outfile = argv[2]
else:
    dir_path = os.path.dirname(infile)
    filebody, ext = os.path.splitext(infile)
    outfile = os.path.join(dir_path, f"{filebody}-symmetrized.cif")


def get_member_variables(self):
        class_type = self.__class__
        return [method for method in dir(class_type) if not callable(getattr(class_type, method))]

def get_members(self):
        class_type = self.__class__
        funcs = [method for method in dir(class_type) if callable(getattr(class_type, method))]
        vars = [method for method in dir(class_type) if not callable(getattr(class_type, method))]
        return funcs, vars, getattr(self, "__dict__", None)

#https://github.com/materialsproject/pymatgen/blob/v2023.10.4/pymatgen/symmetry/structure.py#L18-L151
def __str__(self) -> str:
        outs = [
            "SymmetrizedStructure",
            f"Full Formula ({self.composition.formula})",
            f"Reduced Formula: {self.composition.reduced_formula}",
            f"Spacegroup: {self.spacegroup.int_symbol} ({self.spacegroup.int_number})",
        ]

        def to_str(x):
            return f"{x:>10.6f}"

        outs.append(f"abc   : {' '.join(to_str(val) for val in self.lattice.abc)}")
        outs.append(f"angles: {' '.join(to_str(val) for val in self.lattice.angles)}")
        if self._charge:
            outs.append(f"Overall Charge: {self._charge:+}")
        outs.append(f"Sites ({len(self)})")
        data = []
        props = self.site_properties
        keys = sorted(props)
        for idx, sites in enumerate(self.equivalent_sites):
            site = sites[0]
            row = [str(idx), site.species_string]
            row.extend([to_str(j) for j in site.frac_coords])
            row.append(self.wyckoff_symbols[idx])
            for k in keys:
                row.append(props[k][idx])
            data.append(row)
        outs.append(tabulate(data, headers=["#", "SP", "a", "b", "c", "Wyckoff", *keys]))
        return "\n".join(outs)

print()
print(f"infile: {infile}")
print(f"outfile: {outfile}")

struct = Structure.from_file(infile, primitive = False)
print()
print("Input structure:")
print(struct)
#print("matrix:", struct.lattice.matrix)
#print("coord:", struct.cart_coords)
#print("atomic num:", struct.atomic_numbers)

symmetry_analyzer = SpacegroupAnalyzer(struct, symprec = prec)
print()
print("symmetry_analyzer=", symmetry_analyzer)
print("  get_symmetry_dataset:", symmetry_analyzer.get_symmetry_dataset())
print("  get_space_group_symbol:", symmetry_analyzer.get_space_group_symbol())
print("  get_space_group_number:", symmetry_analyzer.get_space_group_number())
print("  is_laue:", symmetry_analyzer.is_laue())
print("  get_hall:", symmetry_analyzer.get_hall())
print("  get_symmetry_operations:", symmetry_analyzer.get_symmetry_operations(cartesian = False))
print()
print("  get_point_group_symbol:", symmetry_analyzer.get_point_group_symbol())
#print("  get_point_group_operations:", symmetry_analyzer.get_point_group_operations())


symmetrized_structure = symmetry_analyzer.get_symmetrized_structure()
print()
variables = get_member_variables(symmetrized_structure)
members = get_members(symmetrized_structure)[0]
print("variables:")
for m in variables:
    print(f"  variable: {m}")
print("members:")
for m in members:
    if m in variables:
        print(f"  variable: {m}")
    else:
        print(f"  method: {m}")

print()
print("Symmetrized structure:")
print("__str__:", __str__(symmetrized_structure))
exit()

print(symmetrized_structure)
print("lattice:", symmetrized_structure.lattice)
print("distance_matrix:", symmetrized_structure.distance_matrix)
print("total charge:", symmetrized_structure.charge)
print("composition:", symmetrized_structure.composition)
#print("formula:", symmetrized_structure.formula)
print("volume:", symmetrized_structure.volume)
print("density:", symmetrized_structure.density)
print("properties:", symmetrized_structure.properties)
print("matrix:", symmetrized_structure.lattice.matrix)
print("cart coord:", symmetrized_structure.cart_coords)
print("frac coord:", symmetrized_structure.frac_coords)
#print("atomic num:", symmetrized_structure.atomic_numbers)

print()
print("sites:")
sites = symmetrized_structure.sites
site_variables = get_member_variables(sites[0])
site_members = get_members(symmetrized_structure.sites[0])[0]
print("variables:")
for m in site_variables:
    print(f"  variable: {m}")
for s in sites:
    print(f"  {s.species_string}: {s.label} ({s.frac_coords[0]}, {s.frac_coords[1]}, {s.frac_coords[2]})")

print("find_equivalent_sites:")
eq_sites = []
for s in sites:
    ls = symmetrized_structure.find_equivalent_sites(s)
#    print("  eq:", ls)
    if s is ls[0]:
        print("  asymmetric:", ls[0])

    for l in ls:
        if l is not s:
            eq_sites.append(l)

print()
print("species_and_occu")
spec_occ = symmetrized_structure.species_and_occu
spec_variables = get_member_variables(spec_occ[0])
sepc_members = get_members(spec_occ[0])[0]
print("variables:")
for m in spec_variables:
    print(f"  variable: {m}")
for s in spec_occ:
    print(f"  {s.alphabetical_formula} ({s.formula}) ({s.reduced_formula}) (w={s.weight}): z={s.charge} {s.element_composition} {s.elements}, f={s.fractional_composition}")

composition = spec_occ[0].element_composition
print("composition=", composition, type(composition))
print("dict:", composition.to_data_dict)
print("composition:", composition.to_data_dict['unit_cell_composition'])
c_variables = get_member_variables(composition)
c_members = get_members(composition)[0]
print("variables:")
for m in c_variables:
    print(f"  variable: {m}")
print("fractional_composition: ", composition.fractional_composition)
exit()

print()
#print("site_properties:", site_properties.sites)
#print("species:", symmetrized_structure.species)
#print("elements:", symmetrized_structure.elements)
print("types_of_species:", symmetrized_structure.types_of_species)
#print("types_of_specie:", symmetrized_structure.types_of_specie)
species = symmetrized_structure.types_of_species
for s, l in zip(species, symmetrized_structure.labels):
    print(f"  {str(s)}: {l}")

#print("equivalent_sites:", symmetrized_structure.equivalent_sites)
#for site in symmetrized_structure.equivalent_sites[0]:
#    print("site:", site)
print()
print(" dict:", symmetrized_structure.as_dict())
exit()



symmetrized_structure.to(fmt = 'cif', filename = outfile)