# from plotting import kin_dataplotter, plot_alpha_t0, dataplotter
from pyctp import cubic, cpa, pcsaft, extended_csp, multiparameter, saftvrmie
import numpy as np
import matplotlib.pyplot as plt
# import matplotlib.gridspec as gs
# import time
from models.kempers01 import Kempers01
# import pandas as pd
# import warnings
# import file_handling as fh
# import plotting
from matplotlib.cm import get_cmap
# import tools
import ternary
from models.test_KineticGas import KineticGas
from scipy.optimize import curve_fit
import pandas as pd
from scipy import constants as const

#eos = cubic.cubic()
#print('Loaded EOS')
#eos.init('H2O,O2,N2', 'SW')
#print('Initialized EOS')
#
#model = Kempers01('H2O,O2,N2', eos=eos, x=[0.1, 0.9*0.2, 0.9*0.8], temp=450, pres=1e5, phase=2)
#print('Initialized Kempers01')
#temps = np.linspace(25, 300, 3) + 273
#
#soret, kin = model.get_soret_temp(temps, mode='cov', kin=True)

print('Running main')

comps = 'AR,KR'

mie_df = pd.read_excel('models/mie_dev.xlsx')
sigma = np.array([mie_df.loc[mie_df['comp'] == comp]['sigma'].iloc[0] for comp in comps.split(',')]) * 1e-10
epsilon = np.array([mie_df.loc[mie_df['comp'] == comp]['epsilon'].iloc[0] for comp in comps.split(',')]) * const.Boltzmann

sigma_12 = 0.5 * sum(sigma)
epsilon_12 = np.sqrt(np.product(epsilon))

reduced_rho = np.array([0.797, 0.81])
reduced_T = np.array([0.805, 0.85])

rho_func = lambda r_rho: r_rho / (sigma_12 ** 3)
T_func = lambda r_T: r_T * epsilon_12 / const.Boltzmann

rho = reduced_rho / (sigma_12 ** 3)
T = reduced_T * epsilon_12 / const.Boltzmann
id_gas_p = rho * const.Boltzmann * T

n = 1500 / const.Avogadro
V = 1500 / rho

eos = saftvrmie.saftvrmie()
eos.init(comps)

pres1, = eos.pressure_tv(T[0], V[0], [n/2, n/2])
pres2, = eos.pressure_tv(T[1], V[1], [n/2, n/2])

print('Ideal gas pressure : ', id_gas_p/1e5)
print('Real pressure (bar): ', np.array([pres1, pres2]) / 1e5)

model = Kempers01(comps, eos, x=[0.5,0.5], pres=pres1, temp=T[0], particle_density=rho[0])

print('Reported values: ', '10.5 ± 1.3, 14.4 ± 1.2')

print('Using Thermopack computed pressure')

print('T = ', round(T[0], 0), ', p = ', round(pres1/1e5, 0), sep='')
print('COV :', [x*1e3 for x in model.get_soret_cov(kin=True)])
print('COM :', [x*1e3 for x in model.get_soret_com(kin=True)])

model.set_temp(T[1])
model.set_pres(pres2, particle_density=rho[1])

print('T = ', round(T[1], 0), ', p = ', round(pres2/1e5, 0), sep='')
print('COV :', [x*1e3 for x in model.get_soret_cov(kin=True)])
print('COM :', [x*1e3 for x in model.get_soret_com(kin=True)])

print('\nUsing ideal gas pressure :')

model.set_temp(T[0])
model.set_pres(id_gas_p[0], particle_density=rho[0])

print('T = ', round(T[0], 0), ', p = ', round(id_gas_p[0]/1e5, 0), sep='')
print('COV :', [x*1e3 for x in model.get_soret_cov(kin=True)])
print('COM :', [x*1e3 for x in model.get_soret_com(kin=True)])

model.set_temp(T[1])
model.set_pres(id_gas_p[1], particle_density=rho[1])

