source_current/synthetic_task/plot_results.py

import os, re, glob
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
import matplotlib.ticker as mticker
import warnings
warnings.filterwarnings("ignore")


sns.set_theme(style="whitegrid", context="talk")
palette = sns.color_palette()

batch_size = 8
BASE_DIR = f"../synthetic_results_general_{batch_size}"
# BASE_DIR = f"../synthetic_results_1_compare_SCS_OSQP_dim200_debug"
# METHODS = [
#     "cvxpylayer",
#     "lpgd",
#     "bpqp",
#     "ffocp_eq",
#     "dqp",
#     "qpth",
#     "ffoqp_eq",   
# ]
METHODS = [
    "cvxpylayer",
    "qpth",
    "lpgd",
    "bpqp",
    "dqp",
    "ffocp_eq",
    "ffoqp_eq_schur"  
]
METHODS_LEGEND = {
    "cvxpylayer": "CvxpyLayer",
    "lpgd": "LPGD",
    "bpqp": "BPQP",
    "ffocp_eq": "FFOCP",
    "dqp": "dQP",
    "qpth": "qpth",
    "ffoqp_eq_schur": "FFOQP",
}

METHODS_STEPS = [method+"_steps" for method in METHODS]

method_order = [METHODS_LEGEND[m] for m in METHODS]

markers = ["o", "s", "D", "^", "v", "x", "P", "s"]
markers_dict = {method: markers[i] for i, method in enumerate(method_order)}

LINEWIDTH = 1.5

def is_subsequence(a, b):
    it = iter(b)
    return all(x in it for x in a)

def load_results_CP(base_dir=BASE_DIR, methods=METHODS, methods_legend=None):
    if methods_legend is None:
        methods_legend = METHODS_LEGEND
    dfs = []
    for m in methods:
        pattern = os.path.join(base_dir, m, "*.csv")
        for fp in sorted(glob.glob(pattern)):
            m = m.removesuffix("_steps")
            
            df = pd.read_csv(fp)
            df["method"] = methods_legend[m]

            fname = os.path.basename(fp)
            def grab(pat, cast=float):
                mo = re.search(pat, fname)
                return cast(mo.group(1)) if mo else np.nan

            df["seed"] = grab(r"_seed(\d+)", int)
            df["n"] = grab(r"n(\d+)", int)
            df["lr"]   = grab(r"lr([0-9eE\.\-]+)", float)
            df["ydim"] = grab(r"ydim(\d+)", int)
            df["backwardTol"] = grab(r"backwardTol([0-9eE\+\-\.]+?)(?:_|\.csv)", float)
            dfs.append(df)
        
        print("method: ", m)
        # print(df)
        
    if not dfs:
        raise FileNotFoundError(f"No CSVs found under {base_dir}.")
    return pd.concat(dfs, ignore_index=True, sort=False)

def load_results_QP(base_dir=BASE_DIR, methods=METHODS, methods_legend=None):
    if methods_legend is None:
        methods_legend = METHODS_LEGEND
    dfs = []
    for m in methods:
        pattern = os.path.join(base_dir, m, "*.csv")
        for fp in sorted(glob.glob(pattern)):
            m = m.removesuffix("_steps")
            if m == "ffocp_eq" and "backwardTol" not in os.path.basename(fp):
                continue
            
            df = pd.read_csv(fp)
            df["method"] = methods_legend[m]

            fname = os.path.basename(fp)
            def grab(pat, cast=float):
                mo = re.search(pat, fname)
                return cast(mo.group(1)) if mo else np.nan

            df["seed"] = grab(r"_seed(\d+)", int)
            df["n"] = grab(r"n(\d+)", int)
            df["lr"]   = grab(r"lr([0-9eE\.\-]+)", float)
            df["ydim"] = grab(r"ydim(\d+)", int)
            df["backwardTol"] = grab(r"backwardTol([0-9eE\+\-\.]+?)(?:_|\.csv)", float)
            dfs.append(df)
        
        print("method: ", m)
        # print(df)
        
    if not dfs:
        raise FileNotFoundError(f"No CSVs found under {base_dir}.")
    return pd.concat(dfs, ignore_index=True, sort=False)

def plot_time_vs_method(df, time_names=['forward_time', 'backward_time'], plot_path=BASE_DIR, plot_name_tag=""):
    df_avg_method = df.groupby('method')[time_names].mean().reset_index()

    # Convert wide → long format so Seaborn can handle grouped bars
    df_long = df_avg_method.melt(id_vars='method', 
                                value_vars=time_names,
                                var_name='Metrics', 
                                value_name='Time')

    plt.figure(figsize=(8,5))
    sns.barplot(data=df_long, x='method', y='Time', hue='Metrics')
    plt.ylabel("Time")
    plt.title("Forward and Backward Time")
    plt.savefig(f"{plot_path}/{plot_name_tag}_time_vs_method.png", dpi=300, bbox_inches='tight')
    plt.close()

