print('Running')
#from pyctp import cubic
print('#', end='')
import numpy as np
# print('#', end='')
import matplotlib.pyplot as plt
# print('#', end='')
from scipy.constants import Avogadro, Boltzmann, Planck
from scipy.constants import gas_constant as R
from scipy.optimize import root
from scipy import linalg as lin
# print('#', end='')
import os
print('\nImports finished')

z = np.linspace(1, 10, 1000)
n1 = np.sqrt(z)
n2 = 3 + 2 * np.sin(z * np.pi / 5)
n3 = 5 / z
n4 = np.log(z)
nT = n1 + n2 + n3 + n4
T = z**2

dn1 = np.diff(n1)
dn2 = np.diff(n2)
dn3 = np.diff(n3)
dn4 = np.diff(n4)
dnT = np.diff(nT)
dT = np.diff(T)

x1 = n1/nT
x2 = n2/nT
x3 = n3/nT
x4 = n4/nT

dx1 = np.diff(x1)
dx2 = np.diff(x2)
dx3 = np.diff(x3)
dx4 = np.diff(x4)

x1 = x1[:-1] + dx1
x2 = x2[:-1] + dx2
x3 = x3[:-1] + dx3
x4 = x4[:-1] + dx4

nT = nT[:-1] + dnT
n1 = n1[:-1] + dn1
n2 = n2[:-1] + dn2
n3 = n3[:-1] + dn3
n4 = n4[:-1] + dn4

T = T[:-1] + dT

Dx1 = (- x1 / nT) * (dn2 + dn3 + dn4) + (1 - x1) * dn1 / nT

fig, axs = plt.subplots(4, 1, sharex='all')
ax1, ax2, ax3, ax4 = axs
twn4 = ax4.twinx()
ax1.plot(z[1:], df1)
ax1.plot(z[1:], Df1, linestyle='--')
ax2.plot(z[1:], df2)
ax2.plot(z[1:], Df2)
ax3.plot(z[1:], dx1)
ax3.plot(z[1:], Dx1, linestyle='--')

ax4.plot(z[1:], n1)
ax4.plot(z[1:], n2)
ax4.plot(z[1:], n3)
ax4.plot(z[1:], n4)
twn4.plot(z[1:], T)

ax1.set_ylabel('df1')
ax2.set_ylabel('df2')
ax3.set_ylabel('f(T,n)')
ax4.set_ylabel(r'$n$')
twn4.set_ylabel('T')
plt.show()
exit(0)

n1_list = np.concatenate((np.linspace(1, 2, 100), np.ones(100) * 2))
n2_list = np.concatenate((np.ones(100) * 2, np.linspace(2, 3, 100)))
nT = n1_list + n2_list

eos = cubic.cubic()
eos.init('HE', 'VdW')

mu1 = np.empty_like(n1_list)
mu2 = np.empty_like(n2_list)

V = 6000
T = 1e-3

m1 = 4.0026 * 1e-3 / Avogadro
m2 = 20.1797 * 1e-3 / Avogadro

q_el_He = 2

q = lambda m, q_el: q_el * (((2 * np.pi * m * Boltzmann * T) / Planck**2) ** (3/2)) * V

mu1_id = -R * T * np.log(q(m1, q_el_He) / (n1_list * Avogadro))
mu2_id = -R * T * np.log(q(m2, 2) / (n2_list * Avogadro))

p = np.empty_like(nT)

for i in range(len(n2_list)):
    n1, n2 = n1_list[i], n2_list[i]
    mu, = eos.chemical_potential_tv(T, V, [n1])
    #mu1[i], mu2[i] = mu[0], mu[1]
    mu1[i] = mu[0]
    p[i], = eos.pressure_tv(T, V, [n1])

p_id = (n1_list * R * T) / V

plt.plot(n1_list, p, label='tpc')
plt.plot(n1_list, p_id, label='id', linestyle=':')
plt.legend()
plt.savefig('ideal_gas_test')
plt.clf()

fig, ax = plt.subplots(figsize=(10,5))
twn = ax.twinx()

delta_mu = mu1 - mu1_id

ax.plot(n1_list, mu1, label=r'$\mu_1$ (tpc)', color='b')
#ax.plot(nT, mu2, label=r'$\mu_2$ (tpc)', color='r')
ax.plot(n1_list, mu1_id, label=r'$\mu_1$ (id)', color='g', linestyle='--')
#ax.plot(nT, -R * T * np.log(q(m2) / (n2_list * Avogadro)), label=r'$\mu_2$ (id)', color='orange', linestyle=':')
twn.plot(n1_list, delta_mu, label=r'$\Delta \mu_1$', color='g')
#twn.plot(nT, mu2 / (-R * T * np.log(q(m2) / (n2_list * Avogadro))), label=r'$\Delta \mu_2$', color='orange')
#twn.plot(nT, mu1 / (-R * T * np.log(q(m1) / (n1_list * Avogadro))) + mu2 / (-R * T * np.log(q(m2) / (n2_list * Avogadro))), color='black')
plt.figlegend()
#twn.set_ylim(delta_mu[0] - 0.5, delta_mu[0] + 0.5)
plt.savefig('ideal_gas_mu')

