import sys
import os
import yaml
import matplotlib.pyplot as plt
import numpy as np


def total_boxplot(DATA, creator = 'any', tester = 'any', simulation = "dummy_simulation", savepath = "./"):

    fig, ax = plt.subplots()

    data_to_plot = []
    labels = []
    title = f'Boxplot of scores in {simulation},'
    xlabel = ""

    colors_list = plt.colormaps['tab10'].colors

    if creator == tester == 'any':
        print(f"Warning: setting both creator and tester as 'any' plots only self testing!")
        title += f" self testing"
        xlabel = "self tester"
        for creator_name, data in DATA.items():
            if creator_name in data:
                data_to_plot.append(data[creator_name])
                labels.append(creator_name)

    elif creator == 'any':
        title += f" tester: {tester}"
        xlabel = "creator"
        for creator_name, data in DATA.items():
            if tester in data:
                data_to_plot.append(data[tester])
                #labels.append(f"{creator_name}-{tester}")
                labels.append(creator_name)

    elif tester == 'any':
        title += f" creator: {creator}"
        xlabel = "tester"
        if creator in DATA:
            for tester_name, data in DATA[creator].items():
                data_to_plot.append(data)
                #labels.append(f"{creator}-{tester}")
                labels.append(tester_name)


    if len(data_to_plot) == 0:
        print("No appropriate data!")
        return


    #print(len(data_to_plot), len(labels))
    elements = ax.boxplot(data_to_plot, patch_artist=True, tick_labels=labels)
    for i, patch in enumerate(elements['boxes']):
        patch.set_facecolor(colors_list[(i)%10])

    ax.grid()
    #ax.legend()
    ax.set(xlabel=xlabel, ylabel='Score')
    ax.set_title(title)

    ax.figure.savefig(f'{savepath}/{creator}_{tester}_{simulation}.png', bbox_inches='tight')


def heat_map(DATA, criteria = np.median, savepath = "./"):

    names_to_nums = {name: num for num, name in enumerate(list(DATA.keys()))}
    values = np.zeros( (len(names_to_nums)+1, len(names_to_nums)+1) )

    testers = {}
    for creator, Data in DATA.items():
        total_creator = []
        for tester, data in Data.items():
            if not tester in testers:
                testers[tester] = []
            if creator in names_to_nums and tester in names_to_nums:
                creator_idx = names_to_nums[creator]
                tester_idx = names_to_nums[tester]
                values[creator_idx, tester_idx] = criteria(data)
                total_creator += data
                testers[tester] +=data
            values[creator_idx, len(names_to_nums)] = criteria(total_creator)

    for tester, data in testers.items():
        tester_idx = names_to_nums[tester]
        values[len(names_to_nums), tester_idx] = criteria(data)

    fig, ax = plt.subplots()
    im = ax.imshow(values, cmap='coolwarm')

    for i in range(values.shape[0]):
        for j in range(values.shape[1]):
            text = ax.text(j, i, round(values[i, j],1), ha="center", va="center", color="k")

    ax.set_xticks(range(len(DATA.keys())+1), labels=list(DATA.keys())+['mean'] )
    ax.set_yticks(range(len(DATA.keys())+1), labels=list(DATA.keys())+['mean'] )

    ax.set(xlabel='tester', ylabel='creator')
    ax.set_title(f"Cross-validation heatmap ({criteria.__name__})")

    ax.figure.savefig(f"{savepath}/cross_validation_{criteria.__name__}.png")


def convert_score(data):

    score = data['score']
    Nsof = len(data['exted_sofs'])
    Ncol = len(data['collisions'])
    Ksof = data['k_params']['_k_sof']
    Khp = data['k_params']['_k_hp']
    Kch = data['k_params']['_k_charge']
    Khp_max = 100
    Kch_max = 100

    Kch_r = (score - Ksof*Nsof-Khp*(Khp_max-Ncol)/Khp_max) * Kch_max/Kch

    score_new = Ksof*Nsof*(Kch*Kch_r/Khp_max+Khp*(Khp_max-Ncol)/Khp_max)
    #print(type(score_new))

    return score_new

if __name__ == '__main__':

    global_path = "/home/gestures6/Easy_sim/dummy_simulation/stats/"
    # 1. Name (Dima)
    # 2. Name-Name (Dima-Elisey)

    DATA = {} # Name of creator - {} - Name of tester

    for dirname in sorted(os.listdir(global_path)):
        dirpath = os.path.join(global_path, dirname)
        if os.path.isdir(dirpath) and not dirname.startswith('!'):
            data = []
            if not '-' in dirname:
                name = dirname
                if not name in DATA:
                    DATA[name] = {}
                DATA[name][name] = data
            else:
                names = dirname.split('-')
                name1 = names[0]
                name2 = names[1]
                if not name1 in DATA:
                    DATA[name1] = {}
                DATA[name1][name2] = data
            for filename in sorted(os.listdir(dirpath)):
                file_path = os.path.join(dirpath, filename)
                if file_path.endswith('.yaml'):
                    with open(file_path, 'r', encoding='utf-8') as file:
                        datafile = yaml.load(file, Loader = yaml.Loader)
                        data.append(convert_score(datafile)) # TODO change score to new formula
                        #data.append(datafile['score'])


    #for cr, data in DATA.items():
        #print(f"{cr}: {data}")
    savepath = global_path
    total_boxplot(DATA, creator = 'any', tester = 'any', savepath = savepath)

    for name in DATA.keys():

        total_boxplot(DATA, creator = name, tester = 'any', savepath = savepath)
        total_boxplot(DATA, creator = 'any', tester = name, savepath = savepath)

    heat_map(DATA, savepath = savepath)

    plt.show()