print('Starter opp...')
from pyctp import saftvrmie
from models.kempers01 import Kempers01
import numpy as np
import matplotlib.pyplot as plt
print('Nu kjør vi!!')
comps = 'XE,AR'

temp = 300
pres = -518.9e5
phase = 1

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


def dmudn_TP(x, total_moles, temp=300, pres=1e5):
    '''
    Calculate chemical potential derivative with respect to number of moles at constant temperature and pressure
    :return: ndarray, dmudn[i,j] = dmu_idn_j
    '''
    global eos, phase

    x = np.array(x)

    v, dvdn = eos.specific_volume(temp, pres, x, phase, dvdn=True)
    mu, dmudn_TV = eos.chemical_potential_tv(temp, v * total_moles,
                                                  x * total_moles, dmudn=True)
    pres, dpdn = eos.pressure_tv(temp, v * total_moles, x * total_moles, dpdn=True)

    return dmudn_TV - np.tensordot(dpdn, dvdn, axes=0)


def dmudx_TP(x, total_moles, temp=300, pres=1e5):
    '''
    Calculate chemical potential derivative with respect to mole fraction of components
    at constant temperature and pressure
    :return: ndarray, dmudx[i,j] = dmu_idn_j
    '''

    x = np.array(x)

    dmudn = dmudn_TP(x, total_moles, temp=temp, pres=pres)

    M1 = (np.tensordot(np.ones(len(x)), -1 / x, axes=0) +
          (np.identity(len(x)) * (1 / x + 1 / (1 - x)))) * total_moles / len(x)

    return np.dot(dmudn, M1)/2

x = [0.44574155, 0.55425845]

T_list = np.linspace(120, 123)
mu1_list = np.empty_like(T_list)
mu2_list = np.empty_like(T_list)
rho_list = np.empty_like(T_list)

dmu1dx1_list = np.empty_like(T_list)
dmu2dx1_list = np.empty_like(T_list)

for i, T in enumerate(T_list):
    phase = eos.guess_phase(T, pres, x)
    V, = eos.specific_volume(T, pres, x, phase)
    mu, = eos.chemical_potential_tv(T, V, x)
    dmudx = dmudx_TP(x, 1, T, pres)
    rho_list[i] = 1 / (V * 1e3)
    mu1_list[i] = mu[0]
    mu2_list[i] = mu[1]
    dmu1dx1_list[i] = dmudx[0, 0]
    dmu2dx1_list[i] = dmudx[1, 0]

fig, axs = plt.subplots(3, 1, sharex='all', figsize=(10,5))
ax1, ax2, ax3 = axs
tw1 = ax1.twinx()
tw2 = ax2.twinx()

ax1.plot(T_list, mu1_list, color='r')
tw1.plot(T_list, mu2_list, color='b')
ax1.set_ylabel(r'$\mu_1$', color='r')
tw1.set_ylabel(r'$\mu_2$', color='b')

ax2.plot(T_list, dmu1dx1_list, color='r')
tw2.plot(T_list, dmu2dx1_list, color='b')
ax2.set_ylabel(r'$d\mu_1 / dx_1$', color='r')
tw2.set_ylabel(r'$d\mu_2 / dx_1$', color='b')

ax3.plot(T_list, rho_list)
ax3.set_ylabel(r'$\rho$ [kmol/m$^3$]')
plt.suptitle(comps)
plt.savefig('test_dmudx_'+comps.replace(',', '_'))

exit(0)
N = 100
t = np.linspace(0, 5, N)
n1 = 3 + t
n2 = 1 + 2 * t
nT = n1 + n2

x1 = n1 / nT
x2 = n2 / nT

dx1 = np.diff(x1)
dn1 = np.diff(n1)
dx2 = np.diff(x2)
dn2 = np.diff(n2)

mu1_list = np.empty_like(x1)
mu2_list = np.empty_like(x2)

dmu1dn_list = np.empty_like(dx1)
dmu2dn_list = np.empty_like(dx1)

dmu1dx_list = np.empty_like(dx1)
dmu2dx_list = np.empty_like(dx1)

for i in range(N-1):
    V = eos.specific_volume(temp, pres, [x1[i], x2[i]], 2)[0] * nT[i]
    mu, = eos.chemical_potential_tv(temp, V, [n1[i], n2[i]])
    mu1_list[i], mu2_list[i] = mu

    dmudn = dmudn_TP([x1[i], x2[i]], nT[i])
    print(dmudn[0, 0] * dn1[i])
    dmu1dn_list[i] = dmudn[0, 0] * dn1[i] + dmudn[0, 1] * dn2[i]
    dmu2dn_list[i] = dmudn[1, 0] * dn1[i] + dmudn[1, 1] * dn2[i]

    dmudx = dmudx_TP([x1[i], x2[i]], nT[i])
    dmu1dx_list[i] = dmudx[0, 0] * dx1[i] + dmudx[0, 1] * dx2[i]
    dmu2dx_list[i] = dmudx[1, 0] * dx1[i] + dmudx[1, 1] * dx2[i]

V = eos.specific_volume(temp, pres, [x1[-1], x2[-1]], 2)[0] * nT[-1]
mu, = eos.chemical_potential_tv(temp, V, [n1[-1], n2[-1]])
mu1_list[-1], mu2_list[-1] = mu

dmu1 = np.diff(mu1_list)
dmu2 = np.diff(mu2_list)

plt.plot(t[:-1], dmu1, color='r', label='1')
plt.plot(t[:-1], dmu1dx_list, color='r', linestyle='--', marker='v', markevery=5)
plt.plot(t[:-1], dmu2, color='b', label='2')
plt.plot(t[:-1], dmu2dx_list, color='b', linestyle='--', marker='x', markevery=5)
plt.legend()
plt.savefig('dmudn')