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

runid = "Run ID"
x_label = "Copy Type"
y_label = "Throughput in GiB/s"
var_label = "Configuration"
types = ["intersock-n0ton4", "internode-n0ton1"]
types_nice = ["Inter-Socket Copy", "Inter-Node Copy"]
copy_methods = ["dstcopy", "srccopy", "xcopy"]
copy_methods_nice = [ "Engine on DST-Node", "Engine on SRC-Node", "Cross-Copy / Both Engines" ]
title = "Performance of Engine Location - Copy Operation on DDR with Size 1 MiB and 1 Engine per WQ"

index = [runid, x_label, var_label]
data = []

def calc_throughput(size_bytes,time_microseconds):
    time_seconds = time_microseconds * 1e-9
    size_gib = size_bytes / (1024 ** 3)
    throughput_gibs = size_gib / time_seconds
    return throughput_gibs


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


def load_and_process_copy_json(file_path,method_label):
    with open(file_path, 'r') as file:
        data = json.load(file)
        
        # Extracting time from JSON structure
        if method_label == "xcopy":
            # For xcopy method, add times from two entries and divide by 3
            time0 = data["list"][0]["report"]["time"]
            time1 = data["list"][1]["report"]["time"]

            return {
                "combined" : [sum(x) / 4 for x in zip(time0["combined"], time1["combined"])],
                "submission" : [sum(x) / 4 for x in zip(time0["completion"], time1["completion"])],
                "completion" : [sum(x) / 4 for x in zip(time0["submission"], time1["submission"])]
            }

        else:
            return data["list"][0]["report"]["time"]

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

    method_index = index_from_element(method_label,copy_methods)
    method_nice = copy_methods_nice[method_index]

    idx = 0
    for file_path in file_paths:
        time = load_and_process_copy_json(file_path,method_label)
        times.append(time["combined"])
        idx = idx + 1

    throughput = [[calc_throughput(1024*1024,time) for time in t] for t in times]

    idx = 0
    for run_set in throughput:
        run_idx = 0
        for run in run_set:
            data.append({ runid : run_idx, x_label: types_nice[idx], var_label : method_nice, y_label : throughput[idx][run_idx]})
            run_idx = run_idx + 1
        idx = idx + 1

# Main function to iterate over files and create plots for the new benchmark
def main():
    folder_path = "benchmark-results/"

    for method_label in copy_methods:
        copy_file_paths = [os.path.join(folder_path, f"{method_label}-{type_label}-1mib-1e.json") for type_label in types]
        create_copy_dataset(copy_file_paths, method_label)

    df = pd.DataFrame(data)
    df.set_index(index, inplace=True)
    df = df.sort_values(y_label)

    sns.barplot(x=x_label, y=y_label, hue=var_label, data=df, palette="rocket", errorbar="sd")

    plt.title(title)
    plt.savefig(os.path.join(folder_path, "plot-perf-enginelocation.png"), bbox_inches='tight')
    plt.show()

if __name__ == "__main__":
    main()