"""
Classes to compute vibrational properties with phonopy and ML potentials.
"""
from __future__ import annotations

import numpy as np
import json
import abipy.core.abinit_units as abu

from monty.dev import requires
from monty.string import list_strings, marquee
from monty.termcolor import cprint
from ase.calculators.calculator import Calculator
from ase.atoms import Atoms
from pymatgen.io.phonopy import get_phonopy_structure
from abipy.core.structure import Structure
from abipy.dfpt.ddb import DdbFile
from abipy.tools.context_managers import Timer
from abipy.ml.aseml import RX_MODE, CalcBuilder, AseResults, MlBase, relax_atoms, dataframe_from_results_list
try:
    import phonopy
    from phonopy import Phonopy
    from phonopy.structure.atoms import PhonopyAtoms
except ImportError:
    phonopy = None
    Phonopy = None


def cprint_traceback(color="red") -> None:
    """Print traceback."""
    import traceback
    from monty.termcolor import cprint
    cprint(traceback.format_exc(), color=color)


@requires(phonopy, "phonopy should be installed to calculate phonons")
def get_phonopy(structure: Structure,
                supercell_matrix,
                calculator: Calculator,
                distance=0.01,
                primitive_matrix=None,
                remove_drift=True,
                ) -> Phonopy:
    """
    Args:
        structure: Structure object.
        supercell_matrix: Supercell matrix.
        calculator: ASE calculator to be attached to the atoms.
        distance: Distance of finite displacements in Angstrom.
        primitive_matrix
        remove_drift: True if the drift in the forces should be removed.

    Based on a similar implementation available at: https://github.com/modl-uclouvain/randomcarbon/blob/main/randomcarbon/run/phonon.py
    """
    unitcell = Structure.as_structure(structure).get_phonopy_atoms()

    phonon = Phonopy(unitcell, supercell_matrix=supercell_matrix, primitive_matrix=primitive_matrix)
    phonon.generate_displacements(distance=distance)

    forces_list = []
    nsc = len(phonon.supercells_with_displacements)
    print(f"Calculating forces for {nsc} supercell configurations ...")

    for i, sc in enumerate(phonon.supercells_with_displacements):
        a = Atoms(symbols=sc.symbols,
                  positions=sc.positions,
                  masses=sc.masses,
                  cell=sc.cell,
                  pbc=True,
                  calculator=calculator)


        forces = a.get_forces()
        if remove_drift:
            drift_force = forces.sum(axis=0)
            for force in forces:
                force -= drift_force / forces.shape[0]

        forces_list.append(forces)
        print(f"\t{i+1} of {nsc} completed ...")

    phonon.produce_force_constants(forces_list)

    return phonon


