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

from common import calc_throughput, load_time_mesurements, get_task_count

result_path = "benchmark-results/"
output_path = "benchmark-plots/"

runid = "Run ID"
x_label = "Destination Node"
y_label = "Throughput in GiB/s"


title_allnodes = \
    """Copy Throughput in GiB/s tested for 1GiB Elements\n
    Using all 8 DSA Chiplets available on the System"""
title_smartnodes = \
    """Copy Throughput in GiB/s tested for 1GiB Elements\n
    Using Cross-Copy for Intersocket and all 4 Chiplets of Socket for Intrasocket"""
title_difference = \
    """Gain in Copy Throughput in GiB/s of All-DSA vs. Smart Assignment"""

description_smartnodes = \
    """Copy Throughput in GiB/s tested for 1GiB Elements\n
    Nodes of {8...15} are HBM accessors for their counterparts (minus 8)\n
    Using all 4 DSA Chiplets of a Socket for Intra-Socket Operation\n
    And using only the Source and Destination Nodes DSA for Inter-Socket"""
description_allnodes = \
    """Copy Throughput in GiB/s tested for 1GiB Elements\n
    Nodes of {8...15} are HBM accessors for their counterparts (minus 8)\n
    Using all 8 DSA Chiplets available on the System"""

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

data_peakbench_andre = [
    { runid : 0, x_label : 8, y_label : 64 },
    { runid : 0, x_label : 11, y_label : 63 },
    { runid : 0, x_label : 12, y_label : 40 },
    { runid : 0, x_label : 15, y_label : 54 }
]


# 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):
    divisor = get_task_count(file_path)
    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, src_node, dst_node):
    size = 1024*1024*1024

    timing = get_timing(file_path)
    run_idx = 0
    for t in timing:
        tp = calc_throughput(size, t)
        data.append({ runid : run_idx, x_label : dst_node, y_label : tp})
        run_idx = run_idx + 1


def plot_bar(table,title,node_config):
    plt.figure(figsize=(2, 3))

    sns.barplot(x=x_label, y=y_label, data=table, palette="mako", errorbar="sd")

    plt.ylim(0, 75)

    plt.savefig(os.path.join(output_path, f"plot-{node_config}-cpu-throughput.pdf"), bbox_inches='tight')
    plt.show()


# 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(node_config,title,ext):
    src_node = 0
    for dst_node in {8,11,12,15}:
        file = os.path.join(result_path, f"copy-n{src_node}ton{dst_node}-1gib-{node_config}{ext}.json")
        process_file_to_dataset(file, src_node, dst_node)

    df = pd.DataFrame(data)

    data.clear()
    df.set_index(index, inplace=True)

    if ext == "brute": node_config = ext
    plot_bar(df, title, node_config)

    return df


def plotandre():
    df = pd.DataFrame(data_peakbench_andre)
    df.set_index(index, inplace=True)
    plot_bar(df, title_allnodes, "andrepeak")
    return df


if __name__ == "__main__":
    dandr = plotandre()
    dall = main("allnodes", title_allnodes, "-cpu")
    dbrt = main("brute", title_allnodes, "-cpu")