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

RESULTS_DIR = "results/2025-11-12-dual-floats-v2"
FIGURES_DIR = "figures/2025-11-12-dual-floats-v2"
if not os.path.exists(FIGURES_DIR):
    os.mkdir(FIGURES_DIR)
NUM_BINS = 10
df_baselines_disaggregated = pd.read_json(os.path.join(RESULTS_DIR, "baselines.json"))

df_baselines_disaggregated.replace({"compiler": {"gcomp": "grev"}}, inplace=True)
df_baselines_disaggregated["implementation"] = (
    df_baselines_disaggregated["compiler"]
    + " ("
    + df_baselines_disaggregated["typing"]
    + ")"
)
df_baselines_disaggregated.sort_values(
    by=["bench", "compiler", "typing", "cflags"], inplace=True, ignore_index=True
)

df_baselines = (
    df_baselines_disaggregated.groupby(["bench", "compiler", "typing", "cflags"])
    .agg(
        mean_runtime=("runtime", "mean"),
        std_runtime=("runtime", "std"),
        runs=("run", "count"),
    )
    .reset_index()
)
df_baselines["implementation"] = (
    df_baselines["compiler"] + " (" + df_baselines["typing"] + ")"
)


def preprocess_grev(filename: str, use_hex=False):
    df = pd.read_json(filename)
    df_baselines_disaggregated.replace({"compiler": {"gcomp": "grev"}}, inplace=True)

    split_df = df["basename"].str.split("_", n=3, expand=True)

    df["bucket"] = split_df[1].astype(int)
    df["sample"] = split_df[2].astype(int)
    if use_hex:
        df["hex_string"] = split_df[3]
        df["binary_string"] = df["hex_string"].apply(lambda x: bin(int(x, 16))[2:])
    else:
        df["binary_string"] = split_df[3]

    df.loc[:, "group"] = df["bucket"].astype(str) + "_" + df["sample"].astype(str)

    df = df[["group", "bucket", "sample", "run", "binary_string", "runtime"]]

    df["n_ones"] = df["binary_string"].str.count("1")
    df["total_length"] = df["binary_string"].str.len()
    df["precision"] = df["n_ones"] / df["total_length"]

    df["configuration"] = df["binary_string"]
    df = df.drop(
        columns=["n_ones", "total_length", "bucket", "sample", "group", "binary_string"]
    )

    df.sort_values(by="precision", ascending=True, inplace=True)
    df.reset_index(inplace=True, drop=True)
    return df


def preprocess_grift(filename: str):
    df = pd.read_json(filename)
    df.rename(columns={"basename": "configuration"}, inplace=True)
    df.sort_values(by="precision", ascending=True, inplace=True)
    df.reset_index(inplace=True, drop=True)
    max_precision = max(df["precision"])
    if max_precision > 1.0:
        pass
    min_precision = min(df["precision"])
    if min_precision > 0.0:
        pass
    if (max_precision > 1.0) or (min_precision > 0.0):
        df["precision"] = (df["precision"] - min_precision) / (
            max_precision - min_precision
        )
    return df


def load_data(dir: str):
    df = pd.DataFrame()
    for file in os.listdir(RESULTS_DIR):
        path = os.path.join(RESULTS_DIR, file)
        if os.path.isdir(path):
            continue
        basename = os.path.splitext(file)[0]
        parts = basename.split("_")
        compiler = parts[0]
        if compiler.startswith("baselines"):
            continue
        bench = parts[1]
        typing = parts[2]
        if typing == "mpairs":
            bench += "_mpairs"
            typing = parts[3]
            parts = parts[1:]
        elif typing == "body":
            bench += "_body"
            typing = parts[3]
            parts = parts[1:]
        cflags = parts[3:]
        if len(cflags) > 1 and cflags[0] == "succinct":
            typing += "_succinct"
            cflags = cflags[1:]
        elif len(cflags) > 1 and cflags[0] == "noop":
            typing += "_noop"
            cflags = cflags[1:]
        elif len(cflags) > 1 and cflags[0] == "proxy":
            typing += "_proxy"
            cflags = cflags[1:]
        elif len(cflags) > 1 and cflags[0] == "mono":
            typing += "_mono"
            cflags = cflags[1:]
        cflags = " ".join(cflags)

        if compiler == "gcomp":
            bench_df = preprocess_grev(path, use_hex=True)
            compiler = "grev"
        elif compiler == "grift":
            bench_df = preprocess_grift(path)
        else:
            raise ValueError("unexepcted compiler")
        df = pd.concat(
            [
                df,
                bench_df.assign(
                    compiler=compiler, bench=bench, typing=typing, cflags=cflags
                ),
            ],
            ignore_index=True,
        )
    return df


df_disaggregated = load_data(RESULTS_DIR)

df = (
    df_disaggregated.groupby(
        ["compiler", "bench", "typing", "cflags", "configuration", "precision"]
    )
    .agg(
        mean_runtime=("runtime", "mean"),
        std_runtime=("runtime", "std"),
        runs=("run", "count"),
    )
    .reset_index()
)

sns.color_palette("tab20")
palette = sns.color_palette("tab20")
colors = {
    "grev": {
        "static": palette[0],
        "dynamic": palette[2],
        "gradual": palette[8],
        "gradual_proxy": palette[12],
    },
    "grift": {
        "static": palette[10],
        "gradual": palette[4],
        "gradual_mono": palette[16],
    },
    "racket": {
        "dynamic": palette[18],
    },
}

