import numpy as np
from fnmatch import fnmatch
import os
import argparse
import setup_paths
from nomadcore.unit_conversion.unit_conversion import convert_unit_function
from nomadcore.parser_backend import *
from nomadcore.local_meta_info import loadJsonFile, InfoKindEl
from phonopy.units import EvTokJmol, Kb as kBoltzmann

parser_info = {"name": "parser_quasi_harmonic_properties", "version": "1.0"}

path = "../../../../nomad-meta-info/meta_info/nomad_meta_info/quasi-harmonic-properties.nomadmetainfo.json"
metaInfoPath = os.path.normpath(
    os.path.join(os.path.dirname(os.path.abspath(__file__)), path))
metaInfoEnv, warns = loadJsonFile(filePath=metaInfoPath,
                                  dependencyLoader=None,
                                  extraArgsHandling=InfoKindEl.ADD_EXTRA_ARGS,
                                  uri=None)

def prep_ref(ref_listP, ref_listS, Emit, kindP, kindS):
#    sCalc = Emit.openSection("section_calculation_to_calculation_refs")
    for ref in ref_listP:
        sCalc = Emit.openSection("section_calculation_to_calculation_refs")

        Emit.addValue("calculation_to_calculation_kind", kindP)
        Emit.addValue("calculation_to_calculation_external_url", ref)
        Emit.closeSection("section_calculation_to_calculation_refs", sCalc)
    sCalc = Emit.openSection("section_calculation_to_calculation_refs")
    for ref in ref_listS:
        sCalc = Emit.openSection("section_calculation_to_calculation_refs")
        Emit.addValue("calculation_to_calculation_kind", kindS)
        Emit.addValue("calculation_to_calculation_external_url", ref)
        Emit.closeSection("section_calculation_to_calculation_refs", sCalc)


def parse(name):
    Parse = JsonParseEventsWriterBackend(metaInfoEnv)
    Parse.startedParsingSession(name, parser_info)
    sRun = Parse.openSection("section_run")
    system = Parse.openSection("section_system")
    Parse.addValue("spacegroup_3D_number", spacegroup)
    Parse.addValue("x_qhp_formula", material)
    Parse.closeSection("section_system", system)
    results = Parse.openSection("section_single_configuration_calculation")
    #fref = Parse.openSection("section_calculation_to_calculation_refs")
    #Parse.addValue("calculation_to_calculation_kind", "source_folder")
    #Parse.addValue("calculation_to_calculation_external_url", refname)
    #Parse.closeSection("section_calculation_to_calculation_refs", fref)
    prep_ref(Phonon_path, Static_path, Parse, "source_phonon_calculations", "source_static_calculations")
    frameSeq = Parse.openSection("section_frame_sequence")
    Parse.addArrayValues("frame_sequence_local_frames_ref", np.array([results]))
    sTD = Parse.openSection("section_thermodynamical_properties")
    Parse.addArrayValues("thermodynamical_property_temperature", T)
    Parse.addArrayValues("thermodynamical_property_heat_capacity_C_v", cV)
    Parse.addArrayValues("x_qhp_quasi_harmonic_free_energy", Free)
    Parse.addArrayValues("x_qhp_quasi_harmonic_bulk_modulus", B)
    Parse.addArrayValues("x_qhp_quasi_harmonic_volume", V)
    Parse.addArrayValues("x_qhp_quasi_harmonic_thermal_expansion", alpha)
    Parse.closeSection("section_thermodynamical_properties", sTD)
    Parse.closeSection("section_frame_sequence", frameSeq)
    Parse.closeSection("section_single_configuration_calculation", results)
    Parse.closeSection("section_run", sRun)
    Parse.finishedParsingSession("ParseSuccess", None)

if __name__ == '__main__':
    import sys
    parser = argparse.ArgumentParser(description='Parses a txt file containing quasi harmonic data.')
    parser.add_argument('mainFileUri',
                        help='The uri of the main file associated with this calculation.')
    parser.add_argument('mainFilePath', default = None,
                        help='The path to the main file associated with this calculation.')

    args = parser.parse_args()
    if args.mainFilePath:
        mainDir = os.path.dirname(args.mainFilePath)
        os.chdir(mainDir)
    with open("Quasi-harmonic_properties.txt") as qh:
        for line in qh:
            if "Quasi-harmonic" in line:
                spacegroup = int(line.split()[-1])
                material = str(line.split()[-3])
    name = args.mainFilePath
    refname = name.split("/Quasi-harmonic_properties.txt")[0]
    #cwd_ref = os.getcwd(refname)
    con_ref = os.listdir(mainDir)
    Phonon_path = []
    Static_path = []
    for con in con_ref:
        if fnmatch(con, "Phonon*"):
            for sub in os.listdir("%s/%s" % (refname, con)):
                if fnmatch(sub, "phonopy-FHI-aims-displacement*"):
                    Phonon_path.append("%s/%s/%s/%s.out" % (refname,con,sub,sub))
        elif fnmatch(con, "Murn_data"):
            for sub in os.listdir("%s/%s" % (refname, con)):
                if fnmatch(sub, "aims*"):
                    Static_path.append("%s/%s/%s" % (refname,con,sub))
    T, V, alpha, B, Free, cV = np.loadtxt("Quasi-harmonic_properties.txt").T
    cV = kBoltzmann * cV
    cubAtocubm = convert_unit_function('angstrom**3', 'meter**3')
    eVpercubAtoJoulespercubm = convert_unit_function('eV*angstrom**-3', 'joules*meter**-3')
    eVtoJoules = convert_unit_function('eV', 'joules')
    eVperKtoJoules = convert_unit_function('eV*K**-1', 'joules*K**-1')
    Free = eVtoJoules(Free)     
    V = cubAtocubm(V)
    B = eVpercubAtoJoulespercubm(B)
    cV = eVperKtoJoules(cV)
    parse(name)