import numpy as np
import matplotlib.pyplot as plt
from pyctp import saftvrmie


def mid(lst):
    return lst[1:] - 0.5*np.diff(lst), np.diff(lst)


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

T, p = 300, 1e5

t = np.linspace(1, 50)
n1 = np.sqrt(t)
n2 = 1.5 * t + np.sin(np.pi * t)
nT = n1 + n2

x1 = n1/nT
x2 = n2/nT

f1tv = np.empty_like(nT)
f1tp = np.empty_like(nT)
z = np.empty_like(nT)
V_list = np.empty_like(nT)
for i in range(len(nT)):
    ph = eos.guess_phase(T, p, [x1[i], x2[i]])
    V_list[i] = eos.specific_volume(T, p, [x1[i], x2[i]], ph)[0]
    V = V_list[i] * nT[i]
    f1tv[i], _ = eos.fugacity_tv(T, V, [n1[i], n2[i]])[0]
    f1tp[i], _ = eos.thermo(T, p, [x1[i], x2[i]], 2)[0]
    z[i], = eos.zfac(T, p, [x1[i], x2[i]], 2)


f1tv, df1tv = mid(f1tv)
f1tp, df1tp = mid(f1tp)
z, dz = mid(z)

x1, dx1 = mid(x1)
x2, dx2 = mid(x2)
n1, dn1 = mid(n1)
n2, dn2 = mid(n2)
nT, dnT = mid(nT)
V, _ = mid(V_list)
t, dt = mid(t)


fig, axs = plt.subplots(3, 1, sharex='all')
ax1, ax2, ax3 = axs
ax1.plot(t, f1tv - np.log(x1 * p)) # f1tv = fugacity_tv(T, V, n)
ax1.plot(t, np.exp(f1tp))          # f1tp = thermo(T, p, x, guess_phase)
ax1.legend()


ax2.plot(t, V / nT)
ax2.set_ylabel(r'$V_m$')

ax3.plot(t, z)
ax3.set_ylabel('Z')
plt.savefig('test_fug_2')