import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import sys

# shortpath = sys.argv[1]
# shortpath = "unit_q_${q}_k_${k}_beta_${b}"

# arguments_mu=[0.0001, 0.0002, 0.0004, 0.0006, 0.0008]
arguments_mu=[0.001, 0.01, 0.1]
betas = [0.5]
arguments_q = [5, 10, 50, 100, 200]
# arguments_q = [50, 100, 200]
arguments_k = [1, 2, 3, 4, 5]


# alg_names = ["szoht", "svrgzoht",  "sarah-zht", "b-svrgzoht"]
alg_names = ["b-svrgzoht"]
# alg_names = ["sarah-zht"]
# alg_names = ["sarah-zht", "b-svrgzoht"]
true_names = {
            'szoht': 'SZOHT',
            'svrgzoht': 'VR-SZHT',
            'sarah-zht': 'SARAH-ZHT',
            'b-svrgzoht' : 'BVR-SZHT(ours)'
            }


flag = True
for k in arguments_k:
    for q in arguments_q:
        plt.figure()
        # plt.ylim(0.35, 0.5)
        for mu in arguments_mu:
            for b in betas:
                # path = f'./results/{shortpath}' +f'_mu_0.0001'
                path = f'./results/unit_q_{q}_k_{k}_beta_{b}_mu_{mu}'
                print(path)
                results_mean = pd.read_csv(f'{path}/results_mean.csv', header=None)
                results_std = pd.read_csv(f'{path}/results_std.csv', header=None)
                it_count = pd.read_csv(f'{path}/it_count.csv', header=None)
                nizo = pd.read_csv(f'{path}/nizo.csv', header=None)
                nht = pd.read_csv(f'{path}/nht.csv', header=None)

                for alg_name in alg_names:
                    curve_mean = np.array(eval(results_mean[(results_mean[1] == alg_name)].iloc[:, -1].iloc[0]))
                    print(len(np.array(eval(results_mean[(results_mean[1] == alg_name)].iloc[:, -1].iloc[0]))))  
                    curve_std = np.array(eval((results_std[(results_std[1] == alg_name)].iloc[:, -1].iloc[0])))
                    nizo_data = np.array(eval(nizo[(nizo[1] == alg_name)].iloc[:, -1].iloc[0]))
                    plt.plot(nizo_data, curve_mean, label=f'{true_names[alg_name]}_q={q}_k={k}_mu={mu}', linestyle='-', marker='^', markersize=1)
                    plt.fill_between(nizo_data, curve_mean-curve_std, curve_mean+curve_std, alpha=0.3)
                    plt.legend()
        plt.ylabel(r'$f(\omega)$', fontweight='bold', fontsize=20)
        plt.xlabel('IZO', fontsize=20)
        # plt.savefig(f'{path}/izo_{shortpath}.png')
        # plt.show()
        plt.savefig(f'sen_reg_q_{q}_k_{k}_mu_{mu}.png')
        plt.close()

