import os
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

from common import calc_throughput, index_from_element, load_time_mesurements

runid = "Run ID"
x_label = "Thread Count"
y_label = "Throughput in GiB/s"
var_label = "Transfer Size"
thread_counts = ["1t", "2t", "12t"]
thread_counts_nice = ["1 Thread", "2 Threads", "12 Threads"]
size_labels = ["1mib", "1gib"]
size_labels_nice = ["1 MiB", "1 GiB"]

title = \
    """Total Throughput showing cost of MT Submit\n
    Copying 120x split on n Threads Intra-Node on DDR\n
    """

description =  \
    """Total Throughput showing cost of MT Submit\n
    Running 120 Copy Operations split on n Threads\n
    Copying Intra-Node on DDR performed for multiple Configurations\n
    """

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


# loads the measurements from a given file and processes them
# so that they are normalized, meaning that the timings returned
# are nanoseconds per element transfered
def get_timing(file_path,thread_count):
    divisor = 0

    if thread_count == "1t": divisor = 1
    elif thread_count == "2t" : divisor = 2
    elif thread_count == "12t" : divisor = 12

    return [ x / divisor for x in load_time_mesurements(file_path)]


# procceses a single file and appends the desired timings
# to the global data-array, handles multiple runs with a runid
# and ignores if the given file is not found as some
# configurations may not be benchmarked
def process_file_to_dataset(file_path, size_label, thread_count):
    size_index = index_from_element(size_label,size_labels)
    size_nice = size_labels_nice[size_index]
    threadc_index = index_from_element(thread_count, thread_counts)
    thread_count_nice = thread_counts_nice[threadc_index]
    data_size = 0

    if size_label == "1gib" : data_size = 1024*1024*1024
    elif size_label == "1mib" : data_size = 1024*1024

    timing = get_timing(file_path, thread_count)
    run_idx = 0
    for t in timing:
        data.append({ runid : run_idx, x_label: thread_count_nice, var_label : size_nice, y_label : calc_throughput(data_size, t)})
        run_idx = run_idx + 1



# loops over all possible configuration combinations and calls
# process_file_to_dataset for them in order to build a dataframe
# which is then displayed and saved
def main():
    result_path = "benchmark-results/"
    output_path = "benchmark-plots/"

    for size in size_labels:
        for thread_count in thread_counts:
            file = os.path.join(result_path, f"mtsubmit-{thread_count}-{size}.json")
            process_file_to_dataset(file, size, thread_count)

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

    plt.figure(figsize=(4, 4))
    plt.ylim(0, 30)
    sns.barplot(x=x_label, y=y_label, hue=var_label, data=df, palette="mako", errorbar="sd")

    plt.savefig(os.path.join(output_path, "plot-mtsubmit.pdf"), bbox_inches='tight')
    plt.show()


if __name__ == "__main__":
    main()