'''
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, Boltzmann
from models.kempers import Kempers
from scipy.optimize import root
from models.alpha_t0_empirical import Alpha_T0_empirical
from models.py_KineticGas import KineticGas
import warnings, platform
from scipy.constants import Avogadro


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,
                 mole_weights=None, particle_density=None, sigma=None, epsilon=None):
        '''
        :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.mole_weights = mole_weights
            self.particle_density = particle_density
            self.sigma = sigma
            self.epsilon = epsilon
            if mole_weights is not None and particle_density is not None:
                id_gas_pressure = particle_density * Boltzmann * self.temp
                self.kinetic_gas = KineticGas(comps, eos=None, mole_fracs=x, N=alpha_t0_N, p=id_gas_pressure, T=self.temp, mole_weights=mole_weights, sigma=sigma, epsilon=epsilon)
            elif mole_weights is not None:
                self.kinetic_gas = KineticGas(comps, eos=None, mole_fracs=x, N=alpha_t0_N, p=self.pres,
                                              T=self.temp, mole_weights=mole_weights, sigma=sigma, epsilon=epsilon)
            elif particle_density is not None:
                id_gas_pressure = particle_density * Boltzmann * self.temp
                self.kinetic_gas = KineticGas(comps, eos=self.eos, mole_fracs=x, N=alpha_t0_N, p=id_gas_pressure, T=self.temp, sigma=sigma, epsilon=epsilon)
            else:
                self.kinetic_gas = KineticGas(comps, eos=self.eos, mole_fracs=x, N=alpha_t0_N, p=self.pres, T=self.temp, sigma=sigma, epsilon=epsilon)

    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:
            if self.mole_weights is not None and self.particle_density is not None:
                id_gas_pressure = self.particle_density * Boltzmann * self.temp
                self.kinetic_gas = KineticGas(self.comps, eos=None, mole_fracs=self.x, N=self.alpha_T0_N, p=id_gas_pressure,
                                              T=self.temp, mole_weights=self.mole_weights, sigma=self.sigma, epsilon=self.epsilon)
            elif self.mole_weights is not None:
                self.kinetic_gas = KineticGas(self.comps, eos=None, mole_fracs=self.x, N=self.alpha_T0_N, p=self.pres,
                                              T=self.temp, mole_weights=self.mole_weights, sigma=self.sigma, epsilon=self.epsilon)
            elif self.particle_density is not None:
                id_gas_pressure = self.particle_density * Boltzmann * self.temp
                self.kinetic_gas = KineticGas(self.comps, eos=self.eos, mole_fracs=self.x, N=self.alpha_T0_N, p=id_gas_pressure,
                                              T=self.temp, sigma=self.sigma, epsilon=self.epsilon)
            else:
                self.kinetic_gas = KineticGas(self.comps, eos=self.eos, mole_fracs=self.x, N=self.alpha_T0_N, p=self.pres, T=self.temp, sigma=self.sigma, epsilon=self.epsilon)

    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

        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:
            warnings.warn('Solution did not converge for composition :' + str(self.x) + ', Temperature :'+ str(self.temp) + ', Pressure : '+str(self.pres/1e5))
            alpha = np.full_like(solved.x, np.nan)
            alpha = solved.x
        else:
            alpha = solved.x

        soret = alpha / self.temp
        if kin is True:
            kin_contrib = alpha_T0 / self.temp #* 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

        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:
            warnings.warn('Solution did not converge for composition :'+str(self.x)+ ', Temperature :' + str(self.temp))
            alpha = np.full_like(solved.x, np.nan)
        else:
            alpha = solved.x

        soret = alpha / self.temp

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