from __future__ import absolute_import

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

import numpy as np
import scipy.optimize as opt
from . import equation_of_state as eos
from ..utils.math import bracket
import warnings


def bulk_modulus_fourth(volume, params):
    """
    compute the bulk modulus as per the fourth order
    birch-murnaghan equation of state.  Returns bulk
    modulus in the same units as the reference bulk
    modulus.  Pressure must be in :math:`[Pa]`.
    """

    invVrel = params["V_0"] / volume
    f = 0.5 * (pow(invVrel, 2.0 / 3.0) - 1.0)

    Xi = (3.0 / 4.0) * (4.0 - params["Kprime_0"])
    Zeta = (3.0 / 8.0) * (
        (params["K_0"] * params["Kprime_prime_0"])
        + params["Kprime_0"] * (params["Kprime_0"] - 7.0)
        + 143.0 / 9.0
    )

    K = (
        5.0
        * f
        * pow((1.0 + 2.0 * f), 5.0 / 2.0)
        * params["K_0"]
        * (1.0 - (2.0 * Xi * f) + (4.0 * Zeta * pow(f, 2.0)))
    ) + (
        pow(1.0 + (2.0 * f), 7.0 / 2.0)
        * params["K_0"]
        * (1.0 - (4.0 * Xi * f) + (12.0 * Zeta * pow(f, 2.0)))
    )

    return K


