import os
import json
import pandas as pd
from itertools import chain
import seaborn as sns
import matplotlib.pyplot as plt

runid = "Run ID"
x_label = "Thread Count"
y_label = "Throughput in GiB/s"
var_label = "Thread Counts"
thread_counts = ["1t", "2t", "4t", "8t", "12t"]
thread_counts_nice = ["1 Thread", "2 Threads", "4 Threads", "8 Threads", "12 Threads"]
engine_counts = ["1mib-1e_PREVENT_FROM_DISPLAYING", "1mib-4e_PREVENT_FROM_DISPLAYING", "1gib-1e", "1gib-4e"]
engine_counts_nice = ["1 E/WQ and Tasksize 1 MiB", "4 E/WQ and Tasksize 1 MiB", "1 E/WQ and Tasksize 1 GiB", "4 E/WQ and Tasksize 1 GiB"]
title = "Total Throughput - 120 Copy Operations split on Threads Intra-Node on DDR"

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

def calc_throughput(size_bytes,time_ns):
    time_seconds = time_ns * 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_time_mesurements(file_path):
    with open(file_path, 'r') as file:
        data = json.load(file)
     
        count = data["count"]
        iterations = data["list"][0]["task"]["iterations"]

        # work queue size is 120 which is split over all available threads
        # therefore we divide the result by 120/n_threads to get the per-element speed

        return {
            "total" : sum([x / (iterations * 120) for x in list(chain([data["list"][i]["report"]["time"]["total"] for i in range(count)]))]),
            "combined" : [x / 120 for x in list(chain(*[data["list"][i]["report"]["time"]["combined"] for i in range(count)]))],
            "submission" : [x / 120 for x in list(chain(*[data["list"][i]["report"]["time"]["submission"] for i in range(count)]))],
            "completion" : [x / 120 for x in list(chain(*[data["list"][i]["report"]["time"]["completion"] for i in range(count)]))]
        }

def process_file_to_dataset(file_path, engine_label, thread_count):
    engine_index = index_from_element(engine_label,engine_counts)
    engine_nice = engine_counts_nice[engine_index]
    threadc_index = index_from_element(thread_count, thread_counts)
    thread_count_nice = thread_counts_nice[threadc_index]
    data_size = 0

    if engine_label in ["1gib-1e", "1gib-4e"]: data_size = 1024*1024*1024
    elif engine_label in ["1mib-1e", "1mib-4e"]: data_size = 1024*1024
    else: data_size = 0

    try:
        time = load_time_mesurements(file_path)["combined"]
        run_idx = 0
        for t in time:
            data.append({ runid : run_idx, x_label: thread_count_nice, var_label : engine_nice, y_label : calc_throughput(data_size, t)})
            run_idx = run_idx + 1
    except FileNotFoundError:
        return


def main():
    folder_path = "benchmark-results/"

    for engine_label in engine_counts:
        for thread_count in thread_counts:
            file = os.path.join(folder_path, f"mtsubmit-{thread_count}-{engine_label}.json")
            process_file_to_dataset(file, engine_label, thread_count)

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

    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-mtsubmit.png"), bbox_inches='tight')
    plt.show()


if __name__ == "__main__":
    main()