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

def draw_graph():
    benchmarks = [
        "vBasicMath",
        "vCrc",
        "vFFT",
        "vSha",
        "vStringSearch",
        "vMatMul",
        "vConv2d",
        "vAes",
    ]
    rows = []
    for benchmark in benchmarks:
        try:
            with open(f"output/adaptive/{benchmark}.pickle", "rb") as f:
                df = pickle.load(f)
        except:
            continue

        columns = ["bench_name", "init", "getCSC", "adjust", "total cycles"]
        df = df[columns]

        df["init (%)"] = df["init"] / df["total cycles"] * 100
        df["getCSC (%)"] = df["getCSC"] / df["total cycles"] * 100
        df["adjust (%)"] = df["adjust"] / df["total cycles"] * 100

        columns = ["bench_name", "init (%)", "getCSC (%)", "adjust (%)"]

        df = df[columns]
        converge_index = 4
        columns = ["init (%)", "getCSC (%)", "adjust (%)"]
        row = df.loc[0:converge_index][columns].mean()
        row["converged"] = False
        row["benchmark"] = benchmark[1:].lower()
        rows.append(row)

        row = df.loc[converge_index:][columns].mean()
        row["converged"] = True
        row["benchmark"] = benchmark[1:].lower()
        rows.append(row)

    df = pd.DataFrame(rows)
    df = df.sort_values(by=["benchmark", "converged"])
    df = df[["benchmark", "converged", "init (%)", "getCSC (%)", "adjust (%)"]]

    print(df[df["converged"] == False][columns].mean())
    print(df[df["converged"] == True][columns].mean())

    plot_utils.set_theme_seaborn()

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

    df = df.drop(columns=["init (%)"])
    df = df.rename(
        columns={
            "getCSC (%)": "csc",
            "adjust (%)": "adjust",
        }
    )

    df = df.pivot(
        index="benchmark",
        columns=["converged"],
    )

    df["adjust", False] += df["csc", False]
    df["adjust", True] += df["csc", True]

    df = df.stack(level="converged")
    df = df.reset_index(level="converged")

    print(df)

    fig, ax = plt.subplots(figsize=(12, 3.5))

    sns.barplot(
        data=df,
        x="benchmark",
        y="adjust",
        hue="converged",
        order=bench_order,
        ax=ax,
    )

    for bar in ax.patches:
        bar.set_hatch("///")

    sns.barplot(
        data=df,
        x="benchmark",
        y="csc",
        hue="converged",
        order=bench_order,
        ax=ax
    )

    labels = [
        "adjust",
        "adjust",
        "getCSC (run #1~5)",
        "getCSC (run #5~20)",
    ]

    new_labels = ["Before convergence (run 1~5)", "After convergence (run 6~20)"]

    handles, previous_labels = ax.get_legend_handles_labels()
    new_handles = [handles[2], handles[3]]
    ax.legend(handles=new_handles, labels=new_labels, ncols=1, columnspacing=0.8)

    ax.set_ylabel("Overhead (%)")
    ax.set_xlabel("Benchmark")
    ax.get_yaxis().set_minor_locator(matplotlib.ticker.AutoMinorLocator())
    ax.grid(visible=True, which="both", axis="y")
    ax.tick_params(left=True, length=8, width=2)

    return fig