def volume_fourth_order(pressure, params):
    """
    Volume according to the fourth-order
    Birch-Murnaghan equation of state.

    :param pressure: Pressure in the same units that are supplied for the reference bulk
    modulus (params['K_0'])
    :type pressure: float
    :param params: parameter dictionary
    :type params: dictionary
    :return: V/V_0
    :rtype: float
    """

    def delta_pressure(x):
        return pressure_birch_murnaghan_fourth(params["V_0"] / x, params) - pressure

    try:
        sol = bracket(delta_pressure, params["V_0"], 1.0e-2 * params["V_0"])
    except ValueError:
        raise ValueError(
            "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])


def pressure_birch_murnaghan_fourth(invVrel, params):
    """
    Pressure according to the fourth-order
    Birch-Murnaghan equation of state.

    :param invVrel: V/V_0
    :type invVrel: float or numpy array
    :param params: parameter dictionary
    :type params: dictionary
    :return: Pressure in the same units that are supplied for the reference bulk
    modulus (params['K_0'])
    :rtype: float or numpy array
    """

    f = 0.5 * (pow(invVrel, 2.0 / 3.0) - 1.0)
    Xi = (3.0 / 4.0) * (4.0 - params["Kprime_0"])
    Zeta = (3.0 / 8.0) * (
        (params["K_0"] * params["Kprime_prime_0"])
        + params["Kprime_0"] * (params["Kprime_0"] - 7.0)
        + 143.0 / 9.0
    )

    return (
        3.0
        * f
        * pow(1.0 + 2.0 * f, 5.0 / 2.0)
        * params["K_0"]
        * (1.0 - (2.0 * Xi * f) + (4.0 * Zeta * pow(f, 2.0)))
        + params["P_0"]
    )


class BM4(eos.EquationOfState):
    """
    Base class for the isothermal Birch Murnaghan equation of state.  This is fourth order in strain, and
    has no temperature dependence.
    """

    def volume(self, pressure, temperature, params):
        """
        Returns volume :math:`[m^3]` as a function of pressure :math:`[Pa]`.
        """
        return volume_fourth_order(pressure, params)

    def pressure(self, temperature, volume, params):
        return pressure_birch_murnaghan_fourth(params["V_0"] / volume, params)

    def isothermal_bulk_modulus_reuss(self, pressure, temperature, volume, params):
        """
        Returns isothermal bulk modulus :math:`K_T` :math:`[Pa]` as a function of pressure :math:`[Pa]`,
        temperature :math:`[K]` and volume :math:`[m^3]`.
        """
        return bulk_modulus_fourth(volume, params)

    def isentropic_bulk_modulus_reuss(self, pressure, temperature, volume, params):
        """
        Returns adiabatic bulk modulus :math:`K_s` of the mineral. :math:`[Pa]`.
        """
        return bulk_modulus_fourth(volume, params)

    def shear_modulus(self, pressure, temperature, volume, params):
        """
        Returns shear modulus :math:`G` of the mineral. :math:`[Pa]`
        """
        return 0.0

    def entropy(self, pressure, temperature, volume, params):
        """
        Returns the molar entropy :math:`\\mathcal{S}` of the mineral. :math:`[J/K/mol]`
        """
        return 0.0

    def molar_internal_energy(self, pressure, temperature, volume, params):
        """
        Returns the internal energy :math:`\\mathcal{E}` of the mineral. :math:`[J/mol]`
        """
        x = np.power(volume / params["V_0"], -1.0 / 3.0)
        x2 = x * x
        x4 = x2 * x2
        x6 = x4 * x2
        x8 = x4 * x4

        xi1 = 3.0 * (4.0 - params["Kprime_0"]) / 4.0
        xi2 = (
            3.0
            / 8.0
            * (
                params["K_0"] * params["Kprime_prime_0"]
                + params["Kprime_0"] * (params["Kprime_0"] - 7.0)
            )
            + 143.0 / 24.0
        )

        intPdV = (
            -9.0
            / 2.0
            * params["V_0"]
            * params["K_0"]
            * (
                (xi1 + 1.0) * (x4 / 4.0 - x2 / 2.0 + 1.0 / 4.0)
                - xi1 * (x6 / 6.0 - x4 / 4.0 + 1.0 / 12.0)
                + xi2 * (x8 / 8 - x6 / 2 + 3.0 * x4 / 4.0 - x2 / 2.0 + 1.0 / 8.0)
            )
        )

        return -intPdV + params["E_0"]

    def gibbs_free_energy(self, pressure, temperature, volume, params):
        """
        Returns the Gibbs free energy :math:`\\mathcal{G}` of the mineral. :math:`[J/mol]`
        """
        # G = int VdP = [PV] - int PdV = E + PV

        return (
            self.molar_internal_energy(pressure, temperature, volume, params)
            + volume * pressure
        )

    def molar_heat_capacity_v(self, pressure, temperature, volume, params):
        """
        Since this equation of state does not contain temperature effects, simply return a very large number. :math:`[J/K/mol]`
        """
        return 1.0e99

    def molar_heat_capacity_p(self, pressure, temperature, volume, params):
        """
        Since this equation of state does not contain temperature effects, simply return a very large number. :math:`[J/K/mol]`
        """
        return 1.0e99

    def thermal_expansivity(self, pressure, temperature, volume, params):
        """
        Since this equation of state does not contain temperature effects, simply return zero. :math:`[1/K]`
        """
        return 0.0

    def grueneisen_parameter(self, pressure, temperature, volume, params):
        """
        Since this equation of state does not contain temperature effects, simply return zero. :math:`[unitless]`
        """
        return 0.0

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

        if "E_0" not in params:
            params["E_0"] = 0.0
        if "P_0" not in params:
            params["P_0"] = 0.0

        # If G and Gprime are not included this is presumably deliberate,
        # as we can model density and bulk modulus just fine without them,
        # so just add them to the dictionary as nans
        if "G_0" not in params:
            params["G_0"] = float("nan")
        if "Gprime_0" not in params:
            params["Gprime_0"] = float("nan")

        # Check that all the required keys are in the dictionary
        expected_keys = ["V_0", "K_0", "Kprime_0"]
        for k in expected_keys:
            if k not in params:
                raise KeyError("params object missing parameter : " + k)

        # Finally, check that the values are reasonable.
        if params["P_0"] < 0.0:
            warnings.warn("Unusual value for P_0", stacklevel=2)
        if params["V_0"] < 1.0e-7 or params["V_0"] > 1.0e-3:
            warnings.warn("Unusual value for V_0", stacklevel=2)
        if params["K_0"] < 1.0e9 or params["K_0"] > 1.0e13:
            warnings.warn("Unusual value for K_0", stacklevel=2)
        if params["Kprime_0"] < 0.0 or params["Kprime_0"] > 10.0:
            warnings.warn("Unusual value for Kprime_0", stacklevel=2)
        if params["Kprime_prime_0"] > 0.0 or params["Kprime_prime_0"] < -10.0:
            warnings.warn("Unusual value for Kprime_prime_0", stacklevel=2)
