import os
import sys
from tabulate import tabulate
from pymatgen.core.structure import Structure
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.io.cif import CifParser, CifWriter


from tklib.tkcrystal.tkpymatgen import tkPymatgen


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")


pm = tkPymatgen()

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

struct = pm.read(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)

inf = pm.get_symmetry_inf(symprec = prec)
print()
print("symmetry:")
print(f"  space group #{inf['spg_num']} {inf['spg_name']}")
print(f"  symmetry operations:", inf["sym_op"])
 
symmetrized_structure = pm.get_symmetrized_structure()
print()
pm.print_variables(symmetrized_structure)
#pm.print_methods(symmetrized_structure)

print()
print("Symmetrized structure:")
#print(pm.__str__(symmetrized_structure))
sinf = pm.get_structure_inf()
print(symmetrized_structure)
#print("dict:", sinf['dict'])
print("lattice parameters:", sinf["lattice_parameters"])
print("lattice vectors:", sinf["lattice_vectors"])
print("gij:", sinf["gij"])
#print("formula:", sinf["formula"])
print("volume:", sinf["volume"])
print("density:", sinf["density"])
#print("frac coord:", sinf["frac_coord"])

print()
print("sites:")
sites = pm.get_sites(symmetrized_structure)
#pm.print_variables(sites[0])
for s in sites:
    site_inf = pm.get_site_inf(s)
    print(f"   {site_inf['composition']} ({site_inf['label']}): ", end = '')
    print(f"({site_inf['x']}, {site_inf['y']}, {site_inf['z']})")
    print(f"       [{site_inf['species']}] ")

print("asymmetric sites:")
asites = pm.get_asymmetric_sites(symmetrized_structure)
for asites in asites:
#    print(f"   {asites}")
    site_inf = pm.get_site_inf(asites)
    print(f"   {site_inf['composition']} ({site_inf['label']}): ", end = '')
    print(f"({site_inf['x']}, {site_inf['y']}, {site_inf['z']})")
    print(f"       [{site_inf['species']}] ")

print()
print("species_and_occu")
site_occ = pm.get_species_and_occu(symmetrized_structure)
#pm.print_variables(site_occ)
for s in site_occ:
    inf = pm.get_site_occ_inf(s)
    composition = s.element_composition
    cinf = pm.get_composition_inf(composition)
    print(f"  {inf['site_name']}:  w={inf['weight']} natoms={cinf['natoms']}  nspecies={cinf['nspecies']}")
    for atom_name, occ in cinf['composition'].items():
        print(f"     {atom_name} occ={occ}")

print("")
print(f"Save the symmetrized structure to [{outfile}]")
#pm.to(filename = outfile)
#symmetrized_structure.to(fmt = 'cif', filename = outfile)
writer = CifWriter(symmetrized_structure, symprec = prec)
writer.write_file(outfile)