print('Starter opp...')
from pyctp import saftvrmie
from models.kempers01 import Kempers01
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import root
from scipy.constants import Avogadro
from scipy.constants import gas_constant as R
print('Nu kjør vi!')


class IdealGas:

    def __init__(self):
        pass

    def pressure(self, T, V, n):
        nT = sum(n)
        return R * T * nT / V

    def specific_volume(self, T, p):
        return R * T / p

comps = 'KR,AR'

T_red = 0.85
rho_red = 0.81
x = np.array([0.5, 0.5])

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

m1, sig_1, eps_1_div_k, la1, lr1 = eos.get_pure_fluid_param(1)
m2, sig_2, eps_2_div_k, la2, lr2 = eos.get_pure_fluid_param(2)
eos.set_pure_fluid_param(1, m1, sig_1, eps_1_div_k, 6, 12)
eos.set_pure_fluid_param(2, m2, sig_2, eps_2_div_k, 6, 12)
eos.redefine_critical_parameters()

sig_12 = 0.5 * (sig_1 + sig_2)
eps_12_div_k = np.sqrt(eps_1_div_k * eps_2_div_k)

T = T_red * eps_12_div_k
rho = rho_red / (sig_12 ** 3)

pres = eos.pressure_tv(T, Avogadro / rho, x)[0]

nT = 1
V = Avogadro / rho

delta_T = 0.3 * T
TA = T - 0.5 * delta_T
TB = TA + delta_T

kempers = Kempers01(eos, comps, temp=T, pres=pres, x=x)

_, ST0 = kempers.get_soret(mode='cov', kin=True)

def model(n):
    n1A, n1B, n2A, n2B, n1A0, n1B0, n2A0, n2B0 = n

    nA = n1A + n2A
    nB = n1B + n2B

    n1 = n1A + n1B
    n2 = n2A + n2B

    nT = n1 + n2

    x1 = n1 / nT
    x2 = n2 / nT

    x1A = n1A / nA
    x2A = n2A / nA
    x1B = n1B / nB
    x2B = n2B / nB

    nA0 = n1A0 + n2A0
    nB0 = n1B0 + n2B0

    x1A0 = n1A0 / nA
    x2A0 = n2A0 / nA
    x1B0 = n1B0 / nB
    x2B0 = n2B0 / nB

    vA, dvAdn = eos.specific_volume(TA, pres, [x1A, x2A])
    vB, dvBdn = eos.specific_volume(TB, pres, [x1B, x2B])

    v1A, v2A = dvAdn
    v1B, v2B = dvBdn

    eq1 = pres - eos.pressure_tv(TA, V/2, [n1A, n2A], 1)
    eq2 = pres - eos.pressure_tv(TB, V/2, [n1B, n2B], 1)
    eq3 = n1A + n1B - x1[0] * nT
    eq4 = 0

    eqs = [eq1, eq2, eq3, eq4]
    for i in range(len(n)):
        if n[i] < 0:
            eqs[i] = n[i] * 1e6