import os
import csv
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

output_path = "./plots"
prefetch_result = "./evaluation-results/qdp-xeonmax-prefetch-tca2-tcb1-tcj1-tmul8-wl4294967296-cs16777216.csv"
dram_result = "./evaluation-results/qdp-xeonmax-dram-tca2-tcb0-tcj1-tmul8-wl4294967296-cs2097152.csv"

tt_name = "rt-ns"
function_names = [ "scana-run", "scanb-run", "aggrj-run" ]
fn_nice = [ "Scan A, Filter", "Scan B, Prefetch", "Aggregate, Project + Sum" ]

def read_timings_from_csv(fname) -> tuple[list[float], list[str]]:
    t = {}
    total_time = 0

    # Read data from CSV file
    with open(fname, newline='') as csvfile:
        reader = csv.DictReader(csvfile, delimiter=';')
        for row in reader:
            total_time += int(row[tt_name])
            for i in range(len(function_names)):
                t[fn_nice[i]] = t.get(fn_nice[i], 0) + int(row[function_names[i]])

    t = {key: value * 100 / total_time for key, value in t.items() if value != 0}

    total = sum(list(t.values()))
    if total < 100.0:
        t["Waiting / Other"] = 100.0 - total

    return list(t.values()), list(t.keys())


def get_data_prefetch_cache_access() -> tuple[list[float], list[str]]:
    total = 0.3
    data = [ 0.07, 0.19, 0.04 ]
    data = [ x * 100 / total for x in data ]
    keys = ["numa_alloc_onnode", "dml::make_mem_move_task", "dml::hardware_device::submit"]

    return data,keys

def get_data_prefetch_total() -> tuple[list[float], list[str]]:
    return read_timings_from_csv(prefetch_result)

def get_data_dram_total() -> tuple[list[float], list[str]]:
    return read_timings_from_csv(dram_result)

# 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(data: tuple[list[float], list[str]], fname):
    palette_color = sns.color_palette('mako') 
    fig, ax = plt.subplots(figsize=(6, 3), subplot_kw=dict(aspect="equal"))

    wedges, texts = ax.pie(data[0], wedgeprops=dict(width=0.5), startangle=-40, colors=palette_color)

    bbox_props = dict(boxstyle="square,pad=0.3", fc="w", ec="k", lw=0.72)
    kw = dict(arrowprops=dict(arrowstyle="-"), bbox=bbox_props, zorder=0, va="center")

    for i, p in enumerate(wedges):
        ang = (p.theta2 - p.theta1)/2. + p.theta1
        y = np.sin(np.deg2rad(ang))
        x = np.cos(np.deg2rad(ang))
        horizontalalignment = {-1: "right", 1: "left"}[int(np.sign(x))]
        connectionstyle = f"angle,angleA=0,angleB={ang}"
        kw["arrowprops"].update({"connectionstyle": connectionstyle})
        ax.annotate(f"{data[1][i]} - {data[0][i]:2.1f}%", xy=(x, y), xytext=(1.35*np.sign(x), 1.4*y), horizontalalignment=horizontalalignment, **kw)    
    
    fig.savefig(os.path.join(output_path, fname), bbox_inches='tight')


if __name__ == "__main__":
    main(get_data_prefetch_cache_access(), "plot-timing-prefetch-cacheaccess.pdf")
    main(get_data_prefetch_total(), "plot-timing-prefetch-totalexec.pdf")
    main(get_data_dram_total(), "plot-timing-dram-totalexec.pdf")