# -*- coding: utf-8 -*-
"""
Soave Redlich Kwong cubic equation of state model.
"""

import numpy as np
import numpy.linalg as la
from clausius._interfaces import Model
from clausius.utils import ERROR_TOL, phase_struct, volume_completion
from clausius.eos.basicmixing import vanDerWaalsMixing

__all__ = ['SoaveRedlichKwong']


class SoaveRedlichKwong(Model):
    """The Soave Redlich Kwong equation of state.
    Requires that the mixture have valid critical temperature, critical
    pressure, and acentric factor for all compounds.

    References:
    -----------
    Soave, G.: Chem. Eng. Sci., 27: 1197 (1972)."""

    def __init__(self, mixingrule=vanDerWaalsMixing()):
        """Unless specified a van Der Waals mixing rule with zero binary
        interaction parameters is used."""
        self.mixingrule = mixingrule
        self._companionmatrix = np.diag(np.ones((2,), dtype=np.float), -1)

    def __call__(self, mixture):
        """Takes a Mixture with a set temperature and pressure and returns the
        Mixture with the compressibility and departure functions completed"""

        Ai, TdAdTi, Bi = self._eosparams(mixture)

        # Vapor Phase
        eosargsV = self.mixingrule(mixture.composition['vapor'], Ai, TdAdTi, Bi)
        ZV = self._compressibility(*eosargsV)['vapor']

        # Liquid Phase
        eosargsL = self.mixingrule(mixture.composition['liquid'], Ai, TdAdTi, Bi)
        ZL = self._compressibility(*eosargsL)['liquid']

        # Completing state info
        mixture.compressibility = (ZV, ZL)
        mixture.volume = (volume_completion(mixture.pressure, mixture.temperature, ZV), volume_completion(mixture.pressure, mixture.temperature, ZL))
        mixture.departure['logfugacity'] = (self._logfugacity(ZV, *eosargsV), self._logfugacity(ZL, *eosargsL))
        mixture.departure['enthalpy'] = (self._enthalpydeparture(ZV, *eosargsV), self._enthalpydeparture(ZL, *eosargsL))
        mixture.departure['entropy'] = (self._entropydeparture(ZV, *eosargsV), self._entropydeparture(ZL, *eosargsL))
        mixture.departure['gibbs'] = (self._gibbsdeparture(ZV, *eosargsV), self._gibbsdeparture(ZL, *eosargsL))
        mixture.departure['helmholtz'] = (self._helmholtzdeparture(ZV, *eosargsV), self._helmholtzdeparture(ZL, *eosargsL))

        return mixture

    def _eosparams(self, mixture):
        """Calculate the EOS parameters A, B, and TdA/dT"""
        fw = 0.48 + mixture.critical['acentric'] * (1.574 - 0.176 * mixture.critical['acentric'])
        alpha = 1 + fw * (1 - np.sqrt(mixture.reduced['temperature']))
        A = 0.42748 * alpha**2 * mixture.reduced['pressure'] / mixture.reduced['temperature']**2
        B = 0.08664 * mixture.reduced['pressure'] / mixture.reduced['temperature']
        TdAdT = -0.42748 * mixture.reduced['pressure'] * mixture.reduced['temperature']**-1.5 * alpha * fw

        return A, TdAdT, B

    def _compressibility(self, A, TdAdT, dAdn, B, dBdn):
        """Calculate the compressibility from the parameters A and B"""
        c = np.array([- 1, A - B * (1 + B), -(A * B)], dtype=np.float)
        self._companionmatrix[0, :] = -c
        result = la.eigvals(self._companionmatrix)
        result = result[result.imag < ERROR_TOL].real
        return np.array((result.max(), result.min()), dtype=phase_struct)

    def _logfugacity(self, Z, A, TdAdT, dAdn, B, dBdn):
        """Calculates the natural logarithm of the fugacity coefficient from parameters A, B, and partial derivatives"""

        return -np.log(np.fabs(Z - B)) + (dBdn / B) * (Z - 1) - (A / B) * ((dAdn / A) - (dBdn / B)) * np.log(1 + (B / Z))

    def _gibbsdeparture(self, Z, A, TdAdT, dAdn, B, dBdn):
        """Calculates the gibbs free energy departure function from parameters A, B, and dA/dT"""

        return (A / B) * np.log(1 + (B / Z)) + np.log(np.fabs(Z - B)) + 1 - Z

    def _helmholtzdeparture(self, Z, A, TdAdT, dAdn, B, dBdn):
        """Calculates the helmholtz free energy departure function from parameters A, B, and dA/dT"""

        return (A / B) * np.log(1 + (B / Z)) + np.log(np.fabs(Z - B))

    def _enthalpydeparture(self, Z, A, TdAdT, dAdn, B, dBdn):
        """Calculates the enthalpy departure function from parameters A, B, and dA/dT"""

        return ((A - TdAdT) / B) * np.log(1 + (B / Z)) + 1 - Z

    def _entropydeparture(self, Z, A, TdAdT, dAdn, B, dBdn):
        """Calculates the entropy departure function from parameters A, B, and dA/dT"""

        return -(TdAdT / B) * np.log(1 + (B / Z)) - np.log(np.fabs(Z - B))

    def logfug(self, mixture, phase):
        """Method called by the flash calculation methods, should return the natural logarithm of the fugacity coefficient"""

        Ai, TdAdTi, Bi = self._eosparams(mixture)
        eosargs = self.mixingrule(mixture.composition[phase], Ai, TdAdTi, Bi)
        Z = self._compressibility(*eosargs)[phase]
        return self._logfugacity(Z, *eosargs)
