import pickle
import os
import pandas as pd
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",
    ]
    duplication_ratios = list(range(10, 101, 10))
    config_name = "adaptive"

    all_dfs = []

    for benchmark in benchmarks:
        for duplication_ratio in duplication_ratios:
            try:
                df = pd.read_pickle(f"output/{benchmark}/duplication_ratio_{duplication_ratio}.pickle")
                all_dfs.append(df)
            except FileNotFoundError:
                continue
    
    all_df = pd.concat(all_dfs, ignore_index=True)
    print(all_df)

    rc = {
        "lines.linewidth": 1.5,
        "axes.titlepad": 10,
        "ytick.major.pad": -4,
        "ytick.labelsize": 20,
        "legend.fontsize": 23,
        "axes.labelsize": 25,
    }
    plot_utils.set_theme_seaborn(kind="bar", rc_custom=rc)

    fig_size = (18, 6)
    n_rows = 2
    n_cols = 4
    hspace = 0.05
    fig = plt.figure(figsize=fig_size)
    axes = fig.subplots(n_rows, n_cols, sharex=True, gridspec_kw={"hspace": hspace})

    for i, benchmark in enumerate(benchmarks):
        ax = axes.reshape(-1)[i]
        mask = (all_df["benchmark"] == benchmark)
        df = all_df[mask]
        if df.empty:
            continue
        sns.barplot(data=df, x="duplication_ratio", y="func", ax=ax)
        ax.set_title(benchmark)
        ax.set_xlabel("Duplication Ratio (%)")
        ax.set_ylabel("Func Size (bytes)")
        ax.grid(True)

        ax.set_title(benchmark)
        ax.set_xlabel("")
        ax.set_ylabel("")
        if i / 4 >= 1:
            ax.set_xlabel("Execution #")
        if i in [0, 4]:
            ax.set_ylabel("Func Size (kB)")

        ax.yaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(lambda x, pos: f"{x/1000:.1f}"))
        ax.yaxis.grid(visible=True, which="both")
        # ax.yaxis.set_major_locator(matplotlib.ticker.MaxNLocator(3))
        ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator(3))

        ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator(2))
        ax.xaxis.grid(visible=True, which="both")


        if i == 0:
            ax.legend(
                ncol=1, loc="upper left", bbox_to_anchor=(0.33, 0.995), labelspacing=0.3,
                handlelength=1,
            )
            # ax.get_legend().remove()
        else:
            if ax.get_legend() is not None:
                ax.get_legend().remove()
    
    return fig