class MlPhonopyWithDDB(MlBase):
    """
    Compute phonons with phonopy and a ML potential starting from a DDB file and compare the results.
    """
    def __init__(self, ddb_filepath,
                 distance, asr, dipdip, line_density, qppa,
                 relax_mode, fmax, pressure, steps, optimizer, nn_names,
                 verbose, workdir, prefix=None,
                 supercell=None):
        """
        Args:
            ddb_filepath: DDB filepath.
            distance: Distance of finite displacement in Angstrom.
            asr: Enforce acoustic sum-rule. Abinit variable.
            dipdip: Treatment of dipole-dipole term. Abinit variable.
            line_density: Defines the a density of k-points per reciprocal atom to plot the phonon dispersion.
            qppa: Defines the homogeneous q-mesh used for the DOS in units of q-points per atom.
            relax_mode: "ions" to relax ions only, "cell" for ions + cell, "no" for no relaxation.
            fmax: tolerance for relaxation convergence. Here fmax is a sum of force and stress forces.
            pressure: Target pressure.
            steps: max number of steps for relaxation.
            optimizer: name of the ASE optimizer to use for relaxation.
            nn_names: String or list of strings defining the NN potential. See also CalcBuilder.
            verbose: Verbosity level.
            workdir: Working directory, None to generate temporary directory automatically.
            prefix: Prefix for workdir
            supercell: with supercell dimensions. None to use the supercell from the DDB file.
        """
        super().__init__(workdir, prefix)

        self.distance = float(distance)
        self.asr = asr
        self.dipdip = dipdip
        self.relax_mode = relax_mode
        RX_MODE.validate(self.relax_mode)
        self.fmax = fmax
        self.pressure = pressure
        self.steps = steps
        self.optimizer = optimizer
        self.nn_names = list_strings(nn_names)
        self.verbose = verbose
        self.supercell = supercell
        self.line_density = line_density
        self.qppa = qppa

        self.ddb_filepath = ddb_filepath
        with DdbFile(self.ddb_filepath) as ddb:
            self.initial_atoms = ddb.structure.to_ase_atoms()
            if self.supercell is None:
                # Take supercell from the q-mesh used in the DDB.
                self.supercell = np.eye(3) * ddb.guessed_ngqpt

        if self.supercell is None:
            raise ValueError("Supercell must be specified in input!")

        self.abi_nac_params = {}

    def to_string(self, verbose=0):
        """String representation with verbosity level `verbose`."""
        supercell = self.supercell.tolist()
        s = f"""\

{self.__class__.__name__} parameters:

     ddb_path      = {self.ddb_filepath}
     supercell     = {supercell}
     distance      = {self.distance}
     qppa          = {self.qppa}
     asr           = {self.asr}
     dipdip        = {self.dipdip}
     line_density  = {self.line_density}
     relax_mode    = {self.relax_mode}
     fmax          = {self.fmax}
     steps         = {self.steps}
     optimizer     = {self.optimizer}
     pressure      = {self.pressure}
     nn_names      = {self.nn_names}
     workdir       = {self.workdir}
     verbose       = {self.verbose}
"""
        return s

    def run(self) -> None:
        """
        Run MlPhonopyWithDDB.
        """
        # Run anaddb computation and save results in self.
        with DdbFile(self.ddb_filepath) as ddb, Timer(header="Starting anaddb ph-bands computation...") as timer:
            if ddb.has_lo_to_data and self.dipdip != 0:
                # according to the phonopy website 14.399652 is not the coefficient for abinit
                # probably it relies on the other conventions in the output.
                out = ddb.anaget_epsinf_and_becs()
                self.abi_nac_params = {"born": out.becs.values, "dielectric": out.epsinf, "factor": 14.399652}

            # ab-initio phonons from the DDB.
            with ddb.anaget_phbst_and_phdos_files(
                        qppa=self.qppa, line_density=self.line_density,
                        asr=self.asr, dipdip=self.dipdip, verbose=self.verbose) as g:

                phbst_file, phdos_file = g[0], g[1]
                self.abi_phbands = phbst_file.phbands
                # The q-points passed to phonopy.
                self.py_qpoints = [[qpt.frac_coords for qpt in self.abi_phbands.qpoints]]

        data = {}
        data["ddb"] = self.abi_phbands.get_phfreqs_stats_dict()
        d_nn = data["nn_names"] = {}
        for nn_name in self.nn_names:
            try:
                d = self._run_nn_name(nn_name)
                d_nn[nn_name] = d
                d_nn[nn_name]["exception"] = None
            except Exception as exc:
                #raise
                cprint_traceback()
                d_nn[nn_name] = {}
                d_nn[nn_name]["exception"] = str(exc)

        self.write_json("data.json", data, info="JSON file with final results.")
        self._finalize()

    def _run_nn_name(self, nn_name: str) -> None:
        """
        Run calculation for a single NN potential.
        """
        workdir = self.workdir

        calculator = CalcBuilder(nn_name).get_calculator()
        atoms = self.initial_atoms.copy()
        atoms.calc = calculator
        natom = len(atoms)

        if self.relax_mode == RX_MODE.cell:
            raise RuntimeError("One should not relax the cell when comparing ML phonons with a DDB file!")

        if self.relax_mode != RX_MODE.no:
            print(f"Relaxing input DDB atoms with relax mode: {self.relax_mode}.")
            relax_kws = dict(optimizer=self.optimizer,
                             relax_mode=self.relax_mode,
                             fmax=self.fmax,
                             pressure=self.pressure,
                             steps=self.steps,
                             traj_path=self.get_path(f"{nn_name}_relax.traj", "ASE relax trajectory"),
                             verbose=self.verbose,
                            )

            relax = relax_atoms(atoms, **relax_kws)
            relax.summarize(["DDB_initial", "DDB_relaxed"])
            #print(relax)

        # Call phonopy to compute phonons with finite difference and ML potential.
        # Include non-analytical term if dipoles are available in the DDB file.
        with Timer(header=f"Calling get_phonopy with {nn_name=}", footer=""):
            phonon = get_phonopy(atoms, self.supercell, calculator,
                                 distance=self.distance, primitive_matrix=None, remove_drift=True)
            if self.abi_nac_params:
                print("Including dipolar term in phonopy using BECS and eps_inf taken from DDB.")
                phonon.nac_params = self.abi_nac_params

        # Save phonopy object in Yaml format.
        phonon.save(filename=workdir / f"phonopy_params.yaml", settings={'force_constants': True})

        from abipy.dfpt.phonons import PhononBands, PhononBandsPlotter
        with Timer(header="Starting phonopy ph-band computation...", footer=""):
            phonon.run_band_structure(self.py_qpoints, with_eigenvectors=True)

        plt = phonon.plot_band_structure()
        plt.savefig(workdir / f"phonopy_{nn_name}_phbands{self.fig_ext}")
        plt.close()

        bands_dict = phonon.get_band_structure_dict()
        nqpt = 0
        py_phfreqs, py_displ_cart = [], []
        for q_list, w_list, eig_list in zip(bands_dict['qpoints'], bands_dict['frequencies'], bands_dict['eigenvectors']):
            nqpt += len(q_list)
            py_phfreqs.extend(w_list)
            py_displ_cart.extend(eig_list)

        py_phfreqs = np.reshape(py_phfreqs, (nqpt, 3*natom)) / abu.eV_to_THz
        py_displ_cart = np.reshape(py_displ_cart, (nqpt, 3*natom, 3*natom))

        # Build abipy phonon bands from phonopy results.
        py_phbands = PhononBands(self.abi_phbands.structure, self.abi_phbands.qpoints, py_phfreqs,
                                 # FIXME: Use phononopy displacement
                                 self.abi_phbands.phdispl_cart,
                                 non_anal_ph=None,
                                 amu=self.abi_phbands.amu,
                                 epsinf=self.abi_phbands.epsinf,
                                 zcart=self.abi_phbands.zcart,
                                 )

        # Compute diff stats.
        mabs_wdiff_ev = np.abs(py_phbands.phfreqs - self.abi_phbands.phfreqs).mean()

        ph_plotter = PhononBandsPlotter(key_phbands=[
            (f"phonopy with {nn_name}", py_phbands),
            ("ABINIT DDB", self.abi_phbands),
        ])
        mae_str = f"MAE {1000 * mabs_wdiff_ev:.3f} meV"
        print(mae_str)
        latex_formula = self.abi_phbands.structure.latex_formula
        ph_plotter.combiplot(show=False, title=f"{latex_formula}: {mae_str}", units="meV",
                             savefig=str(workdir / f"combiplot_{nn_name}.pdf"))

        data = dict(
            mabs_wdiff_ev=mabs_wdiff_ev,
            **py_phbands.get_phfreqs_stats_dict()
        )

        # Compute phonon DOS and generate file with figure.
        #phonon.auto_total_dos(plot=True)
        #plt.savefig(workdir / f"phonopy_{nn_name}_phdos{self.fig_ext}")
        #plt.close()

        # Compute site-project phonon DOS and generate file with figure.
        #phonon.auto_projected_dos(plot=True)
        #plt.savefig(workdir / f"phonopy_{nn_name}_pjdos{self.fig_ext}")
        #plt.close()

        #phonon.run_thermal_properties(t_step=10, t_max=1000, t_min=0)
        #phonon.write_yaml_thermal_properties(filename=workdir / f"phonopy_{nn_name}_thermal_properties.yaml")

        return data


