import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from models.test_KineticGas import KineticGas
from models.py_KineticGas_11_10 import KineticGas as real_KineticGas
from pycThermopack.pyctp import cubic
from scipy.optimize import curve_fit
from matplotlib.cm import get_cmap
from matplotlib.colors import Normalize
from scipy import constants as const
import pandas as pd

#Ar, He
mole_weights = np.array([39.95, 4])
sigma_i = 3.405 * 1e-10 # m
epsilon_i = 122.1 * const.Boltzmann
sigma_j = 2.64 * 1e-10 # m
epsilon_j = 10.9 * const.Boltzmann

sigma_ij = [[sigma_i, 0.5 * (sigma_i + sigma_j)], [0.5 * (sigma_i + sigma_j), sigma_j]]
epsilon = [epsilon_i, epsilon_j]

comps = 'AR,HE'

N = 5
diffusion_volume = [16.1, 2.88]

def torr_2_Pa(p_torr):
    return p_torr * (1e5/760)

def atm_2_Pa(p_atm):
    return p_atm * 1.01325 * 1e5

def Fuller(T, p, mole_weights, volumes): #https://pubs.acs.org/doi/pdf/10.1021/ie50677a007?casa_token=i4bNazNv9UAAAAAA:Ie6wz8medLTHGkONA75DQ8N1LHh7utmhcIIhbDTU0LdrLPqaCC81I-hy83BB1UBM1GMi0KH1Dc3ghg
    mA, mB = mole_weights
    vA, vB = volumes
    return 1e-3 * T**1.75 * np.sqrt((1/mA) + (1/mB)) / ( p * ((vA**(1/3) + vB**(1/3))**2))

def Chapman_Enskog(T, p, mole_weights, sigma_ij):
    erg_Boltzmann = const.Boltzmann * 1e7
    p *= 10
    n = p / (erg_Boltzmann * T) # molecules / cm^3
    sigma = 0.5 * (sigma_ij[0][0] + sigma_ij[1][1]) * 1e2 #cm
    mole_weights = mole_weights / const.Avogadro
    return (3/(32 * n * sigma**2)) * np.sqrt((8 * erg_Boltzmann * T / np.pi) \
           * ((1 / (mole_weights[0])) + (1/ (mole_weights[1]))))


def kinetic_x(x, sigma_i, sigma_j):
    global mole_weights
    kin_list = [KineticGas(comps, mole_weights, [xi, 1-xi], 300, 1e5, 5) for xi in x]
    return [k.interdiffusion(300) for k in kin_list]

def kinetic_T(T, p, sigma_i, sigma_j):
    global mole_weights
    kin_list = [KineticGas(comps, mole_weights, [0.5, 0.5], Ti, p, 5) for Ti in T]
    return np.array([k.interdiffusion(Ti) for k, Ti in zip(kin_list, T)])

def zero_order_kinetic(T, p, sigma_ij, mole_weights):

    kin_list = [KineticGas(comps, mole_weights, [0.5, 0.5], Ti, p, 0) for Ti in T]
    return np.array([k.interdiffusion(Ti) for k, Ti in zip(kin_list, T)])


def fit_T_dependence(T, A, B):
    return A * T ** B

def test_T():
    T, p = 300, 1
    T_list = np.linspace(250, 500)

    D12_fuller = Fuller(T_list, p, mole_weights, diffusion_volume)
    D12_chapman_enskog = Chapman_Enskog(T_list, atm_2_Pa(p), mole_weights, sigma_ij)
    D12_kinetic = kinetic_T(T_list, atm_2_Pa(p), sigma_i, sigma_j)
    D12_zero_order = zero_order_kinetic(T_list, atm_2_Pa(p), sigma_ij, mole_weights)

    coeff, cov = curve_fit(fit_T_dependence, T_list, D12_kinetic)

    plt.plot(T_list, D12_fuller * 1e-4, label='Fuller')
    plt.plot(T_list, D12_kinetic, label='5. order Kinetic, $d_{BH}$')
    plt.plot(T_list, D12_zero_order, label=r'0. order Kinetic, $d_{BH}$')
    plt.plot(T_list, fit_T_dependence(T_list, coeff[0], coeff[1]), linestyle='--', label='%.1E' % coeff[0] + r' $\cdot T^{'+str(round(coeff[1], 2)) + '}$')


    plt.ticklabel_format(axis='Y', style='sci', scilimits=(0,0))
    plt.ylabel(r'D$_{AB}$ [m$^2$s$^{-1}$]')
    plt.xlabel('T [K]')
    plt.legend()
    plt.title(r'$x_{Ar} = x_{He} = 0.5$')
    plt.show()

    x1_vals = np.arange(0.1, 0.91, step=0.2)
    cmap = get_cmap('viridis')
    norm = Normalize(0.1, 0.9)
    for x1 in x1_vals:
        kin_list = [KineticGas(comps, mole_weights, [x1, 1-x1], Ti, atm_2_Pa(p), 5) for Ti in T_list]
        plt.plot(T_list, [k.interdiffusion(Ti) for k, Ti in zip(kin_list, T_list)], label=round(x1,1), color=cmap(norm(x1)))

    plt.plot(T_list, D12_fuller * 1e-4, label='Fuller', color='r')
    plt.legend()
    plt.show()


    kin_list = [KineticGas(comps, mole_weights, [0.5, 0.5], Ti, p, 0) for Ti in T_list]
    d_BH_1_list = np.array([k.get_hard_sphere_radius(comps)[0, 0] for k in kin_list])
    d_BH_2_list = np.array([k.get_hard_sphere_radius(comps)[1, 1] for k in kin_list])
    sigma_1_list = np.array([k.get_sigma(comps)[0, 0] for k in kin_list])
    sigma_2_list = np.array([k.get_sigma(comps)[1, 1] for k in kin_list])

    plt.plot(T_list, d_BH_1_list / sigma_1_list, label=r'$d^{BH}_{1} / \sigma_{1}$')
    plt.plot(T_list, d_BH_2_list / sigma_2_list, label=r'$d^{BH}_{2} / \sigma_{2}$')
    plt.plot(T_list, (d_BH_1_list + d_BH_2_list) / (sigma_1_list + sigma_2_list), label=r'$d^{BH}_{12} / \sigma_{12}$')
    plt.xlabel('T [K]')
    plt.legend()
    plt.title(r'$\epsilon_{1} = 122.1 k_B$,  $\epsilon_{2} = 10.9 k_B$')
    plt.show()

