# coding: utf-8
from __future__ import annotations

import os
import abc
import numpy as np
import abipy.core.abinit_units as abu

from scipy.interpolate import UnivariateSpline
from monty.collections import dict2namedtuple
from monty.functools import lazy_property
from monty.termcolor import cprint
from pymatgen.analysis.eos import EOS
from abipy.core.func1d import Function1D
from abipy.tools.plotting import add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt
from abipy.tools.typing import Figure
from abipy.electrons.gsr import GsrFile
from abipy.dfpt.ddb import DdbFile
from abipy.dfpt.phonons import PhononBandsPlotter, PhononDos, PhdosFile
from abipy.dfpt.gruneisen import GrunsNcFile


class AbstractQHA(metaclass=abc.ABCMeta):
    """
    Abstract class for the quasi-harmonic approximation analysis.
    Provides some basic methods and plotting utils, plus a converter to write input files for phonopy-qha or to
    generate an instance of phonopy.qha.core.QHA. These can be used to obtain other quantities and plots.
    Does not include electronic entropic contributions for metals.
    """

    def __init__(self, structures, energies, eos_name='vinet', pressure=0):
        """
        Args:
            structures: list of structures at different volumes.
            energies: list of SCF energies for the structures in eV.
            eos_name: string indicating the expression used to fit the energies. See pymatgen.analysis.eos.EOS.
            pressure: value of the pressure in GPa that will be considered in the p*V contribution to the energy.
        """
        self.structures = structures
        self.energies = np.array(energies)
        self.eos = EOS(eos_name)
        self.pressure = pressure

        self.volumes = np.array([s.volume for s in structures])
        self.iv0 = np.argmin(energies)

    def fit_energies(self, tstart=0, tstop=800, num=100):
        """
        Performs a fit of the energies as a function of the volume at different temperatures.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns:
            `namedtuple` with the following attributes::

                tot_en: numpy array with shape (nvols, num) with the energies used for the fit
                fits: list of subclasses of pymatgen.analysis.eos.EOSBase, depending on the type of
                    eos chosen. Contains the fit for the energies at the different temperatures.
                min_en: numpy array with the minimum energies for the list of temperatures
                min_vol: numpy array with the minimum volumes for the list of temperatures
                temp: numpy array with the temperatures considered
        """
        tmesh = np.linspace(tstart, tstop, num)

        # array with phonon energies and shape (n_vol, n_temp)
        ph_energies = self.get_vib_free_energies(tstart, tstop, num)

        tot_en = self.energies[np.newaxis, :].T + ph_energies + self.volumes[np.newaxis, :].T * self.pressure / abu.eVA3_GPa

        # list of fits objects, one for each temperature
        #fits = [self.eos.fit(self.volumes, e) for e in tot_en.T]

        fits = []
        for kt, e in zip(tmesh, tot_en.T):
            try:
                f = self.eos.fit(self.volumes, e)
                fits.append(f)
            except Exception as exc:
                msg = f"""
EOS fit failed for T={kt} with exception:

    {str(exc)}

Very likely the minimum volume is not in the input range.
Try to change the temperature range with the `tstart`, `tstop` optional arguments
"""
                cprint(msg, color="red")

        # list of minimum volumes and energies, one for each temperature
        min_volumes = np.array([fit.v0 for fit in fits])
        min_energies = np.array([fit.e0 for fit in fits])

        return dict2namedtuple(tot_en=tot_en, fits=fits, min_en=min_energies, min_vol=min_volumes, temp=tmesh)

    @abc.abstractmethod
    def get_vib_free_energies(self, tstart=0, tstop=800, num=100):
        """
        Generates the vibrational free energy corresponding to all the structures.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns: A numpy array of `num` values of the vibrational contribution to the free energy
        """

    @abc.abstractmethod
    def get_thermodynamic_properties(self, tstart=0, tstop=800, num=100):
        """
        Generates all the thermodynamic properties corresponding to all the volumes.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns:
            `namedtuple` with the following attributes for all the volumes:

                tmesh: numpy array with the list of temperatures. Shape (num).
                cv: constant-volume specific heat, in eV/K. Shape (nvols, num).
                free_energy: free energy, in eV. Shape (nvols, num).
                entropy: entropy, in eV/K. Shape (nvols, num).
                zpe: zero point energy in eV. Shape (nvols).
        """

    @property
    def nvols(self) -> int:
        """Number of volumes"""
        return len(self.volumes)

    @property
    def natoms(self) -> int:
        """Number of atoms in the unit cell."""
        return len(self.structures[0])

    def set_eos(self, eos_name: str) -> None:
        """
        Set the EOS model used for the fit.

        Args:
            eos_name: string indicating the expression used to fit the energies. See pymatgen.analysis.eos.EOS.
        """
        self.eos = EOS(eos_name)

    @add_fig_kwargs
    def plot_energies(self, tstart=0, tstop=800, num=10, ax=None, **kwargs) -> Figure:
        """
        Plots the energies as a function of volume at different temperatures.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 10.
            ax: |matplotlib-Axes| or None if a new figure should be created.

        Returns: |matplotlib-Figure|
        """
        f = self.fit_energies(tstart, tstop, num)

        ax, fig, plt = get_ax_fig_plt(ax)
        xmin, xmax = np.floor(self.volumes.min() * 0.97), np.ceil(self.volumes.max() * 1.03)
        x = np.linspace(xmin, xmax, 100)

        for fit, e, t in zip(f.fits, f.tot_en.T - self.energies[self.iv0], f.temp):
            ax.scatter(self.volumes, e, label=t, color='b', marker='x', s=5)
            ax.plot(x, fit.func(x) - self.energies[self.iv0], color='b', lw=1)

        ax.plot(f.min_vol, f.min_en - self.energies[self.iv0], color='r', lw=1, marker='x', ms=5)

        ax.set_xlabel(r'V (${\AA}^3$)')
        ax.set_ylabel('E (eV)')
        #ax.grid(True)

        return fig

    def get_thermal_expansion_coeff(self, tstart=0, tstop=800, num=100) -> Function1D:
        """
        Calculates the thermal expansion coefficient as a function of temperature, using
        finite difference on the fitted values of the volume as a function of temperature.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns: |Function1D|
        """
        f = self.fit_energies(tstart, tstop, num)

        dt = f.temp[1] - f.temp[0]
        alpha = (f.min_vol[2:] - f.min_vol[:-2]) / (2 * dt) / f.min_vol[1:-1]

        return Function1D(f.temp[1:-1], alpha)

    @add_fig_kwargs
    def plot_thermal_expansion_coeff(self, tstart=0, tstop=800, num=100, ax=None, **kwargs) -> Figure:
        """
        Plots the thermal expansion coefficient as a function of the temperature.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.
            ax: |matplotlib-Axes| or None if a new figure should be created.

        Returns: |matplotlib-Figure|
        """
        ax, fig, plt = get_ax_fig_plt(ax)

        if 'linewidth' not in kwargs and 'lw' not in kwargs:
            kwargs['linewidth'] = 2

        if "color" not in kwargs:
            kwargs["color"] = "b"

        alpha = self.get_thermal_expansion_coeff(tstart, tstop, num)

        ax.plot(alpha.mesh, alpha.values, **kwargs)
        ax.set_xlabel(r'T (K)')
        ax.set_ylabel(r'$\alpha$ (K$^{-1}$)')
        ax.grid(True)

        ax.set_xlim(tstart, tstop)
        ax.get_yaxis().get_major_formatter().set_powerlimits((0, 0))

        return fig

    @add_fig_kwargs
    def plot_vol_vs_t(self, tstart=0, tstop=800, num=100, ax=None, **kwargs) -> Figure:
        """
        Plot the volume as a function of the temperature.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.
            ax: |matplotlib-Axes| or None if a new figure should be created.

        Returns: |matplotlib-Figure|
        """
        ax, fig, plt = get_ax_fig_plt(ax)
        f = self.fit_energies(tstart, tstop, num)

        if 'linewidth' not in kwargs and 'lw' not in kwargs:
            kwargs['linewidth'] = 2

        if "color" not in kwargs:
            kwargs["color"] = "b"

        ax.plot(f.temp, f.min_vol, **kwargs)
        ax.set_xlabel('T (K)')
        ax.set_ylabel(r'V (${\AA}^3$)')
        ax.set_xlim(tstart, tstop)
        ax.grid(True)

        return fig

    @add_fig_kwargs
    def plot_phbs(self, phbands, temperatures=None, t_max=1000, colormap="plasma", **kwargs) -> Figure:
        """
        Given a list of |PhononBands| plots the band structures with a color depending on
        the temperature using a |PhononBandsPlotter|.
        If temperatures are not given they will be deduced inverting the dependence of volume
        with respect to the temperature. If a unique volume could not be identified an error will
        be raised.

        Args:
            phbands: list of |PhononBands| objects.
            temperatures: list of temperatures.
            t_max: maximum temperature considered for plotting.
            colormap: matplotlib color map.

        Returns: |matplotlib-Figure|
        """
        if temperatures is None:
            tv = self.get_t_for_vols([b.structure.volume for b in phbands], t_max=t_max)
            temperatures = []
            for b, t in zip(phbands, tv):
                if len(t) != 1:
                    raise ValueError("Couldn't find a single temperature for structure with "
                                     "volume {}. Found {}: {}".format(b.structure.volume, len(t), list(t)))
                temperatures.append(t[0])

        temperatures_str = ["{:.0f} K".format(t) for t in temperatures]

        import matplotlib.pyplot as plt
        cmap = plt.get_cmap(colormap)
        colors = [cmap(t / max(temperatures)) for t in temperatures]
        labels_phbs = zip(temperatures_str, phbands)

        pbp = PhononBandsPlotter(labels_phbs)
        pbp._LINE_COLORS = colors
        pbp._LINE_STYLES = ['-']

        fig = pbp.combiplot(show=False, **kwargs)

        return fig

    def get_vol_at_t(self, t) -> float:
        """
        Calculates the volume corresponding to a specific temperature.

        Args:
            t: temperature in K

        Returns: The volume in Ang^3.
        """
        f = self.fit_energies(t, t, 1)

        return f.min_vol[0]

    def get_t_for_vols(self, vols, t_max=1000) -> list[float]:
        """
        Find the temperatures corresponding to a specific volume.
        The search is performed interpolating the V(T) dependence with a spline and
        finding the roots with of V(t) - v.
        It may return more than one temperature for a volume in case of non monotonic behavior.

        Args:
            vols: list of volumes
            t_max: maximum temperature considered for the fit

        Returns:
            A list of lists of temperatures. For each volume more than one temperature can
            be identified.
        """
        if not isinstance(vols, (list, tuple, np.ndarray)):
            vols = [vols]

        f = self.fit_energies(0, t_max, t_max+1)

        temps = []
        for v in vols:
            spline = UnivariateSpline(f.temp, f.min_vol - v, s=0)
            temps.append(spline.roots())

        return temps

    def write_phonopy_qha_inputs(self, tstart=0, tstop=2100, num=211, path=None) -> None:
        """
        Writes nvols thermal_properties-i.yaml files that can be used as inputs for phonopy-qha.
        Notice that phonopy apparently requires the value of the 300 K temperature to be present
        in the list. Choose the values of tstart, tstop and num to satisfy this condition.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 211.
            path: a path to a folder where the files will be stored
        """
        if path is None:
            path = os.getcwd()

        thermo = self.get_thermodynamic_properties(tstart=tstart, tstop=tstop, num=num)

        np.savetxt(os.path.join(path, 'e-v.dat'), np.array([self.volumes, self.energies]).T, fmt='%.10f')

        # generator for thermal_properties.yaml extracted from phonopy:
        # phonopy.phonon.thermal_properties.ThermalProperties._get_tp_yaml_lines
        for j in range(self.nvols):
            lines = []
            lines.append("# Thermal properties / unit cell (natom)")
            lines.append("")
            lines.append("unit:")
            lines.append("  temperature:   K")
            lines.append("  free_energy:   kJ/mol")
            lines.append("  entropy:       J/K/mol")
            lines.append("  heat_capacity: J/K/mol")
            lines.append("")
            lines.append("natom: %5d" % (self.natoms))

            lines.append("num_modes: %d" % (3 * self.natoms))
            lines.append("num_integrated_modes: %d" % (3 * self.natoms))

            lines.append("")
            lines.append("zero_point_energy: %15.7f" % (thermo.zpe[j] * abu.e_Cb * abu.Avogadro))
            lines.append("high_T_entropy:    %15.7f" % 0) # high_T_entropy is not used in QHA
            lines.append("")
            lines.append("thermal_properties:")
            fe = thermo.free_energy[j] * abu.e_Cb * abu.Avogadro
            entropy = thermo.entropy[j] * abu.e_Cb * abu.Avogadro
            cv = thermo.cv[j] * abu.e_Cb * abu.Avogadro
            temperatures = thermo.tmesh
            for i, t in enumerate(temperatures):
                lines.append("- temperature:   %15.7f" % t)
                lines.append("  free_energy:   %15.7f" % (fe[i] / 1000))
                lines.append("  entropy:       %15.7f" % entropy[i])
                # Sometimes 'nan' of C_V is returned at low temperature.
                if np.isnan(cv[i]):
                    lines.append("  heat_capacity: %15.7f" % 0)
                else:
                    lines.append("  heat_capacity: %15.7f" % cv[i])
                lines.append("  energy:        %15.7f" %
                             (fe[i] / 1000 + entropy[i] * t / 1000))

                lines.append("")

            with open(os.path.join(path, "thermal_properties-{}.yaml".format(j)), 'wt') as f:
                f.write("\n".join(lines))

    def get_phonopy_qha(self, tstart=0, tstop=2100, num=211, eos='vinet', t_max=None,
                        energy_plot_factor=None, pressure=None):
        """
        Creates an instance of phonopy.qha.core.QHA that can be used generate further plots and output data.
        The object is returned right after the construction. The "run()" method should be executed
        before getting results and plots.
        Notice that phonopy apparently requires the value of the 300 K temperature to be present
        in the list. Choose the values of tstart, tstop and num to satisfy this condition.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 211.
            eos: the expression used to fit the energies in phonopy. Possible values are "vinet",
                "murnaghan" and "birch_murnaghan". Passed to phonopy's QHA.
            t_max: maximum temperature. Passed to phonopy's QHA.
            energy_plot_factor: factor multiplying the energies. Passed to phonopy's QHA.
            pressure: pressure value, passed to phonopy.

        Returns: An instance of phonopy.qha.core.QHA
        """
        try:
            from phonopy.qha.core import QHA as QHA_phonopy
        except ImportError as exc:
            print("Phonopy is required to generate the QHA phonopy object")
            raise exc

        thermo = self.get_thermodynamic_properties(tstart=tstart, tstop=tstop, num=num)

        fe = thermo.free_energy.T * abu.e_Cb * abu.Avogadro / 1000
        entropy = thermo.entropy.T * abu.e_Cb * abu.Avogadro
        cv = thermo.cv.T * abu.e_Cb * abu.Avogadro
        temperatures = thermo.tmesh

        en = self.energies + self.volumes * self.pressure / abu.eVA3_GPa

        qha_p = QHA_phonopy(
            volumes=self.volumes,
            electronic_energies=en,
            temperatures=temperatures,
            cv=cv,
            entropy=entropy,
            fe_phonon=fe,
            pressure=pressure,
            eos=eos,
            t_max=t_max,
            energy_plot_factor=energy_plot_factor
        )

        return qha_p