def plot_time_vs_epoch(df, time_names=['forward_time', 'backward_time'], iteration_name='epoch', plot_path=BASE_DIR, plot_name_tag=""):
    df_avg_epoch = df.groupby(['method', iteration_name])[time_names].mean().reset_index()

    # --- Forward Time Figure ---
    plt.figure(figsize=(8,5))
    sns.lineplot(data=df_avg_epoch, x=iteration_name, y=time_names[0], hue='method', marker=None, dashes=False, linewidth=LINEWIDTH)
    plt.ylabel("Forward Time")
    plt.title(f"Forward Time vs {iteration_name}")
    
    plt.savefig(f"{plot_path}/{plot_name_tag}_forward_time_vs_{iteration_name}.png", dpi=300, bbox_inches='tight')
    plt.close()

    # --- Backward Time Figure ---
    plt.figure(figsize=(8,5))
    sns.lineplot(data=df_avg_epoch, x=iteration_name, y=time_names[1], hue='method', marker=None, dashes=False, linewidth=LINEWIDTH)
    plt.ylabel("Backward Time")
    plt.title(f"Backward Time vs {iteration_name}")
    
    plt.savefig(f"{plot_path}/{plot_name_tag}_backward_time_vs_{iteration_name}.png", dpi=300, bbox_inches='tight')
    plt.close()

# def plot_total_time_vs_method(df, time_names=['forward_time', 'backward_time'], plot_path=BASE_DIR, plot_name_tag=""):
#     # Group by method, average over epochs and seeds
#     df_avg_method = df.groupby('method')[time_names].mean().reset_index()

#     # --- Stacked Bar Chart ---
#     methods = df_avg_method['method']
#     forward = df_avg_method[time_names[0]]
#     backward = df_avg_method[time_names[1]]

#     plt.figure(figsize=(8,5))
#     plt.bar(methods, forward, label=time_names[0], color=palette[0])
#     plt.bar(methods, backward, bottom=forward, label=time_names[1], color=palette[1])
#     plt.ylabel("Time")
#     plt.title("Total Time vs Method")
#     plt.legend()
#     plt.savefig(f"{plot_path}/{plot_name_tag}_total_time_vs_method.pdf", dpi=300, bbox_inches='tight')
#     # plt.savefig(f"{plot_path}/{plot_name_tag}_total_time_vs_method.png", dpi=300, bbox_inches='tight')
#     plt.close()

def plot_total_time_vs_method(
    df,
    task_title_tag=None,
    time_names=["forward_time", "backward_time"],
    plot_path=BASE_DIR,
    plot_name_tag="",
):
    available_methods = set(df["method"].unique())
    # print("hhhhhhh\n", df[df['method']=='FFOQP'])
    # assert("FFOQP" in available_methods)
    filtered_method_order = [m for m in method_order if m in available_methods]
    print("filtered_method_order:", filtered_method_order)
    print("available method: ", available_methods)

    df_avg_method = df.groupby("method")[time_names].mean().reindex(filtered_method_order)

    methods = df_avg_method.index.tolist()
    forward = df_avg_method[time_names[0]].to_numpy()
    backward = df_avg_method[time_names[1]].to_numpy()

    total_times = forward + backward
    max_time = np.nanmax(total_times)
    min_time = 0
    y_max = max_time * 1.05
    
    #### get groups of method to plot them close together
    groups = [
        ["CvxpyLayer", "qpth"], # KKT based
        ["LPGD","BPQP", "dQP"], # optimization based
        ["FFOCP", "FFOQP"], # our methods
    ]
    
    # assert groups respect the filtered method order
    flat_groups = [m for g in groups for m in g]
   
    groups = [[m for m in g if m in filtered_method_order] for g in groups]
    groups = [g for g in groups if len(g) > 0]
    
    flat_groups = [m for g in groups for m in g]
    print("groups:", flat_groups)
    print("methods:", methods)
    assert is_subsequence(flat_groups, methods), "Method groups do not respect the method order."

    # dashed_methods = {"CvxpyLayer", "LPGD", "BPQP", "FFOCP"}  # CP
    
    ######### set bar positions with custom gaps
    # inner_gap = 0.25 #1.5     # spacing between bars inside a group
    # group_gap = len(methods)*0.5/7   ##so that when num methods is 7, group gap is 3.0   # extra spacing between groups
    width = 0.08 #0.75 # bar width
    inner_gap = 0.08 #1.5     # spacing between bars inside a group
    group_gap = len(methods)*0.125/7   ##so that when num methods is 7, group gap is 3.0   # extra spacing between groups


    x = []
    # labels = []
    pos = 0.0

    for g in groups:
        for m in g:
            x.append(pos)
            # labels.append(m)
            pos += inner_gap
        pos += group_gap  # extra space after each group
        
    ########## set bar label positions with custom gaps
    # inner_gap = 1.5     # spacing between bars inside a group
    # group_gap = len(methods)*2/7   ##so that when num methods is 7, group gap is 3.0   # extra spacing between groups
    # label_x = []
    # pos = 0.0
    # for g in groups:
    #     for m in g:
    #         label_x.append(pos)
    #         pos += inner_gap
    #     pos += group_gap  # extra space after each group
    label_x = x

    
    # x = np.arange(len(methods))
    # width = 0.25 #0.75

    fig, ax = plt.subplots(figsize=(14, 7))
    sf = mticker.ScalarFormatter(useOffset=False)
    sf.set_scientific(False)
    ax.yaxis.set_major_formatter(sf)
    ax.yaxis.get_offset_text().set_visible(False)

    for i, m in enumerate(methods):
        # is_dashed = (m in dashed_methods)
        # ls = (0, (4, 2)) if is_dashed else "solid"
        # lw = 2.0 if is_dashed else 1.5

        # ax.bar(x[i], forward[i], width=width,
        #        color=palette[0], edgecolor="black", linewidth=lw, linestyle=ls)
        # ax.bar(x[i], backward[i], bottom=forward[i], width=width,
        #        color=palette[1], edgecolor="black", linewidth=lw, linestyle=ls)
        
        ax.bar(x[i], forward[i], width=width,
               color=palette[0], edgecolor="black")
        ax.bar(x[i], backward[i], bottom=forward[i], width=width,
               color=palette[1], edgecolor="black")


    ax.set_xticks(label_x)
    methods = ["Cvxpy\nLayer" if m == "CvxpyLayer" else m for m in methods]
    ax.set_xticklabels(methods)
    ax.set_ylabel("Time")
    if task_title_tag is not None:
        ax.set_title(f"{task_title_tag}: Total Time vs Methods")
    else:
        ax.set_title("Total Time vs Methods")
    ax.set_ylim(min_time, y_max)
    min_x = min([min(x), min(label_x)])
    max_x = max([max(x), max(label_x)])
    ax.set_xlim(min_x - width, max_x + width)

    handles = [
        Line2D([], [], linestyle="none", label="Comp. Time"),
        Patch(facecolor=palette[0], edgecolor="black", label="Forward"),
        Patch(facecolor=palette[1], edgecolor="black", label="Backward"),
        # Line2D([], [], linestyle="none", label=""),
        # Line2D([], [], linestyle="none", label="Solvers"),
        # Line2D([0], [0], color="black", linewidth=2, linestyle=(0, (4, 2)), label="CP methods"),
        # Line2D([0], [0], color="black", linewidth=2, linestyle="solid", label="QP methods"),
    ]

    leg = ax.legend(
        handles=handles,
        loc="upper right",
        frameon=True,
        handlelength=2.2,
        handletextpad=0.8,
        borderpad=0.7,
        labelspacing=0.1,
    )
    leg._legend_box.align = "left"

    for h, t in zip(leg.legend_handles, leg.get_texts()):
        txt = t.get_text()
        if txt in ("Comp. Time", "Solvers"):
            t.set_weight("bold")
            if hasattr(h, "set_visible"):
                h.set_visible(False)
        if txt == "":
            if hasattr(h, "set_visible"):
                h.set_visible(False)
            t.set_color((0, 0, 0, 0))

    fig.savefig(f"{plot_path}/{plot_name_tag}_total_time_vs_method.pdf", dpi=300, bbox_inches="tight")
    plt.close(fig)
    