print('T = ', round(T[1], 0), ', p = ', round(id_gas_p[1]/1e5, 0), sep='')
print('COV :', [x*1e3 for x in model.get_soret_cov(kin=True)])
print('COM :', [x*1e3 for x in model.get_soret_com(kin=True)])

'''
for comp, vals in zip(['H2O', 'O2', 'N2'], kin):
    plt.plot(temps, vals, label=comp)

x = np.array([[0.1], [0.9*0.2], [0.9*0.8]])
plt.plot(temps, np.sum(soret * x * (1 - x), axis=0), color='black', label='Sum')
plt.legend()
plt.title(r'$x_{H_2O} = 0.1$')
plt.savefig('kinetic_multicomponent')

x_H2O = np.linspace(0, 1)
x_N2 = (1 - x_H2O) * 0.8
x_O2 = (1 - x_H2O) * 0.2

x_arr = np.array([x_H2O, x_O2, x_N2]).transpose()

cmap = get_cmap('cool')
for T in temps:
    model.set_temp(T)
    kin, _ = model.get_soret_comp(x_arr, mode='cov', kin=True)
    plt.plot(x_H2O, kin[0], color=cmap(T/max(temps)), label=r'H$_2$O ' + str(round(T,0)) + 'K')
    plt.plot(x_H2O, kin[1], color=cmap(T/max(temps)), label=r'O$_2$ ' + str(round(T,0)) + 'K', linestyle='--')
    plt.plot(x_H2O, kin[2], color=cmap(T / max(temps)), label=r'N$_2$ ' + str(round(T, 0)) + 'K', linestyle=':')
    print(T)

plt.xlabel(r'$x_{H_2O}$')
plt.ylabel(r'$S_T$ [K$^{-1}$]')
plt.legend(loc='lower right')
plt.savefig('multicomp_soret')


cmap_tri = get_cmap('bwr')

x_H2O = np.linspace(0, 1)
x_N2 = (1 - x_H2O) * 0.8
x_O2 = (1 - x_H2O) * 0.2

x_arr = np.array([x_H2O, x_O2, x_N2]).transpose()

soret_H2O_vals = np.loadtxt('ternary_kin.txt')
soret_values = []
kin_values = []
i = 0
def soret_H2O(x):
    global i
    print(x)
    #soret, kin = model.get_soret_comp(x, mode='cov', kin=True)
    #soret_values.append(soret.flatten())
    #kin_values.append(kin.flatten())
    soret = soret_H2O_vals[i]
    #soret_values.append(soret[0])
    i += 1
    return soret[0]


def test_tax_corners(x):
    return (x[0] - x[1]) * (1 - x[2])
scale = 60

fig, tax = ternary.figure(scale=scale)
tax.set_axis_limits({'b' : [0,1], 'l' : [0, 1], 'r' : [0, 1]})
tax.heatmapf(soret_H2O, boundary=False,
             style='triangular', cmap=cmap_tri,
             vmin=-0.26, vmax=0.26)


tax.plot(x_arr*scale, color='lime', label=r'H$_2$O in Air')

offset = 0.14
tax.boundary(linewidth=1.0)
tax.set_title(r'S$_{T, H_2O}^{0}$'+'\nT = 450 K', loc='left')
tax.top_corner_label(r'O$_2$')
tax.right_axis_label(r'x$_{O_2}$', offset=offset)
tax.right_corner_label(r'H$_2$O')
tax.bottom_axis_label(r'x$_{H_2O}$', offset=offset)
tax.left_corner_label(r'N$_2$')
tax.left_axis_label(r'x$_{N_2}$', offset=offset)
tax._redraw_labels()
tax.get_ticks_from_axis_limits(multiple=11)
tax.set_custom_ticks(tick_formats='%.1f', offset=0.025)
tax.get_axes().axis('off')
tax.clear_matplotlib_ticks()
tax.legend(loc='upper left')
plt.savefig('ternary_kin_450K')
tax.show()

#np.savetxt('ternary_soret_450K.txt', soret_values)
#np.savetxt('ternary_kin.txt', kin_values)
'''