class QHA(AbstractQHA):
    """
    Object to extract results in the quasi-harmonic approximation from electronic and phonon calculations
    at different volumes.
    Provides some basic methods and plotting utils, plus a converter to write input files for phonopy-qha or to
    generate an instance of phonopy.qha.core.QHA. These can be used to obtain other quantities and plots.
    Does not include electronic entropic contributions for metals.

    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: QHA
    """

    @classmethod
    def from_files(cls, gsr_paths, phdos_paths):
        """
        Creates an instance of QHA from a list of GSR files and a list of PHDOS.nc files.
        The list should have the same size and the volumes should match.

        Args:
            gsr_paths: list of paths to GSR files.
            phdos_paths: list of paths to PHDOS.nc files.

        Returns: A new instance of QHA
        """
        energies = []
        structures = []
        for gp in gsr_paths:
            with GsrFile.from_file(gp) as g:
                energies.append(g.energy)
                structures.append(g.structure)

        #doses = [PhononDos.as_phdos(dp) for dp in phdos_paths]

        doses = []
        structures_from_phdos = []
        for path in phdos_paths:
            with PhdosFile(path) as p:
                doses.append(p.phdos)
                structures_from_phdos.append(p.structure)

        cls._check_volumes(structures, structures_from_phdos)

        return cls(structures, doses, energies)

    def __init__(self, structures, doses, energies, eos_name='vinet', pressure=0):
        """
        Args:
            structures: list of structures at different volumes.
            doses: list of |PhononDos| at volumes corresponding to the structures.
            energies: list of SCF energies for the structures in eV.
            eos_name: string indicating the expression used to fit the energies. See pymatgen.analysis.eos.EOS.
            pressure: value of the pressure in GPa that will be considered in the p * V contribution to the energy.
        """
        super().__init__(structures=structures, energies=energies, eos_name=eos_name, pressure=pressure)
        self.doses = doses

    @staticmethod
    def _check_volumes(struct_list1, struct_list2):

        lens = (len(struct_list1), len(struct_list2))
        if lens[0] != lens[1]:
            raise RuntimeError("Expecting lists with same number of structures. Got %s" % str(lens))
        vols1 = [s.volume for s in struct_list1]
        vols2 = [s.volume for s in struct_list2]
        ierr = 0
        for v1, v2 in zip(vols1, vols2):
            if abs(v1 - v2) > 1e-3:
                ierr += 1
                print("Volume1: %s != Volume2: %s" % (v1, v2))

        if ierr != 0:
            raise RuntimeError("Expecting lists with same volumes!")

    def get_vib_free_energies(self, tstart=0, tstop=800, num=100) -> np.ndarray:
        """
        Generates the vibrational free energy from the phonon DOS.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns: A numpy array of `num` values of the vibrational contribution to the free energy
        """
        f = np.zeros((self.nvols, num))

        for i, dos in enumerate(self.doses):
            f[i] = dos.get_free_energy(tstart, tstop, num).values

        return f

    def get_thermodynamic_properties(self, tstart=0, tstop=800, num=100):
        """
        Generates all the thermodynamic properties corresponding to all the volumes using the phonon DOS.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns:
            `namedtuple` with the following attributes for all the volumes:

                tmesh: numpy array with the list of temperatures. Shape (num).
                cv: constant-volume specific heat, in eV/K. Shape (nvols, num).
                free_energy: free energy, in eV. Shape (nvols, num).
                entropy: entropy, in eV/K. Shape (nvols, num).
                zpe: zero point energy in eV. Shape (nvols).
        """
        tmesh = np.linspace(tstart, tstop, num)
        cv = np.zeros((self.nvols, num))
        free_energy = np.zeros((self.nvols, num))
        entropy = np.zeros((self.nvols, num))
        internal_energy = np.zeros((self.nvols, num))
        zpe = np.zeros(self.nvols)

        for i, d in enumerate(self.doses):
            cv[i] = d.get_cv(tstart, tstop, num).values
            free_energy[i] = d.get_free_energy(tstart, tstop, num).values
            entropy[i] = d.get_entropy(tstart, tstop, num).values
            zpe[i] = d.zero_point_energy

        return dict2namedtuple(tmesh=tmesh, cv=cv, free_energy=free_energy, entropy=entropy,
                               zpe=zpe)