def plot_loss_vs_epoch(df, loss_metric_name, task_title_tag=None, iteration_name='epoch', plot_path=BASE_DIR, plot_name_tag="", loss_range=None, stride=50):
    df_avg_epoch = df.groupby(['method', iteration_name])[[loss_metric_name]].mean().reset_index()
    
    # print(df_avg_epoch)
    # epoch_0_loss = -0.00950829166918993
    # df_avg_epoch.loc[df_avg_epoch[iteration_name] == 0, loss_metric_name] = epoch_0_loss
    # df_avg_epoch = df_avg_epoch[df_avg_epoch[iteration_name] % stride == 0]

    # --- Forward Time Figure ---
    plt.figure(figsize=(8,5))
    ax = sns.lineplot(data=df_avg_epoch, x=iteration_name, y=loss_metric_name, hue='method', dashes=False, linewidth=LINEWIDTH, hue_order=method_order)
    plt.ylabel("loss")
    if task_title_tag is not None:
        plt.title(f"{task_title_tag}: Loss vs iteration")
    else:
        plt.title("Loss vs iteration")
    
    # plt.legend(
    #     title=None,
    #     loc="lower center",
    #     bbox_to_anchor=(0.5, -0.25),
    #     ncol=(df["method"].nunique()),           
    #     frameon=False
    # )
    
    if loss_range is not None:
        ax.set_ylim(loss_range)
    
    plt.savefig(f"{plot_path}/{plot_name_tag}_{loss_metric_name}_vs_{iteration_name}.pdf", dpi=300, bbox_inches='tight')
    plt.close()

def plot_total_time_vs_method_by_ydim(df, plot_path, time_names=('forward_time','backward_time'), tag="syn"):
    os.makedirs(plot_path, exist_ok=True)
    ydims = sorted(df["ydim"].dropna().unique().astype(int))
    for y in ydims:
        d = df[df["ydim"] == y]
        plot_total_time_vs_method(
            d,
            time_names=list(time_names),
            plot_path=plot_path,
            plot_name_tag=f"{tag}_ydim{y}"
        )

def plot_total_time_vs_ydim_for_methods(
    df,
    methods_subset,
    markers_dict=markers_dict,
    time_names=('forward_time', 'backward_time'),
    plot_path=BASE_DIR,
    plot_name_tag="syn_methods_vs_ydim"
):
    os.makedirs(plot_path, exist_ok=True)
    df_sub = df[df["method"].isin(methods_subset)].copy()
    df_sub = df_sub.dropna(subset=["ydim"])
    df_sub["ydim"] = df_sub["ydim"].astype(int)

    df_group = (
        df_sub
        .groupby(["method", "ydim"])[list(time_names)]
        .mean()
        .reset_index()
    )

    df_group["total_time"] = df_group[time_names[0]] + df_group[time_names[1]]

    plt.figure(figsize=(8, 5))

    for idx, method in enumerate(methods_subset):
        d = df_group[df_group["method"] == method].copy()
        d = d.sort_values("ydim")
        marker = markers_dict.get(method, markers[idx % len(markers)])
        plt.plot(
            d["ydim"],
            d["total_time"],
            label=method,
            marker=marker,
            linewidth=LINEWIDTH,
        )

    plt.xlabel("y_dim")
    plt.ylabel("Total Time")
    plt.title("Total Time vs y_dim (selected methods)")
    
    # Reorder legend items to: CvxpyLayer, qpth, LPGD, BPQP, dQP, FFOCP, FFOQP
    legend_order = ["CvxpyLayer", "qpth", "LPGD", "BPQP", "dQP", "FFOCP", "FFOQP"]
    handles, labels = plt.gca().get_legend_handles_labels()
    label_to_handle = dict(zip(labels, handles))
    ordered_handles = [label_to_handle[l] for l in legend_order if l in label_to_handle]
    ordered_labels = [l for l in legend_order if l in label_to_handle]
    plt.legend(ordered_handles, ordered_labels, ncol=len(ordered_labels))
    
    plt.tight_layout()
    plt.savefig(
        f"{plot_path}/{plot_name_tag}_total_time_vs_ydim.pdf",
        dpi=300,
        bbox_inches="tight",
    )
    plt.close()


