import matplotlib.pyplot as plt
from tools.compute_constants import *
from tools import file_handler

def p_co2(T,a):
    #returnerer p_co2, gitt Temperatur og loading

    K = K2(T)       #bruker exponensialregresjon for å finne likevektskonstanten
    K_h = henry_constant(T,a)       #regner ut Henrykonstanten

    p = (K_h * (a**2))/(K*((1-2*a)**2))
    return p

def plot_pco2_a(T=350.0):
    #plotter pco2 som funksjon av loading

    a_values = np.linspace(0.1,0.49,50)
    p_values = [p_co2(float(T),a) for a in a_values]

    plt.plot(a_values,p_values, label='T = '+ str(T))

    da = (a_values[-1] - a_values[0])/50
    plt.xlim(a_values[0] - da, a_values[-1] + da)
    plt.yscale('log')
    plt.xlabel('Loading')
    plt.ylabel('p_co2/[bar]')
    plt.legend(loc='upper left')

def plot_pco2_T(a=1.0):
    # plotter pco2 som funksjon av Temperatur
    T_values = np.linspace(298, 360, 150)
    p_values = [p_co2(T, a) for T in T_values]

    dt = (T_values[-1] - T_values[0]) / 50
    plt.xlim(T_values[0] - dt, T_values[-1] + dt)
    plt.plot(T_values, p_values, label='a = '+str(a))
    plt.xlabel('T/[K]')
    plt.ylabel('p_co2/[bar]')
    plt.yscale('log')

def plot_loading_pco2(T=350.0):
    #plotter loading som funksjon av p_co2

    a_values = np.linspace(0, 0.49, 50)
    p_values = [p_co2(float(T), a) for a in a_values]

    plt.plot(p_values, a_values, label='T = ' + str(T))

    dp = (p_values[-1] - p_values[0]) / 50
    plt.xlim(p_values[0] - dp, p_values[-1] + dp)
    plt.xscale('log')
    plt.xlabel('p_co2/[bar]')
    plt.ylabel('loading')
    plt.legend(loc='lower right')

def plot_K(style = ''):
    #henter K2(T) fra regresjon
    T = np.linspace(270,430,150)
    K = [K2(t) for t in T]

    #henter K2(t) fra datafil
    data = file_handler.file_to_list_tuples('K2(T).txt')
    T_data, K_data = [[None for i in range(len(data))] for x in range(2)]
    i = 0
    for temp,val in data:
        T_data[i], K_data[i] = float(temp),float(val)
        i += 1

    plt.figure()

    #genererer plot
    plt.plot(T,K,label = 'Regresjon : 4.83E16 * exp(-0.094)')
    plt.scatter(T_data, K_data, label = 'Gitt data for K2(T)', marker = 'o')

    dk = (K_data[0] - K_data[-1])/10
    #pynter på plotten
    plt.xlim(T[0]-5,T[-1]+5)
    #plt.ylim(K_data[-1] - dk, K_data[0] + dk)
    plt.xlabel('T/[K]')
    plt.ylabel('K_2')
    plt.legend()
    plt.yscale('log')

    plt.savefig('plots/K_2(T)_log')
    plt.show()

def gen_figure_pco2(var=None, space=[0,0,0]):
    #'var' er enten 'T' eller 'a' og er den varierende variabelen.
    #space inneholder [start,stop,num] for np.linspace, er de verdiene av den varante variabelen
    #det skal plottes for.
    konst = 'a'
    if var == 'a':
        konst = 'T'

    plt.figure()
    for x in np.linspace(space[0],space[1],space[2]):
        print(konst + ' = ', x)
        if var == 'a':
            plot_pco2_a(x)
        elif var == 'T':
            plot_pco2_T(x)

    plt.legend(loc='upper left')
    plt.savefig('plots/p_co2('+ var + ')')
    plt.show()

def gen_plot_loading(T_values = [270,430,5]):
    plt.figure()
    for T in np.linspace(T_values[0],T_values[1],T_values[2]):
        plot_loading_pco2(T)
    plt.show()
    plt.savefig('plots/loading(p_co2)2')