class QHA3PF(AbstractQHA):
    """
    Object to extract results in the quasi-harmonic approximation from several electronic energies at different
    volumes and three phonon calculations.
    Provides some basic methods and plotting utils, plus a converter to write input files for phonopy-qha or to
    generate an instance of phonopy.qha.core.QHA. These can be used to obtain other quantities and plots.
    Does not include electronic entropic contributions for metals.

    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: QHA3PF
    """

    @classmethod
    def from_files(cls, gsr_paths, phdos_paths, ind_doses) -> QHA3PF:
        """
        Creates an instance of QHA from a list of GSR files and a list of PHDOS.nc files.
        The list should have the same size and the volumes should match.

        Args:
            gsr_paths: list of paths to GSR files.
            phdos_paths: list of paths to three PHDOS.nc files.
            ind_doses: list of three values indicating, for each of the three doses, the index of the
                corresponding gsr_file in "gsr_paths".

        Returns: A new instance of QHA3PF
        """
        energies = []
        structures = []
        for gp in gsr_paths:
            with GsrFile.from_file(gp) as g:
                energies.append(g.energy)
                structures.append(g.structure)

        doses = [PhononDos.as_phdos(dp) for dp in phdos_paths]

        return cls(structures, doses, energies, ind_doses)

    def __init__(self, structures, doses, energies, ind_doses, eos_name='vinet', pressure=0, fit_degree=2):
        """
        Args:
            structures: list of structures at different volumes.
            doses: list of three |PhononDos| at different volumes corresponding to the some of the structures.
            energies: list of SCF energies for the structures in eV.
            ind_doses: list of three values indicating, for each of the three doses, the index of the
                corresponding structure in "structures".
            eos_name: string indicating the expression used to fit the energies. See pymatgen.analysis.eos.EOS.
            pressure: value of the pressure in GPa that will be considered in the p*V contribution to the energy.
        """
        super().__init__(structures=structures, energies=energies, eos_name=eos_name, pressure=pressure)
        self.doses = doses
        self.ind_doses = ind_doses
        self.fit_degree = fit_degree
        self._ind_energy_only = [i for i in range(len(structures)) if i not in ind_doses]

    def get_thermodynamic_properties(self, tstart=0, tstop=800, num=100):
        """
        Generates all the thermodynamic properties corresponding to all the volumes using the phonon DOS.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns:
            `namedtuple` with the following attributes for all the volumes:

                tmesh: numpy array with the list of temperatures. Shape (num).
                cv: constant-volume specific heat, in eV/K. Shape (nvols, num).
                free_energy: free energy, in eV. Shape (nvols, num).
                entropy: entropy, in eV/K. Shape (nvols, num).
                zpe: zero point energy in eV. Shape (nvols).
        """
        tmesh = np.linspace(tstart, tstop, num)
        cv = self._get_thermodynamic_prop("cv", tstart, tstop, num)
        free_energy = self._get_thermodynamic_prop("free_energy", tstart, tstop, num)
        entropy = self._get_thermodynamic_prop("entropy", tstart, tstop, num)
        zpe = np.zeros(self.nvols)

        for i, dos in zip(self.ind_doses, self.doses):
            zpe[i] = dos.zero_point_energy

        dos_vols = self.volumes[self.ind_doses]
        missing_vols = self.volumes[self._ind_energy_only]

        fit_params = np.polyfit(dos_vols, zpe[self.ind_doses], self.fit_degree)
        zpe[self._ind_energy_only] = np.poly1d(fit_params)(missing_vols)

        return dict2namedtuple(tmesh=tmesh, cv=cv, free_energy=free_energy, entropy=entropy,
                               zpe=zpe)

    def _get_thermodynamic_prop(self, name, tstart, tstop, num):
        """
        Helper function to get a generic thermodynamic property for all the volumes.

        Args:
            name: name of the property to calculate. Possible values in "internal_energy",
                "free_energy", "entropy", "c_v".
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns:
            Numpy array with the values of the thermodynamic properties at the different
            volumes with size (nvols, num).
        """
        prop_doses = np.array([getattr(dos, "get_" + name)(tstart, tstop, num).values for dos in self.doses])

        p = np.zeros((self.nvols, num))

        for i, prop_dos in zip(self.ind_doses, prop_doses):
            p[i] = prop_dos

        dos_vols = self.volumes[self.ind_doses]
        missing_vols = self.volumes[self._ind_energy_only]

        # generate fit objects from the known dos values
        for i_temp in range(num):
            fit_params = np.polyfit(dos_vols, prop_doses[:, i_temp], self.fit_degree)
            p[self._ind_energy_only, i_temp] = np.poly1d(fit_params)(missing_vols)

        return p

    def get_vib_free_energies(self, tstart=0, tstop=800, num=100):
        """
        Generates the vibrational free energy corresponding to all the structures, either from the phonon DOS
        or from a fit of the known values.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns:
            A numpy array of `num` values of the vibrational contribution to the free energy
        """

        return self._get_thermodynamic_prop("free_energy", tstart, tstop, num)