colormap = {
    f"{compiler} ({typing})": colors[compiler][typing]
    for compiler in colors
    for typing in colors[compiler]
}
bench_order = sorted(df["bench"].unique())

np.random.seed(42)
jitter_amount = (1 / NUM_BINS) * 0.75

##############################
#####     BAR PLOTS      #####
##############################

legend_changes = {
    "grev (gradual_proxy)": "GrEv/G",
    "grev (gradual)": "GrEv/MV",
    "grev (static)": "StaticGrEv",
    "grift (gradual)": "Grift/G",
    "grift (gradual_mono)": "Grift/MS",
    "grift (static)": "StaticGrift",
    "racket (dynamic)": "Racket",
}

colormap_new = {
    new_name: colormap[old_name]
    for old_name, new_name in legend_changes.items()
    if old_name in colormap
}

float_bench = ["blackscholes", "fft", "n_body", "ray"]

general_bench = ["array", "matmult", "qsort_mpairs", "quicksort", "sieve", "tak"]

df_bar = df_disaggregated[["compiler", "bench", "typing", "cflags", "runtime"]].copy()
df_bar = df_bar[df_bar["typing"] != "gradual_succinct"]
df_bar["implementation"] = df_bar["compiler"] + " (" + df_bar["typing"] + ")"

df_baselines_bar = df_baselines_disaggregated[
    ["compiler", "bench", "typing", "cflags", "runtime"]
].copy()

df_baselines_bar["implementation"] = (
    df_baselines_bar["compiler"] + " (" + df_baselines_bar["typing"] + ")"
)

df_merge = pd.merge(df_bar, df_baselines_bar, on=["bench"], suffixes=("", "_base"))


def create_barplot(
    df,
    imp,
    legend=True,
    legend_handle="upper left",
    legend_anchor=(0.8, 1),
    rotation=45,
    height=4.5,
):
    WIDTH_PER_CATEGORY = 1.2
    PLOT_HEIGHT = height
    FONTSIZE = 12
    FONTFAMILY = "serif"
    FONTWEIGHT = 530
    n = len(df["bench"].unique())
    aspect = (n * WIDTH_PER_CATEGORY) / PLOT_HEIGHT

    df["bench"] = df["bench"].apply(lambda x: x.replace("_", "-"))

    bar_plot = sns.catplot(
        df,
        kind="bar",
        x="bench",
        y="speedup",
        hue="implementation",
        hue_order=[
            impl
            for impl in colormap_new.keys()
            if impl in df["implementation"].unique()
        ],
        palette=colormap_new,
        legend=legend,
        height=PLOT_HEIGHT,
        aspect=aspect,
    )

    for ax in bar_plot.axes.flat:
        ax.axhline(y=1.0, color="gray", linestyle="--", linewidth=1.5)

    bar_plot.tick_params(axis="x", rotation=rotation)
    bar_plot.set_axis_labels(
        "",
        "Speedup with respect to " + legend_changes[imp],
        fontsize=FONTSIZE,
        fontweight=FONTWEIGHT,
        fontfamily=FONTFAMILY,
    )
    bar_plot.set_xticklabels(
        fontsize=FONTSIZE, fontfamily=FONTFAMILY, fontweight="bold"
    )
    bar_plot.set_yticklabels(
        fontsize=FONTSIZE, fontfamily=FONTFAMILY, fontweight=FONTWEIGHT
    )
    if legend:
        bar_plot.legend.set_title(None)
        sns.move_legend(
            bar_plot,
            legend_handle,
            bbox_to_anchor=legend_anchor,
            prop={"family": FONTFAMILY, "size": FONTSIZE, "weight": FONTWEIGHT},
            frameon=True,
        )
        frame = bar_plot.legend.get_frame()
        frame.set_facecolor("white")
        frame.set_alpha(1)
        frame.set_boxstyle("square")
        bar_plot.legend.set_zorder(100)
    plt.tight_layout()

    df["bench"] = df["bench"].apply(lambda x: x.replace("-", "_"))
    return bar_plot


base_implementations = ["grift (static)", "racket (dynamic)"]

for imp in base_implementations:

    df = df_merge[df_merge["implementation_base"] == imp].copy()
    df["speedup"] = df["runtime_base"] / df["runtime"]

    df_float = df[df["bench"].isin(float_bench)]
    df_general = df[df["bench"].isin(general_bench)]
    df = df.replace(legend_changes)
    df_float = df_float.replace(legend_changes)
    df_general = df_general.replace(legend_changes)

    order = legend_changes.values()

    bar_plot = create_barplot(df, imp)
    bar_plot.savefig(
        os.path.join(
            FIGURES_DIR,
            "GrEv-G-MV_Grift-G-MS_vs_" + legend_changes[imp].replace("/", "-") + "_all",
        )
    )

    bar_plot = create_barplot(df_float, imp, legend=imp != "racket (dynamic)")
    bar_plot.savefig(
        os.path.join(
            FIGURES_DIR,
            "GrEv-G-MV_Grift-G-MS_vs_"
            + legend_changes[imp].replace("/", "-")
            + "_float",
        )
    )

    bar_plot = create_barplot(df_general, imp, legend=imp != "racket (dynamic)")
    bar_plot.savefig(
        os.path.join(
            FIGURES_DIR,
            "GrEv-G-MV_Grift-G-MS_vs_"
            + legend_changes[imp].replace("/", "-")
            + "_general",
        )
    )


