import matplotlib.pyplot as plt
import numpy as np
from itertools import cycle
import argparse
import pickle
import yaml


#import matplotlib.font_manager
from matplotlib import rc
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
rc('text', usetex=True)





def plot_parameters():
    ''' Plot the parameters in separate subplots with uncertainties.

    Args:
        dat (dict):  data dictionary
        deparameterize (bool):  flag indicating if parameters should be
          deparameterized via 2**theta
        ref:    reference value to be plotted with parameters
    '''
 

    
    name_file = ['SNR12V120_Pf','SNR12V70_Pf','SNR12V30_Pf']
    #name_file = ['SNR12V120_Pf_MAG','SNR12V70_Pf_MAG','SNR12V30_Pf_MAG']
    #name_file = ['SNR12V120_Pb_MAG','SNR12V70_Pb_MAG','SNR12V30_Pb_MAG']

    name_file = ['slice2.3_Pa']
    vencs = ['180','105','45']
    path0 = '/home/yeye/Desktop/kalman/results/'


    fig, axes = plt.subplots(1,1,figsize=(12,7))
    col = cycle(['orangered', 'dodgerblue', 'limegreen', 'C3','C4'])

    true_values = {
            3: 4800,
            4: 7200,
            5: 11520,
            6: 11520,
            2: 75
            }



    for nn,name in enumerate(name_file):



        path1 = path0 + name + '/'
        inputfile_path = path1 + 'input.yaml'
        dat = np.load(path1 + 'theta_stats.npz')


        with open(inputfile_path) as file:
            inputfile = yaml.full_load(file)

        col_ = next(col)
    

        dim = dat['theta'].shape[-1]

        current_val = []
        ids_type = []
        labels = []
        ids = []

        for bnd_c in inputfile['estimation']['boundary_conditions']:
            
            if 'windkessel' in bnd_c['type']:
                for bnd_set in inputfile['boundary_conditions']:
                    if bnd_c['id'] == bnd_set['id']:
                        ids.append(bnd_c['id'])
                        ids_type.append('windkessel')
                        current_val.append(bnd_set['parameters']['R_d'])
                        


            elif 'dirichlet' in bnd_c['type']:
                current_val.append(inputfile['boundary_conditions'][0]['parameters']['U'])
                ids.append(bnd_c['id'])
                ids_type.append('dirichlet')
                labels.append('$U')




        t = dat['times']
        theta = dat['theta']
        P = dat['P_theta']

        legends = cycle(labels)

        legends_=next(legends)

        if dim == 1:
            theta = theta.reshape((-1, 1))
            P = P.reshape((-1, 1, 1))
        
        idx = 0

        for i in range(len(ids)):

            cur_key = ids[i]
            rec_value = np.round(2**theta[-1, idx]*current_val[i],2)
            curve = 2**theta[:, idx]*current_val[i]
            std_down =  2**(-np.sqrt(P[:, idx, idx]))*curve
            std_up =  2**np.sqrt(P[:, idx, idx])*curve
            dash_curve = true_values[ids[i]]  + t*0


            if ids_type[i] == 'dirichlet':
                axes.plot(t, curve , '-', color=col_,label= '$(venc \ '+ vencs[nn] + ' \ cm/s) \ U = ' + str(rec_value) + '/' + str(true_values[cur_key])  + '$', linewidth = 5)
                #axes.plot(t, curve , '-', color=col_,label= legends_ + vencs[nn] + ' \ cm/s$', linewidth = 5)
                axes.fill_between(t, std_down, std_up, alpha=0.3, color=col_)
                legends_=next(legends)
                axes.plot(t, dash_curve , color='black',ls='--' , linewidth = 3)


            idx +=1



    axes.set_ylabel(r'$U$',fontsize=36)
    axes.legend(fontsize=30,loc='upper right')
    axes.set_xlim([0,0.35])
    axes.set_ylim([10,160])
    axes.set_xlabel(r'$t  (s)$',fontsize=36)
    axes.set_box_aspect(1/2)
    plt.xticks(fontsize=28)
    plt.yticks(fontsize=28)
    plt.savefig('U.png')   
    plt.show()



    #path_paper = '/home/yeye/A_aliasing_kalman/latex/0_preprint/Figures/'
    #fig1.savefig('Rd_'+ name_file +'.png')
    





if __name__ == '__main__':

    plot_parameters()