def plot_total_loss_vs_method_by_ydim(df, plot_path):
    os.makedirs(plot_path, exist_ok=True)
    ydims = sorted(df["ydim"].dropna().unique().astype(int))
    for y in ydims:
        d = df[df["ydim"] == y]
        plot_loss_vs_epoch(
            d,
            loss_metric_name="train_df_loss",
            iteration_name="iter",
            plot_path=plot_path,
            plot_name_tag=f"syn_ydim{y}"
        )


def plot_total_time_vs_method_by_backwardTol(df, time_names=['forward_time', 'backward_time'],
                                             plot_path=BASE_DIR, plot_name_tag=""):
    """
    Plots stacked bar chart of total time vs backward tolerance for each method.
    Each method is plotted separately in a for-loop to ensure only one method at a time.
    """
    os.makedirs(plot_path, exist_ok=True)

    df = df.dropna(subset=['method', 'backwardTol']).copy()

    # Get unique methods
    methods = df['method'].unique()
    assert(len(methods)==1)

    for method in methods:
        d = df[df['method'] == method].copy()

        # Remove duplicate backwardTol per method (take mean)
        d = d.groupby('backwardTol')[time_names].mean().reset_index()

        plt.figure(figsize=(8,5))
        plt.bar(d['backwardTol'].astype(str), d[time_names[0]], label=time_names[0], color=palette[0])
        plt.bar(d['backwardTol'].astype(str), d[time_names[1]], bottom=d[time_names[0]],
                label=time_names[1], color=palette[1])
        
        plt.xlabel("backward tolerance")
        plt.ylabel("Time")
        plt.title(f"Total Time vs Backward Tolerance ({method})")
        plt.legend()
        plt.tight_layout()
        plt.savefig(f"{plot_path}/{plot_name_tag}_total_time_vs_backwardTol_{method}.pdf", 
                    dpi=300, bbox_inches='tight')
        plt.close()


def plot_loss_vs_epoch_method_tol(df, loss_metric='train_df_loss', iteration='epoch',
                                  plot_path=BASE_DIR, plot_name_tag=""):
    """
    Simple: plot loss vs epoch, each line = one (method, backwardTol) tuple.
    """
    import matplotlib.pyplot as plt
    import os

    os.makedirs(plot_path, exist_ok=True)
    df = df.dropna(subset=['method', 'backwardTol', iteration, loss_metric])

    plt.figure(figsize=(8,5))

    for (method, tol), d in df.groupby(['method', 'backwardTol']):
        d = d.sort_values(iteration)
        plt.plot(d[iteration], d[loss_metric], label=f"{method}-tol{tol:.0e}")

    plt.xlabel(iteration)
    plt.ylabel("loss")
    plt.title(f"Loss vs {iteration}")
    plt.legend()
    plt.tight_layout()
    plt.savefig(f"{plot_path}/{plot_name_tag}_{loss_metric}_vs_{iteration}.pdf", dpi=300, bbox_inches='tight')
    plt.close()