exit(0)

nT = n1 + n2
x1 = n1 / nT
x2 = n2 / nT

V = 1
T = 300

m1 = 2
m2 = 9
m3 = 5
q = lambda m: ((2 * np.pi * m * Boltzmann * T) / Planck) ** (3/2) * V

q1 = q(m1)
q2 = q(m2)
q3 = q(m3)

Q1 = (np.e * q1 / n1) ** (n1)
Q2 = (np.e * q2 / n2) ** (n2)
Q3 = (np.e * q3 / n3) ** (n3)

Q = Q1 * Q2 * Q3

F = - R * T * (np.log(Q1) + np.log(Q2) + np.log(Q3)) / 1e3

# mu1 = (R / nT) * (np.log(x1) + x1 * np.log(x1) + x2 * np.log(x2) + x3 * np.log(x3))
# mu2 = (R / nT) * (np.log(x2) + x1 * np.log(x1) + x2 * np.log(x2) + x3 * np.log(x3))
# mu3 = (R / nT) * (np.log(x3) + x1 * np.log(x1) + x2 * np.log(x2) + x3 * np.log(x3))

mu1 = lambda n1: - R * T * (np.log(q1 / n1)) / 1e3
mu2 = lambda n2: - R * T * (np.log(q2 / n2)) / 1e3
mu3 = lambda n3: - R * T * (np.log(q3 / n3)) / 1e3

dF = np.diff(F)
dn1 = np.diff(n1)
dn2 = np.diff(n2)
dn3 = np.diff(n3)

nT_mid = nT[:-1] + np.diff(nT)
n1_mid = n1[:-1] + np.diff(n1)
n2_mid = n2[:-1] + np.diff(n2)
n3_mid = n3[:-1] + np.diff(n3)

fig, ax = plt.subplots()
twn = ax.twinx()
#ax.plot(nT, F, color='r')
twn.plot(nT, F, color='g')
ax.plot(nT_mid, mu1(n1_mid)*dn1 + mu2(n2_mid)*dn2 + mu3(n3_mid)*dn3)
ax.plot(nT_mid, dF, color='r', linestyle='--')
ax.set_ylabel('dF')
twn.set_ylabel('F')
plt.show()

'''

s = 1.9
e = 1

rho1_list = np.linspace(20,40)*1e3
T1 = 130

eos = saftvrmie.saftvrmie()
eos.init('AR')

params = eos.get_pure_fluid_param(1)
_, sigma1, eps1_div_k, _, _ = params
print('Params 1 :', params)

red_rho = rho1_list * (sigma1**3) * Avogadro
red_T = T1 / eps1_div_k

pres1 = np.array([eos.pressure_tv(130, 1, [rho])[0] for rho in rho1_list])
red_pres1 = pres1 * (sigma1**3) / (eps1_div_k * Boltzmann)

eos.set_pure_fluid_param(1, params[0], s * params[1], e * params[2], params[3], params[4])
eos.redefine_critical_parameters()

params = eos.get_pure_fluid_param(1)
_, sigma1, eps1_div_k, _, _ = params
print('Params 2 :', params)

rho2_list = red_rho / (Avogadro * (sigma1)**3)
T2 = red_T * eps1_div_k

pres2 = np.array([eos.pressure_tv(T2, 1, [rho])[0] for rho in rho2_list])
red_pres2 = pres1 * (sigma1**3) / (eps1_div_k * Boltzmann)

fig, axs = plt.subplots(2, 1)
ax1, ax2 = axs
ax1.plot(rho1_list/1e3, pres1/1e5, label='1')
ax1.plot(rho2_list/1e3, pres2/1e5, label='2')
ax1.set_xlabel(r'$\rho$ [kmol/m$^3$]')
ax1.set_ylabel(r'p [bar]')

ax2.plot(red_rho, red_pres1, label='1')
ax2.plot(red_rho, red_pres2, label='2')
ax2.set_xlabel(r'$\rho^*$ [-]')
ax2.set_ylabel(r'p$^*$ [-]')

plt.suptitle('AR, e = '+str(e)+', s = '+str(s))

save_path = 'CSP_e_'+str(e).replace('.','_')+'_s_'+str(s).replace('.','_')
plt.savefig(save_path)
print('Saved figure to :', save_path)
'''