
""" This code is used generate the plot in Figure 2.
Before executing this file, please execute the following to generate the data in the plot:
python WMC_synthetic/240519_compute_geometric_LambdavN.py """



import os
import sys
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rcParams
import seaborn as sns

rcParams['font.family'] = 'Times New Roman'
rcParams['font.size'] = 28
rcParams['mathtext.fontset'] = 'stix'

sys.path.insert(0, './')

# Load saved variables
data = np.load('WMC_synthetic/240519-variables.npz')
inradius = data['inradius']
gamma_W = data['gamma_W']
ineq = data['ineq']
LL = data['LL']
LU = data['LU']
lambd_list = data['lambd_list']
N_list = data['N_list']
alpha_list = data['alpha_list']
seed_list = data['seed_list']


# Calculate mean and standard deviations of output metrics
r_mean = np.zeros((len(N_list), len(alpha_list)))
r_std = np.zeros((len(N_list), len(alpha_list)))
LL_mean = np.zeros((len(N_list), len(alpha_list)))
LL_std = np.zeros((len(N_list), len(alpha_list)))
LU_mean = np.zeros((len(N_list), len(alpha_list)))
LU_std = np.zeros((len(N_list), len(alpha_list)))
gamma_W_mean = np.zeros((len(N_list), len(alpha_list)))
gamma_W_std = np.zeros((len(N_list), len(alpha_list)))
ineq_mean = np.zeros((len(N_list), len(alpha_list)))
ineq_std = np.zeros((len(N_list), len(alpha_list)))
lambd_mean = np.zeros((len(N_list), len(alpha_list)))
lambd_std = np.zeros((len(N_list), len(alpha_list)))

for alph in range(len(alpha_list)):
    for Ni in range(len(N_list)):
        r_mean[Ni, alph] = np.mean(inradius[Ni, :, alph])
        r_std[Ni, alph] = np.std(inradius[Ni, :, alph])
        LL_mean[Ni, alph] = np.mean(LL[Ni, :, alph])
        LL_std[Ni, alph] = np.std(LL[Ni, :, alph])
        LU_mean[Ni, alph] = np.mean(LU[Ni, :, alph])
        LU_std[Ni, alph] = np.std(LU[Ni, :, alph])
        gamma_W_mean[Ni, alph] = np.mean(gamma_W[Ni, :, alph])
        gamma_W_std[Ni, alph] = np.std(gamma_W[Ni, :, alph])
        ineq_mean[Ni, alph] = np.mean(ineq[Ni, :, alph])
        ineq_std[Ni, alph] = np.std(ineq[Ni, :, alph])
        lambd_mean[Ni, alph] = np.mean(lambd_list[Ni, :, alph])
        lambd_std[Ni, alph] = np.std(lambd_list[Ni, :, alph])

# Generate plots
linestyle = [{"linestyle":"solid", "linewidth":2, "markeredgewidth":1, "elinewidth":1, "capsize":4},
             {"linestyle":"dotted", "linewidth":2, "markeredgewidth":1, "elinewidth":1, "capsize":4},
             {"linestyle":(0, (3, 10, 1, 10)),  "markeredgewidth":1, "elinewidth":1, "capsize":4}]
# font = {#'weight' : 'bold',
        # 'size'   : 22}
fig, axes = plt.subplots(nrows=1, ncols=len(alpha_list), figsize=(24, 8))
fig.tight_layout(pad=4)
# plt.rc('font', **font)
N_list = [i*2 for i in N_list]

colors = sns.color_palette("deep")


for alph in range(len(alpha_list)):
    axes[alph].tick_params(axis='both', which='major')
    axes[alph].errorbar(N_list,lambd_mean[:,alph],lambd_std[:,alph],color=colors[3],**linestyle[0])
    axes[alph].errorbar(N_list,LL_mean[:,alph],LL_std[:,alph],color=colors[2],**linestyle[0])
    axes[alph].errorbar(N_list,LU_mean[:,alph],LU_std[:,alph],color=colors[0],**linestyle[0])
    axes[alph].set_xlabel("$N$")
    if alph == 0:
        axes[alph].legend(["$\lambda$","$\lambda^l$","$\lambda^u$"])
    axes[alph].set_title(f"$\lambda={alpha_list[alph]}\cdot \lambda^l$")


#Save plots
filename = f'./output_WMCsynthetic/Lambda_r{len(seed_list)}.pdf'
if not os.path.exists(os.path.dirname(filename)):
    os.makedirs(os.path.dirname(filename))
plt.savefig(filename, bbox_inches='tight')
plt.close()