def test_x():
    p = 1
    x1_list = np.linspace(0.001, 0.999, 50)
    x2_list = 1 - x1_list

    D12_fuller = Fuller(np.full_like(x1_list, 300), p, mole_weights, diffusion_volume)
    #D12_chapman_enskog = Chapman_Enskog(np.full_like(x1_list, 300), atm_2_Pa(p), mole_weights, sigma_ij)
    D12_kinetic = kinetic_x(x1_list, sigma_i, sigma_j)
    D12_zero_order = zero_order_kinetic(np.full_like(x1_list, 300), atm_2_Pa(p), sigma_ij, mole_weights)

    plt.plot(x1_list, D12_fuller * 1e-4, label='Fuller')
    #plt.plot(x1_list, D12_chapman_enskog * 1e-4, label='Chapman-Enskog')
    plt.plot(x1_list, D12_kinetic, label=r'5. order Kinetic, $d_{BH}$')
    plt.plot(x1_list, D12_zero_order, label=r'0. order Kinetic, $d_{BH}$')

    plt.ticklabel_format(axis='Y', style='sci', scilimits=(0,0))
    plt.ylabel(r'D$_{AB}$ [m$^2$s$^{-1}$]')
    plt.xlabel(r'x$_{Ar}$ [-]')
    plt.legend()
    plt.title('T = 300 K')
    plt.show()

def test_p():
    p_list = np.linspace(0.5, 1.5)
    D12_fuller = Fuller(300, p_list, mole_weights, diffusion_volume) * 1e-4

    kin5_list = [KineticGas(comps, mole_weights, [0.5, 0.5], 300, atm_2_Pa(pi), 5) for pi in p_list]
    D12_5_order = [k.interdiffusion(300) for k in kin5_list]

    kin0_list = [KineticGas(comps, mole_weights, [0.5, 0.5], 300, atm_2_Pa(pi), 0) for pi in p_list]
    D12_0_order = [k.interdiffusion(300) for k in kin0_list]

    plt.plot(p_list, D12_fuller, label='Fuller')
    plt.plot(p_list, D12_5_order, label='5. order')
    plt.plot(p_list, D12_0_order, label='0. order')

    plt.ticklabel_format(axis='Y', style='sci', scilimits=(0,0))
    plt.ylabel(r'D$_{AB}$ [m$^2$s$^{-1}$]')
    plt.xlabel(r'p [atm]')
    plt.legend()
    plt.title(r'T = 300 K, $x_{Ar} = x_{He} = 0.5$')
    plt.show()

def test_multi():
    p_list = np.linspace(0.5, 1.5)
    D12_fuller = Fuller(300, p_list, mole_weights, diffusion_volume) * 1e-4

    kin5_list = [KineticGas(comps, mole_weights, [0.5, 0.5], 300, atm_2_Pa(pi), 5) for pi in p_list]
    D12_5_order = [k.interdiffusion(300) for k in kin5_list]

    kin_real_list = [real_KineticGas(comps, None, [0.5, 0.5], 300, atm_2_Pa(pi), 5, mole_weights=mole_weights) for pi in p_list]
    D12_real = [k.interdiffusion()[0] for k in kin_real_list]

    plt.plot(p_list, D12_fuller, label='Fuller')
    plt.plot(p_list, D12_5_order, label='5. order')
    plt.plot(p_list, D12_real, label='multicomp', linestyle='--')

    plt.ticklabel_format(axis='Y', style='sci', scilimits=(0,0))
    plt.ylabel(r'D$_{AB}$ [m$^2$s$^{-1}$]')
    plt.xlabel(r'p [atm]')
    plt.legend()
    plt.title(r'T = 300 K, $x_{Ar} = x_{He} = 0.5$')
    plt.show()

