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

output_path = "./plots"
hbm_result = "./evaluation-results/current/qdp-xeonmax-hbm-tca4-tcb0-tcj1-tmul32-wl4294967296-cs2097152.csv"
dram_result = "./evaluation-results/current/qdp-xeonmax-dram-tca2-tcb0-tcj1-tmul32-wl4294967296-cs2097152.csv"
prefetch_result = "./evaluation-results/current/qdp-xeonmax-prefetch-tca1-tcb1-tcj1-tmul32-wl4294967296-cs8388608.csv"
distprefetch_result = "./evaluation-results/current/qdp-xeonmax-distprefetch-tca1-tcb1-tcj1-tmul32-wl4294967296-cs8388608.csv"

tt_name = "rt-ns"
function_names = [ "scana-run", "scanb-run", "aggrj-run" ]
fn_nice = [ "Scan A", "Scan B", "Aggregate" ]

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

    row_count = 0
    
    with open(fname, newline='') as csvfile:
        reader = csv.DictReader(csvfile, delimiter=';')
        for row in reader:
            row_count = row_count + 1
            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 / (1000 * 1000 * row_count) for key, value in t.items() if value != 0}

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


def read_total_time_from_csv(fname) -> float:
    time = 0
    row_count = 0
    
    with open(fname, newline='') as csvfile:
        reader = csv.DictReader(csvfile, delimiter=';')
        for row in reader:
            row_count = row_count + 1
            time += int(row["rt-ns"])

    return time / (1000 * 1000 * row_count)


def read_cache_hitrate_from_csv(fname) -> float:
    hitrate = 0
    row_count = 0
    
    with open(fname, newline='') as csvfile:
        reader = csv.DictReader(csvfile, delimiter=';')
        for row in reader:
            row_count = row_count + 1
            hitrate += float(row["cache-hr"])

    return (hitrate * 100) / row_count


def generate_speedup_table():
    baseline = read_total_time_from_csv(dram_result)
    columns = [ "Configuration", "Speedup", "Cache Hitrate", "Raw Time" ]

    names = [
        "DDR-SDRAM (Baseline)",
        "HBM (Upper Limit)",
        "Prefetching",
        "Prefetching, Distributed Columns"
    ]

    rawtime = [
        read_total_time_from_csv(dram_result),
        read_total_time_from_csv(hbm_result),
        read_total_time_from_csv(prefetch_result),
        read_total_time_from_csv(distprefetch_result),
    ]

    speedup = [
        baseline / rawtime[0],
        baseline / rawtime[1],
        baseline / rawtime[2],
        baseline / rawtime[3]
    ]

    cachehr = [
        0,
        0,
        read_cache_hitrate_from_csv(prefetch_result),
        read_cache_hitrate_from_csv(distprefetch_result)
    ]

    data = [
        [ names[0], f"x{speedup[0]:1.2f}", r" \textemdash ", f"{rawtime[0]:.2f} ms" ],
        [ names[1], f"x{speedup[1]:1.2f}", r" \textemdash ", f"{rawtime[1]:.2f} ms" ],
        [ names[2], f"x{speedup[2]:1.2f}", f"{cachehr[2]:2.2f} \%", f"{rawtime[2]:.2f} ms" ],
        [ names[3], f"x{speedup[3]:1.2f}", f"{cachehr[3]:2.2f} \%", f"{rawtime[3]:.2f} ms" ]
    ]

    return pd.DataFrame(data, columns=columns)


def generate_rawtime_base_table():
    baseline = read_total_time_from_csv(dram_result)
    columns = [ "Configuration", "Raw Time" ]

    names = [
        "DDR-SDRAM (Baseline)",
        "HBM (Upper Limit)"
    ]

    rawtime = [
        read_total_time_from_csv(dram_result),
        read_total_time_from_csv(hbm_result)
    ]

    data = [
        [ names[0], f"{rawtime[0]:.2f} ms" ],
        [ names[1], f"{rawtime[1]:.2f} ms" ]
    ]

    return pd.DataFrame(data, columns=columns)


def tex_table(df, fname):
    with open(os.path.join(output_path, fname), "w") as of:
        of.write(df.to_latex(index=False))


# 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 donut_plot(data: tuple[list[float], list[str]], fname):
    palette_color = sns.color_palette('mako_r') 
    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.2f} ms", xy=(x, y), xytext=(1.35*np.sign(x), 1.4*y), horizontalalignment=horizontalalignment, **kw)    
    
    plt.rcParams.update({'font.size': 18})
    fig.savefig(os.path.join(output_path, fname), bbox_inches='tight')


def main():
    donut_plot(read_timings_from_csv(prefetch_result), "plot-timing-prefetch.pdf")
    donut_plot(read_timings_from_csv(distprefetch_result), "plot-timing-distprefetch.pdf")
    donut_plot(read_timings_from_csv(dram_result), "plot-timing-dram.pdf")
    donut_plot(read_timings_from_csv(hbm_result), "plot-timing-hbm.pdf")
    donut_plot(read_timings_from_csv(prefetch_result), "plot-timing-prefetch.pdf")
    tex_table(generate_speedup_table(), "table-qdp-speedup.tex")
    tex_table(generate_rawtime_base_table(), "table-qdp-baseline.tex")


if __name__ == "__main__":
    main()