class MlPhonopy(MlBase):
    """
    Compute phonons with phonopy and ML potential.
    """
    def __init__(self, structure, supercell, distance, line_density, qppa,
                 relax_mode, fmax, pressure, steps, optimizer, nn_names,
                 verbose, workdir, prefix=None):
        """
        Args:
            structure: Structure object
            supercell: Supercell dimensions.
            distance: Distance of finite displacements in Angstrom.
            line_density: Defines the a density of k-points per reciprocal atom to plot the phonon dispersion.
            qppa: Defines the homogeneous q-mesh used for the DOS in units of q-points per atom.
            relax_mode: "ions" to relax ions only, "cell" for ions + cell, "no" for no relaxation.
            fmax: tolerance for relaxation convergence. Here fmax is a sum of force and stress forces.
            pressure: Target pressure.
            steps: max number of steps for relaxation.
            optimizer: name of the ASE optimizer to use for relaxation.
            nn_names: String or list of strings defining the NN potential. See also CalcBuilder.
            verbose: Verbosity level.
            workdir: Working directory, None to generate temporary directory automatically.
            prefix: Prefix for workdir.
        """
        super().__init__(workdir, prefix)

        self.initial_atoms = structure.to_ase_atoms()

        self.distance = float(distance)
        #self.asr = asr
        #self.dipdip = dipdip
        self.relax_mode = relax_mode
        RX_MODE.validate(self.relax_mode)
        self.fmax = fmax
        self.pressure = pressure
        self.steps = steps
        self.optimizer = optimizer
        self.nn_names = list_strings(nn_names)
        self.verbose = verbose
        self.supercell = supercell
        self.line_density = line_density
        self.qppa = qppa
        #self.abi_nac_params = {}

    def to_string(self, verbose=0):
        """String representation with verbosity level `verbose`."""
        supercell = self.supercell.tolist()
        s = f"""\

{self.__class__.__name__} parameters:

     supercell     = {supercell}
     distance      = {self.distance}
     qppa          = {self.qppa}
     line_density  = {self.line_density}
     relax_mode    = {self.relax_mode}
     fmax          = {self.fmax}
     steps         = {self.steps}
     optimizer     = {self.optimizer}
     pressure      = {self.pressure}
     nn_names      = {self.nn_names}
     workdir       = {self.workdir}
     verbose       = {self.verbose}
"""
        return s

    def run(self) -> None:
        """Run calculation."""
        data = {}
        d_nn = data["nn_names"] = {}
        for nn_name in self.nn_names:
            try:
                d = self._run_nn_name(nn_name)
                d_nn[nn_name] = d
                d_nn[nn_name]["exception"] = None
            except Exception as exc:
                cprint_traceback()
                d_nn[nn_name] = {}
                d_nn[nn_name]["exception"] = str(exc)

        self.write_json("data.json", data, info="JSON file with final results.")
        self._finalize()

    def _run_nn_name(self, nn_name: str) -> None:
        """
        Run calculation for a single NN potential.
        """
        workdir = self.workdir

        calculator = CalcBuilder(nn_name).get_calculator()
        atoms = self.initial_atoms.copy()
        atoms.calc = calculator
        natom = len(atoms)

        if self.relax_mode != RX_MODE.no:
            print(f"Relaxing input atoms with relax mode: {self.relax_mode}.")
            relax_kws = dict(optimizer=self.optimizer,
                             relax_mode=self.relax_mode,
                             fmax=self.fmax,
                             pressure=self.pressure,
                             steps=self.steps,
                             traj_path=self.get_path(f"{nn_name}_relax.traj", "ASE relax trajectory"),
                             verbose=self.verbose,
                            )

            relax = relax_atoms(atoms, **relax_kws)
            relax.summarize(["initial", "relaxed"])
            print(relax)

        # Call phonopy to compute phonons with finite difference and ML potential.
        # Include non-analytical term if dipoles are available in the DDB file.
        with Timer(header=f"Calling get_phonopy with {nn_name=}", footer=""):
            phonon = get_phonopy(atoms, self.supercell, calculator,
                                 distance=self.distance, primitive_matrix=None, remove_drift=True)

        plt = phonon.auto_band_structure(
              npoints=101,
              with_eigenvectors=False,
              with_group_velocities=False,
              plot=True,
              write_yaml=True,
              filename=workdir / f"{nn_name}_band.yml",
        )
        plt.savefig(workdir / f"phonopy_{nn_name}_phbands{self.fig_ext}")
        plt.close()

        # Save phonopy object in Yaml format.
        phonon.save(filename=workdir / f"phonopy_params.yaml", settings={'force_constants': True})

        # Compute phonon DOS and generate file with figure.
        phonon.auto_total_dos(plot=True)
        plt.savefig(workdir / f"phonopy_{nn_name}_phdos{self.fig_ext}")
        plt.close()

        # Compute site-project phonon DOS and generate file with figure.
        #phonon.auto_projected_dos(plot=True)
        #plt.savefig(workdir / f"phonopy_{nn_name}_pjdos{self.fig_ext}")
        #plt.close()

        # Compute thermal properties.
        phonon.run_thermal_properties(t_step=10, t_max=1000, t_min=0)
        phonon.write_yaml_thermal_properties(filename=workdir / f"phonopy_{nn_name}_thermal_properties.yaml")

        phonon.plot_thermal_properties()
        plt.savefig(workdir / f"phonopy_{nn_name}_thermal_properties{self.fig_ext}")
        plt.close()

        return {}