def plot_time_scaling_vs_ydim(
    df,
    time_names=("forward_time", "backward_time"),
    plot_path=BASE_DIR,
    plot_name_tag="syn",
    methods_order=None,
    markers_dict=markers_dict,
    dashed_methods=("CvxpyLayer", "LPGD", "BPQP", "FFOCP"),  # CP
    agg="mean",
    logy=False,
    y_min=None,
    y_max=None,
    filter_backwardTol=None,
):
    os.makedirs(plot_path, exist_ok=True)

    d = df.copy().dropna(subset=["ydim", "method"])
    d["ydim"] = d["ydim"].astype(int)

    if "backwardTol" in d.columns and filter_backwardTol is not None:
        d = d[np.isclose(d["backwardTol"].astype(float), float(filter_backwardTol))]

    gfunc = "median" if agg == "median" else "mean"
    g = getattr(d.groupby(["method", "ydim"], as_index=False)[list(time_names)], gfunc)()
    g["total_time"] = g[time_names[0]] + g[time_names[1]]

    if methods_order is None:
        methods_order = (
            [m for m in d["method"].cat.categories if m in g["method"].unique()]
            if isinstance(d["method"].dtype, pd.CategoricalDtype)
            else sorted(g["method"].unique())
        )
    else:
        methods_order = [m for m in methods_order if m in g["method"].unique()]

    plt.rcParams.update({
        "font.size": 18, "axes.titlesize": 20, "axes.labelsize": 18,
        "legend.fontsize": 16, "xtick.labelsize": 16, "ytick.labelsize": 16,
    })

    fig, axes = plt.subplots(1, 3, figsize=(15.5, 4.2), sharex=True)
    panels = [
        ("total_time", "Total time"),
        (time_names[0], "Forward time"),
        (time_names[1], "Backward time"),
    ]

    for ax, (col, title) in zip(axes, panels):
        for i, method in enumerate(methods_order):
            gg = g[g["method"] == method].sort_values("ydim")
            if gg.empty:
                continue

            ax.plot(
                gg["ydim"], gg[col],
                label=method,
                marker=markers_dict.get(method, markers[i % len(markers)]),
                markersize=5.5,
                linewidth=2.0,
                linestyle="solid",
            )

        ax.set_title(title)
        ax.set_xlabel("y dim")
        ax.grid(True, which="major", alpha=0.25)
        ax.grid(True, which="minor", alpha=0.12)
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)

        if logy:
            ax.set_yscale("log")
        if y_min is not None or y_max is not None:
            ax.set_ylim(bottom=y_min, top=y_max)

        ax.yaxis.set_major_formatter(mticker.ScalarFormatter(useOffset=False))
        ax.yaxis.get_offset_text().set_visible(False)

    axes[0].set_ylabel("time (s)")

    # Reorder legend items to: CvxpyLayer, qpth, LPGD, BPQP, dQP, FFOCP, FFOQP (single row)
    legend_order = ["CvxpyLayer", "qpth", "LPGD", "BPQP", "dQP", "FFOCP", "FFOQP"]
    handles, labels = axes[0].get_legend_handles_labels()
    label_to_handle = dict(zip(labels, handles))
    ordered_handles = [label_to_handle[l] for l in legend_order if l in label_to_handle]
    ordered_labels = [l for l in legend_order if l in label_to_handle]
    fig.legend(
        ordered_handles, ordered_labels,
        loc="lower center",
        bbox_to_anchor=(0.5, 0.02),
        fontsize=18,
        ncol=len(ordered_labels),
        frameon=False,
        handlelength=2.4,
        handletextpad=0.6,
        columnspacing=1.0,
    )

    # leave space at bottom for legend
    fig.tight_layout(rect=[0, 0.10, 1, 1])

    out = os.path.join(plot_path, f"{plot_name_tag}_time_scaling_vs_ydim.pdf")
    fig.savefig(out, dpi=300, bbox_inches="tight")
    plt.close(fig)

def discover_batch_dirs(root_dir="..", prefix="synthetic_results_"):
    cand = glob.glob(os.path.join(root_dir, f"{prefix}*"))
    out = []
    for p in cand:
        if not os.path.isdir(p):
            continue
        name = os.path.basename(p.rstrip("/"))
        mo = re.match(rf"^{re.escape(prefix)}(\d+)$", name)
        if mo:
            out.append((int(mo.group(1)), p))
    out.sort(key=lambda x: x[0])
    if not out:
        raise FileNotFoundError(f"No dirs like {prefix}{{int}} under {root_dir}")
    return out


def load_results_across_batches(
    root_dir="..",
    batch_sizes=None,
    loader_fn=None,
    methods=None,
    methods_legend=None,
    prefix="synthetic_results_",
):
    if loader_fn is None:
        raise ValueError("Please pass loader_fn=load_results_CP or load_results_QP")
    if methods is None:
        methods = METHODS
    if methods_legend is None:
        methods_legend = METHODS_LEGEND

    if batch_sizes is None:
        batch_dirs = discover_batch_dirs(root_dir=root_dir, prefix=prefix)
    else:
        batch_dirs = [(int(bs), os.path.join(root_dir, f"{prefix}{int(bs)}")) for bs in batch_sizes]

    dfs = []
    for bs, dpath in batch_dirs:
        if not os.path.isdir(dpath):
            print(f"[WARN] missing dir: {dpath}, skip")
            continue

        df = loader_fn(base_dir=dpath, methods=methods, methods_legend=methods_legend)
        df = df.rename(columns=lambda c: c.strip() if isinstance(c, str) else c)
        df["batch_size"] = int(bs)
        dfs.append(df)

    if not dfs:
        raise FileNotFoundError("No results loaded across batches. Check dirs / patterns.")
    out = pd.concat(dfs, ignore_index=True, sort=False)
    return out


# def plot_time_scaling_vs_batch(
#     df,
#     time_names=("forward_time", "backward_time"),
#     plot_path=".",
#     plot_name_tag="syn",
#     methods_order=None,
#     dashed_methods=("CvxpyLayer", "LPGD", "BPQP", "FFOCP"),  # CP
#     agg="mean",
#     logx=True,
#     logy=False,
#     y_min=None,
#     y_max=None,
#     filter_ydim=None,
#     filter_backwardTol=None,
# ):
#     os.makedirs(plot_path, exist_ok=True)
#     markers = ["o", "s", "D", "^", "v", "x", "P", "s"]
#     markers_dict = {method: markers[i] for i, method in enumerate(methods_order)}

#     d = df.copy()

#     if "batch_size" not in d.columns:
#         raise KeyError("df must contain column 'batch_size' (use load_results_across_batches).")
#     if "method" not in d.columns:
#         raise KeyError("df must contain column 'method'.")

#     if filter_ydim is not None:
#         if "ydim" not in d.columns:
#             raise KeyError("filter_ydim is set but df has no 'ydim' column.")
#         d = d.dropna(subset=["ydim"])
#         d = d[d["ydim"].astype(int) == int(filter_ydim)]

#     if "backwardTol" in d.columns and filter_backwardTol is not None:
#         d = d.dropna(subset=["backwardTol"])
#         d = d[np.isclose(d["backwardTol"].astype(float), float(filter_backwardTol))]

