import numpy as np
import os
from glob import glob
import matplotlib.pyplot as plt
import matplotlib as mpl


path = "/home/gongjinlan/BO_tuning_results/main_1560002003-BO-spider9-finished/"
dirs = glob(path + "*/")

# Base settings for alph, sigma_sq, l
base_names = ["Sampling method", "UCB alpha", "Kernel sigma_sq", "Kernel l"]
base_setting = ["LHS", "0.5", "1", "0.1"]
parameter_file = "dynamic_parameters.txt"
max_dirs = 20
model_name = "Spider9"

# Set matplotlib font
font = {'font.size' : 60}

# Initiate groups with base
zero = ["/".join(x.split("/")[:-2]) + "/" for x in dirs][0] + "0/"
groups = []
for k in range(len(base_setting)):
    groups += [[[zero, base_setting]]]

for ix, my_dir in enumerate(dirs):
    my_sub_dirs = [float(path.split("/")[-2]) for path in glob(my_dir + "*/")]
    my_sub_dirs.sort()
    my_sub_dirs = [str(x) for x in my_sub_dirs]

    # Save all fitnesses in a list. This can help you to later construct the
    sd_list = []
    for sd in my_sub_dirs:
        if os.path.isfile(my_dir + sd + "/" + parameter_file):
            sd_list += [sd]

    # Get dynamic parameters
    params_all = [(line.rstrip('\n')).split(",") for line in open(my_dir + sd_list[0] + "/" + parameter_file)]
    params = params_all[0]
    real_param = params_all[0]

    # Check dynamic parameters for thread-safety:
    for sd in sd_list:
        temp_param = [(line.rstrip('\n')).split(",") for line in open(my_dir + sd + "/" + parameter_file)]
        if not all(x == real_param for x in temp_param):
            differ = [x  for x in temp_param if  not x == real_param]
            differ = [d for d in differ if len(d) != 3]
            if len(differ) > 0:
                None
                #print(f"WARNING: DYNAMIC PARAMETERS DIFFER FOR {sd}")
            for d in differ:
                print(d)

    # Group folders together based on which param they differ in from
    equality = [0]*len(base_setting)
    c= 0
    for k in range(len(base_setting)):
        if not base_setting[k] == params[k]:
            groups[k] += [[my_dir, params]]
            print(f"Added {my_dir} to group {k} with params {params}")
            c += 1
    if c == 0:
        print(f"Didn't add {my_dir}")


for ix, group in enumerate(groups):
    print("\n\n")
    # Sort group
    group = sorted(group, key=lambda x: x[1][ix])
    labels = []
    fitnesses_list_mean = []
    fitnesses_list_error = []

    for i, e in enumerate(group):
        print(e)
        labels += [e[1][ix]]

        # Get fitnesses that are generated by gridsearchanalysis
        fitnesses = np.loadtxt(e[0] + "experiment_fitnesses.txt", delimiter =",")
        print(f"{ix}, {i} uses {len(fitnesses)} runs")

        fitnesses = fitnesses[:,0]
        error = 1.96*np.std(fitnesses)/np.sqrt(len(fitnesses))
        mean = np.mean(fitnesses)
        fitnesses_list_mean += [mean]
        fitnesses_list_error += [error]

    f, ax = plt.subplots(1,1,figsize=(10,7.5))
    plt.rcParams["font.size"] = "32"
    plt.xlabel(base_names[ix], fontsize=34)
    plt.ylabel("fitness", fontsize=34)
    plt.title(model_name)
    plt.xlim((-0.5, len(group) -0.5))
    plt.errorbar([i for i in range(len(group))],
                 fitnesses_list_mean,
                 yerr = fitnesses_list_error,
                 linewidth =3.5,
                 capsize = 6,
                 capthick = 3.0,
                 marker="o",
                 ms = 10,
                 ls ="none")
    ax.set_xticks([i for i in range(len(group))])
    ax.set_xticklabels(labels)
    ax.tick_params(labelsize=34)

    plt.tight_layout()
    plt.savefig(path + model_name + "-" + str(ix) + ".pdf")
    plt.clf()
    #plt.show()




