'''
Author: Vegard G. Jervell
Date: December 2020
Purpose: Implementation of L. J. T. M. Kempers' 2001 model for prediction of Soret coefficients in multicomponent
        mixtures, derived in "A comprehensive thermodynamic theory of the Soret effect in a multicomponent gas,
        liquid, or solid", 2001, Journal of Chemical Physics
        doi : http://dx.doi.org/10.1063/1.1398315
Requires : numpy, scipy, ThermoPack, KineticGas
Note : KineticGas is only 2-component compatible, replacing the KineticGas module with a module capable of predicting
        thermal diffusion coefficients for multicomponent ideal-gas mixtures will make this module multicomponent-compatible.
'''

import numpy as np
from scipy.constants import gas_constant
from pycThermopack.pyctp import cubic
from pycThermopack.pyctp import cpa
from models.kempers import Kempers
from scipy.optimize import root
from models.alpha_t0_empirical import Alpha_T0_empirical
from models.kineticgas import KineticGas

class Kempers01(Kempers):
    def __init__(self, comps, eos, x=[0.5, 0.5], temp=300, pres=1e5, phase=1,
                 alpha_t0_empirical=False, alpha_t0_method='wheights', alpha_t0_N = 7):
        '''
        :param comps (str): comma separated list of components
                :param eos (ThermoPack): Initialized equation of state, with same components as Kempers
        :param x (1darray): list of mole fractions
        :param temp (float > 0): Temperature [K]
        :param pres (float > 0): Pressure [Pa]
        :param phase: Phase of mixture, used for calculating dmudn_TP, see thermo.thermopack for phase identifiers
        :param alpha_t0_empirical (bool): True if using empirical alpha_T0 values
        :param alpha_t0_method (str): what method to use for empirical alpha_T0 values (see Alpha_T0_empirical)
        :param alpha_t0_N (int > 0): Order of approximation when using analytical alpha_T0 values
        '''
        super().__init__(comps, eos, x=x, temp=temp, pres=pres, phase=phase)

        self.alpha_T0_empirical = alpha_t0_empirical
        if alpha_t0_empirical:
            self.alpha_T0_method = alpha_t0_method
            self.kinetic_gas = Alpha_T0_empirical(comps, method=alpha_t0_method)
        else:
            self.alpha_T0_N = alpha_t0_N
            self.kinetic_gas = KineticGas(comps, self.eos, mole_fracs=x, N=alpha_t0_N)

    def reset_alpha_t0(self):
        #reset alpha_t0_getter with current mole fractions
        if self.alpha_T0_empirical:
            self.kinetic_gas = Alpha_T0_empirical(self.comps, mole_fracs=self.x, method=self.alpha_T0_method)
        else:
            self.kinetic_gas = KineticGas(self.comps, self.eos, mole_fracs=self.x, N=self.alpha_T0_N)

    def get_soret_cov(self, kin=False):
        '''
        Get soret coefficients at current settings, center of volume frame of reference
        :return: (ndarray) soret coefficients
        '''

        R = gas_constant
        v, dvdn = self.eos.specific_volume(self.temp, self.pres, self.x, self.phase, dvdn=True)
        h, dhdn = self.eos.enthalpy(self.temp, self.pres, self.x, self.phase, dhdn=True)
        h0, dh0dn = self.eos.enthalpy(self.temp, 1e-5, self.x, 2, dhdn=True)

        dmudx = self.dmudx_TP()
        alpha_T0 = self.kinetic_gas.alpha_T0(self.temp)

        #using alpha_T0 as initial guess for root solver
        initial_guess = alpha_T0 * self.temp

        N = len(self.x)
        #Defining the set of equations
        def eq_set(alpha):
            eqs = np.zeros(N)
            for i in range(N-1):
                eqs[i] = ((dhdn[-1] - dh0dn[-1])/dvdn[-1]) - ((dhdn[i] - dh0dn[i])/dvdn[i])\
                         + R * self.temp * ((alpha_T0[i] * (1 - self.x[i]) / dvdn[i])
                                            - (alpha_T0[-1] * (1 - self.x[-1])/dvdn[-1]))\
                         - sum((dmudx[i, j]/dvdn[i] - dmudx[-1, j]/dvdn[-1]) * self.x[j] * (1 - self.x[j]) * alpha[j]
                               for j in range(N - 1))

            eqs[N-1] = sum(self.x * (1 - self.x) * alpha)
            return eqs

        #Solve the set of equations, warn if non-convergent
        solved = root(eq_set, initial_guess)
        if solved.success is False:
            print('Solution did not converge for composition :', self.x, ', Temperature :', self.temp)
        alpha = solved.x

        soret = alpha / self.temp

        if kin is True:
            kin_contrib = alpha_T0 * R /(self.x * dmudx.diagonal())
            return soret, kin_contrib
        else:
            return soret

    def get_soret_com(self, kin=False):
        '''
        Get soret coefficients at current settings, center of mass frame of reference
        :return: (ndarray) soret coefficients
        '''

        R = gas_constant
        M = self.mole_weights = np.array([self.eos.compmoleweight(self.eos.getcompindex(comp))
                                          for comp in self.comps.split(',')])
        h, dhdn = self.eos.enthalpy(self.temp, self.pres, self.x, self.phase, dhdn=True)
        h0, dh0dn = self.eos.enthalpy(self.temp, 1e-5, self.x, 2, dhdn=True)

        dmudx = self.dmudx_TP()
        alpha_T0 = self.kinetic_gas.alpha_T0(self.temp)

        # using alpha_T0 as initial guess for root solver
        initial_guess = alpha_T0 * self.temp

        N = len(self.x)

        # Defining the set of equations
        def eq_set(alpha):
            eqs = np.zeros(N)
            for i in range(N - 1):
                eqs[i] = ((dhdn[-1] - dh0dn[-1]) / M[-1]) - ((dhdn[i] - dh0dn[i]) / M[i]) \
                         + R * self.temp * ((alpha_T0[i] * (1 - self.x[i]) / M[i])
                                            - (alpha_T0[-1] * (1 - self.x[-1]) / M[-1])) \
                         - sum(
                    (dmudx[i, j] / M[i] - dmudx[-1, j] / M[-1]) * self.x[j] * (1 - self.x[j]) * alpha[j]
                    for j in range(N - 1))

            eqs[N - 1] = sum(self.x * (1 - self.x) * alpha)
            return eqs

        # Solve the set of equations, warn if non-convergent
        solved = root(eq_set, initial_guess)
        if solved.success is False:
            print('Solution did not converge for composition :', self.x, ', Temperature :', self.temp)
        alpha = solved.x

        soret = alpha / self.temp

        if kin is True:
            kin_contrib = alpha_T0 * R /(self.x * dmudx.diagonal())
            return soret, kin_contrib
        else:
            return soret