racket_mean_runtime = df_baselines_bar[
    df_baselines_bar["implementation"] == "racket (dynamic)"
]["runtime"].mean()

df_baselines_bar["speedup"] = racket_mean_runtime / df_baselines_bar["runtime"]
df_baselines_bar = df_baselines_bar.replace(legend_changes)
df_baselines_bar = df_baselines_bar[df_baselines_bar["compiler"] != "racket"]

df_float = df_baselines_bar[df_baselines_bar["bench"].isin(float_bench)]
df_general = df_baselines_bar[df_baselines_bar["bench"].isin(general_bench)]

bar_plot = create_barplot(
    df_general, "racket (dynamic)", legend_anchor=(0.8, 1), rotation=0
)
bar_plot.savefig(
    os.path.join(FIGURES_DIR, "GrEv-Static_Grift-Static_vs_Racket_general")
)

bar_plot = create_barplot(df_float, "racket (dynamic)", legend=False, rotation=0)
bar_plot.savefig(os.path.join(FIGURES_DIR, "GrEv-Static_Grift-Static_vs_Racket_float"))

df_baselines_bar = df_baselines_bar[
    df_baselines_bar["bench"].isin(float_bench + general_bench)
]

bar_plot = create_barplot(df_baselines_bar, "racket (dynamic)")
bar_plot.savefig(os.path.join(FIGURES_DIR, "GrEv-Static_Grift-Static_vs_Racket_all"))

############################

