import os
import json
import pandas as pd
from pandas.core.ops import methods
from typing import List
import seaborn as sns
import matplotlib.pyplot as plt

x_label = "Size of Submitted Task"
y_label = "Time to Copy 1 KiB in Microseconds"
var_label = "Submission Type"
sizes = ["1kib", "4kib", "1mib", "1gib"]
sizes_nice = ["1 KiB", "4 KiB", "1 MiB", "1 GiB"]
types = ["bs10", "bs50", "ms10", "ms50", "ssaw"]
types_nice = ["Batch, Size 10", "Batch, Size 50", "Multi-Submit, Count 10", "Multi Submit, Count 50", "Single Submit"]

data = {
    x_label : sizes_nice,
    types_nice[0] : [],
    types_nice[1] : [],
    types_nice[2] : [],
    types_nice[3] : [],
    types_nice[4] : []
}

stdev = {}


def index_from_element(value,array):
    for (idx,val) in enumerate(array):
        if val == value: return idx
    return 0


def load_and_process_submit_json(file_path,s,t):
    with open(file_path, 'r') as file:
        data = json.load(file)
        time_microseconds = data["list"][0]["report"]["time"]["combined_avg"]
        if t not in stdev: stdev[t] = dict()
        stdev[t][s] = data["list"][0]["report"]["time"]["combined_stdev"]
        return time_microseconds


def stdev_functor(values):
    v = values[0]
    sd = stdev[v]
    return (v - sd, v + sd)


# Function to plot the graph for the new benchmark
def plot_submit_graph(file_paths, type_label):
    times = []

    type_index = index_from_element(type_label,types)
    type_nice = types_nice[type_index]

    idx = 0
    for file_path in file_paths:
        time_microseconds = load_and_process_submit_json(file_path,sizes_nice[idx],type_nice)
        times.append(time_microseconds)
        idx = idx + 1

    # Adjust time measurements based on type
    # which can contain multiple submissions
    if type_label in {"bs10", "ms10"}:
        times = [time / 10 for time in times]
    elif type_label in {"ms50", "bs50"}:
        times = [time / 50 for time in times]

    times[0] = times[0] / 1
    times[1] = times[1] / 4
    times[2] = times[2] / 1024
    times[3] = times[3] / (1024 * 1024)

    data[type_nice] = times


# Main function to iterate over files and create plots for the new benchmark
def main():
    folder_path = "benchmark-results/submit-bench/"  # Replace with the actual path to your folder

    for type_label in types:
        file_paths = [os.path.join(folder_path, f"submit-{type_label}-{size}-1e.json") for size in sizes]
        plot_submit_graph(file_paths, type_label)

    df = pd.DataFrame(data)
    dfm = pd.melt(df, id_vars=x_label, var_name=var_label, value_name=y_label)

    error_values: List[float] = []
    for index,row in dfm.iterrows():
        s = dfm[x_label][index]
        t = dfm[var_label][index]
        error_values.append(stdev[t][s])

    dfm["Stdev"] = error_values

    print(dfm)

    sns.catplot(x=x_label, y=y_label, hue=var_label, data=dfm, kind='bar', height=5, aspect=1, palette="viridis", errorbar=("ci", 100))
    plt.title("Performance of Submission Methods - Copy Operatione tested Intra-Node on DDR")
    plt.savefig(os.path.join(folder_path, "plot-perf-submitmethod.png"))
    plt.show()

if __name__ == "__main__":
    main()