
import numpy as np
import scipy.linalg as lin
import scipy.constants as const
from scipy.integrate import quad
import pandas as pd
import os, sys, platform
sys.path.append(os.path.dirname(os.getcwd())+'/soret_model/cpp/release')
sys.path.append(os.path.dirname(os.getcwd())+'/soret_model')
sys.path.append(os.path.dirname(os.getcwd())+'/soret_model/cpp')

if (platform.system() == 'Linux'):
    from release_ubuntu.KineticGas import cpp_KineticGas
else:
    from release_mac.KineticGas import cpp_KineticGas


class KineticGas:

    computed_points = {}

    def __init__(self, comps, eos=None, mole_fracs=None, T=300, p=1e5, N=7, mole_weights=None):
        '''
        :param comps (str): Comma-separated list of components, following Thermopack-convention
        :param eos (thermopack): An initialized equation of state, only used to get mole weights
        :param mole_fracs (array-like): list of mole fractions
        :param N (int > 0): Order of approximation.
                            Be aware that N > 10 can be detrimental to runtime. This should be implemented in C++ or Fortran.
        '''

        self.T = T
        self.p = p

        #Packing out variables in an n-component compatible way
        complist = comps.split(',')
        if mole_weights is None:
            self.mole_weights = np.array([eos.compmoleweight(eos.getcompindex(comp)) for comp in complist]) * 1e-3 / const.Avogadro
        else:
            self.mole_weights = np.array(mole_weights) * 1e-3 / const.Avogadro
        self.mole_fracs = np.array(mole_fracs)
        self.m0 = np.sum(self.mole_weights)
        self.M = self.mole_weights/self.m0

        self.sigmaij = self.get_hard_sphere_radius(comps)
        self.sigma = np.diag(self.sigmaij)

        #Calculate the (temperature independent) soret-coefficient at the ideal gas state
        self.N = N

        self.binary_mole_frac_matr = np.zeros_like(self.sigmaij)
        for i in range(len(self.binary_mole_frac_matr)):
            for j in range(len(self.binary_mole_frac_matr)):
                if (mole_fracs[i] + mole_fracs[j] < 1e-10):
                    self.binary_mole_frac_matr[i, j] = 0
                else:
                    self.binary_mole_frac_matr[i,j] = mole_fracs[i] / (mole_fracs[i] + mole_fracs[j])

        binary_KineticGas_matr = np.empty_like(self.sigmaij, dtype=cpp_KineticGas)
        for i in range(1, len(binary_KineticGas_matr)):
            for j in range(i):
                binary_mole_weights = [self.mole_weights[i], self.mole_weights[j]]
                binary_sigmaij = [[self.sigmaij[i,i], self.sigmaij[i,j]], [self.sigmaij[j,i], self.sigmaij[j,j]]]
                binary_mole_fracs = [self.binary_mole_frac_matr[i,j], self.binary_mole_frac_matr[j,i]]
                binary_KineticGas_matr[i,j] = cpp_KineticGas(binary_mole_weights, binary_sigmaij, binary_mole_fracs, T, p, N)

        binary_alpha_T0_matr = np.zeros_like(self.sigmaij)
        for i in range(1, len(binary_KineticGas_matr)):
            for j in range(i):
                A = binary_KineticGas_matr[i, j].A_matrix
                A = np.array(A)
                delta = binary_KineticGas_matr[i, j].delta_vector

                try:
                    d = lin.solve(A, delta)
                except (ValueError, np.linalg.LinAlgError):
                    binary_alpha_T0_matr[i,j] = np.nan
                    binary_alpha_T0_matr[j,i] = np.nan
                    continue

                d_1, d0, d1 = d[N - 1], d[N], d[N + 1]
                binary_alpha_T0_matr[i,j] = - (5 / (2 * d0)) * ((self.binary_mole_frac_matr[i,j] * d1 / np.sqrt(self.M[i])) + (self.binary_mole_frac_matr[j,i] * d_1 / np.sqrt(self.M[j])))
                binary_alpha_T0_matr[j,i] = - (self.binary_mole_frac_matr[i,j] / self.binary_mole_frac_matr[j,i]) * binary_alpha_T0_matr[i,j]

        self.binary_d_matr = np.zeros((len(self.sigmaij), len(self.sigmaij), 3))
        for i in range(1, len(binary_KineticGas_matr)):
            for j in range(i):
                A = binary_KineticGas_matr[i, j].A_matrix
                A = np.array(A)
                delta = binary_KineticGas_matr[i, j].delta_vector

                try:
                    d = lin.solve(A, delta)
                except (ValueError, np.linalg.LinAlgError):
                    binary_alpha_T0_matr[i,j] = np.nan
                    binary_alpha_T0_matr[j,i] = np.nan
                    continue

                self.binary_d_matr[i, j, 0] = d[N - 1]
                self.binary_d_matr[i, j, 1] = d[N]
                self.binary_d_matr[i, j, 2] = d[N + 1]
                self.binary_d_matr[j, i, 0] = d[N - 1]
                self.binary_d_matr[j, i, 1] = d[N]
                self.binary_d_matr[j, i, 2] = d[N + 1]

        self.kT_vec = np.zeros_like(self.mole_weights)
        for i in range(len(binary_alpha_T0_matr)):
            for j in range(len(binary_alpha_T0_matr)):
                self.kT_vec[i] += (mole_fracs[i] + mole_fracs[j]) * binary_alpha_T0_matr[i, j]

    def alpha_T0(self, T):
        return self.kT_vec / (np.array(self.mole_fracs) * (1 - np.array(self.mole_fracs)))

    def interdiffusion(self):
        Dij_matr = np.zeros_like(self.sigmaij)
        print('x : ', self.mole_fracs, 'x_tilde :')
        for line in self.binary_mole_frac_matr:
            for x in line:
                print(x, end=' ')
            print()
        for i in range(1, len(Dij_matr)):
            for j in range(i):
                print(i, j, ':', self.binary_mole_frac_matr[i,j], self.binary_d_matr[i,j,1], self.m0,
                      (0.5 * self.binary_mole_frac_matr[i,j] * self.binary_mole_frac_matr[j,i]
                       * np.sqrt(2 * const.Boltzmann * self.T / self.m0) * self.binary_d_matr[i,j,1]))

                if 0 < self.binary_mole_frac_matr[i,j] < 1:
                    Dij_matr[i,j] = 1 / (0.5 * self.binary_mole_frac_matr[i,j] * self.binary_mole_frac_matr[j,i] * np.sqrt(2 * const.Boltzmann * self.T / self.m0) * self.binary_d_matr[i,j,1])
                    Dij_matr[j,i] = Dij_matr[i, j]
                else:
                    Dij_matr[i,j] = 0 # Logically, this should be NaN (the diffusion coefficient between two species where one of them has zero particles is undefined)
                    Dij_matr[j,i] = 0 # but for consistency in further computation must be set to 0, se further comments

        S = np.dot(Dij_matr, self.mole_fracs) # if D[i,j] = NaN => S[i] = NaN, when in reality NaN values should simply not be included in the sum
        print(Dij_matr)
        print('#'*50)
        #print(1/Dij_matr, S)
        return (1 - self.mole_fracs) / S

    def binary_diffusion(self):
        Dij_matr = np.zeros_like(self.sigmaij)
        for i in range(len(Dij_matr)):
            for j in range(i):
                if 0 < self.binary_mole_frac_matr[i, j] < 1:
                    Dij_matr[i, j] = 0.5 * self.binary_mole_frac_matr[i, j] * self.binary_mole_frac_matr[j, i] * np.sqrt(
                            2 * const.Boltzmann * self.T / self.m0) * self.binary_d_matr[i, j, 1]
                    Dij_matr[j, i] = Dij_matr[i, j]
                else:
                    Dij_matr[i, j] = np.nan
                    Dij_matr[j, i] = np.nan

        return Dij_matr

    def get_hard_sphere_radius(self, comps):
        '''
        Get Barker-Henderson diameters for each pair of particles.
        Using Lorentz-Berthleot rules for combining Mie-parameters for each pair of particles
        Gets Mie-parameters from the file mie_dev.xlsx.

        :param comps: Comma seperated list of components
        :return: N x N matrix of hard sphere diameters, where sigma_ij = 0.5 * (sigma_i + sigma_j),
                such that the diagonal is the radius of each component, and off-diagonals are the average diameter of
                component i and j.
        '''
        df = pd.read_excel(os.path.dirname(__file__) + '/mie_dev.xlsx')
        sigma_i = np.array([df.loc[df['comp'] == comp]['sigma'].iloc[0] for comp in comps.split(',')])
        epsilon_i = np.array([df.loc[df['comp'] == comp]['epsilon'].iloc[0] for comp in comps.split(',')])
        sigma_ij = 0.5 * np.sum(np.meshgrid(sigma_i, np.vstack(sigma_i)), axis=0)
        epsilon_ij = np.sqrt(epsilon_i * np.vstack(epsilon_i))
        epsilon_ij *= const.Boltzmann
        sigma_ij *= 1e-10

        return np.array([[quad(self.BH_integrand, 0, sigma_ij[i, j], args=(sigma_ij[i, j], epsilon_ij[i, j]))[0]
                          for i in range(len(sigma_ij))]
                         for j in range(len(sigma_ij))])

    def get_sigma(self, comps):
        df = pd.read_excel(os.path.dirname(__file__) + '/mie_dev.xlsx')
        sigma_i = np.array([df.loc[df['comp'] == comp]['sigma'].iloc[0] for comp in comps.split(',')]) * 1e-10

        sigma_ij = 0.5 * np.sum(np.meshgrid(sigma_i, np.vstack(sigma_i)), axis=0)
        return sigma_ij

    def BH_integrand(self, r, sigma, epsilon):
        lambda_r = 12
        lambda_a = 6
        return (1 - np.exp(-self.u_Mie(r, sigma, epsilon, lambda_r, lambda_a) / (self.T * const.Boltzmann)))

    def u_Mie(self, r, sigma, epsilon, lambda_r, lambda_a):
        C = lambda_r / (lambda_r - lambda_a) * (lambda_r / lambda_a) ** (lambda_a / (lambda_r - lambda_a))
        return C * epsilon * ((sigma / r) ** lambda_r - (sigma / r) ** lambda_a)