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

T = 300
p = 1e3
V = 10

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

fig, axs = plt.subplots(2, 1)
ax1, ax2 = axs

for N in [5, 10, 20, 50]:
    t = np.linspace(1, 10, N)
    x1, x2 = 0.3, 0.7
    n1 = 1.5 * t
    n2 = 2.5 * t
    nT = n1 + n2
    nT0 = nT
    n10 = n1
    n20 = n2

    V, = eos.specific_volume(T, p, [x1, x2], 2)

    dn1 = np.diff(n1)
    dn2 = np.diff(n2)
    dnT = np.diff(nT)

    x10 = n1 / nT
    x20 = n2 / nT

    f1 = np.empty_like(nT0)
    f2 = np.empty_like(nT0)
    for i in range(len(nT0)):
        fug, = eos.fugacity_tv(T, V*nT[i], [n1[i], n2[i]])
        f1[i], f2[i] = fug

    dx1 = np.diff(x10)
    dx2 = np.diff(x20)

    n1 = n10[1:] + 0.5 * dn1
    n2 = n20[1:] + 0.5 * dn2
    nT = nT[1:] + 0.5 * dnT
    x1 = x10[1:] + 0.5 * dx1
    x2 = x20[1:] + 0.5 * dx2

    Df1 = np.empty_like(nT)
    Df2 = np.empty_like(nT)
    for i in range(len(nT)):
        _, dfdn = eos.fugacity_tv(T, V*nT[i], [n1[i], n2[i]], dlnphidn=True)
        Df1[i] = (dfdn[0, 0] * dn1[i] + dfdn[0, 1] * dn2[i])
        Df2[i] = (dfdn[1, 1] * dn2[i] + dfdn[1, 0] * dn1[i])

    #ax1.plot(nT0, np.concatenate(([f1[0]], f1[0] + np.array([sum(Df1[:i+1]) for i in range(len(Df1))]))), label=N, linestyle='--')
    #ax2.plot(nT0, np.concatenate(([f2[0]], f2[0] + np.array([sum(Df2[:i+1]) for i in range(len(Df1))]))), linestyle='--')

f1 = np.empty_like(nT0)
f2 = np.empty_like(nT0)
for i in range(len(nT0)):
    V, = eos.specific_volume(T, p, [0.3, 0.7], 2)
    fug, = eos.fugacity_tv(T, V*nT0[i], [n10[i], n20[i]])
    f1[i], f2[i] = fug

df1, df2 = np.diff(f1), np.diff(f2)
dn1 = np.diff(n10)
nT = nT0[1:] - 0.5 * dnT
n1 = 0.3 * nT

ax1.plot(nT, df1/dn1, color='b')
ax2.plot(nT, n1 * df1/dn1, color='b')

plt.figlegend()
plt.savefig('test_fug')