from pprint import pprint
import re

from pymatgen.io.cif import CifParser
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.core.structure import Structure
from pymatgen.core.lattice import Lattice
from pymatgen.core.periodic_table import Element


from tklib.tkapplication import tkApplication
from tklib.tkutils import terminate, pint, pfloat, getarg, getintarg, getfloatarg


infile = 'SrTiO3.cif'


app    = tkApplication()
infile      = getarg( 1, infile)


#==========================================
# Main prgram
#==========================================
def print_matrix(m):
    print(f"| {m[0][0]:8.4f} {m[0][1]:8.4f} {m[0][2]:8.4f} |")
    print(f"| {m[1][0]:8.4f} {m[1][1]:8.4f} {m[1][2]:8.4f} |")
    print(f"| {m[2][0]:8.4f} {m[2][1]:8.4f} {m[2][2]:8.4f} |")

def print_inf(structure):
    print(structure)

    structure_dict = structure.as_dict()
    print("structure keys:", structure_dict.keys())
    structure_inf  = structure_dict.get('structure', None)
    spacegroup_inf = structure_dict.get('spacegroup', None)
    charge_inf     = structure_dict.get('charge', None)
#    print("  structure: ", structure_inf)
#    print("    keys:", structure_inf.keys())
#    print("  spacegroup: ", spacegroup_inf)
#    print("  charge : ", charge_inf)

    analyzer = SpacegroupAnalyzer(structure)
#    spacegroup = analyzer.get_space_group_symbol()
    spg_name, spg_num = structure.get_space_group_info()
    symmetry_ops = analyzer.get_symmetry_operations()
    symmetry_matrix = [op.as_dict()['matrix'] for op in symmetry_ops]
    nsym = len(symmetry_matrix)
    print("")
    print(f"Space group: {spg_name} #{spg_num}")
    print(f"Symmetry operatoins: {nsym}")
    for i in range(nsym):
        print(f"op #{i+1}")
        print_matrix(symmetry_matrix[i])

    print("")
    print("Lattice: ")
    lattice_inf = structure.lattice
    a, b, c, alpha, beta, gamma = lattice_inf.parameters
    volume = lattice_inf.volume
    aij    = lattice_inf.matrix
    print("  Lattice parameters:", a, b, c, alpha, beta, gamma)
    print("  Lattice vectors:")
    pprint(aij)
    print("  Volume: ", volume)
#    lattice_inf = structure_dict['structure']['lattice']
#    a, b, c            = structure.lattice.abc
#    alpha, beta, gamma = structure.lattice.angles

    """
    print("")
    print("Species  : ")
    species_inf = structure.species
    print("dir(spec)=", dir(species_inf))
#    print("vars(spec)=", vars(species_inf))
    print("species_inf=", species_inf)
#    print("dir(species_inf)=", dir(species_inf))
#    print("vars(species_inf)=", vars(species_inf))
    mat_spec_occ = symmetrized_structure.species_and_occu
    print("mat_spec_occ:")
    print("mat_spec_occ=", mat_spec_occ)
#    for occ in mat_spec_occ:
#        print("  occ:", occ.to_data_dict)
#        print("  occ:", dir(occ))
    exit()
    """

    print("")
    eq_sites = getattr(structure, "equivalent_sites", None)
    if eq_sites is None:
        print("No information for equivalent sites")
    else:
        print("Independent sites:")
        for sites in eq_sites:
            print("  ", sites[0].species, sites[0].frac_coords)
# 以降のデータは等価位置なので表示しない        
#        print("dir(sites[0])=", dir(sites[0]))
#        for site in sites:
#            print("  ", site)

    print("")
    print("All sites:")
    sites_inf = structure.sites
    for site in sites_inf:
#        print("  ", site)
#    print(sites_inf[0].coords)
        composition = site.species
        for atom_name in composition.keys():
            print("  ", atom_name, site.frac_coords, "  occ=", composition[atom_name])


def main():
    logfile = app.replace_path(infile, template = ["{dirname}", "{filebody}-out.txt"])
    print(f"Open logfile [{logfile}]")
    app.redirect(targets = ["stdout", logfile], mode = 'w')

    convCIFfile = app.replace_path(infile, template = ["{dirname}", "{filebody}-symmetrized.cif"])

    print("")
    print(f"input: {infile}")   
    print(f"log file: {logfile}")
    print(f"output symmetrized CIF file: {convCIFfile}")

    print("")
    print(f"Read [{infile}]")
    structure = Structure.from_file(infile)
    
    print("")
    print("Structure as original input:")
    print_inf(structure)

    print("")
    print("Symmetrized structure:")
    analyzer = SpacegroupAnalyzer(structure)
    symmetrized_structure = analyzer.get_symmetrized_structure()
    print_inf(symmetrized_structure)

    app.terminate(pause = True)


if __name__ == "__main__":
    main()