class QHA3P(AbstractQHA):
    """
    Object to extract results in the quasi-harmonic approximation from several electronic energies at different
    volumes and three phonon calculations.
    Provides some basic methods and plotting utils, plus a converter to write input files for phonopy-qha or to
    generate an instance of phonopy.qha.core.QHA. These can be used to obtain other quantities and plots.
    Does not include electronic entropic contributions for metals.

    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: QHA3P
    """

    @classmethod
    def from_files(cls, gsr_paths, grun_file_path, ind_doses) -> QHA3P:
        """
        Creates an instance of QHA from a list of GSR files and a list PHDOS.nc files.
        The list should have the same size and the volumes should match.

        Args:
            gsr_paths: list of paths to GSR files.
            phdos_paths: list of paths to three PHDOS.nc files.
            ind_doses: list of three values indicating, for each of the three doses, the index of the
                corresponding gsr_file in "gsr_paths".

        Returns: A new instance of QHA3P
        """
        energies = []
        structures = []
        for gp in gsr_paths:
            with GsrFile.from_file(gp) as g:
                energies.append(g.energy)
                structures.append(g.structure)

        gruns = GrunsNcFile(grun_file_path)

        return cls(structures, gruns, energies, ind_doses)

    def __init__(self, structures, gruns, energies, ind_grun, eos_name='vinet', pressure=0):
        """
        Args:
            structures: list of structures at different volumes.
            doses: list of three |PhononDos| at different volumes corresponding to the some of the structures.
            energies: list of SCF energies for the structures in eV.
            ind_grun: list of three values indicating, for each of the three doses, the index of the
                corresponding structure in "structures".
            eos_name: string indicating the expression used to fit the energies. See pymatgen.analysis.eos.EOS.
            pressure: value of the pressure in GPa that will be considered in the p*V contribution to the energy.
        """

        super().__init__(structures=structures, energies=energies, eos_name=eos_name, pressure=pressure)
        self.grun = gruns
        self.ind_grun = ind_grun
        self._ind_energy_only = [i for i in range(len(structures)) if i not in ind_grun]

    def close(self) -> None:
        """Close files."""
        self.grun.close()

    def get_thermodynamic_properties(self, tstart=0, tstop=800, num=100):
        """
        Generates the thermodynamic properties corresponding to all the structures, either from the phonon
        frequencies or from a fit of the know values.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns:
            `namedtuple` with the following attributes for all the volumes:

                tmesh: numpy array with the list of temperatures. Shape (num).
                cv: constant-volume specific heat, in eV/K. Shape (nvols, num).
                free_energy: free energy, in eV. Shape (nvols, num).
                entropy: entropy, in eV/K. Shape (nvols, num).
                zpe: zero point energy in eV. Shape (nvols).
        """
        w = self.fitted_frequencies

        tmesh = np.linspace(tstart, tstop, num)
        weights = self.grun.doses['qpoints'].weights

        free_energy = np.zeros((self.nvols, num))
        cv = np.zeros((self.nvols, num))
        entropy = np.zeros((self.nvols, num))
        zpe = np.zeros(self.nvols)

        for i in range(self.nvols):
            free_energy[i] = [get_free_energy(w[i], weights, t) for t in tmesh]
            cv[i] = [get_cv(w[i], weights, t) for t in tmesh]
            entropy[i] = [get_entropy(w[i], weights, t) for t in tmesh]
            zpe[i] = get_zero_point_energy(w, weights)

        return dict2namedtuple(tmesh=tmesh, cv=cv, free_energy=free_energy, entropy=entropy,
                               zpe=zpe)

    @lazy_property
    def fitted_frequencies(self) -> np.ndarray:
        """
        A numpy array with size (nvols, nqpts_ibz, 3*natoms) containing the phonon frequencies
        for all the volumes, either from the original phonon calculations or fitted
        """
        w = self.grun.wvols_qibz
        dv = np.abs(self.volumes[self.ind_grun[2]] - self.volumes[self.ind_grun[1]])

        w0 = w[:, 0, :]
        w2 = w[:, 2, :]
        d1 = (w2 - w0) / (2 * dv)
        w1 = w[:, 1, :]
        d2 = (w0 - 2 * w1 + w2) / dv ** 2
        v0 = self.volumes[self.ind_grun[1]]

        w_full = np.zeros((self.nvols, w.shape[0], w.shape[2]))

        for i, ig in enumerate(self.ind_grun):
            w_full[ig] = w[:, i, :]

        for i in self._ind_energy_only:
            dv = self.volumes[i] - v0
            w_full[i] = w1 + d1 * dv + d2 * dv ** 2 / 2

        return w_full

    def get_vib_free_energies(self, tstart=0, tstop=800, num=100) -> np.ndarray:
        """
        Generates the vibrational free energy corresponding to all the structures, either from the phonon DOS
        or from a fit of the know values.

        Args:
            tstart: The starting value (in Kelvin) of the temperature mesh.
            tstop: The end value (in Kelvin) of the mesh.
            num: int, optional Number of samples to generate. Default is 100.

        Returns: A numpy array of `num` values of the vibrational contribution to the free energy
        """
        w = self.fitted_frequencies

        tmesh = np.linspace(tstart, tstop, num)
        weights = self.grun.doses['qpoints'].weights

        f = np.zeros((self.nvols, num))

        for i in range(self.nvols):
            f[i] = [get_free_energy(w[i], weights, t) for t in tmesh]

        return f