def test_ternary():
    t_comps = 'AR,KR,NE'
    df = pd.read_csv('data/Ar_Kr_Ne_interdiff.csv', sep=', ')
    mie_df = pd.read_excel('models/mie_dev.xlsx')

    sigma_i = np.array([mie_df.loc[mie_df['comp'] == comp]['sigma'].iloc[0] for comp in t_comps.split(',')]) * 1e-10
    epsilon_i = np.array([mie_df.loc[mie_df['comp'] == comp]['epsilon'].iloc[0] for comp in t_comps.split(',')]) * const.Boltzmann

    x_Ar = np.array(df['x_Ar'].tolist())
    x_Kr = np.array(df['x_Kr'].tolist())
    x_Ne = np.array(df['x_Ne'].tolist())
    D12_data = np.array(df['D12'])
    D31_data = np.array(df['D31'])
    D23_data = np.array(df['D23'])
    reduced_T = 0.8
    reduced_rho = 0.7901
    reduced_p = np.array(df['p'])

    red_eps = epsilon_i[0]
    red_sig = sigma_i[0]

    T = reduced_T * red_eps / const.Boltzmann
    rho = reduced_rho / red_sig**3
    p = reduced_p * (red_eps / red_sig**3)
    id_gas_p = rho * const.Boltzmann * T
    N = 300

    M_Ar = 39.948 / const.Avogadro
    M_Kr = 83.798 / const.Avogadro
    M_Ne = 20.1797 / const.Avogadro

    m = N * (x_Ar * M_Ar + x_Ne * M_Ne + x_Kr * M_Kr)

    V = m / rho
    x = [x_Ar, x_Kr, x_Ne]

    print('p : ', p / 1e5)
    print('T : ', T)
    print('rho : ', rho)
    print('V :', V)
    print('filled V :', sum([4/3 * np.pi * (sigma_i[i]/2)**3 * N * x[i] for i in range(3)]))

    dimless_D_factor = np.sqrt(m/red_eps) / red_sig

    plt.scatter(x_Ar, D12_data, label=r'$D_{Ar,Kr}$')
    plt.scatter(x_Ar, D31_data, label=r'$D_{Ar,Ne}$')
    plt.scatter(x_Ar, D23_data, label=r'$D_{Kr,Ne}$')
    plt.title(r'$x_{Ne} = 0.1$')
    plt.xlabel(r'$x_{Ar}$')
    plt.ylabel(r'D [-]')
    plt.legend()
    #plt.show()

    plt.clf()

    fig, axs = plt.subplots(2,1)

    plt.sca(axs[0])

    plt.scatter(x_Ar, D12_data/dimless_D_factor, label=r'$D_{Ar,Kr}$', color='r')
    plt.scatter(x_Ar, D31_data/dimless_D_factor, label=r'$D_{Ar,Ne}$', color='g')
    plt.scatter(x_Ar, D23_data/dimless_D_factor, label=r'$D_{Kr,Ne}$', color='b')
    #plt.title(r'$x_{Ne} = 0.1$')
    plt.xlabel(r'$x_{Ar}$')
    plt.ylabel(r'D [m$^2$s$^{-1}$]')
    plt.legend()


    kin_list = [real_KineticGas('AR,KR,NE', mole_fracs=[x_Ar[i], x_Kr[i], x_Ne[i]], T=T, p=id_gas_p, mole_weights=[M_Ar * const.Avogadro, M_Kr * const.Avogadro, M_Ne * const.Avogadro]) for i in range(len(x_Ar))]
    pred_diff = np.array([k.binary_diffusion() for k in kin_list])

    plt.sca(axs[1])

    plt.plot(x_Ar, pred_diff[:, 0, 1], label=r'$D_{Ar,Kr}$', color='r')
    plt.plot(x_Ar, pred_diff[:, 0, 2], label=r'$D_{Ar,Ne}$', color='g')
    plt.plot(x_Ar, pred_diff[:, 1, 2], label=r'$D_{Kr,Ne}$', color='b')
    #plt.title(r'$x_{Ne} = 0.1$')
    plt.xlabel(r'$x_{Ar}$')
    plt.ylabel(r'D [m$^2$s$^{-1}$]')
    plt.legend()

    plt.suptitle(r'$x_{Ne} = 0.1$')
    plt.show()

if __name__ == '__main__':
    #k = real_KineticGas('AR,AR', mole_fracs=[0.5, 0.5], T = 300, p = 1e5, mole_weights=[39.948, 39.948])
    #print(k.binary_diffusion())
    test_ternary()