#     d = d.dropna(subset=["batch_size", "method"])
#     d["batch_size"] = d["batch_size"].astype(int)

#     gfunc = "median" if agg == "median" else "mean"
#     g = getattr(d.groupby(["method", "batch_size"], as_index=False)[list(time_names)], gfunc)()
#     g["total_time"] = g[time_names[0]] + g[time_names[1]]

#     if methods_order is None:
#         if isinstance(d["method"].dtype, pd.CategoricalDtype):
#             methods_order = [m for m in d["method"].cat.categories if m in g["method"].unique()]
#         else:
#             methods_order = sorted(g["method"].unique())
#     else:
#         methods_order = [m for m in methods_order if m in g["method"].unique()]

#     plt.rcParams.update({
#         "font.size": 13, "axes.titlesize": 14, "axes.labelsize": 13,
#         "legend.fontsize": 11, "xtick.labelsize": 12, "ytick.labelsize": 12,
#     })

#     fig, axes = plt.subplots(1, 3, figsize=(15.5, 4.2), sharex=True)
#     panels = [
#         ("total_time", "Total time"),
#         (time_names[0], "Forward time"),
#         (time_names[1], "Backward time"),
#     ]
#     dashed_methods = set(dashed_methods)

#     for ax, (col, title) in zip(axes, panels):
#         for i, method in enumerate(methods_order):
#             gg = g[g["method"] == method].sort_values("batch_size")
#             if gg.empty:
#                 continue

#             is_dashed = method in dashed_methods
#             ax.plot(
#                 gg["batch_size"], gg[col],
#                 label=method,
#                 marker=markers_dict.get(method, markers[i % len(markers)]),
#                 markersize=5.5,
#                 linewidth=2.2 if is_dashed else 2.0,
#                 linestyle=(0, (4, 2)) if is_dashed else "solid",
#             )

#         ax.set_title(title)
#         ax.set_xlabel("batch size")
#         ax.grid(True, which="major", alpha=0.25)
#         ax.grid(True, which="minor", alpha=0.12)
#         ax.spines["top"].set_visible(False)
#         ax.spines["right"].set_visible(False)

#         if logx:
#             ax.set_xscale("log", base=2)
#             ax.xaxis.set_major_formatter(mticker.ScalarFormatter())
#             ax.set_xticks(sorted(g["batch_size"].unique()))
#             ax.get_xaxis().set_major_formatter(mticker.ScalarFormatter())

#         if logy:
#             ax.set_yscale("log")
#         if y_min is not None or y_max is not None:
#             ax.set_ylim(bottom=y_min, top=y_max)

#         ax.yaxis.set_major_formatter(mticker.ScalarFormatter(useOffset=False))
#         ax.yaxis.get_offset_text().set_visible(False)

#     axes[0].set_ylabel("time (s)")

#     handles, labels = axes[0].get_legend_handles_labels()
#     ncol = min(len(labels), 4)
#     fig.legend(
#         handles, labels,
#         loc="lower center",
#         bbox_to_anchor=(0.5, -0.02),
#         fontsize=14,
#         ncol=ncol,
#         frameon=False,
#         handlelength=2.4,
#         handletextpad=0.6,
#         columnspacing=1.2,
#     )

#     fig.tight_layout(rect=[0, 0.12, 1, 1])

#     suffix = ""
#     if filter_ydim is not None:
#         suffix += f"_ydim{int(filter_ydim)}"
#     if filter_backwardTol is not None:
#         suffix += f"_bwdTol{filter_backwardTol:g}"

#     out = os.path.join(plot_path, f"{plot_name_tag}_time_scaling_vs_batch{suffix}.pdf")
#     fig.savefig(out, dpi=300, bbox_inches="tight")
#     plt.close(fig)

#     print(f"[saved] {out}")