def get_free_energy(w, weights, t):
    """
    Calculates the free energy in eV from the phonon frequencies on a regular grid.

    Args:
         w: the phonon frequencies
         weights: the weights of the q-points
         t: the temperature
    """
    wdkt = w / (abu.kb_eVK * t)

    # if w=0 set fe=0
    fe = np.choose(w > 0, (0, w / 2 + abu.kb_eVK * t * np.log(1 - np.exp(-wdkt))))
    ind = np.where(w >= 0)

    return np.dot(weights[ind[0]], fe[ind]).sum()


def get_cv(w, weights, t):
    """
    Calculates the constant-volume specific heat in eV/K from the phonon frequencies on a regular grid.

    Args:
         w: the phonon frequencies
         weights: the weights of the q-points
         t: the temperature
    """
    wdkt = w / (abu.kb_eVK * t)

    # if w=0 set cv=0
    cv = np.choose(w > 0, (0, abu.kb_eVK * wdkt ** 2 * np.exp(wdkt) / (np.exp(wdkt) - 1) ** 2))
    ind = np.where(w >= 0)

    return np.dot(weights[ind[0]], cv[ind]).sum()


def get_zero_point_energy(w, weights):
    """
    Calculates the zero point energy in eV from the phonon frequencies on a regular grid.

    Args:
         w: the phonon frequencies
         weights: the weights of the q-points
    """
    zpe = np.choose(w > 0, (0, w / 2))
    ind = np.where(w >= 0)

    return np.dot(weights[ind[0]], zpe[ind]).sum()


