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

import plot_utils


def validate_and_cleanup_df(df):
    # rename_dict = {
    #     "checkpoint_enabled": "ckpt",
    #     "voltage_check_enabled": "VC",
    #     "pass_counter_enabled": "PC"
    # }
    # df = df.rename(columns=rename_dict)

    mask = df["is_correct"] == False
    if mask.any():
        print("drop rows with incorrect output")
        print(df[mask])

    df = df[~mask]
    assert df["is_correct"].all()

    # def category_mapper(row):
    #     if row["ckpt"] == False:
    #         return "Baseline"
    #     else:
    #         if row["VC"] == False:
    #             assert(row["PC"] == True)
    #             return "PC"
    #         else:
    #             assert(row["PC"] == False)
    #             return "VC"

    # category_column_name = "config"
    # df[category_column_name] = df.apply(category_mapper, axis=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",
        "stats",
        # "ckpt",
        # "VC",
        # "PC",
        "is_correct",
    ]
    df = df.drop(columns=drop_columns)
    return df


def get_num_cloned_loops(benchmark):
    d = {
        "vBasicMath": 5,
        "vCrc": 4,
        "vFFT": 9,
        "vSha": 6,
        "vStringSearch": 6,
        "vMatMul": 4,
        "vConv2d": 5,
        "vAes": 15,
    }
    return d[benchmark]


def get_base_df():
    benchmarks = [
        "vBasicMath",
        "vCrc",
        "vFFT",
        "vSha",
        "vStringSearch",
        "vMatMul",
        "vConv2d",
        "vAes",
    ]
    # benchmarks = [
    #     "vStringSearch"
    # ]
    output_dir = "/home/ybkim/workspace/imc/imc_freertos_app_m33/imc/exprs/date2025/5_size_overhead/output"

    all_dfs = []

    for benchmark in benchmarks:
        pickle_filename = f"{output_dir}/{benchmark}.pickle"
        if not os.path.exists(pickle_filename):
            print(f"pass loading {benchmark}")
            continue
        with open(pickle_filename, "rb") as f:
            df = pickle.load(f)
            mask = df["config"] == "original"
            base = int(df[mask]["text"].iloc[0])
            df["normalized_text"] = df["text"] / base
            # df = df[~mask]
            all_dfs.append(df)

    orig_df = pd.concat(all_dfs)
    return orig_df


def draw_graph():
    df = get_base_df()
    print(df)

    rows = []

    benchmarks = df["benchmark"].unique()
    for benchmark in benchmarks:
        bench_name = benchmark.lower()[1:]
        print(bench_name)

        n_cloned_loops = get_num_cloned_loops(benchmark)

        mask = (df["benchmark"] == benchmark) & (df["config"] == "original")
        func_size_original = int(df[mask]["func"])

        mask = (df["benchmark"] == benchmark) & (df["config"] == "adaptive")
        func_size_after = int(df[mask]["func"])

        func_size_overhead = func_size_after / func_size_original
        func_size_original = f"{func_size_original:,}"
        func_size_after = f"{func_size_after:,}"

        text_overhead = float(df[mask]["normalized_text"] * 100 - 100)

        rows.append({
            "benchmark": bench_name,
            "cloned loops (N)": n_cloned_loops,
            "function size (orig)": func_size_original,
            "function size (after)": func_size_after,
            "function size overhead": func_size_overhead,
            "\\texttt{.text} increased (%)": text_overhead,
        })

    d = pd.DataFrame(rows)
    print(d)
    print("average function size overhead: ")
    print(d["function size overhead"].mean())
    print("average code size overhead: ")
    print(d["\\texttt{.text} increased (%)"].mean())
    d = d.drop(columns=["function size overhead"])
    latex = d.to_latex(
        index=False,
        escape=True,
        float_format="%.2f",
        column_format="ccccc",
    )
    print(latex)


if __name__ == "__main__":
    draw_graph()
