import numpy as np
import matplotlib.pyplot as plt
import pickle
from matplotlib import cm

def calculate_F(m_AB, w_AB):
    cmap = cm.get_cmap('cool')
    states, multiplicity = np.unique(m_AB, return_counts=True)
    temp = np.linspace(0,10,5) #Sets the temperature range

    helmholtz = lambda state, w, T: w_AB * state - T * np.log(w) #Helmholtz energi

    #Regn ut og plot F(states) for hver temperatur
    for T in temp:
        F = helmholtz(states, multiplicity, T)
        plt.plot(states, F, label=T, color=cmap(T/max(temp)))
        
    plt.legend(title='Temp [-]')
    plt.title(r'$w_{AB} =$ '+str(w_AB))
    plt.savefig('Helmholtz_wab_'+str(w_AB))
    plt.show()
    plt.close()

with open('state_files/m_AB_24.pkl', 'rb') as f: #Opens up the list of all microstates.
    m_AB = pickle.load(f)[0] #m_AB is the list.

calculate_F(m_AB, 1)