def get_entropy(w, weights, t):
    """
    Calculates the entropy in eV/K from the phonon frequencies on a regular grid.

    Args:
         w: the phonon frequencies
         weights: the weights of the q-points
         t: the temperature
    """
    wd2kt = w / (2 * abu.kb_eVK * t)
    coth = lambda x: 1.0 / np.tanh(x)

    s = np.choose(w > 0, (0, w*coth(wd2kt) / 2 / t - abu.kb_eVK * np.log(2 * np.sinh(wd2kt))))
    ind = np.where(w >= 0)

    return np.dot(weights[ind[0]], s[ind]).sum()


class AbstractQmeshAnalyzer(metaclass=abc.ABCMeta):
    """
    Abstract class for the analysis of the convergence wrt to the q-mesh used to compute the phonon DOS. 
    Relies on abstract methods implemented in AbstractQHA.
    """

    fontsize = 8

    colormap = "viridis"

    def _consistency_check(self):
        if not hasattr(self, "qha_list") or not self.qha_list:
            raise RuntimeError("Please call the run method to compute the QHA!")

    def set_eos(self, eos_name: str):
        """
        Set the EOS model used for the fit.

        Args:
            eos_name: string indicating the expression used to fit the energies. See pymatgen.analysis.eos.EOS.
        """
        self._consistency_check()
        for qha in self.qha_list:
            qha.set_eos(eos_name)

    @add_fig_kwargs
    def plot_energies(self, **kwargs) -> Figure:
        """
        Plots the energies as a function of volume at different temperatures.
        kwargs are propagated to the analogous method of QHA.
        """
        self._consistency_check()
        # Build grid of plots.
        ax_list, fig, plt = get_axarray_fig_plt(None, nrows=self.num_qmeshes, ncols=1,
                                                sharex=True, sharey=True, squeeze=False)
        ax_list = ax_list.ravel()

        for qha, ngqpt, ax in zip(self.qha_list, self.ngqpt_list, ax_list):
            qha.plot_energies(ax=ax, show=False, **kwargs)
            ax.set_title("ngpqt: %s" % str(ngqpt), fontsize=self.fontsize)
        return fig

    @add_fig_kwargs
    def plot_thermal_expansion_coeff(self, **kwargs) -> Figure:
        """
        Plots the thermal expansion coefficient as a function of the temperature.
        kwargs are propagated to the analogous method of QHA.
        """
        self._consistency_check()
        ax, fig, plt = get_ax_fig_plt(None)
        cmap = plt.get_cmap(self.colormap)
        for iq, (qha, ngqpt) in enumerate(zip(self.qha_list, self.ngqpt_list)):
            qha.plot_thermal_expansion_coeff(ax=ax,
                                             color=cmap(float(iq) / self.num_qmeshes),
                                             label="ngqpt: %s" % str(ngqpt), show=False, **kwargs)
        ax.legend(loc="best", fontsize=self.fontsize, shadow=True)
        return fig

    @add_fig_kwargs
    def plot_vol_vs_t(self, **kwargs) -> Figure:
        """
        Plot the volume as a function of the temperature.
        kwargs are propagated to the analogous method of QHA.
        """
        self._consistency_check()
        ax, fig, plt = get_ax_fig_plt(None)
        cmap = plt.get_cmap(self.colormap)
        for iq, (qha, ngqpt) in enumerate(zip(self.qha_list, self.ngqpt_list)):
            qha.plot_vol_vs_t(ax=ax,
                              color=cmap(float(iq) / self.num_qmeshes),
                              label="ngqpt: %s" % str(ngqpt), show=False, **kwargs)
        ax.legend(loc="best", fontsize=self.fontsize, shadow=True)
        return fig