def plot_time_scaling_vs_batch(
    df,
    plot_path=".",
    plot_name_tag="syn",
    methods_order=None,
    markers_dict=None,
    dashed_methods=("CvxpyLayer", "LPGD", "BPQP", "FFOCP"),
    agg="mean",                 # "mean" or "median"
    use_log2_x=True,
    logy=False,
    y_min=None,
    y_max=None,
    filter_backwardTol=None,    # optional if df has backwardTol
):
    os.makedirs(plot_path, exist_ok=True)
    d = df.copy().dropna(subset=["batch_size", "method"])
    d["batch_size"] = d["batch_size"].astype(int)

    if "backwardTol" in d.columns and filter_backwardTol is not None:
        d = d[np.isclose(d["backwardTol"].astype(float), float(filter_backwardTol))]

    for c in ["forward_time", "backward_time"]:
        if c not in d.columns:
            raise ValueError(f"missing column: {c}")
        d[c] = d[c].astype(float)

    # derive step/overhead
    if "overhead_time" in d.columns:
        d["overhead_time"] = d["overhead_time"].astype(float)
        d["step_time"] = d["forward_time"] + d["backward_time"] + d["overhead_time"]
    else:
        if "step_time" in d.columns:
            d["step_time"] = d["step_time"].astype(float)
        elif "total_time" in d.columns:
            d["step_time"] = d["total_time"].astype(float)
        else:
            d["step_time"] = d["forward_time"] + d["backward_time"]
        d["overhead_time"] = (d["step_time"] - d["forward_time"] - d["backward_time"]).clip(lower=0.0)

    gfunc = "median" if agg == "median" else "mean"
    g = getattr(
        d.groupby(["method", "batch_size"], as_index=False)[
            ["forward_time", "backward_time", "overhead_time", "step_time"]
        ],
        gfunc,
    )()

    eps = 1e-12
    g["throughput"] = g["batch_size"] / g["step_time"].clip(lower=eps)
    g["per_sample"] = g["step_time"] / g["batch_size"].clip(lower=1)       # sec/sample

    # method order (match your style)
    if methods_order is None:
        methods_order = (
            [m for m in d["method"].cat.categories if m in g["method"].unique()]
            if isinstance(d["method"].dtype, pd.CategoricalDtype)
            else sorted(g["method"].unique())
        )
    else:
        methods_order = [m for m in methods_order if m in g["method"].unique()]

    # style (copy yours)
    plt.rcParams.update({
        "font.size": 13, "axes.titlesize": 14, "axes.labelsize": 13,
        "legend.fontsize": 11, "xtick.labelsize": 12, "ytick.labelsize": 12,
    })

    dashed_methods = set(dashed_methods)
    markers = ["o", "s", "D", "^", "v", "P", "X", "*", "<", ">"]

    xvals = np.array(sorted(g["batch_size"].unique().tolist()), dtype=int)
    xplot = np.log2(xvals) if use_log2_x else xvals
    xticks = xplot
    xtick_labels = [str(v) for v in xvals]
    xlabel = "batch size" if use_log2_x else "batch size"

    fig, axes = plt.subplots(1, 3, figsize=(15.5, 4.2), sharex=True)
    panels = [
        ("throughput", "Throughput (samples/sec)"),
        ("step_time", "Time per step (sec)"),
        ("per_sample", "Per-sample cost (sec/sample)"),
    ]

    for ax, (col, title) in zip(axes, panels):
        for i, method in enumerate(methods_order):
            gg = g[g["method"] == method].sort_values("batch_size")
            if gg.empty:
                continue

            xx = np.log2(gg["batch_size"].to_numpy()) if use_log2_x else gg["batch_size"].to_numpy()
            is_dashed = method in dashed_methods
            mk = (markers_dict.get(method) if markers_dict else None) or markers[i % len(markers)]

            ax.plot(
                xx, gg[col].to_numpy(),
                label=method,
                marker=mk,
                markersize=5.5,
                linewidth=2.2 if is_dashed else 2.0,
                linestyle=(0, (4, 2)) if is_dashed else "solid",
            )

            # decomposition shading only on panel 2 (step_time)
            if col == "step_time":
                fwd = gg["forward_time"].to_numpy()
                bwd = gg["backward_time"].to_numpy()
                ovh = gg["overhead_time"].to_numpy()
                ax.fill_between(xx, 0, fwd, alpha=0.06)
                ax.fill_between(xx, fwd, fwd + bwd, alpha=0.06)
                ax.fill_between(xx, fwd + bwd, fwd + bwd + ovh, alpha=0.06)

        ax.set_title(title)
        ax.set_xlabel(xlabel)
        ax.set_xticks(xticks)
        ax.set_xticklabels(xtick_labels)

        ax.grid(True, which="major", alpha=0.25)
        ax.grid(True, which="minor", alpha=0.12)
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)

        if logy:
            ax.set_yscale("log")
        if y_min is not None or y_max is not None:
            ax.set_ylim(bottom=y_min, top=y_max)

        ax.yaxis.set_major_formatter(mticker.ScalarFormatter(useOffset=False))
        ax.yaxis.get_offset_text().set_visible(False)

    # shared legend (same as your function)
    handles, labels = axes[0].get_legend_handles_labels()
    ncol = min(len(labels), 4)
    fig.legend(
        handles, labels,
        loc="lower center",
        bbox_to_anchor=(0.5, -0.02),
        fontsize=14,
        ncol=ncol,
        frameon=False,
        handlelength=2.4,
        handletextpad=0.6,
        columnspacing=1.2,
    )
    fig.tight_layout(rect=[0, 0.12, 1, 1])

    out = os.path.join(plot_path, f"{plot_name_tag}_batch_scaling_3panels.pdf")
    fig.savefig(out, dpi=300, bbox_inches="tight")
    plt.close(fig)
    return out


def load_comp_grad_results_metrics(
    base_dir,
    methods,
    methods_legend,
    parse_backwardTol=True,
):
    dfs = []
    for mm in methods:
        pattern = os.path.join(base_dir, mm, "*.csv")
        for fp in sorted(glob.glob(pattern)):
            m = mm.removesuffix("_steps")
            df = pd.read_csv(fp)
            df["method"] = methods_legend[m]

            fname = os.path.basename(fp)
            def grab(pat, cast=float):
                mo = re.search(pat, fname)
                return cast(mo.group(1)) if mo else np.nan

            df["seed"] = grab(r"_seed(\d+)", int)
            df["ydim"] = grab(r"ydim(\d+)", int)
            df["lr"]   = grab(r"lr([0-9eE\.\-]+)", float)

            if parse_backwardTol:
                df["backwardTol"] = grab(r"backwardTol([0-9eE\+\-\.]+?)(?:_|\.csv)", float)

            dfs.append(df)

    if not dfs:
        raise FileNotFoundError(f"No CSVs found under {base_dir}.")
    out = pd.concat(dfs, ignore_index=True, sort=False)
    out = out.rename(columns=lambda c: c.strip() if isinstance(c, str) else c)
    return out


