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

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


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

N = 100

t = np.linspace(1, 10, 100)
n1 = t + np.sqrt(t)
n2 = 2 * t  + np.sin(t)
nT = n1 + n2

x1 = n1/nT
x2 = n2/nT

x1, dx1 = mid(x1)
x2, dx2 = mid(x2)
n1, dn1 = mid(n1)
n2, dn2 = mid(n2)
nT, dnT = mid(nT)

T, p = 300, 1e5

f1 = np.empty_like(nT)
for i in range(len(nT)):
    fug, dfugdn = eos.thermo(T, p, [x1[i], x2[i]], 2, dlnfugdn=True)
    f1[i] = fug[0]
    Df1 = dfugdn[0, 0] * dx1 + dfugdn[0, 1] * dx2
    Df2 = dfugdn[0, 0] * dn1 + dfugdn[0, 1] * dn2

t, _ = mid(t)

f1_test = np.array([f1[0] + sum(Df1[:i]) for i in range(len(Df1))])
f2_test = np.array([f1[0] + sum(Df2[:i]) for i in range(len(Df2))])

plt.plot(t, f1)
plt.plot(t, f1_test, linestyle='--')
plt.plot(t, f2_test, linestyle='--')
plt.savefig('test_fug_3')