from pycThermopack.pyctp import cubic, cpa, extended_csp
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
from matplotlib.colors import Normalize

phases = ['Two-phase', 'Liquid', 'Vapour', 'Minimum Gibbs', 'Single', 'Solid', 'Fake']

def round_lst(x, n):
    return [round(xi, n) for xi in x]


def humidity_2_mole_frac(rh, T):
    eos = cubic.cubic()
    eos.init('H2O,O2,N2', 'PR')
    x = np.array([1-1e-3, 1e-3 * 0.2, 1e-3 * 0.8])
    x, y, betaV, betaL, ph = eos.two_phase_tpflash(T, 101325, x)

    return y[0] * rh

def mole_frac_2_humidity(x, T):
    eos = cubic.cubic()
    eos.init('H2O,O2,N2', 'PR')
    x_tot = np.array([1-1e-3, 1e-3 * 0.2, 1e-3 * 0.8])
    x_liq, y, betaV, betaL, ph = eos.two_phase_tpflash(T, 101325, x)

    return x / y[0]

def print_arr(arr):
    for line in arr:
        for x in line:
            print(x, end=' ')
        print()


if __name__ == '__main__':

    eos_list = [None] * 5
    names = ['SRK', 'PR', 'PT', 'SW', 'SPUNG', 'CPA']
    comps = 'H2O,O2,N2'
    temp_list = np.linspace(80, 120, 50) + 273
    pres_list = np.zeros_like(temp_list)
    n = 0

    eq = cubic.cubic()
    eq.init(comps, 'SRK')
    for i, temp in enumerate(temp_list):
        pres_list[i] = humidity_2_mole_frac(1, temp, eq)
    plt.plot(temp_list-273, pres_list, label=names[n])
    n+=1

    eq = cubic.cubic()
    eq.init(comps, 'PR')
    for i, temp in enumerate(temp_list):
        pres_list[i] = humidity_2_mole_frac(1, temp, eq)
    plt.plot(temp_list-273, pres_list, label=names[n])
    n+=1

    eq = cubic.cubic()
    eq.init('H2O', 'PT')
    for i, temp in enumerate(temp_list):
        pres_list[i] = humidity_2_mole_frac(1, temp, eq)
    plt.plot(temp_list-273, pres_list, label=names[n])
    n+=1

    eq = cubic.cubic()
    eq.init(comps, 'SW')
    for i, temp in enumerate(temp_list):
        pres_list[i] = humidity_2_mole_frac(1, temp, eq)
    plt.plot(temp_list-273, pres_list, label=names[n])
    n+=1

    eq = extended_csp.ext_csp()
    eq.init(comps, 'SRK', 'Classic', 'SW', 'MBWR32', 'N2')
    for i, temp in enumerate(temp_list):
        pres_list[i] = humidity_2_mole_frac(1, temp, eq)
    plt.plot(temp_list-273, pres_list, label=names[n])
    n+=1

    #eq = cpa.cpa()
    #eq.init(comps, 'SRK')
    #for i, temp in enumerate(temp_list):
    #    pres_list[i] = humidity_2_mole_frac(1, temp, eq)
    #plt.plot(temp_list-273, pres_list, label=names[n])

    plt.vlines(100, 0, 1, color='black')
    plt.hlines(1, 90, 110, color='black')
    plt.legend()
    plt.savefig('Vapour pressure')