def plot_comp_grad_metrics_vs_iter_by_ydim(
    df,
    metrics=("cosine_sim_val", "l2_dist_val", "l_inf_dist_val"),
    plot_path=".",
    plot_name_tag="syn",
    filter_method=None,
    filter_backwardTol=None,
    agg="mean",
    logy_dist=True,                     
    cosine_ylim=None,                   
    ydim_order=None,                    
    markers=("o","s","D","^","v","x","P","*"),
    linewidth=2.0,
    legend_ncol=5,
    legend_y=-0.06,
    legend_fontsize=11,
):
    os.makedirs(plot_path, exist_ok=True)

    d = df.copy().dropna(subset=["iter", "ydim"])
    d["iter"] = d["iter"].astype(int)
    d["ydim"] = d["ydim"].astype(int)

    if filter_method is not None:
        d = d[d["method"] == filter_method]

    if filter_backwardTol is not None and "backwardTol" in d.columns:
        d = d[np.isclose(d["backwardTol"].astype(float), float(filter_backwardTol))]

    if "backwardTol" in d.columns and filter_backwardTol is None and d["backwardTol"].nunique() > 1:
        bt = d["backwardTol"].mode().iloc[0]
        d = d[np.isclose(d["backwardTol"].astype(float), float(bt))]

    gfunc = "median" if agg == "median" else "mean"
    g = getattr(d.groupby(["ydim", "iter"], as_index=False)[list(metrics)], gfunc)()

    if ydim_order is None:
        ydim_order = sorted(g["ydim"].unique())

    titles = {
        "cosine_sim_val": "Cosine similarity",
        "l2_dist_val":    r"$\ell_2$ distance",
        "l_inf_dist_val": r"$\ell_\infty$ distance",
    }

    plt.rcParams.update({
        "font.size": 13,
        "axes.titlesize": 14,
        "axes.labelsize": 13,
        "xtick.labelsize": 12,
        "ytick.labelsize": 12,
    })

    fig, axes = plt.subplots(1, 3, figsize=(15.5, 4.2), sharex=True)

    for ax, metric in zip(axes, metrics):
        for i, ydim in enumerate(ydim_order):
            gg = g[g["ydim"] == ydim].sort_values("iter")
            if gg.empty:
                continue
            ax.plot(
                gg["iter"], gg[metric],
                label=f"ydim={ydim}",
                marker=markers[i % len(markers)],
                markersize=5.0,
                linewidth=linewidth,
            )

        ax.set_title(titles.get(metric, metric))
        ax.set_xlabel("iterations")
        ax.grid(True, which="major", alpha=0.25)
        ax.grid(True, which="minor", alpha=0.12)
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.yaxis.set_major_formatter(mticker.ScalarFormatter(useOffset=False))
        ax.yaxis.get_offset_text().set_visible(False)

        if metric == "cosine_sim_val" and cosine_ylim is not None:
            ax.set_ylim(*cosine_ylim)
        if metric != "cosine_sim_val" and logy_dist:
            ax.set_yscale("log")

    axes[0].set_ylabel("value")

    handles, labels = axes[0].get_legend_handles_labels()
    fig.legend(
        handles, labels,
        loc="lower center",
        bbox_to_anchor=(0.5, legend_y),
        ncol=min(len(labels), legend_ncol),
        frameon=False,
        fontsize=legend_fontsize,
        handlelength=2.4,
        handletextpad=0.6,
        columnspacing=1.2,
        borderaxespad=0.0,
    )

    fig.tight_layout(rect=[0, 0.14, 1, 1])

    tag = plot_name_tag
    if filter_method is not None:
        tag += f"_{filter_method}"
    out = os.path.join(plot_path, f"{tag}_metrics_vs_iter_by_ydim.pdf")
    fig.savefig(out, dpi=300, bbox_inches="tight")
    plt.close(fig)



if __name__=="__main__":
    task_title = "SOCP"
    df = load_results_CP()
    df = df.rename(columns=lambda c: c.strip() if isinstance(c, str) else c)
    df["method"] = pd.Categorical(df["method"], categories=method_order, ordered=True)
    df = df[df["ydim"] == 800]

    print("loaded df")
    print(df)
    # assert(1==0)

    plot_total_time_vs_method(df, task_title_tag=task_title, time_names=['forward_time', 'backward_time'], plot_path=BASE_DIR, plot_name_tag="syn_soc_ydim800")
    # plot_total_time_vs_method_by_backwardTol(df, time_names=['forward_time', 'backward_time'],
    #                                          plot_path=BASE_DIR, plot_name_tag="syn")
   
   #########################################
    df = load_results_CP(methods=METHODS_STEPS)
    df = df.rename(columns=lambda c: c.strip() if isinstance(c, str) else c)
    df["method"] = pd.Categorical(df["method"], categories=method_order, ordered=True)
    df = df[df["ydim"] == 800]

    print("loaded df steps")

    # plot_total_time_vs_method(df, time_names=['forward_solve_time', 'backward_solve_time'], plot_path=BASE_DIR, plot_name_tag="syn_steps_solve")
    # plot_total_time_vs_method(df, time_names=['forward_setup_time', 'backward_setup_time'], plot_path=BASE_DIR, plot_name_tag="syn_steps_setup")
    
    plot_loss_vs_epoch(df, "train_df_loss", task_title_tag=task_title, iteration_name='iter', plot_path=BASE_DIR, plot_name_tag="syn_soc_steps_ydim800")
    # plot_loss_vs_epoch_method_tol(df, loss_metric='train_df_loss', iteration='iter',
    #                               plot_path=BASE_DIR, plot_name_tag="syn_steps")