import matplotlib.pyplot as plt
import numpy as np 
import pandas as pd
import os
from data_utils import data_mean,bootstrapping_CI,computing_confidence_intervals   

os.chdir('/home/vsansi01/ACC_R_values')

#for T2M only
xx_ticks = ["2 weeks","3 weeks","4 weeks"]
xx = [0,1,2]
xx_graphop = [0,1]

def generating_data(week):

    files_S2S = f'ACC_S2S_t2m_week{week}.npy'
    files_S2S_ensemble = f'ACC_S2S_ensemble_mean_t2m_week{week}.npy'
    files_GraphFT = f'ACC_GraphFT_t2m_week{week}.npy'
    
    if week == 4:
        files_GraphOP = f'ACC_GraphOP_t2m_week{week-1}.npy'
    else:
        files_GraphOP = f'ACC_GraphOP_t2m_week{week}.npy'
    
    return [files_S2S, files_S2S_ensemble, files_GraphFT, files_GraphOP]

#["2 weeks","3 weeks","4 weeks"] in this order
acc_values_graphFT = [data_mean(generating_data(2))[2],data_mean(generating_data(3))[2],data_mean(generating_data(4))[2]]
acc_values_S2S = [data_mean(generating_data(2))[0],data_mean(generating_data(3))[0],data_mean(generating_data(4))[0]]
acc_values_S2S_ensemble_mean = [data_mean(generating_data(2))[1],data_mean(generating_data(3))[1],data_mean(generating_data(4))[1]]
acc_values_graphOP = [data_mean(generating_data(2))[3],data_mean(generating_data(3))[3]]

errors_graphFT = [bootstrapping_CI(generating_data(2))[2],bootstrapping_CI(generating_data(3))[2],bootstrapping_CI(generating_data(4))[2]]
errors_graphOP = [bootstrapping_CI(generating_data(2))[3],bootstrapping_CI(generating_data(3))[3]]
errors_S2S = [bootstrapping_CI(generating_data(2))[0],bootstrapping_CI(generating_data(3))[0],bootstrapping_CI(generating_data(4))[0]]
errors_S2S_ensemble_mean = [bootstrapping_CI(generating_data(2))[1],bootstrapping_CI(generating_data(3))[1],bootstrapping_CI(generating_data(4))[1]]


plt.figure()
plt.plot(xx,acc_values_graphFT,'r',label='GraphFT')
plt.plot(xx,acc_values_S2S,'b',label='S2S')
plt.plot(xx_graphop,acc_values_graphOP,'k',label='GraphOP')
plt.plot(xx,acc_values_S2S_ensemble_mean,'c--',label='S2S ensemble mean')

#computing confidence intervals 
plt.errorbar(xx,acc_values_S2S_ensemble_mean, yerr=errors_S2S_ensemble_mean, ecolor='c',fmt = 'None',capsize=3, capthick=0.5)
plt.errorbar(xx,acc_values_graphFT, yerr=errors_graphFT, ecolor='r',fmt = 'None',capsize=3, capthick=0.5)
plt.errorbar(xx_graphop,acc_values_graphOP, yerr=errors_graphOP, ecolor='k',fmt = 'None',capsize=3, capthick=0.5)
plt.errorbar(xx,acc_values_S2S, yerr=errors_S2S, ecolor='b',fmt = 'None',capsize=3, capthick=0.5)

plt.legend()
plt.xlabel('Lead-time')
plt.ylabel('ACC for T2M')
plt.xticks(ticks=xx,labels=xx_ticks)
plt.tight_layout()
plt.savefig('/home/vsansi01/acc_vs_time_bootstrapping_test.jpeg',dpi=200)