# This file is part of BurnMan - a thermoelastic and thermodynamic toolkit for the Earth and Planetary Sciences
# Copyright (C) 2012 - 2015 by the BurnMan team, released under the GNU
# GPL v2 or later.

from __future__ import absolute_import

from os import path
import numpy as np
import scipy.optimize as opt

try:
    # scipy's factorial was moved to special in scipy 1.3.0+
    from scipy.special import factorial
except ImportError:
    from scipy.misc import factorial

import warnings

from . import equation_of_state as eos
from .. import constants as constants
from ..utils.chemistry import read_masses
from ..utils.math import bracket

atomic_masses = read_masses()
# energy_states should provide the energies and degeneracies of each electronic level in a variety of elements


class DKS_L(eos.EquationOfState):
    """
    Base class for the finite strain liquid equation of state detailed
    in :cite:`deKoker2013` (supplementary materials).
    """

    """
    Ideal gas contributions (translational and electronic)
    to thermodynamic properties
    """

    def _ln_partition_function(self, mass, temperature):
        """
        Calculates the natural log of the partition function
        """
        return 3.0 / 2.0 * np.log(temperature) + 3.0 / 2.0 * np.log(
            mass * constants.Boltzmann / (2 * np.pi * constants.Dirac * constants.Dirac)
        )

    def _F_ig(self, temperature, volume, params):
        """
        The ideal gas contribution to the helmholtz free energy
        Eq. S6, see also eq. 16.72 of Callen., 1985; p. 373
        """

        V = volume / constants.Avogadro
        figoverRT = 0.0
        for element, N in params["formula"].items():  # N is a.p.f.u
            if N > 1.0e-5:
                mass = atomic_masses[element] / constants.Avogadro
                figoverRT += -N * (
                    np.log(V) + self._ln_partition_function(mass, temperature) + 1.0
                ) + N * np.log(N)
        return constants.gas_constant * temperature * figoverRT

    def _S_ig(self, temperature, volume, params):
        """
        The ideal gas contribution to the entropy
        """

        V = volume / constants.Avogadro
        entropy_sum = 0.0
        for element, N in params["formula"].items():  # N is a.p.f.u
            if N > 1.0e-5:
                mass = atomic_masses[element] / constants.Avogadro
                entropy_sum -= -N * (
                    np.log(V)
                    + self._ln_partition_function(mass, temperature)
                    + 5.0 / 2.0
                ) + N * np.log(N)
        return constants.gas_constant * entropy_sum

    def _C_v_ig(self, temperature, volume, params):
        """
        The ideal gas contribution to the heat capacity
        """

        n_atoms = 0
        for element, N in params["formula"].items():
            n_atoms += N
        return 1.5 * constants.gas_constant * n_atoms

    def _P_ig(self, temperature, volume, params):
        """
        The ideal gas contribution to the pressure
        PV = nRT
        """

        n_atoms = 0
        for element, N in params["formula"].items():
            n_atoms += N
        return n_atoms * constants.gas_constant * temperature / volume

    def _K_T_ig(self, temperature, volume, params):
        """
        The ideal gas contribution to the isothermal bulk modulus
        V * d/dV(-nRT/V) = V*nRT/V^2
        """
        n_atoms = 0
        for element, N in params["formula"].items():
            n_atoms += N
        return n_atoms * constants.gas_constant * temperature / volume

    def _alphaK_T_ig(self, temperature, volume, params):
        """
        The ideal gas contribution to the product of the
        thermal expansivity and isothermal bulk modulus
        d/dT(nRT/V) = nR/V
        """

        n_atoms = 0
        for element, N in params["formula"].items():
            n_atoms += N
        return n_atoms * constants.gas_constant / volume

    """
    Electronic contributions to thermodynamic properties
    """

    def _zeta(
        self, temperature, volume, params
    ):  # eq. S5a, beta in deKoker thesis (3.34)
        return params["zeta_0"] * (np.power(volume / params["el_V_0"], params["xi"]))

    def _dzetadV(self, temperature, volume, params):
        return (
            params["zeta_0"]
            * params["xi"]
            * (np.power(volume / params["el_V_0"], params["xi"]))
            / volume
        )

    def _d2zetadV2(self, temperature, volume, params):
        return (
            params["zeta_0"]
            * params["xi"]
            * (params["xi"] - 1.0)
            * (np.power(volume / params["el_V_0"], params["xi"]))
            / volume
            / volume
        )

    def _Tel(self, temperature, volume, params):  # eq. S5b
        return params["Tel_0"] * (np.power(volume / params["el_V_0"], params["eta"]))

    def _dTeldV(self, temperature, volume, params):
        return (
            params["Tel_0"]
            * params["eta"]
            * (np.power(volume / params["el_V_0"], params["eta"]))
            / volume
        )

    def _d2TeldV2(self, temperature, volume, params):
        return (
            params["Tel_0"]
            * params["eta"]
            * (params["eta"] - 1.0)
            * (np.power(volume / params["el_V_0"], params["eta"]))
            / volume
            / volume
        )

    def _gimel(
        self, temperature_el, temperature, volume, params
    ):  # -F_el/zeta, 3.30 in de Koker thesis
        return 0.5 * (
            temperature * temperature - temperature_el * temperature_el
        ) - temperature * temperature_el * np.log(temperature / temperature_el)

    def _dgimeldTel(self, temperature_el, temperature, volume, params):
        return (temperature - temperature_el) - temperature * np.log(
            temperature / temperature_el
        )

    def _dgimeldT(self, temperature_el, temperature, volume, params):
        return (temperature - temperature_el) - temperature_el * np.log(
            temperature / temperature_el
        )

    def _d2gimeldTdTel(self, temperature_el, temperature, volume, params):
        return -np.log(temperature / temperature_el)

    def _d2gimeldTel2(self, temperature_el, temperature, volume, params):
        return (temperature / temperature_el) - 1.0

    def _F_el(self, temperature, volume, params):  # F_el
        temperature_el = self._Tel(temperature, volume, params)
        if temperature < temperature_el:
            F_el = 0
        else:
            F_el = -self._zeta(temperature, volume, params) * self._gimel(
                temperature_el, temperature, volume, params
            )
        return F_el

    def _S_el(self, temperature, volume, params):  # S_el
        temperature_el = self._Tel(temperature, volume, params)
        if temperature < temperature_el:
            S_el = 0
        else:
            S_el = self._zeta(temperature, volume, params) * self._dgimeldT(
                temperature_el, temperature, volume, params
            )
        return S_el

    def _P_el(self, temperature, volume, params):  # P_el
        temperature_el = self._Tel(temperature, volume, params)
        if temperature < temperature_el:
            P_el = 0
        else:
            P_el = self._dzetadV(temperature, volume, params) * self._gimel(
                temperature_el, temperature, volume, params
            ) + self._zeta(temperature, volume, params) * self._dTeldV(
                temperature, volume, params
            ) * self._dgimeldTel(
                temperature_el, temperature, volume, params
            )
        return P_el

    def _K_T_el(self, temperature, volume, params):  # K_T_el
        temperature_el = self._Tel(temperature, volume, params)
        if temperature < temperature_el:
            K_T_el = 0
        else:
            K_T_el = -volume * (
                self._d2zetadV2(temperature, volume, params)
                * self._gimel(temperature_el, temperature, volume, params)
                + 2.0
                * self._dzetadV(temperature, volume, params)
                * self._dgimeldTel(temperature_el, temperature, volume, params)
                * self._dTeldV(temperature, volume, params)
                + self._zeta(temperature, volume, params)
                * (
                    self._d2TeldV2(temperature, volume, params)
                    * self._dgimeldTel(temperature_el, temperature, volume, params)
                    + self._dTeldV(temperature, volume, params)
                    * self._dTeldV(temperature, volume, params)
                    * self._d2gimeldTel2(temperature_el, temperature, volume, params)
                )
            )
        return K_T_el

    def _alphaK_T_el(self, temperature, volume, params):  # (alphaK_T)_el
        temperature_el = self._Tel(temperature, volume, params)
        if temperature < temperature_el:
            alphaK_T_el = 0
        else:
            alphaK_T_el = self._dzetadV(temperature, volume, params) * self._dgimeldT(
                temperature_el, temperature, volume, params
            ) + self._zeta(temperature, volume, params) * self._d2gimeldTdTel(
                temperature_el, temperature, volume, params
            ) * self._dTeldV(
                temperature, volume, params
            )
        return alphaK_T_el

    def _C_v_el(self, temperature, volume, params):  # C_el, eq. 3.28 of de Koker thesis
        temperature_el = self._Tel(temperature, volume, params)
        zeta = self._zeta(temperature, volume, params)

        if temperature > temperature_el:
            Cv_el = zeta * (temperature - temperature_el)
        else:
            Cv_el = 0.0
        return Cv_el

    """
    Excess (bonding) contributions to thermodynamic properties
    """

    # Finite strain
    def _finite_strain(self, temperature, volume, params):  # f(V), eq. S3a
        return (1.0 / 2.0) * (np.power(params["V_0"] / volume, 2.0 / 3.0) - 1.0)

    def _dfdV(self, temperature, volume, params):  # f(V), eq. S3a
        return (-1.0 / 3.0) * np.power(params["V_0"] / volume, 2.0 / 3.0) / volume

    def _d2fdV2(self, temperature, volume, params):
        return (
            (5.0 / 9.0) * np.power(params["V_0"] / volume, 2.0 / 3.0) / volume / volume
        )

    # Temperature
    def _theta(self, temperature, volume, params):  # theta, eq. S3b
        return np.power(temperature / params["T_0"], params["m"]) - 1.0

    def _dthetadT(self, temperature, volume, params):
        return (
            params["m"]
            * np.power(temperature / params["T_0"], params["m"])
            / temperature
        )

    def _d2thetadT2(self, temperature, volume, params):
        return (
            params["m"]
            * (params["m"] - 1.0)
            * np.power(temperature / params["T_0"], params["m"])
            / temperature
            / temperature
        )

    def _F_xs(self, temperature, volume, params):  # F_xs, eq. S2
        f = self._finite_strain(temperature, volume, params)
        theta = self._theta(temperature, volume, params)
        energy = 0.0
        for i in range(len(params["a"])):
            ifact = factorial(i, exact=False)
            for j in range(len(params["a"][0])):
                jfact = factorial(j, exact=False)
                energy += (
                    params["a"][i][j]
                    * np.power(f, i)
                    * np.power(theta, j)
                    / ifact
                    / jfact
                )
        return energy

    def _S_xs(self, temperature, volume, params):  # F_xs, eq. 3.18
        f = self._finite_strain(temperature, volume, params)
        theta = self._theta(temperature, volume, params)
        entropy = 0.0
        for i in range(len(params["a"])):
            ifact = factorial(i, exact=False)
            for j in range(len(params["a"][0])):
                if j > 0:
                    jfact = factorial(j, exact=False)
                    entropy += (
                        j
                        * params["a"][i][j]
                        * np.power(f, i)
                        * np.power(theta, j - 1.0)
                        / ifact
                        / jfact
                    )
        return -self._dthetadT(temperature, volume, params) * entropy

    def _P_xs(self, temperature, volume, params):  # P_xs, eq. 3.17 of de Koker thesis
        f = self._finite_strain(temperature, volume, params)
        theta = self._theta(temperature, volume, params)
        pressure = 0.0
        for i in range(len(params["a"])):
            ifact = factorial(i, exact=False)
            if i > 0:
                for j in range(len(params["a"][0])):
                    jfact = factorial(j, exact=False)
                    pressure += (
                        float(i)
                        * params["a"][i][j]
                        * np.power(f, float(i) - 1.0)
                        * np.power(theta, float(j))
                        / ifact
                        / jfact
                    )
        return -self._dfdV(temperature, volume, params) * pressure

    def _K_T_xs(
        self, temperature, volume, params
    ):  # K_T_xs, eq. 3.20 of de Koker thesis
        f = self._finite_strain(temperature, volume, params)
        theta = self._theta(temperature, volume, params)
        K_ToverV = 0.0
        for i in range(len(params["a"])):
            ifact = factorial(i, exact=False)
            for j in range(len(params["a"][0])):
                if i > 0:
                    jfact = factorial(j, exact=False)
                    prefactor = (
                        float(i)
                        * params["a"][i][j]
                        * np.power(theta, float(j))
                        / ifact
                        / jfact
                    )
                    K_ToverV += (
                        prefactor
                        * self._d2fdV2(temperature, volume, params)
                        * np.power(f, float(i - 1))
                    )
                if i > 1:
                    dfdV = self._dfdV(temperature, volume, params)
                    K_ToverV += (
                        prefactor
                        * dfdV
                        * dfdV
                        * float(i - 1)
                        * np.power(f, float(i - 2))
                    )
        return volume * K_ToverV

    def _alphaK_T_xs(self, temperature, volume, params):  # eq. 3.21 of de Koker thesis
        f = self._finite_strain(temperature, volume, params)
        theta = self._theta(temperature, volume, params)
        sum_factors = 0.0
        for i in range(len(params["a"])):
            ifact = factorial(i, exact=False)
            if i > 0:
                for j in range(len(params["a"][0])):
                    if j > 0:
                        jfact = factorial(j, exact=False)
                        sum_factors += (
                            float(i)
                            * float(j)
                            * params["a"][i][j]
                            * np.power(f, float(i - 1))
                            * np.power(theta, float(j - 1))
                            / ifact
                            / jfact
                        )

        return (
            -self._dfdV(temperature, volume, params)
            * self._dthetadT(temperature, volume, params)
            * sum_factors
        )

    def _C_v_xs(
        self, temperature, volume, params
    ):  # Cv_xs, eq. 3.22 of de Koker thesis
        f = self._finite_strain(temperature, volume, params)
        theta = self._theta(temperature, volume, params)
        C_voverT = 0.0
        for i in range(len(params["a"])):
            ifact = factorial(i, exact=False)
            for j in range(len(params["a"][0])):
                if j > 0:
                    jfact = factorial(j, exact=False)
                    prefactor = (
                        float(j)
                        * params["a"][i][j]
                        * np.power(f, float(i))
                        / ifact
                        / jfact
                    )
                    C_voverT += (
                        prefactor
                        * self._d2thetadT2(temperature, volume, params)
                        * np.power(theta, float(j - 1))
                    )
                if j > 1:
                    dthetadT = self._dthetadT(temperature, volume, params)
                    C_voverT += (
                        prefactor
                        * dthetadT
                        * dthetadT
                        * float(j - 1)
                        * np.power(theta, float(j - 2))
                    )
        return -temperature * C_voverT

    """
    Magnetic contributions to thermodynamic properties
    (as found in Ramo and Stixrude, 2014)
    """

    def _spin(self, temperature, volume, params):
        S_a = 0.0
        S_b = 0.0
        numerator = 0.0
        numerator_2 = 0.0
        n_atoms = 0.0
        if "spin_a" in params:
            for element, N in params["formula"].items():
                if element == "Fe":
                    n_atoms += N

            VoverVx = volume / params["V_0"]
            S_a = params["spin_a"][0] + params["spin_a"][1] * VoverVx
            S_b = (
                params["spin_b"][0]
                + params["spin_b"][1] / VoverVx
                + params["spin_b"][2] / (np.power(VoverVx, 2.0))
                + params["spin_b"][3] / (np.power(VoverVx, 3.0))
            )

            # S = S_a*T + S_b
            # d(2S + 1)/dV
            numerator = (
                -2.0
                * (
                    -params["spin_a"][1] * temperature
                    + params["spin_b"][1] / (np.power(VoverVx, 2.0))
                    + 2.0 * params["spin_b"][2] / (np.power(VoverVx, 3.0))
                    + 3.0 * params["spin_b"][3] / (np.power(VoverVx, 4.0))
                )
                / params["V_0"]
            )

            # d2S/dV2
            numerator_2 = 2.0 * (
                (
                    2.0 * params["spin_b"][1] / (np.power(VoverVx, 3.0))
                    + 6.0 * params["spin_b"][2] / (np.power(VoverVx, 4.0))
                    + 12.0 * params["spin_b"][3] / (np.power(VoverVx, 5.0))
                )
                / np.power(params["V_0"], 2.0)
            )
        return S_a, S_b, numerator, numerator_2, n_atoms

    def _F_mag(self, temperature, volume, params):
        S_a, S_b, numerator, numerator_2, n_atoms = self._spin(
            temperature, volume, params
        )
        S = S_a * temperature + S_b
        return -n_atoms * constants.gas_constant * temperature * np.log(2.0 * S + 1.0)

    def _S_mag(self, temperature, volume, params):
        S_a, S_b, numerator, numerator_2, n_atoms = self._spin(
            temperature, volume, params
        )
        S = S_a * temperature + S_b
        return (
            n_atoms
            * constants.gas_constant
            * ((2.0 * S_a * temperature / (2.0 * S + 1.0) + np.log(2.0 * S + 1.0)))
        )

    def _P_mag(self, temperature, volume, params):
        S_a, S_b, numerator, numerator_2, n_atoms = self._spin(
            temperature, volume, params
        )
        S = S_a * temperature + S_b
        dFdV = (
            -n_atoms
            * constants.gas_constant
            * temperature
            * numerator
            / (2.0 * S + 1.0)
        )
        return -dFdV

    def _K_T_mag(self, temperature, volume, params):
        S_a, S_b, numerator, numerator_2, n_atoms = self._spin(
            temperature, volume, params
        )
        S = S_a * temperature + S_b
        dFdV = numerator / (2.0 * S + 1.0)
        d2FdV2 = numerator_2 / (2.0 * S + 1.0) - np.power(dFdV, 2.0)

        return -volume * n_atoms * constants.gas_constant * temperature * d2FdV2

    def _alphaK_T_mag(
        self, temperature, volume, params
    ):  # WARNING: numeric differentiation a.t.m.
        return self._P_mag(temperature + 0.5, volume, params) - self._P_mag(
            temperature - 0.5, volume, params
        )

    def _C_v_mag(self, temperature, volume, params):
        S_a, S_b, numerator, numerator_2, n_atoms = self._spin(
            temperature, volume, params
        )
        S = S_a * temperature + S_b
        return (
            n_atoms
            * constants.gas_constant
            * temperature
            * 4.0
            * S_a
            * (S_a * temperature + 2.0 * S_b + 1.0)
            / np.power(2.0 * S + 1.0, 2.0)
        )

    def _aK_T(self, temperature, volume, params):
        aK_T = (
            self._alphaK_T_ig(temperature, volume, params)
            + self._alphaK_T_el(temperature, volume, params)
            + self._alphaK_T_xs(temperature, volume, params)
            + self._alphaK_T_mag(temperature, volume, params)
        )
        return aK_T

    # Pressure
    def pressure(self, temperature, volume, params):
        P = (
            self._P_ig(temperature, volume, params)
            + self._P_el(temperature, volume, params)
            + self._P_xs(temperature, volume, params)
            + self._P_mag(temperature, volume, params)
        )
        return P

    def volume(self, pressure, temperature, params):
        _delta_pressure = (
            lambda x, pressure, temperature, params: pressure
            - self.pressure(temperature, x, params)
        )

        # we need to have a sign change in [a,b] to find a zero. Let us start with a
        # conservative guess:
        args = (pressure, temperature, params)
        try:
            sol = bracket(_delta_pressure, params["V_0"], 1.0e-2 * params["V_0"], args)
        except ValueError:
            raise Exception(
                "Cannot find a volume, perhaps you are outside of the range of validity for the equation of state?"
            )
        return opt.brentq(_delta_pressure, sol[0], sol[1], args=args)

    def isothermal_bulk_modulus_reuss(self, pressure, temperature, volume, params):
        """
        Returns isothermal bulk modulus :math:`[Pa]`
        """
        K_T = (
            self._K_T_ig(temperature, volume, params)
            + self._K_T_el(temperature, volume, params)
            + self._K_T_xs(temperature, volume, params)
            + self._K_T_mag(temperature, volume, params)
        )
        return K_T

    def isentropic_bulk_modulus_reuss(self, pressure, temperature, volume, params):
        """
        Returns adiabatic bulk modulus. :math:`[Pa]`
        """
        K_S = self.isothermal_bulk_modulus_reuss(
            pressure, temperature, volume, params
        ) * (
            1.0
            + temperature
            * self.thermal_expansivity(pressure, temperature, volume, params)
            * self.grueneisen_parameter(pressure, temperature, volume, params)
        )
        return K_S

    def grueneisen_parameter(self, pressure, temperature, volume, params):
        """
        Returns grueneisen parameter. :math:`[unitless]`
        """
        gamma = (
            self._aK_T(temperature, volume, params)
            * volume
            / self.molar_heat_capacity_v(pressure, temperature, volume, params)
        )
        return gamma

    def shear_modulus(self, pressure, temperature, volume, params):
        """
        Returns shear modulus. :math:`[Pa]`
        Zero for fluids
        """
        return 0.0

    def molar_heat_capacity_v(self, pressure, temperature, volume, params):
        """
        Returns heat capacity at constant volume. :math:`[J/K/mol]`
        """
        C_v = (
            self._C_v_ig(temperature, volume, params)
            + self._C_v_el(temperature, volume, params)
            + self._C_v_xs(temperature, volume, params)
            + self._C_v_mag(temperature, volume, params)
        )
        return C_v

    def molar_heat_capacity_p(self, pressure, temperature, volume, params):
        """
        Returns heat capacity at constant pressure. :math:`[J/K/mol]`
        """
        C_p = self.molar_heat_capacity_v(pressure, temperature, volume, params) * (
            1.0
            + temperature
            * self.thermal_expansivity(pressure, temperature, volume, params)
            * self.grueneisen_parameter(pressure, temperature, volume, params)
        )
        return C_p

    def thermal_expansivity(self, pressure, temperature, volume, params):
        """
        Returns thermal expansivity. :math:`[1/K]`
        """
        alpha = self._aK_T(
            temperature, volume, params
        ) / self.isothermal_bulk_modulus_reuss(0.0, temperature, volume, params)
        return alpha

    def gibbs_free_energy(self, pressure, temperature, volume, params):
        """
        Returns the Gibbs free energy at the pressure and temperature of the mineral [J/mol]
        """
        G = (
            self.helmholtz_free_energy(pressure, temperature, volume, params)
            + pressure * volume
        )
        return G

    def entropy(self, pressure, temperature, volume, params):
        """
        Returns the entropy at the pressure and temperature of the mineral [J/K/mol]
        """
        S = (
            self._S_ig(temperature, volume, params)
            + self._S_el(temperature, volume, params)
            + self._S_xs(temperature, volume, params)
            + self._S_mag(temperature, volume, params)
        )
        return S

    def enthalpy(self, pressure, temperature, volume, params):
        """
        Returns the enthalpy at the pressure and temperature of the mineral [J/mol]
        """
        H = (
            self.helmholtz_free_energy(pressure, temperature, volume, params)
            + temperature * self.entropy(pressure, temperature, volume, params)
            + pressure * self.volume(pressure, temperature, params)
        )
        return H

    def helmholtz_free_energy(self, pressure, temperature, volume, params):
        """
        Returns the Helmholtz free energy at the pressure and temperature of the mineral [J/mol]
        """
        F = (
            self._F_ig(temperature, volume, params)
            + self._F_el(temperature, volume, params)
            + self._F_xs(temperature, volume, params)
            + self._F_mag(temperature, volume, params)
        )
        return F

    def molar_internal_energy(self, pressure, temperature, volume, params):
        E = self.helmholtz_free_energy(
            pressure, temperature, volume, params
        ) + temperature * self.entropy(pressure, temperature, volume, params)
        return E

    def validate_parameters(self, params):
        """
        Check for existence and validity of the parameters
        """

        # Check that all the required keys are in the dictionary
        expected_keys = [
            "V_0",
            "T_0",
            "O_theta",
            "O_f",
            "m",
            "a",
            "zeta_0",
            "xi",
            "Tel_0",
            "eta",
        ]
        for k in expected_keys:
            if k not in params:
                raise KeyError("params object missing parameter : " + k)

        # Sometimes the standard electronic volume is different to V_0.
        # If not, make it the same.
        if "el_V_0" not in params:
            params["el_V_0"] = params["V_0"]
