from operator import pos
import xml.etree.ElementTree as ET
from pprint import pprint
from ase import Atoms, symbols
from ase.atom import Atom
from ase.atoms import default
from ase.units import create_units
from ase.utils import iofunction
import numpy as np
from xespresso.utils import modify_text

# Quantum ESPRESSO uses CODATA 2006 internally
units = create_units("2006")

different_names = {"CONTROL": {"forces": "tprnfor", "stress": "tstress"}}


def xml_parser(xmlfile):
    """ """
    pwi_parameters = {}
    # section_names = ['CONTROL', 'SYSTEM', 'ELECTRONS', 'IONS', 'CELL', 'ATOMIC_SPECIES', 'K_POINTS', 'CELL_PARAMETERS', 'CONSTRAINTS', 'OCCUPATIONS', 'ATOMIC_VELECITIES', 'ATOMIC_FORCES']
    section_names = [
        "CONTROL",
        "SYSTEM",
        "ELECTRONS",
        "IONS",
        "CELL",
    ]
    for section_name in section_names:
        pwi_parameters[section_name] = {}
    #
    tree = ET.parse(xmlfile)
    root = tree.getroot()
    input = root.find("./input")
    parameters = {}
    subs = [
        "control_variables",
        "spin",
        "bands",
        "basis",
        "electron_control",
        "ion_control",
        "cell_control",
    ]
    for sub in subs:
        parameters[sub] = {}
        for key in input.find("./%s" % sub):
            parameters[sub][key.tag] = key.text
            parameters[sub].update(key.attrib)
    #
    pwi_parameters["CONTROL"].update(parameters["control_variables"])
    pwi_parameters["SYSTEM"].update(parameters["bands"])
    pwi_parameters["SYSTEM"].update(parameters["basis"])
    pwi_parameters["ELECTRONS"].update(parameters["electron_control"])
    #
    for section_name, paras in different_names.items():
        for key, value in paras.items():
            pwi_parameters[section_name][value] = pwi_parameters[section_name].pop(key)

    # atomic_species
    atomic_species, pseudopotentials, input_ntyp = get_atomic_species(
        input.find("./atomic_species")
    )
    atoms = get_atomic_structure(input.find("./atomic_structure"))
    kpts, koffset = get_kpoint_grid(input.find("./k_points_IBZ"))
    #
    # dftU
    dftU_parameters = get_dftU(input.find("./dft"), input_ntyp)
    pwi_parameters["SYSTEM"].update(dftU_parameters)
    #
    if parameters["spin"]["lsda"].upper() == "TRUE":
        pwi_parameters["SYSTEM"].update({"nspin": 2})
    #
    pwi_parameters = xml2pw(pwi_parameters)
    # Ha to Ry
    for key in ["forc_conv_thr", "etot_conv_thr"]:
        if key in pwi_parameters["CONTROL"]:
            pwi_parameters["CONTROL"][key] = pwi_parameters["CONTROL"][key] * 2
    for key in ["ecutrho", "ecutwfc", "ecutfock"]:
        if key in pwi_parameters["SYSTEM"]:
            pwi_parameters["SYSTEM"][key] = pwi_parameters["SYSTEM"][key] * 2
    #
    pwi_parameters["INPUT_NTYP"] = input_ntyp
    parameters = {
        "input_data": pwi_parameters,
        "pseudopotentials": pseudopotentials,
        "kpts": kpts,
        "koffset": koffset,
        # 'atomic_species': atomic_species,
    }
    return atoms, parameters


def get_atomic_structure(xml):
    """
    Parse atom.

    Parameters
    """
    nat = xml.attrib["nat"]
    alat = xml.attrib["alat"]
    atoms = Atoms()
    positions = []
    symbols = []
    all_species = []
    for item in xml.find("./atomic_positions"):
        species = item.attrib["name"]
        all_species.append(species)
        symbol = species.split("_")[0].split("-")[0]
        if len(symbol) == 3:
            symbol = symbol[0:2]
        symbols.append(symbol)
        index = item.attrib["index"]
        data = item.text.split()
        # These can be fractions and other expressions
        position = [
            float(data[0]) * units["Bohr"],
            float(data[1]) * units["Bohr"],
            float(data[2]) * units["Bohr"],
        ]
        positions.append(position)
    positions = np.array(positions)
    # positions = positions**units['Bohr']
    atoms = Atoms(symbols=symbols, positions=positions)
    atoms.arrays["species"] = all_species
    # cell
    cell = []
    for item in xml.find("./cell"):
        data = item.text.split()
        cell.append([float(data[0]), float(data[1]), float(data[2])])
    cell = np.array(cell)
    cell = cell * units["Bohr"]
    atoms.cell = cell
    atoms.pbc = True
    return atoms


def get_atomic_species(xml):
    """
    Parameters
    """
    default = ["mass", "pseudo_file"]
    atomic_species = []
    pseudopotentials = {}
    input_ntyp = {}
    for sub in xml.findall("./species"):
        species = {}
        name = sub.attrib["name"]
        species["name"] = name
        for item in sub:
            species[item.tag] = item.text
            if item.tag not in default:
                if item.tag not in input_ntyp:
                    input_ntyp[item.tag] = {}
                input_ntyp[item.tag][name] = float(item.text)
        pseudopotentials[name] = species["pseudo_file"]
        atomic_species.append(species)
    return atomic_species, pseudopotentials, input_ntyp


def get_dftU(xml, input_ntyp):
    """
    Parameters
    """
    dftU_parameters = {}
    if not xml.find("./dftU"):
        return dftU_parameters

    for item in xml.find("./dftU"):
        if item.tag == "Hubbard_U":
            dftU_parameters["lda_plus_u"] = True
            if "Hubbard_U" not in input_ntyp:
                input_ntyp["Hubbard_U"] = {}
            input_ntyp["Hubbard_U"][item.attrib["specie"]] = (
                float(item.text) / 0.073498810939358
            )  # Ry to eV
        else:
            dftU_parameters[item.tag] = item.text
    # print(input_ntyp)
    return dftU_parameters


def get_kpoint_grid(xml):
    """
    Parameters
    """
    data = xml[0].attrib
    kpts = (int(data["nk1"]), int(data["nk2"]), int(data["nk3"]))
    koffset = (int(data["k1"]), int(data["k2"]), int(data["k3"]))
    return kpts, koffset


def xml2pw(parameters, package="PW"):
    import copy
    from xespresso.input_parameters import qe_namespace

    new_parameters = copy.deepcopy(parameters)
    for section, paras in parameters.items():
        for key, value in paras.items():
            if key not in qe_namespace[package][section]:
                del new_parameters[section][key]
            else:
                # print(key, parameters[section][key], qe_namespace[package][section][key][0])
                # print(key, parameters[section][key], qe_namespace[package][section][key][1])
                new_parameters[section][key] = modify_text(
                    parameters[section][key],
                    type=qe_namespace[package][section][key][0],
                )
                # if not compare_value(new_parameters[section][key], qe_namespace[package][section][key][1], qe_namespace[package][section][key][0]):
                # del new_parameters[section][key]
    return new_parameters


if __name__ == "__main__":
    xmlfile = "/home/xing/ase/xespresso/restart/scf/fe+u/fe+u.xml"
    xmlfile = "/home/xing/ase/xespresso/restart/relax/h2o/h2o.xml"
    xml_data = xml_parser(xmlfile)
    print(xml_data["atoms"])
    print("pseudopotentials: ", xml_data["pseudopotentials"])
    pprint(xml_data["parameters"])
