import os
import sys

from mp_api.client import MPRester
from pymatgen.core.structure import Structure
from pymatgen.core import Composition
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.io.cif import CifWriter, CifParser


formula = None
output_format = "cif"
prec = 1.0e-3


def usage():
    print()
    print(f"Usage: python {sys.argv[0]} formula [output_format] [prec]\n")

def update_vars():
    global formula, output_formuat, prec

    argv = sys.argv
    narg = len(argv)
    if narg <= 1:
        print("\nError: Chemical formula must be given as the first arg\n")
        input("Press ENTER to terminate>>\n")
        exit()

    formula = argv[1]
    if narg >= 3: output_format = argv[2]
    if narg >= 4: prec = float(argv[3])


def get_structure(formula):
    API_KEY = os.getenv('MP_APIKEY')
    if API_KEY is None:
        print("\nError: Can not get MP API Key from the environment var MP_APIKEY\n")
        return None


    mpr = MPRester(API_KEY)
    if mpr is None:
        print(f"\nError: Can not get MPRester using the given API_KEY [{API_KEY}]\n")
        input("Press ENTER to terminate>>\n")
        exit()

    if ',' in formula:
        elements_in = formula.replace(' ', '').split(',')
        search_results = mpr.materials.summary.search(elements = elements_in)
    else:
        elements_in = None
        search_results = mpr.materials.search(formula = formula)
    if not search_results:
        print(f"No data found for {formula}")
        return None

    structures = []
    for res in search_results:
        material_id = res.material_id

        material_data = mpr.materials.search(material_ids = [material_id])
        if not material_data:
            print(f"No structure data found for material_id {material_id}")
            continue

        structure_dict = material_data[0].structure.as_dict()
        structure   = Structure.from_dict(structure_dict)
        composition = structure.composition
        cformula    = composition.reduced_formula
#        cformula = structure.reduced_formula    #formula.replace(" ", "")
#        composition = Composition(cformula)
        elements_res = [element.symbol for element in composition.elements]
        if elements_in:
#            print("elements_res=", elements_res)
#            print("elements_in=", elements_in)
            not_included = []
            for e in elements_res:
#                print("  e=", e)
                if e not in elements_in: not_included.append(e)

            if len(not_included) >= 1:
                print(f"  ** One or more elements are not included for [{cformula}] in the given condition [{formula}]:", not_included)
                continue
            else:
                print(f"  ** Found [{cformula}] [{material_id}]")

        structure.cmaterial = res
        structure.material_id = material_id
        structure.cformula  = cformula
        structure.celements = elements_res
        structures.append(structure)
        cformula = structure.formula.replace(" ", "")
        print(f"Found material_id for {formula}: {material_id}: {cformula}")

    return structures


def save_structures(structures, output_format, formula):
    print()
    for i, structure in enumerate(structures):
        cformula = structure.cformula
        if output_format == "poscar":
            if i == 0:
                output_path = f"POSCAR"
            else:
                output_path = f"POSCAR_{i+1}"
        else:
            output_path = f"{cformula}_{structure.material_id}.{output_format}"
            '''
            if i == 0:
                output_path = f"{cformula}.{output_format}"
            else:
                output_path = f"{cformula}_{i+1}.{output_format}"
            '''

        symmetry_analyzer = SpacegroupAnalyzer(structure, symprec = prec) #, angle_tolerance = 5.0)
        symmetrized_structure = symmetry_analyzer.get_symmetrized_structure()
#        print()
#        print("symmetrized_structure:")
#        print(symmetrized_structure)
        
        print(f"Saving {cformula} structure to [{output_path}]...")
#        structure.to(filename = output_path, fmt = output_format)
        writer = CifWriter(symmetrized_structure, symprec = prec) #, angle_tolerance = 5.0)
        writer.write_file(output_path)


def main():
    update_vars()

    structures = get_structure(formula)
    if structures:
        save_structures(structures, output_format, formula)

    usage()
    print()
    input("Press ENTER to terminate>>\n")
    exit()


if __name__ == "__main__":
    main()