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

# shortpath = sys.argv[1]
shortpath = 'bodyfat_k_'

# alg_names = ['svrgzoht' , 'sarah-zht',   'bvrht12', 'bvrhtn']
# alg_names = ['sarah-zht']
# alg_names = ['bvrht12']

alg_names = ['bvrht12', 'sarah-zht']
alg1 = 'bvrht12'
alg2 = 'sarah-zht'
true_names = {
              'svrgzoht': 'VR',
              'sarah-zht': 'SARA',
                'bvrht12': 'BVR_1/2',
                'bvrhtn':  'BVR-n'
              }


arguments_k = [8, 9, 10, 11]

# plt.figure()

# for k in arguments_k:

#   path = f'./results/{shortpath}'+f'{k}'
#   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:
      # print(1)
      # curve_mean = 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]}_k={k}', 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(\theta)$', fontweight='bold', fontsize=20)
# plt.xlabel('IFO', fontsize=20)
# plt.savefig(f'{path}/izo_{shortpath}.png')
# plt.show()

plt.figure()

a = 1000

for k in arguments_k:
  path = f'./results/{shortpath}'+f'{k}'
  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)


  curve_mean1 = np.array(eval(results_mean[(results_mean[1] == alg1)].iloc[:, -1].iloc[0]))[a:1991]
  # print(len(np.array(eval(results_mean[(results_mean[1] == alg_name)].iloc[:, -1].iloc[0]))))
  curve_std1 = np.array(eval((results_std[(results_std[1] == alg1)].iloc[:, -1].iloc[0])))[a:1991]
  # nizo_data = np.array(eval(nizo[(nizo[1] == alg1)].iloc[:, -1].iloc[0]))[a:1991]

  curve_mean2 = np.array(eval(results_mean[(results_mean[1] == alg2)].iloc[:, -1].iloc[0]))[a:1991]
  curve_std2 = np.array(eval((results_std[(results_std[1] == alg2)].iloc[:, -1].iloc[0])))[a:1991]
  nizo_data = np.array(eval(nizo[(nizo[1] == alg2)].iloc[:, -1].iloc[0]))[a:1991]

  curve_mean = curve_mean2 - curve_mean1
  curve_std = curve_std2 - curve_std1

  # plt.plot(nizo_data, curve_mean, label=f'{true_names[alg2]}-{true_names[alg1]}_k={k}', linestyle='-', marker='^', markersize=1)
  plt.plot(nizo_data, curve_mean, label=f'k={k}', 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(\theta)$', fontweight='bold', fontsize=20)
plt.xlabel('IFO', fontsize=20)
# plt.savefig(f'{path}/izo_{shortpath}.png')
plt.savefig(f'./izo_sen.png')
plt.show()






  # plt.figure()

  # for alg_name in alg_names:
  #     print(alg_name)
  #     curve_mean = 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])))
  #     nht_data = np.array(eval(nht[(nht[1] == alg_name)].iloc[:, -1].iloc[0]))
  #     plt.plot(nht_data, curve_mean, label=f'{true_names[alg_name]}', linestyle='-', marker='^', markersize=1)
  #     plt.fill_between(nht_data, curve_mean-curve_std, curve_mean+curve_std, alpha=0.3)
  #     plt.legend()
  # plt.ylabel(r'$\mathcal{F}(\theta)$', fontweight='bold', fontsize=20)
  # plt.xlabel('# NHT', fontsize=20)
  # plt.savefig(f'{path}/nht_{shortpath}.png')
  # plt.show()