import matplotlib.pyplot as plt
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 = {} #dict of dicts of points

    def __init__(self, comps, mole_weights, mole_fracs, T, p, N):
        '''
        :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.
        '''

        #Packing out variables in an n-component compatible way
        self.T = T
        self.p = p
        self.mole_weights = np.copy(mole_weights) * 1e-3 / const.Avogadro #kg

        self.mole_fracs = np.copy(mole_fracs)
        self.m0 = np.sum(self.mole_weights)
        self.M = self.mole_weights/self.m0

        self.M1, self.M2 = self.M
        self.x1, self.x2 = self.mole_fracs

        self.sigmaij = self.get_hard_sphere_radius(comps)

        self.n = p / (const.Boltzmann * T) # 1 / m^3

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

        if comps in KineticGas.computed_points.keys():
            if (self.x1, self.x2, self.T, self.p, self.N) in KineticGas.computed_points[comps].keys():
                self.d_1, self.d0, self.d1 = KineticGas.computed_points[comps][(self.x1, self.x2, self.T, self.p, self.N)]
            else:
                self.kin = cpp_KineticGas(self.mole_weights, self.sigmaij, mole_fracs, T, p, N)

                A = self.kin.A_matrix
                A = np.array(A)
                delta = self.kin.delta_vector

                d = lin.solve(A, delta, assume_a='sym')
                if N > 0:
                    self.d_1, self.d0, self.d1 = d[N - 1], d[N], d[N + 1]
                else:
                    self.d_1, self.d0, self.d1 = 0, d[N], 0

                KineticGas.computed_points[comps][(self.x1, self.x2, self.T, self.p, self.N)] = (self.d_1, self.d0, self.d1)

        else:
            self.kin = cpp_KineticGas(self.mole_weights, self.sigmaij, mole_fracs, T, p, N)

            A = self.kin.A_matrix
            A = np.array(A)
            delta = self.kin.delta_vector

            d = lin.solve(A, delta, assume_a='sym')
            if N > 0:
                self.d_1, self.d0, self.d1 = d[N - 1], d[N], d[N + 1]
            else:
                self.d_1, self.d0, self.d1 = 0, d[N], 0

            KineticGas.computed_points[comps] = {(self.x1, self.x2, self.T, self.p, self.N) : (self.d_1, self.d0, self.d1)}

    def alpha_T0(self, T):
        kT = - (5 / (2 * self.d0)) * ((self.x1 * self.d1 / np.sqrt(self.M1)) + (self.x2 * self.d_1 / np.sqrt(self.M2)))
        return T * np.array([kT, - (self.x1 / self.x2) * kT]) / (np.array([self.x2, self.x1]))

    def interdiffusion(self, T):
        return 0.5 * self.x1 * self.x2 * np.sqrt(2 * const.Boltzmann * T / self.m0) * self.d0

    def thermal_diffusion(self, T):
        return - (5/4) * self.x1 * self.x2 * np.sqrt(2 * const.Boltzmann * T / (self.m0 * 1e-3)) * \
               ((self.x1 * self.d1 / np.sqrt(self.M1)) + (self.x2 * self.d_1/np.sqrt(self.M2)))

    def get_hard_sphere_radius(self, comps):
        '''
        Get hard-sphere diameters, assumed to be equal to Mie-potential sigma parameter. 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 = 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 * np.vstack(epsilon))
        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)