"""
Thermodynamic properties
========================
.. moduleauthor:: Benjamin Ye <GitHub: @bbye98>

This module contains classes to evaluate thermodynamic properties of
systems, such as the constant-volume heat capacity.
"""

from io import StringIO
from pathlib import Path
from typing import Union
import warnings

import numpy as np
import pandas as pd

from .. import FOUND_OPENMM, Q_, ureg
from ..algorithm.unit import strip_unit
from .base import Hash

if FOUND_OPENMM:
    from openmm import unit


class ConstantVolumeHeatCapacity:
    """
    A serial implementation to caluclate the constant-volume heat
    capacity :math:`C_V` for a canonical (:math:`NVT`) system.

    The constant-volume heat capacity is defined as

    .. math::

       C_V=\\frac{\\langle U^2\\rangle-\\langle U\\rangle^2}
       {k_\\mathrm{B}T^2}

    where :math:`U` is the total potential energy of the system,
    :math:`\\langle\\cdot\\rangle` denotes the ensemble average,
    :math:`k_\\mathrm{B}` is the Boltzmann constant, and :math:`T` is
    the system temperature.

    Parameters
    ----------
    log_file : `str` or `pathlib.Path`, optional
        Log file generated by the simulation. If not provided, the
        potential energies must be provided directly in `energies`.

    log_format : `str`, optional
        Format of the log file. If not provided, the format will be
        determined automatically (if possible).

        **Valid values**: :code:`"lammps"`, :code:`"openmm"`.

    energies : `numpy.ndarray` or `pint.Quantity`, optional
        Potential or total energies :math:`U`. If not provided, the
        log file must be provided in `log_file`.

    temperature : `float`, `openmm.unit.Quantity`, or `pint.Quantity`, \
    optional
        System temperature :math:`T`. If not provided, the averaged
        temperature from the log file (if available) is used.

    reduced : `bool`, keyword-only, default: :code:`False`
        Specifies whether the data is in reduced units.

    sep : `str`, keyword-only, default: :code:`","`
        Delimiter for OpenMM log files.

    Attributes
    ----------
    results.units : `dict`
        Reference units for the results. For example, to get the
        reference units for `results.temperature`, call
        :code:`results.units["temperature"]`.

    results.energies : `numpy.ndarray`
        Potential or total energies :math:`U`.

        **Reference unit**: :math:`\\mathrm{kcal/mol}` or
        :math:`\\mathrm{kJ/mol}`.

    results.temperature : `float`
        System temperature :math:`T`.

        **Reference unit**: :math:`\\mathrm{K}`.

    results.heat_capacity : `float`
        Constant-volume heat capacity :math:`C_V`.

        **Reference unit**: :math:`\\mathrm{kcal/K}` or
        :math:`\\mathrm{kJ/K}`.
        
    Example
    --------
    First, this analysis class must be imported:

    >>> from mdcraft.analysis.thermodynamics import ConstantVolumeHeatCapacity

    Then, after loading the log files (or providing the energies and temperature directly), the class can be instantiated:
    
    >>> cv = ConstantVolumeHeatCapacity("log.lammps)
    >>> cv.run()

    The results can be obtained under the `results` attribute:
    
    >>> cv.results.energies
    >>> cv.results.heat_capacity
    
    """

    _COLUMNS = {
        "lammps": {
            "energy": [
                "TotEng",
                "KinEng",
                "PotEng",
                "E_angle",
                "E_bond",
                "E_coul",
                "E_dihed",
                "E_impro",
                "E_long",
                "E_vdwl",
            ],
            "temperature": "Temp",
        },
        "openmm": {
            "energy": [
                "Total Energy (kJ/mole)",
                "Kinetic Energy (kJ/mole)",
                "Potential Energy (kJ/mole)",
            ],
            "temperature": "Temperature (K)",
        },
    }

    def __init__(
        self,
        log_file: Union[str, Path] = None,
        log_format: str = None,
        *,
        energies: Union[np.ndarray[float], "unit.Quantity", Q_] = None,
        temperature: Union[float, "unit.Quantity", Q_] = None,
        reduced: bool = False,
        sep: str = ",",
    ) -> None:

        self.results = Hash(units={})
        self._reduced = reduced

        if energies:
            self.results.units["energies"] = ureg.kilojoule / ureg.mole
            self.results.energies = strip_unit(energies, self.results.units["energy"])[
                0
            ]
        elif log_file:
            self._file = log_file if isinstance(log_file, Path) else Path(log_file)
            with open(self._file, "r") as f:
                log = f.read()

            # Determine the simulation toolkit used to write the log
            # file
            if log_format is None:
                for f, cs in self._COLUMNS.items():
                    if any(c in log for c in cs["energy"]):
                        log_format = f
                        break
                else:
                    raise ValueError("Could not determine log file format.")
            self._format = log_format

            # Pre-process LAMMPS log files
            if self._format == "lammps":
                if "minimize" in log:
                    log = log[log.index("Minimization stats:") :]
                log = log.split("\n")
                valid = False
                for i, line in enumerate(log):
                    if "Step" in line:
                        log = log[i:]
                        valid = True
                        break
                if not valid:
                    emsg = "No thermodynamic data found in log file " f"'{log_file}'."
                    raise ValueError(emsg)
                log = "\n".join(log)
                log = log[: log.index("Loop time of ")]
                kwargs = {"sep": r"\s+"}
                self.results.units["energies"] = ureg.kilocalorie / ureg.mole
                self.results.units["heat_capacity"] = ureg.kilocalorie / ureg.kelvin
            elif self._format == "openmm":
                kwargs = {"sep": sep}
                if reduced:
                    warnings.warn("OpenMM simulations always use real units.")
                self.results.units["energies"] = ureg.kilojoule / ureg.mole
                self.results.units["heat_capacity"] = ureg.kilojoule / ureg.kelvin

            # Determine which columns to keep.
            if self._COLUMNS[self._format]["energy"][0] in log:
                cols = self._COLUMNS[self._format]["energy"][:1]
            elif self._COLUMNS[self._format]["energy"][1] in log:
                cols = self._COLUMNS[self._format]["energy"][1:2]
                if self._COLUMNS[self._format]["energy"][2] in log:
                    cols.append(self._COLUMNS[self._format]["energy"][2])
                elif any(e in log for e in self._COLUMNS[self._format]["energy"][3:]):
                    for e in self._COLUMNS[self._format]["energy"][3:]:
                        if e in log:
                            cols.append(e)
                else:
                    raise ValueError("Potential energy column not found.")
            else:
                raise ValueError("Total or kinetic energy column not found.")

            df = pd.read_csv(StringIO(log), **kwargs)
            self.results.energies = df[cols].sum(axis=1).to_numpy()
        else:
            raise ValueError("No log file or energy values provided.")

        if temperature is not None:
            self.temperature, self.results.units["temperature"] = strip_unit(
                temperature
            )
            if self.results.units["temperature"] is None:
                self.results.units["temperature"] = ureg.kelvin
        elif log_file is None:
            raise ValueError("No log file or temperature value provided.")
        else:
            self.temperature = df[self._COLUMNS[self._format]["temperature"]].mean()
            self.results.units["temperature"] = ureg.kelvin

    def run(
        self,
        start: int = None,
        stop: int = None,
        step: int = None,
        frames: Union[np.ndarray[int], slice] = None,
    ) -> None:
        """
        Performs the calculation.

        Parameters
        ----------
        start : `int`, optional
            Starting frame for analysis.

        stop : `int`, optional
            Ending frame for analysis.

        step : `int`, optional
            Number of frames to skip between each analyzed frame.

        frames : `slice` or array-like, optional
            Index or logical array of the desired trajectory frames.

        verbose : `bool`, optional
            Determines whether detailed progress is shown.

        **kwargs
            Additional keyword arguments to pass to
            :class:`MDAnalysis.lib.log.ProgressBar`.

        Returns
        -------
        self : `ConstantVolumeHeatCapacity`
            Analysis object with results.
        """

        if frames is None:
            frames = np.arange(start or 0, stop or len(self.results.energies), step)

        U = self.results.energies[frames]
        if self._reduced:
            self.results.heat_capacity = (
                (U**2).mean() - U.mean() ** 2
            ) / self.temperature**2
        else:
            U *= self.results.units["energies"]
            self.results.heat_capacity = strip_unit(
                ((U**2).mean() - U.mean() ** 2)
                / (
                    ureg.avogadro_constant**2
                    * ureg.boltzmann_constant
                    * (self.temperature * self.results.units["temperature"]) ** 2
                ),
                self.results.units["heat_capacity"],
            )[0]

        return self
