# -*- coding: utf-8 -*-

from functools import cached_property
from pathlib import Path

import numpy
import xmltodict
from scipy.constants import physical_constants

from simmate.toolkit import Structure


class PwscfXml:
    def __init__(self, data: dict, check: bool = False):
        # this class is really just a single dictionary with a bunch of methods
        # attatched to it for reading/picking out certain parts of the dictionary.
        # All methods and properties are parsed "lazily" -- that is they aren't
        # loaded until they are called. Once called, they are cached as well.
        self.data = data

        # TODO: run some prechecks (or just call certain methods) to ensure
        # the dictionary is properly formatted
        if check:
            raise NotImplementedError(
                "checking pwscf.xml data for errors is not yet supported"
            )

    @classmethod
    def from_directory(cls, directory: Path | str = None, **kwargs):
        """
        Loads a `PwscfXml` instance from a output directory that was
        generated by Quantum Espresso's `pw.x` program.

        NOTE: `pwscf.xml` is assumed to be present
        """
        if not directory:
            directory = Path.cwd()
        # TODO: this will load other files in the future (e.g. charge density, bands, etc.)
        return cls.from_file(Path(directory) / "pwscf.xml", **kwargs)

    @classmethod
    def from_file(cls, filename: Path | str = "pwscf.xml", **kwargs):
        """
        Loads a `PwscfXml` instance from an `pwscf.xml` output file that was
        generated by Quantum Espresso's `pw.x` program.
        """

        filename = Path(filename)
        with filename.open() as file:
            xml_data = file.read()
        data_dict = xmltodict.parse(xml_data)

        return cls.from_dict(data_dict, **kwargs)

    @classmethod
    def from_dict(cls, data: dict, **kwargs):
        """
        Loads a `PwscfXml` instance from an `pwscf.xml` output file that was
        generated by Quantum Espresso's `pw.x` program.
        """
        # from_dict is synonmous with the default __init__ method
        return cls(data=data)

    # -------------------------------------------------------------------------

    @cached_property
    def final_structure(self):
        struct_data = self.data["qes:espresso"]["output"]["atomic_structure"]
        site_data = struct_data["atomic_positions"]["atom"]

        # lattice
        lattice_matrix = []
        for vector in ["a1", "a2", "a3"]:
            vector_cleaned = []
            for x in struct_data["cell"][vector].split():
                vector_cleaned.append(float(x))
            lattice_matrix.append(vector_cleaned)

        # species
        species = [s["@name"] for s in site_data]

        # site coords
        site_coords = []
        for s in site_data:
            coords = [float(i) for i in s["#text"].split()]
            site_coords.append(coords)

        # convert everything to Angstroms
        conversion_factor = physical_constants["Bohr radius"][0] * 1e10
        lattice_matrix = numpy.array(lattice_matrix) * conversion_factor
        site_coords = numpy.array(site_coords) * conversion_factor

        return Structure(
            lattice=lattice_matrix,
            species=species,
            coords=site_coords,
            coords_are_cartesian=True,
        )

    @cached_property
    def final_energy(self) -> float:
        rydberg_to_ev = physical_constants["Rydberg constant times hc in eV"][0]
        energy_ry = float(self.data["qes:espresso"]["output"]["total_energy"]["etot"])

        # BUG: something is off here... These values don't match what is in the
        # pwscf.out file. Perhaps these here are per atom?
        natoms = len(self.final_structure)

        # convert to eV & total energy
        return energy_ry * rydberg_to_ev * natoms

    @cached_property
    def site_forces(self) -> numpy.array:
        force_data = self.data["qes:espresso"]["output"]["forces"]["#text"]

        site_forces = []
        sites = force_data.split("\n")
        for site in sites:
            forces = [float(i) for i in site.split()]
            site_forces.append(forces)

        site_forces = numpy.array(site_forces)
        # TODO: convert to differnt units...?

        return site_forces

    @cached_property
    def lattice_stress(self) -> numpy.array:
        stress_data = self.data["qes:espresso"]["output"]["stress"]["#text"]

        stress = []
        vectors = stress_data.split("\n")
        for vector in vectors:
            stress_vector = [float(i) for i in vector.split()]
            stress.append(stress_vector)

        stress = numpy.array(stress)
        # TODO: convert to differnt units...?

        return stress

    # -------------------------------------------------------------------------