class QHAQmeshAnalyzer(AbstractQmeshAnalyzer):

    def __init__(self, gsr_paths, ddb_paths):
        """
        Creates an instance of QHA from a list of GSR files and a list of PHDOS.nc files.
        The list should have the same size and the volumes should match.

        Args:
            gsr_paths: list of paths to GSR files.
            phdos_paths: list of paths to PHDOS.nc files.
        """
        self.gsr_paths = gsr_paths
        self.ddb_paths = ddb_paths

    def run_qlist(self, nqsmall_list, **kwargs):
        """
        """
        self.qha_list = []
        self.ngqpt_list = []

        ddb_list = [DdbFile(p) for p in self.ddb_paths]
        for nqsmall in nqsmall_list:
            phdos_paths = []

            for i, ddb in enumerate(ddb_list):
                phbst_file, phdos_file = ddb.anaget_phbst_and_phdos_files(
                    nqsmall=nqsmall, qppa=None, ndivsm=1, line_density=None, asr=2, chneut=1, dipdip=1,
                    dos_method="tetra", lo_to_splitting="automatic", ngqpt=None, qptbounds=None,
                    anaddb_kwargs=None, verbose=0, spell_check=True,
                    mpi_procs=1, workdir=None, manager=None)

                phdos_paths.append(phdos_file.filepath)
                if i == 0:
                    # These variables added in abinit v8.11. Use nqsmall is not available.
                    ngqpt = 3 * [nqsmall]
                    if "qptrlatt" in phdos_file.reader.rootgrp.variables:
                        ngqpt = np.diagonal(phdos_file.reader.read_value("qptrlatt").T)
                        #shiftq = phdos_file.reader.read_value("shiftq")
                    self.ngqpt_list.append(ngqpt)

                phbst_file.close()
                phdos_file.close()

            qha = QHA.from_files(self.gsr_paths, phdos_paths)
            self.qha_list.append(qha)

        self.ngqpt_list = np.reshape(self.ngqpt_list, (-1, 3))
        self.num_qmeshes = len(self.ngqpt_list)
        for ddb in ddb_list:
            ddb.close()