# bar_plot1 = sns.barplot(
#     df_baselines_disaggregated,
# )
#
# fg = sns.catplot(
#     df_baselines_disaggregated,
#     kind="box",
#     x="implementation",
#     y="runtime",
#     hue="implementation",
#     col="bench",
#     col_wrap=3,
#     sharey=False,
#     col_order=sorted(df_baselines_disaggregated["bench"].unique()),
#     palette=colormap,
#     legend=True,
#     legend_out=False,
#     hue_order=["racket (dynamic)", "grev (dynamic)", "grev (static)", "grift (static)"],
# )
#
# for bench, ax in fg.axes_dict.items():
#     ax.set_ylim(bottom=0)
#     ax.set_ylabel("Runtime (s)", fontweight="bold", fontsize=11)
#     ax.set_xlabel(None)
#
#     ax.set_title(f"{bench}", fontweight="bold", fontsize=12, pad=10)
#     ax.tick_params(labelbottom=True)
#     ax.grid(True, alpha=0.5, linestyle="-", linewidth=1, axis="y")
#     ax.spines["top"].set_visible(False)
#     ax.spines["right"].set_visible(False)
# sns.move_legend(fg, "lower right", title="Implementation")
# fg.savefig(os.path.join(FIGURES_DIR, "baselines"))
#
# fg = sns.catplot(
#     df_baselines_disaggregated,
#     kind="bar",
#     x="typing",
#     y="runtime",
#     hue="compiler",
#     col="bench",
#     col_wrap=4,
#     sharey=False,
#     col_order=sorted(df_baselines_disaggregated["bench"].unique()),
#     legend=True,
#     legend_out=False,
#     height=2,
#     aspect=1,
# )
#
# for bench, ax in fg.axes_dict.items():
#     ax.set_ylim(bottom=0)
#
#     ax.set_title(f"{bench}", pad=12)
#     ax.tick_params(labelbottom=True)
#     ax.grid(True, alpha=0.5, linestyle="-", linewidth=1, axis="y")
#     ax.spines["top"].set_visible(False)
#     ax.spines["right"].set_visible(False)
# fg.set_ylabels("Runtime (s)")
# fg.set_xlabels("Compiler")
# sns.move_legend(fg, "lower right")
#
# fg = sns.catplot(
#     data=df_baselines_disaggregated,
#     kind="box",
#     y="runtime",
#     x="compiler",
#     hue="typing",
#     col="bench",
#     col_wrap=6,
#     height=2,
#     aspect=0.75,
#     legend=True,
#     legend_out=False,
#     sharey=False,
# )
#
# for bench, ax in fg.axes_dict.items():
#     y_max = (
#         df_baselines_disaggregated[df_baselines_disaggregated["bench"] == bench][
#             "runtime"
#         ].max()
#         * 1.15
#     )
#     ax.set_ylim(bottom=0, top=y_max)
#
# fg.set_titles("{col_name}")
# fg.set_ylabels("runtime (s)")
# sns.move_legend(fg, "lower right", title="Implementation")
#
#
# def get_extremes(df):
#     precision_bounds = df.groupby(
#         [
#             "bench",
#             "compiler",
#             "typing",
#             "cflags",
#         ]
#     ).agg(
#         min_precision=("precision", "min"),
#         max_precision=("precision", "max"),
#     )
#
#     df_merged = df.merge(precision_bounds, on=["bench", "compiler", "typing", "cflags"])
#
#     df_min = df_merged[df_merged["precision"] == df_merged["min_precision"]].copy()
#     df_min["precision_type"] = "min"
#
#     df_max = df_merged[df_merged["precision"] == df_merged["max_precision"]].copy()
#     df_max["precision_type"] = "max"
#
#     df_filtered = pd.concat([df_min, df_max], ignore_index=True)
#     df_filtered["implementation"] = (
#         df_filtered["compiler"] + " (" + df_filtered["typing"] + ")"
#     )
#     return df_filtered
#
#
# df_extremes = get_extremes(df_disaggregated)
#
# fg = sns.catplot(
#     data=df_extremes[df_extremes["precision_type"] == "min"],
#     kind="violin",
#     y="runtime",
#     hue="implementation",
#     col="bench",
#     col_wrap=3,
#     sharey=False,
#     hue_order=[
#         "grev (gradual)",
#         "grev (gradual_proxy)",
#         "grift (gradual)",
#         "grift (gradual_mono)",
#     ],
#     col_order=sorted(df_extremes["bench"].unique()),
#     palette=colormap,
#     legend_out=False,
# )
#
# for bench, ax in fg.axes_dict.items():
#     ax.set_ylabel("Runtime (s)", fontweight="bold", fontsize=11)
#     ax.set_xlabel(None)
#
#     ax.set_title(f"{bench} (untyped)", fontweight="bold", fontsize=12, pad=10)
#
#     ax.grid(True, alpha=0.5, linestyle="-", linewidth=1, axis="y")
#     ax.spines["top"].set_visible(False)
#     ax.spines["right"].set_visible(False)
#
#     static_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "static")
#     ].iloc[0]
#
#     dynamic_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "dynamic")
#     ].iloc[0]
#
#     lines = [
#         ax.axhline(
#             y=dynamic_baseline["mean_runtime"],
#             linestyle="--",
#             linewidth=3,
#             label="grev (dynamic)",
#             color=colors["grev"]["dynamic"],
#             alpha=0.8,
#         )
#     ]
#
# handles = fg.legend.legend_handles
# fg._legend.remove()
# fg.add_legend(handles=lines + handles, title="Implementation")
# sns.move_legend(fg, loc="lower right")
# fg.savefig(os.path.join(FIGURES_DIR, "fully_untyped"))
#
# fg = sns.catplot(
#     data=df_extremes[df_extremes["precision_type"] == "max"],
#     kind="violin",
#     y="runtime",
#     hue="implementation",
#     col="bench",
#     col_wrap=3,
#     sharey=False,
#     hue_order=[
#         "grev (gradual)",
#         "grev (gradual_proxy)",
#         "grift (gradual)",
#         "grift (gradual_mono)",
#     ],
#     col_order=sorted(df_extremes["bench"].unique()),
#     palette=colormap,
#     legend_out=False,
# )
#
# for bench, ax in fg.axes_dict.items():
#     ax.set_ylabel("Runtime (s)", fontweight="bold", fontsize=11)
#     ax.set_xlabel(None)
#
#     ax.set_title(f"{bench} (typed)", fontweight="bold", fontsize=12, pad=10)
#
#     ax.grid(True, alpha=0.5, linestyle="-", linewidth=1, axis="y")
#     ax.spines["top"].set_visible(False)
#     ax.spines["right"].set_visible(False)
#
#     static_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "static")
#     ].iloc[0]
#
#     dynamic_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "dynamic")
#     ].iloc[0]
#
#     grift_static_baseline = df_baselines[
#         (df_baselines["compiler"] == "grift")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "static")
#     ].iloc[0]
#
#     lines = [
#         ax.axhline(
#             y=static_baseline["mean_runtime"],
#             linestyle="--",
#             linewidth=3,
#             label="grev (static)",
#             color=colors["grev"]["static"],
#             alpha=0.8,
#         ),
#         ax.axhline(
#             y=grift_static_baseline["mean_runtime"],
#             linestyle="--",
#             linewidth=3,
#             label="grift (static)",
#             color=colors["grift"]["static"],
#             alpha=0.8,
#         ),
#     ]
#
# handles = fg.legend.legend_handles
# fg._legend.remove()
# fg.add_legend(handles=lines + handles, title="Implementation")
# sns.move_legend(fg, loc="lower right")
# fg.savefig(os.path.join(FIGURES_DIR, "fully_typed"))
#
# df_filtered = df[
#     (df["compiler"] == "grev") & (df["typing"] != "gradual_succinct")
# ].copy()
#
# df_filtered["jittered_precision"] = df_filtered["precision"] + np.random.uniform(
#     -jitter_amount, jitter_amount, size=len(df_filtered)
# )
#
# fg = sns.FacetGrid(
#     df_filtered,
#     col="bench",
#     col_wrap=3,
#     height=6,
#     sharey=False,
#     col_order=bench_order,
#     legend_out=False,
# )
#
#
# def plot_facet(data, **kwargs):
#     bench = data["bench"].iloc[0]
#     ax = plt.gca()
#
#     static_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "static")
#     ].iloc[0]
#
#     dynamic_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "dynamic")
#     ].iloc[0]
#
#     ax.axhline(
#         y=static_baseline["mean_runtime"],
#         linestyle="--",
#         linewidth=3,
#         label="grev (static)",
#         color=colors["grev"]["static"],
#         alpha=0.8,
#     )
#     ax.axhline(
#         y=dynamic_baseline["mean_runtime"],
#         linestyle="--",
#         linewidth=3,
#         label="grev (dynamic)",
#         color=colors["grev"]["dynamic"],
#         alpha=0.8,
#     )
#
#     for i, compiler in enumerate(sorted(data["compiler"].unique())):
#         for j, typing in enumerate(
#             data[data["compiler"] == compiler]["typing"].unique()
#         ):
#             subset = data[(data["compiler"] == compiler) & (data["typing"] == typing)]
#             label = f"{compiler} ({typing})" if typing != "grift" else compiler
#
#             ax.errorbar(
#                 x=subset["jittered_precision"],
#                 y=subset["mean_runtime"],
#                 yerr=subset["std_runtime"],
#                 label=label,
#                 fmt="o",
#                 markersize=6,
#                 linestyle="",
#                 color=colors[compiler][typing],
#                 alpha=0.7,
#                 capsize=4,
#                 capthick=1.5,
#                 elinewidth=1.5,
#             )
#
#     ax.set_xlabel("Precision (ratio of static types)")
#     ax.set_ylabel("Runtime (s)")
#     ax.set_title(f"{bench}", fontweight="bold", pad=20)
#     ax.grid(True, alpha=0.5, linestyle="--", linewidth=1)
#
#
# fg.map_dataframe(plot_facet)
# fg.tight_layout()
# fg.add_legend(title="Implementation")
# sns.move_legend(fg, loc="lower right")
# fg.savefig(os.path.join(FIGURES_DIR, "runtimes_grev"))
#
# df_filtered = df[df["typing"] != "gradual_succinct"].copy()
#
# df_filtered["jittered_precision"] = df_filtered["precision"] + np.random.uniform(
#     -jitter_amount, jitter_amount, size=len(df_filtered)
# )
#
# fg = sns.FacetGrid(
#     df_filtered,
#     col="bench",
#     col_wrap=3,
#     height=6,
#     sharey=False,
#     col_order=bench_order,
#     hue_order=[
#         "grev (gradual)",
#         "grev (gradual_proxy)",
#         "grift (gradual)",
#         "grift (gradual_mono)",
#     ],
#     legend_out=False,
# )
#
#
# def plot_facet(data, **kwargs):
#     bench = data["bench"].iloc[0]
#     ax = plt.gca()
#
#     static_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "static")
#     ].iloc[0]
#
#     dynamic_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "dynamic")
#     ].iloc[0]
#
#     static_baseline_grift = df_baselines[
#         (df_baselines["compiler"] == "grift")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "static")
#     ].iloc[0]
#
#     ax.axhline(
#         y=static_baseline["mean_runtime"],
#         linestyle="--",
#         linewidth=3,
#         label="grev (static)",
#         color=colors["grev"]["static"],
#         alpha=0.8,
#     )
#     ax.axhline(
#         y=dynamic_baseline["mean_runtime"],
#         linestyle="--",
#         linewidth=3,
#         label="grev (dynamic)",
#         color=colors["grev"]["dynamic"],
#         alpha=0.8,
#     )
#
#     ax.axhline(
#         y=static_baseline_grift["mean_runtime"],
#         linestyle="--",
#         linewidth=3,
#         label="grift (static)",
#         color=colors["grift"]["static"],
#         alpha=0.8,
#     )
#
#     for i, compiler in enumerate(sorted(data["compiler"].unique())):
#         for j, typing in enumerate(
#             data[data["compiler"] == compiler]["typing"].unique()
#         ):
#             subset = data[(data["compiler"] == compiler) & (data["typing"] == typing)]
#             label = f"{compiler} ({typing})"
#
#             ax.errorbar(
#                 x=subset["jittered_precision"],
#                 y=subset["mean_runtime"],
#                 yerr=subset["std_runtime"],
#                 label=label,
#                 fmt="o",
#                 markersize=6,
#                 linestyle="",
#                 color=colors[compiler][typing],
#                 alpha=0.7,
#                 capsize=4,
#                 capthick=1.5,
#                 elinewidth=1.5,
#             )
#
#     ax.set_xlabel("Precision (ratio of static types)")
#     ax.set_ylabel("Runtime (s)")
#     ax.set_title(f"{bench}", fontweight="bold", pad=20)
#     ax.grid(True, alpha=0.5, linestyle="--", linewidth=1)
#
#
# fg.map_dataframe(plot_facet)
# fg.add_legend(title="Implementation")
# sns.move_legend(fg, loc="lower right")
# fg.savefig(os.path.join(FIGURES_DIR, "runtimes"))
#
# df_filtered = df_disaggregated[df_disaggregated["typing"] != "gradual_succinct"].copy()
#
# df_filtered["jittered_precision"] = df_filtered["precision"] + np.random.uniform(
#     -jitter_amount, jitter_amount, size=len(df_filtered)
# )
#
# fg = sns.FacetGrid(
#     df_filtered,
#     col="bench",
#     col_wrap=3,
#     height=6,
#     sharey=False,
#     col_order=bench_order,
#     legend_out=False,
# )
#
#
# def plot_facet(data, **kwargs):
#     bench = data["bench"].iloc[0]
#     ax = plt.gca()
#
#     static_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "static")
#     ].iloc[0]
#
#     dynamic_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "dynamic")
#     ].iloc[0]
#
#     static_baseline_grift = df_baselines[
#         (df_baselines["compiler"] == "grift")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "static")
#     ].iloc[0]
#
#     ax.axhline(
#         y=static_baseline["mean_runtime"],
#         linestyle="--",
#         linewidth=3,
#         label="grev (static)",
#         color=colors["grev"]["static"],
#         alpha=0.8,
#     )
#     ax.axhline(
#         y=dynamic_baseline["mean_runtime"],
#         linestyle="--",
#         linewidth=3,
#         label="grev (dynamic)",
#         color=colors["grev"]["dynamic"],
#         alpha=0.8,
#     )
#
#     ax.axhline(
#         y=static_baseline_grift["mean_runtime"],
#         linestyle="--",
#         linewidth=3,
#         label="grift (static)",
#         color=colors["grift"]["static"],
#         alpha=0.8,
#     )
#
#     for i, compiler in enumerate(sorted(data["compiler"].unique())):
#         for j, typing in enumerate(
#             data[data["compiler"] == compiler]["typing"].unique()
#         ):
#             subset = data[(data["compiler"] == compiler) & (data["typing"] == typing)]
#             label = f"{compiler} ({typing})"
#
#             ax.scatter(
#                 x=subset["jittered_precision"],
#                 y=subset["runtime"],
#                 label=label,
#                 color=colors[compiler][typing],
#                 s=10,
#             )
#
#     ax.set_xlabel("Precision (ratio of static types)")
#     ax.set_ylabel("Runtime (s)")
#     ax.set_title(f"{bench}", fontweight="bold", pad=20)
#     ax.grid(True, alpha=0.5, linestyle="--", linewidth=1)
#
#
# fg.map_dataframe(plot_facet)
# fg.add_legend(title="Implementation")
# sns.move_legend(fg, loc="lower right")
# fg.savefig(os.path.join(FIGURES_DIR, "runtimes_disaggregated"))
#
# df["slowdown"] = df.apply(
#     lambda row: row["mean_runtime"]
#     / df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["typing"] == "static")
#         & (df_baselines["bench"] == row["bench"])
#     ].iloc[0]["mean_runtime"],
#     axis=1,
# )
# df.sort_values(by=["slowdown"], inplace=True, ignore_index=True)
# df["num_below_slowdown"] = (
#     df.groupby(["compiler", "typing", "bench"])["slowdown"].cumcount() + 1
# )
#
#
# def fmt_slowdown(num):
#     return f"{num:.1f}" if num < 10 else f"{num:.0f}"
#
#
# df_filtered = df[df["typing"] != "gradual_succinct"].copy()
# df_filtered.sort_values(
#     by=["compiler", "typing", "bench"], inplace=True, ignore_index=True
# )
#
# fg = sns.FacetGrid(df_filtered, col="bench", col_wrap=3, height=6, legend_out=False)
#
# x_max = df_filtered["slowdown"].max()
#
#
# def plot_facet(data, **kwargs):
#     bench = data["bench"].iloc[0]
#     ax = plt.gca()
#
#     sns.lineplot(
#         data=data,
#         x="slowdown",
#         y="num_below_slowdown",
#         hue=data["compiler"] + " (" + data["typing"] + ")",
#         linewidth=3,
#         palette=colormap,
#         ax=ax,
#     )
#
#     ax.set_xscale("log")
#     ax.set_ylabel("Number of configurations", fontsize=12)
#     ax.set_xlabel("Slowdown over static grev", fontsize=12)
#     ax.set_title(f"{bench}", fontweight="bold", pad=20)
#
#     for compiler in data["compiler"].unique():
#         for typing in data[data["compiler"] == compiler]["typing"].unique():
#             typing_data = data[
#                 (data["compiler"] == compiler) & (data["typing"] == typing)
#             ]
#             if not typing_data.empty:
#                 last_point = typing_data.iloc[-1]
#                 offset = (-30, 5)
#                 ax.annotate(
#                     f'{fmt_slowdown(last_point["slowdown"])}x',
#                     xy=(last_point["slowdown"], last_point["num_below_slowdown"]),
#                     xytext=offset,
#                     textcoords="offset points",
#                     fontweight="semibold",
#                     fontsize=12,
#                 )
#
#     major_xticks = [1 << n for n in range(int(math.ceil(math.log2(x_max))) + 1)]
#     ax.set_xticks(major_xticks)
#     ax.set_xticklabels([f"{val}" for val in major_xticks])
#
#     static_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "static")
#     ].iloc[0]
#
#     dynamic_baseline = df_baselines[
#         (df_baselines["compiler"] == "grev")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "dynamic")
#     ].iloc[0]
#
#     static_baseline_grift = df_baselines[
#         (df_baselines["compiler"] == "grift")
#         & (df_baselines["bench"] == bench)
#         & (df_baselines["typing"] == "static")
#     ].iloc[0]
#
#     handles, labels = ax.get_legend_handles_labels()
#     lines = [
#         ax.axvline(
#             x=static_baseline["mean_runtime"] / static_baseline["mean_runtime"],
#             linestyle="-",
#             linewidth=3,
#             label="grev (static)",
#             color=colors["grev"]["static"],
#             alpha=0.8,
#         ),
#         ax.axvline(
#             x=dynamic_baseline["mean_runtime"] / static_baseline["mean_runtime"],
#             linestyle="--",
#             linewidth=3,
#             label="grev (dynamic)",
#             color=colors["grev"]["dynamic"],
#             alpha=0.8,
#         ),
#     ]
#     ax.legend(handles=lines + handles, loc="best")
#
#     ax.set_xlim(0.5)
#     ax.margins(0)
#     ax.grid(True, alpha=0.5, linestyle="--", linewidth=1)
#     ax.set_axisbelow(True)
#
#
# fg.map_dataframe(plot_facet)
# fg.add_legend(title="Implementation")
# sns.move_legend(fg, loc="lower right")
# fg.savefig(os.path.join(FIGURES_DIR, "slowdown_over_static_grev"))
#
# df["slowdown_self"] = df.apply(
#     lambda row: row["mean_runtime"]
#     / df_baselines[
#         (df_baselines["compiler"] == row["compiler"])
#         & (df_baselines["typing"] == "static")
#         & (df_baselines["bench"] == row["bench"])
#     ].iloc[0]["mean_runtime"],
#     axis=1,
# )
# df.sort_values(by=["slowdown_self"], inplace=True, ignore_index=True)
# df["num_below_slowdown_self"] = (
#     df.groupby(["compiler", "typing", "bench"])["slowdown_self"].cumcount() + 1
# )
#
# df_filtered = df[df["typing"] != "gradual_succinct"].copy()
# df_filtered.sort_values(
#     by=["compiler", "typing", "bench"], inplace=True, ignore_index=True
# )
#
# fg = sns.FacetGrid(
#     df_filtered,
#     col="bench",
#     col_wrap=3,
#     height=6,
#     legend_out=False,
# )
#
# x_max = df["slowdown_self"].max()
#
#
# def plot_facet(data, **kwargs):
#     bench = data["bench"].iloc[0]
#     ax = plt.gca()
#
#     sns.lineplot(
#         data=data,
#         x="slowdown_self",
#         y="num_below_slowdown_self",
#         hue=data["compiler"] + " (" + data["typing"] + ")",
#         linewidth=3,
#         palette=colormap,
#         ax=ax,
#     )
#
#     ax.set_xscale("log")
#     ax.set_ylabel("Number of configurations", fontsize=12)
#     ax.set_xlabel("Slowdown over self (static)", fontsize=12)
#     ax.set_title(f"{bench}", fontweight="bold", pad=20)
#
#     for compiler in data["compiler"].unique():
#         for typing in data[data["compiler"] == compiler]["typing"].unique():
#             typing_data = data[
#                 (data["compiler"] == compiler) & (data["typing"] == typing)
#             ]
#             if not typing_data.empty:
#                 last_point = typing_data.iloc[-1]
#                 offset = (-30, 5)
#                 ax.annotate(
#                     f'{fmt_slowdown(last_point["slowdown_self"])}x',
#                     xy=(
#                         last_point["slowdown_self"],
#                         last_point["num_below_slowdown_self"],
#                     ),
#                     xytext=offset,
#                     textcoords="offset points",
#                     fontweight="semibold",
#                     fontsize=12,
#                 )
#
#     major_xticks = [1 << n for n in range(int(math.ceil(math.log2(x_max))) + 1)]
#     ax.set_xticks(major_xticks)
#     ax.set_xticklabels([f"{val}" for val in major_xticks])
#
#     handles, labels = ax.get_legend_handles_labels()
#     lines = [
#         ax.axvline(
#             x=1,
#             linestyle="-",
#             linewidth=3,
#             label="self (static)",
#             color="gray",
#             alpha=0.8,
#         ),
#     ]
#     ax.legend(handles=lines + handles, loc="best")
#
#     ax.set_xlim(0.5)
#     ax.margins(0)
#     ax.grid(True, alpha=0.5, linestyle="--", linewidth=1)
#     ax.set_axisbelow(True)
#
#
# fg.map_dataframe(plot_facet)
# fg.add_legend(title="Implementation")
# sns.move_legend(fg, loc="lower right")
# fg.savefig(os.path.join(FIGURES_DIR, "slowdown_over_self"))
#
# df_filtered = df[
#     (df["typing"] != "gradual_succinct") & (df["bench"] != "cps-even-odd")
# ].copy()
# df_filtered.sort_values(
#     by=["compiler", "typing", "bench"], inplace=True, ignore_index=True
# )
#
# df_filtered["slowdown_racket"] = df_filtered.apply(
#     lambda row: row["mean_runtime"]
#     / df_baselines[
#         (df_baselines["compiler"] == "racket")
#         & (df_baselines["typing"] == "dynamic")
#         & (df_baselines["bench"] == row["bench"])
#     ].iloc[0]["mean_runtime"],
#     axis=1,
# )
# df_filtered.sort_values(by=["slowdown_racket"], inplace=True, ignore_index=True)
# df_filtered["num_below_slowdown_racket"] = (
#     df_filtered.groupby(["compiler", "typing", "bench"])["slowdown_racket"].cumcount()
#     + 1
# )
#
# fg = sns.FacetGrid(
#     df_filtered,
#     col="bench",
#     col_wrap=3,
#     height=6,
#     legend_out=False,
#     sharex=False,
#     col_order=sorted(df_filtered["bench"].unique()),
# )
#
# x_max = df_filtered["slowdown_racket"].max()
#
#
# def plot_facet(data, **kwargs):
#     bench = data["bench"].iloc[0]
#     ax = plt.gca()
#
#     sns.lineplot(
#         data=data,
#         x="slowdown_racket",
#         y="num_below_slowdown_racket",
#         hue=data["compiler"] + " (" + data["typing"] + ")",
#         linewidth=3,
#         palette=colormap,
#         ax=ax,
#     )
#
#     ax.set_xscale("log")
#     ax.set_ylabel("Number of configurations", fontsize=12)
#     ax.set_xlabel("Slowdown over Racket", fontsize=12)
#     ax.set_title(f"{bench}", fontweight="bold", pad=20)
#
#     bench_x_max = df_filtered[df_filtered["bench"] == bench]["slowdown_racket"].max()
#     bench_x_min = df_filtered[df_filtered["bench"] == bench]["slowdown_racket"].min()
#
#     upper = int(math.ceil(math.log2(bench_x_max))) + 1
#     lower = int(math.floor(math.log2(bench_x_min))) - 1
#     major_xticks = [2**n for n in range(lower, upper + 1)] + [round(bench_x_max)]
#     ax.set_xticks(major_xticks)
#     ax.set_xticklabels([f"{val}" for val in major_xticks])
#     ax.set_xbound(min(bench_x_min / 2, 1 / 2), bench_x_max * 2)
#     handles, labels = ax.get_legend_handles_labels()
#     lines = [
#         ax.axvline(
#             x=1,
#             linestyle="-",
#             linewidth=3,
#             label="Racket",
#             color="gray",
#             alpha=0.8,
#         ),
#     ]
#     ax.legend(handles=lines + handles, loc="best")
#
#     ax.margins(0)
#     ax.grid(True, alpha=0.5, linestyle="--", linewidth=1)
#     ax.set_axisbelow(True)
#
#
# fg.map_dataframe(plot_facet)
# fg.add_legend(title="Implementation")
# sns.move_legend(fg, loc="lower right")
# fg.savefig(os.path.join(FIGURES_DIR, "slowdown_over_racket"))
#
# summary_df = (
#     df.groupby(["bench", "compiler", "typing"])
#     .agg(
#         best_slowdown=("slowdown", "min"),
#         worst_slowdown=("slowdown", "max"),
#     )
#     .reset_index()
# )
# summary_df["implementation"] = (
#     summary_df["compiler"] + " (" + summary_df["typing"] + ")"
# )
#
# fg = sns.catplot(
#     data=summary_df,
#     kind="bar",
#     y="worst_slowdown",
#     hue="implementation",
#     col="bench",
#     sharey=False,
#     col_wrap=3,
#     height=4,
#     aspect=1.2,
#     alpha=0.9,
#     linewidth=0.5,
#     legend=True,
#     legend_out=False,
#     palette=colormap,
#     margin_titles=True,
#     hue_order=[
#         "grev (gradual)",
#         "grev (gradual_proxy)",
#         "grift (gradual)",
#         "grift (gradual_mono)",
#     ],
# )
#
# for bench, ax in fg.axes_dict.items():
#     ax.set_ylabel("Worst Slowdown (over static grev)", fontweight="bold", fontsize=11)
#     ax.set_xlabel(None)
#
#     ax.set_title(f"{bench}", fontweight="bold", fontsize=12, pad=10)
#
#     for container in ax.containers:
#         ax.bar_label(container, fmt="%.1fx", padding=3, fontsize=12, fontweight="bold")
#
#     ax.grid(True, alpha=0.5, linestyle="-", linewidth=1, axis="y")
#
#     ax.spines["top"].set_visible(False)
#     ax.spines["right"].set_visible(False)
#
# sns.move_legend(fg, loc="lower right", title="Implementation")
#
# fg.savefig(os.path.join(FIGURES_DIR, "slowdown_over_static_grev_worst_case"))
#
# fg = sns.catplot(
#     data=summary_df,
#     kind="bar",
#     y="best_slowdown",
#     hue="implementation",
#     col="bench",
#     sharey=False,
#     col_wrap=3,
#     height=4,
#     aspect=1.2,
#     alpha=0.9,
#     linewidth=0.5,
#     legend=True,
#     legend_out=False,
#     palette=colormap,
#     margin_titles=True,
#     hue_order=[
#         "grev (gradual)",
#         "grev (gradual_proxy)",
#         "grift (gradual)",
#         "grift (gradual_mono)",
#     ],
# )
#
# for bench, ax in fg.axes_dict.items():
#     ax.set_ylabel("Best Slowdown (over static grev)", fontweight="bold", fontsize=11)
#     ax.set_xlabel(None)
#
#     ax.set_title(f"{bench}", fontweight="bold", fontsize=12, pad=10)
#
#     for container in ax.containers:
#         ax.bar_label(container, fmt="%.1fx", padding=3, fontsize=12, fontweight="bold")
#
#     ax.grid(True, alpha=0.5, linestyle="-", linewidth=1, axis="y")
#
#     ax.spines["top"].set_visible(False)
#     ax.spines["right"].set_visible(False)
#
# sns.move_legend(fg, loc="lower right", title="Implementation")
# fg.savefig(os.path.join(FIGURES_DIR, "slowdown_over_static_grev_best_case"))
