import pickle
import os
import pandas as pd
import seaborn as sns
import matplotlib.ticker

import plot_utils

bench_order = ["fft", "basicmath", "crc", "sha", "stringsearch", "aes", "conv2d", "matmul"]

def validate_and_cleanup_df(df):

    assert(df["is_correct"].all())

    df["bench_name"] = df["bench_name"].apply(lambda name: name.lower()[1:])

    def process_stats(row):
        lines = row["stats"]
        result = {}
        for line in lines:
            key, val = line.split(":")
            val = int(val.strip())
            if key.startswith("checkpoint triggered"):
                result["ckpt_trig"] = val
            if key.startswith("checkpoint executed"):
                result["ckpt_exec"] = val
        return result

    df_stats = df.apply(process_stats, axis=1, result_type="expand")
    df = pd.concat([df, df_stats], axis=1)

    drop_columns = [
        "start",
        "end",
        "recovery",
        "outputs",
        "is_correct",
        "stats",
        "debug_outputs"
    ]
    df = df.drop(columns = drop_columns) 
    return df


def draw_graph():
    benchmarks = [ "vBasicMath", "vCrc", "vFFT", "vSha", "vStringSearch", "vMatMul", "vConv2d", "vAes"]
    # benchmarks = ["vBasicMath", "vCrc", "vFFT", "vSha", "vStringSearch"]
    output_dir = "/home/ybkim/workspace/imc/imc_freertos_app_m33/imc/exprs/date2025/1_overhead_stable_power/output"
    configs = ["original", "pass_count", "unroll", "adaptive"]

    all_dfs = []

    for benchmark in benchmarks:
        for config_name in configs:
            pickle_filename = f"{output_dir}/{benchmark}.{config_name}.pickle"
            if not os.path.exists(pickle_filename):
                continue
            with open(pickle_filename, "rb") as f:
                orig_df = pickle.load(f)
                df = validate_and_cleanup_df(orig_df)
                all_dfs.append(df)

    orig_df = pd.concat(all_dfs)

    mask = (orig_df["bench_name"] == "sha") & (orig_df["config"] == "adaptive")
    print(orig_df[mask])

    mask = (orig_df["bench_name"] == "sha") & (orig_df["config"] == "unroll")
    print(orig_df[mask])

    df = orig_df.groupby(["bench_name", "config"]).mean()
    df = df.reset_index()

    df = df.pivot(index="bench_name", columns="config", values="time_taken")

    include_adaptive = True

    df2 = pd.DataFrame()
    df2["counter"] = df["pass_count"] / df["original"]
    df2["unroll"] = df["unroll"] / df["original"]

    if include_adaptive:
        df2["adaptive"] = df["adaptive"] / df["original"]
        n_cols = 3
        config_order = ["counter", "unroll", "adaptive"]  
    else:
        n_cols = 1
        config_order = ["counter", "unroll"]

    df2 = df2.reset_index().melt(id_vars=["bench_name"], var_name="config", value_name="normalized")
    df2 = df2.sort_values(by="bench_name")
    # df2 = df2[df2["config"] == "PC"]
    # df2 = df2[df2["config"] == "VC"]
    # df2 = df2[df2["bench_name"] == "vFFT"]

    plot_utils.set_theme_seaborn()
    g = sns.catplot(
        data=df2,
        kind="bar",
        x="bench_name",
        y = "normalized",
        hue="config",
        aspect=3.0,
        height=4,
        order=bench_order,
        hue_order=config_order,
        legend_out=False
    )

    ax = g.ax
    ax.set_ylim([0, 10])
    ax.set_xlabel("Benchmark")
    ax.set_ylabel("Exec Time (Normalized)")
    ax.set_ylabel("Normalized")
    ax.yaxis.set_major_locator(matplotlib.ticker.MaxNLocator(5))
    ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator(2))
    ax.grid(True, which="both", axis="y")
    sns.despine(g.figure, right=False, top=False)

    for c in ax.containers:
        labels = [f"{(v.get_height()):.1f}" if v.get_height() > 10 else "" for v in c]
        padding = 90
        ax.bar_label(c, labels=labels, label_type="center", padding=padding, size=16)

    num_categories = len(df2.bench_name.unique())
    num_hues = len(df2.config.unique())
    plot_utils.draw_hatch(ax, num_categories, num_hues)

    if include_adaptive:
        ax.legend(ncol=n_cols, loc="upper center", bbox_to_anchor=(0.5, 1.4))
    else:
        ax.legend(ncol=n_cols, loc="upper left")


    return g
