plot_correlation.py

import os
import argparse
import sqlite3

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.colors import is_color_like, to_hex
from scipy import stats


def choose_empty_corner(x_data, y_data, radius=0.22, margin=0.05):
    """
    Choose a relatively empty corner in axes-normalized coordinates.
    This is kept from the numbered-marker style script.
    """
    x = np.asarray(x_data, dtype=float)
    y = np.asarray(y_data, dtype=float)

    x_span = x.max() - x.min()
    y_span = y.max() - y.min()

    if x_span == 0:
        x_norm = np.full_like(x, 0.5, dtype=float)
    else:
        x_norm = (x - x.min()) / x_span

    if y_span == 0:
        y_norm = np.full_like(y, 0.5, dtype=float)
    else:
        y_norm = (y - y.min()) / y_span

    points = np.stack([x_norm, y_norm], axis=1)

    corners = {
        "upper_left":  (margin, 1.0 - margin, "left",  "top"),
        "upper_right": (1.0 - margin, 1.0 - margin, "right", "top"),
        "lower_left":  (margin, margin, "left",  "bottom"),
        "lower_right": (1.0 - margin, margin, "right", "bottom"),
    }

    best_name = None
    best_score = None

    for name, (cx, cy, ha, va) in corners.items():
        corner = np.array([cx, cy], dtype=float)
        dists = np.linalg.norm(points - corner[None, :], axis=1)

        n_near = np.sum(dists < radius)
        min_dist = np.min(dists)
        score = (n_near, -min_dist)

        if best_score is None or score < best_score:
            best_score = score
            best_name = name

    return corners[best_name]


def load_cluster_color_map(color_csv_path, cls_col):
    color_df = pd.read_csv(color_csv_path)

    cluster_col_candidates = [cls_col, "cluster", "Cluster", "cls", "Class"]
    color_col_candidates = ["color", "Color", "colour", "Colour"]
    alias_col_candidates = [
        "alias", "Alias",
        "display_name", "DisplayName", "display", "Display",
        "legend", "Legend",
        "label", "Label",
        "name", "Name",
    ]

    cluster_col = next(
        (col for col in cluster_col_candidates if col is not None and col in color_df.columns),
        None,
    )
    color_col = next(
        (col for col in color_col_candidates if col in color_df.columns),
        None,
    )
    alias_col = next(
        (col for col in alias_col_candidates if col in color_df.columns),
        None,
    )

    if cluster_col is None or color_col is None:
        raise ValueError(
            f"cls_color_csv must contain a cluster column "
            f"({cluster_col_candidates}) and a color column ({color_col_candidates}). "
            f"Found columns: {list(color_df.columns)}"
        )

    if color_df[cluster_col].duplicated().any():
        duplicated = color_df.loc[color_df[cluster_col].duplicated(), cluster_col].tolist()
        raise ValueError(f"Duplicate clusters in cls_color_csv: {duplicated}")

    cluster_keys = color_df[cluster_col].astype(str)
    cluster_color_map = dict(zip(cluster_keys, color_df[color_col].astype(str)))

    if alias_col is not None:
        cluster_alias_map = dict(
            zip(cluster_keys, color_df[alias_col].fillna(color_df[cluster_col]).astype(str))
        )
    else:
        cluster_alias_map = dict(zip(cluster_keys, cluster_keys))

    invalid_colors = [
        color for color in cluster_color_map.values()
        if not is_color_like(color)
    ]
    if invalid_colors:
        raise ValueError(f"Invalid matplotlib colors in cls_color_csv: {invalid_colors}")

    return cluster_color_map, cluster_alias_map


def merged_col_name(col, side, df1, df2, id_col):
    """
    Resolve the column name after pandas.merge(..., suffixes=('_x', '_y')).
    side: 'x' for m1_csv, 'y' for m2_csv.
    """
    if col == id_col:
        return col
    in_df1 = col in df1.columns
    in_df2 = col in df2.columns
    if in_df1 and in_df2:
        return f"{col}_{side}"
    return col


def add_unified_column(df, col):
    """
    If a metadata column exists as col_x/col_y after full merge, expose col as a
    unified alias. Prefer the m1_csv side, matching the old scripts' behavior.
    """
    if col is None or col in df.columns:
        return df
    if f"{col}_x" in df.columns:
        df[col] = df[f"{col}_x"]
    elif f"{col}_y" in df.columns:
        df[col] = df[f"{col}_y"]
    return df


def apply_sql(df, sql):
    """
    Apply an optional SQL filter/projection to the merged dataframe.

    Supported forms:
      --where "Diffusion = 'SiT-B' AND CFG = 1"

    The merged dataframe is registered as table name: df.
    """
    if sql is None:
        return df

    stmt = sql.strip().rstrip(";")
    if not stmt:
        return df

    query = f"SELECT * FROM df WHERE {stmt}"

    conn = sqlite3.connect(":memory:")
    try:
        df.to_sql("df", conn, index=False, if_exists="replace")
        return pd.read_sql_query(query, conn)
    finally:
        conn.close()


def load_and_prepare_data(args):
    df1 = pd.read_csv(args.m1_csv)
    df2 = pd.read_csv(args.m2_csv)

    if args.id_col not in df1.columns or args.m1_col not in df1.columns:
        raise ValueError(
            f"Columns {args.id_col} or {args.m1_col} not found in {args.m1_csv}. "
            f"Found columns: {list(df1.columns)}"
        )
    if args.id_col not in df2.columns or args.m2_col not in df2.columns:
        raise ValueError(
            f"Columns {args.id_col} or {args.m2_col} not found in {args.m2_csv}. "
            f"Found columns: {list(df2.columns)}"
        )

    # Keep the full merged dataframe so --sql can filter by extra columns such as
    # Diffusion, Steps, CFG, n_samples, Arch, Config, etc.
    merge_cols = args.merge_cols if args.merge_cols is not None else [args.id_col]
    missing_in_df1 = [c for c in merge_cols if c not in df1.columns]
    missing_in_df2 = [c for c in merge_cols if c not in df2.columns]
    if missing_in_df1 or missing_in_df2:
        raise ValueError(
            f"Invalid merge_cols={merge_cols}. "
            f"Missing in m1_csv: {missing_in_df1}; "
            f"Missing in m2_csv: {missing_in_df2}"
        )
    df = pd.merge(df1, df2, on=merge_cols, how="inner", suffixes=("_x", "_y"))

    m1_plot_col = merged_col_name(args.m1_col, "x", df1, df2, args.id_col)
    m2_plot_col = merged_col_name(args.m2_col, "y", df1, df2, args.id_col)

    for meta_col in [args.cls_col, "Arch", "Config", "Variant"]:
        df = add_unified_column(df, meta_col)

    df = apply_sql(df, args.where)

    # If the SQL projection renames/preserves the original column names, allow that.
    if m1_plot_col not in df.columns and args.m1_col in df.columns:
        m1_plot_col = args.m1_col
    if m2_plot_col not in df.columns and args.m2_col in df.columns:
        m2_plot_col = args.m2_col

    required_cols = [args.id_col, m1_plot_col, m2_plot_col]
    missing = [col for col in required_cols if col not in df.columns]
    if missing:
        raise ValueError(
            f"Missing required columns after merge/sql: {missing}. "
            f"Available columns: {list(df.columns)}"
        )

    df = df.sort_values(by=args.id_col).reset_index(drop=True)

    if len(df) == 0:
        raise ValueError("No overlapping IDs found between the two CSV files after filtering.")

    df[m1_plot_col] = pd.to_numeric(df[m1_plot_col], errors="coerce")
    df[m2_plot_col] = pd.to_numeric(df[m2_plot_col], errors="coerce")

    before_drop = len(df)
    df = df.dropna(subset=[m1_plot_col, m2_plot_col]).reset_index(drop=True)
    dropped = before_drop - len(df)
    if dropped > 0:
        print(f"Warning: Dropped {dropped} rows with NaN/non-numeric metric values.")

    if len(df) == 0:
        raise ValueError("No valid rows left after dropping NaN/non-numeric metric values.")

    if args.inv_m1:
        df[m1_plot_col] = -df[m1_plot_col]
    if args.inv_m2:
        df[m2_plot_col] = -df[m2_plot_col]

    return df, m1_plot_col, m2_plot_col


def get_style(simplify):
    if simplify:
        return {
            "figsize": (3, 3),
            "right_with_method_legend": None,
            "right_without_method_legend": None,
            "xlabel_fontsize": 14,
            "ylabel_fontsize": 14,
            "tick_labelsize": 12,
            "marker_size": 80,
            "marker_alpha": 0.9,
            "marker_linewidth": 0.8,
            "line_width": 2.0,
            "line_alpha": 0.5,
            "global_line_zorder": 4,
            "global_band_zorder": 3,
            "cluster_line_zorder": 5,
            "cluster_band_zorder": 4,
            "global_band_alpha": 0.15,
            "cluster_band_alpha": 0.2,
            "simple_r_fontsize": 14,
        }

    return {
        "figsize": (3, 2.5),
        "right_with_method_legend": 0.75,
        "right_without_method_legend": 0.95,
        "xlabel_fontsize": 12,
        "ylabel_fontsize": 12,
        "marker_size": 150,
        "marker_text_fontsize": 8,
        "line_width": 2.0,
        "line_alpha": None,
        "global_line_zorder": 3,
        "global_band_zorder": 2,
        "cluster_line_zorder": 5,
        "cluster_band_zorder": 4,
        "global_band_alpha": 0.15,
        "cluster_band_alpha": 0.2,
        "numbered_r_fontsize": 12,
    }


def setup_matplotlib(style, hide_method_legend):
    sns.set_theme(style="whitegrid")
    plt.rcParams["font.family"] = "serif"
    plt.rcParams["font.serif"] = ["Times New Roman"] + plt.rcParams["font.serif"]

    fig, ax = plt.subplots(figsize=style["figsize"])

    if style["right_with_method_legend"] is not None:
        if hide_method_legend:
            plt.subplots_adjust(right=style["right_without_method_legend"])
        else:
            plt.subplots_adjust(right=style["right_with_method_legend"])

    return fig, ax


def prepare_cluster_info(args, df, x_data, y_data):
    cluster_color_map = {}
    cluster_alias_map = {}
    cluster_r_map = {}
    unique_clusters = []

    has_cls = args.cls_col and args.cls_col in df.columns
    if args.cls_color_csv is not None and not has_cls:
        raise ValueError("--cls_color_csv requires a valid --cls_col in the merged dataframe.")

    if has_cls:
        unique_clusters = df[args.cls_col].dropna().unique()

        if args.cls_color_csv is not None:
            cluster_color_map, cluster_alias_map = load_cluster_color_map(
                args.cls_color_csv,
                args.cls_col,
            )
            missing_clusters = sorted(
                [str(cls_val) for cls_val in unique_clusters if str(cls_val) not in cluster_color_map]
            )
            if missing_clusters:
                raise ValueError(
                    "The following clusters are missing in cls_color_csv: "
                    f"{missing_clusters}"
                )
        else:
            palette = sns.color_palette("pastel", len(unique_clusters))
            cluster_color_map = dict(zip([str(cls_val) for cls_val in unique_clusters], palette))
            cluster_alias_map = dict(
                zip(
                    [str(cls_val) for cls_val in unique_clusters],
                    [str(cls_val) for cls_val in unique_clusters],
                )
            )

        print("[Cluster Color Map]")
        for cls_val in unique_clusters:
            key = str(cls_val)
            color = cluster_color_map[key]
            alias = cluster_alias_map.get(key, key)
            if args.simplify:
                print(f"  {key} -> {to_hex(color)}")
            else:
                print(f"  {key} ({alias}) -> {to_hex(color)}")

        for cls_val in unique_clusters:
            cls_mask = df[args.cls_col] == cls_val
            cx = x_data[cls_mask]
            cy = y_data[cls_mask]
            if len(cx) > 1 and cx.nunique() > 1 and cy.nunique() > 1:
                r_val, _ = stats.pearsonr(cx, cy)
                cluster_r_map[cls_val] = r_val
            else:
                cluster_r_map[cls_val] = float("nan")
    else:
        print("[Cluster Color Map] None")

    return unique_clusters, cluster_color_map, cluster_alias_map, cluster_r_map


def safe_pearson(x_data, y_data):
    if len(x_data) > 1 and x_data.nunique() > 1 and y_data.nunique() > 1:
        r, p = stats.pearsonr(x_data, y_data)
        return r, p
    return float("nan"), float("nan")


def fit_line_arrays(x, y):
    if len(x) <= 1 or x.nunique() <= 1:
        return None
    slope, intercept = np.polyfit(x, y, 1)
    x_plot = np.linspace(x.min(), x.max(), 100)
    y_plot = slope * x_plot + intercept
    y_pred = slope * x + intercept
    band = np.std(y - y_pred)
    return x_plot, y_plot, band


def draw_fit(ax, x, y, color, line_width, line_alpha, line_zorder,
             show_band, band_alpha, band_zorder):
    fit_res = fit_line_arrays(x, y)
    if fit_res is None:
        return

    x_plot, y_plot, band = fit_res
    plot_kwargs = {
        "color": color,
        "linewidth": line_width,
        "zorder": line_zorder,
    }
    if line_alpha is not None:
        plot_kwargs["alpha"] = line_alpha

    ax.plot(x_plot, y_plot, **plot_kwargs)

    if show_band:
        ax.fill_between(
            x_plot,
            y_plot - band,
            y_plot + band,
            color=color,
            alpha=band_alpha,
            zorder=band_zorder,
            linewidth=0,
        )


def draw_all_fits(args, ax, df, x_data, y_data, unique_clusters, cluster_color_map, style):
    if args.fit != "linear":
        return

    # Normal style follows script 1: red global line when cluster fit is disabled,
    # otherwise gray global line plus per-cluster lines. Simplified style follows
    # script 2: gray global line, then per-cluster lines when enabled.
    global_color = "gray"
    if (not args.simplify) and args.no_cls_fit:
        global_color = "red"

    draw_fit(
        ax=ax,
        x=x_data,
        y=y_data,
        color=global_color,
        line_width=style["line_width"],
        line_alpha=style["line_alpha"],
        line_zorder=style["global_line_zorder"],
        show_band=args.effective_show_band,
        band_alpha=style["global_band_alpha"],
        band_zorder=style["global_band_zorder"],
    )

    if args.no_cls_fit or not (args.cls_col and args.cls_col in df.columns and cluster_color_map):
        return

    for cls_val in unique_clusters:
        cls_mask = df[args.cls_col] == cls_val
        cx = x_data[cls_mask]
        cy = y_data[cls_mask]
        draw_fit(
            ax=ax,
            x=cx,
            y=cy,
            color=cluster_color_map[str(cls_val)],
            line_width=style["line_width"],
            line_alpha=style["line_alpha"],
            line_zorder=style["cluster_line_zorder"],
            show_band=args.effective_show_band,
            band_alpha=style["cluster_band_alpha"],
            band_zorder=style["cluster_band_zorder"],
        )


def draw_simplified_markers(args, ax, df, x_data, y_data, unique_clusters, cluster_color_map, style):
    if args.cls_col and args.cls_col in df.columns and cluster_color_map:
        for cls_val in unique_clusters:
            cls_mask = df[args.cls_col] == cls_val
            ax.scatter(
                x_data[cls_mask],
                y_data[cls_mask],
                s=style["marker_size"],
                facecolors=cluster_color_map[str(cls_val)],
                edgecolors="black",
                linewidths=style["marker_linewidth"],
                alpha=style["marker_alpha"],
                zorder=10,
            )
    else:
        ax.scatter(
            x_data,
            y_data,
            s=style["marker_size"],
            facecolors="white",
            edgecolors="black",
            linewidths=style["marker_linewidth"],
            alpha=style["marker_alpha"],
            zorder=10,
        )


def format_id_label(value):
    """Format id_col value for circle labels."""
    if pd.isna(value):
        return ""

    if isinstance(value, (int, np.integer)):
        return str(value)

    if isinstance(value, (float, np.floating)) and float(value).is_integer():
        return str(int(value))

    return str(value)


def draw_numbered_markers(args, ax, df, x_data, y_data, ids, cluster_color_map, style):
    mapping_text = []
    seen_legend_ids = set()

    for row_idx, (x, y, name) in enumerate(zip(x_data, y_data, ids)):
        # Directly use the value from args.id_col instead of assigning a new number.
        circle_label = format_id_label(name)

        bg_color = "white"
        if args.cls_col and args.cls_col in df.columns:
            c_val = df[args.cls_col].iloc[row_idx]
            bg_color = cluster_color_map.get(str(c_val), "white")

        ax.scatter(
            x,
            y,
            s=style["marker_size"],
            facecolors=bg_color,
            edgecolors="black",
            zorder=10,
        )
        ax.text(
            x,
            y,
            circle_label,
            ha="center",
            va="center",
            fontsize=style["marker_text_fontsize"],
            fontweight="bold",
            zorder=11,
        )

        if name not in seen_legend_ids:
            mapping_text.append(f"{circle_label}: {name}")
            seen_legend_ids.add(name)

    return mapping_text


def draw_method_legend(fig, mapping_text, id_col):
    if len(mapping_text) == 0:
        return

    total_items = len(mapping_text)
    n_cols = max(1, int(np.ceil(total_items / 30)))
    rows_per_col = int(np.ceil(total_items / n_cols))

    columns = [
        mapping_text[i * rows_per_col:(i + 1) * rows_per_col]
        for i in range(n_cols)
    ]

    fig.text(
        0.77,
        0.85,
        id_col,
        fontsize=12,
        fontweight="bold",
        fontfamily="serif",
    )

    current_x = 0.77
    y_top = 0.82
    text_objects = []

    for col_items in columns:
        legend_text = "\n".join(col_items)
        txt = fig.text(
            current_x,
            y_top,
            legend_text,
            fontsize=9,
            verticalalignment="top",
            fontfamily="serif",
            linespacing=1.4,
        )
        text_objects.append(txt)

    fig.canvas.draw()
    renderer = fig.canvas.get_renderer()

    current_x = 0.77
    for txt in text_objects:
        bbox = txt.get_window_extent(renderer=renderer)
        bbox_fig = bbox.transformed(fig.transFigure.inverted())
        width = bbox_fig.width
        txt.set_position((current_x, y_top))
        current_x += width + 0.01


def draw_cluster_legend(args, ax, unique_clusters, cluster_color_map, cluster_alias_map, cluster_r_map):
    if args.simplify:
        return
    if not (args.cls_col and cluster_color_map and not args.hide_cls_legend):
        return

    import matplotlib.lines as mlines

    handles = [
        mlines.Line2D(
            [0],
            [0],
            marker="o",
            color="w",
            markerfacecolor=cluster_color_map[str(cls_val)],
            markeredgecolor="black",
            markersize=8,
        )
        for cls_val in unique_clusters
    ]

    labels = []
    for k in unique_clusters:
        key = str(k)
        alias = cluster_alias_map.get(key, key)
        r_val = cluster_r_map.get(k, float("nan"))
        if not args.no_cls_fit and not pd.isna(r_val):
            labels.append(f"{alias} ({r_val:.2f})")
        else:
            labels.append(alias)

    n_cols = min(len(labels), 3)
    ax.legend(
        handles,
        labels,
        loc="lower center",
        bbox_to_anchor=(0.5, 1.02),
        ncol=n_cols,
        frameon=False,
        columnspacing=0.6,
        handletextpad=0.3,
        borderaxespad=0.2,
        title_fontproperties={"weight": "bold", "size": 12, "family": "serif"},
        prop={"family": "serif", "size": 10},
    )


def draw_r_box(args, ax, x_data, y_data, style):
    r, p = safe_pearson(x_data, y_data)
    stat_symbol = r"\rho" if args.spearman else "r"

    if args.simplify:
        if r > 0.9:
            tx, ty, ha, va = 0.05, 0.95, "left", "top"
        else:
            tx, ty, ha, va = 0.95, 0.95, "right", "top"
        ax.text(
            tx,
            ty,
            f"${stat_symbol} = {r:.3f}$",
            transform=ax.transAxes,
            fontsize=style["simple_r_fontsize"],
            verticalalignment=va,
            horizontalalignment=ha,
            fontfamily="serif",
            bbox=dict(boxstyle="round", facecolor="none", alpha=0.8, edgecolor="none"),
            zorder=20,
        )
    else:
        tx, ty, ha, va = choose_empty_corner(x_data, y_data)
        ax.text(
            tx,
            ty,
            f"${stat_symbol} = {r:.3f}$",
            transform=ax.transAxes,
            fontsize=style["numbered_r_fontsize"],
            verticalalignment=va,
            horizontalalignment=ha,
            bbox=dict(boxstyle="round", facecolor="white", alpha=0.9, edgecolor="gray"),
        )

    return r


def polish_axes(args, ax, x_data, y_data, x_label, y_label, style):
    ax.set_xlabel(x_label, fontsize=style["xlabel_fontsize"])
    ax.set_ylabel(y_label, fontsize=style["ylabel_fontsize"])

    x_margin = (x_data.max() - x_data.min()) * 0.1
    y_margin = (y_data.max() - y_data.min()) * 0.1
    if x_margin == 0:
        x_margin = 1.0
    if y_margin == 0:
        y_margin = 1.0

    ax.set_xlim(x_data.min() - x_margin, x_data.max() + x_margin)
    ax.set_ylim(y_data.min() - y_margin, y_data.max() + y_margin)

    if args.simplify:
        ax.tick_params(axis="both", labelsize=style["tick_labelsize"])
        plt.grid(False)
        ax.set_xticks([])
        ax.set_yticks([])


def main(args):
    # Different original defaults are preserved here:
    #   normal numbered style: fit defaults to None, band is shown when fit='linear'
    #   simplified style:     fit defaults to linear, band is hidden unless --show_band
    if args.fit is None:
        args.fit = "linear" if args.simplify else None
    args.effective_show_band = args.show_band or (not args.simplify)

    df, m1_plot_col, m2_plot_col = load_and_prepare_data(args)

    y_label = args.m1_name if args.m1_name else args.m1_col
    x_label = args.m2_name if args.m2_name else args.m2_col
    if args.inv_m1:
        y_label = "-" + y_label
    if args.inv_m2:
        x_label = "-" + x_label

    x_data = df[m2_plot_col]
    y_data = df[m1_plot_col]

    if args.spearman:
        x_data = x_data.rank(method="average")
        y_data = y_data.rank(method="average")
        x_label = f"Rank({x_label})"
        y_label = f"Rank({y_label})"

    ids = df[args.id_col]

    style = get_style(args.simplify)
    fig, ax = setup_matplotlib(style, True)

    unique_clusters, cluster_color_map, cluster_alias_map, cluster_r_map = prepare_cluster_info(
        args,
        df,
        x_data,
        y_data,
    )

    # Draw markers and fits through the same data path. The style switch only
    # changes marker form, hard-coded visual parameters, and legend/stat box style.
    if args.simplify:
        draw_simplified_markers(
            args, ax, df, x_data, y_data, unique_clusters, cluster_color_map, style
        )
        mapping_text = []
    else:
        mapping_text = draw_numbered_markers(
            args, ax, df, x_data, y_data, ids, cluster_color_map, style
        )

    draw_all_fits(args, ax, df, x_data, y_data, unique_clusters, cluster_color_map, style)
    r = draw_r_box(args, ax, x_data, y_data, style)

    # if not args.simplify:
        # draw_method_legend(fig, mapping_text, args.id_col)

    draw_cluster_legend(args, ax, unique_clusters, cluster_color_map, cluster_alias_map, cluster_r_map)
    polish_axes(args, ax, x_data, y_data, x_label, y_label, style)

    output_dir = os.path.dirname(args.output_path)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    if args.simplify:
        print(f"Saving simplified figure to {args.output_path} (n={len(df)})")
    else:
        print(f"Saving figure to {args.output_path} (r={r:.3f})")

    plt.savefig(
        args.output_path,
        dpi=300,
        bbox_inches="tight",
        transparent=args.transparent,
    )


def parse_args():
    parser = argparse.ArgumentParser(
        description="Scatter plot with numbered markers or simplified cluster-color style."
    )

    parser.add_argument("--m1_csv", type=str, required=True, help="Path to the first metric CSV")
    parser.add_argument("--m2_csv", type=str, required=True, help="Path to the second metric CSV")
    parser.add_argument("-o", "--output_path", type=str, required=True, help="Output image path")

    parser.add_argument("--m1_col", type=str, required=True, help="Column name for Metric 1, Y-axis")
    parser.add_argument("--m2_col", type=str, required=True, help="Column name for Metric 2, X-axis")
    parser.add_argument("--id_col", type=str, default="VAE", help="Column name for the Method ID/Name")
    parser.add_argument("--cls_col", type=str, default="Cluster", help="Column name to distinguish Method Clusters")

    parser.add_argument("--m1_name", type=str, default=None, help="Label for Y-axis")
    parser.add_argument("--m2_name", type=str, default=None, help="Label for X-axis")
    parser.add_argument("--fit", type=str, default="linear", choices=["linear", "none"], help="Fit type")
    parser.add_argument("--no-cls-fit", action="store_true", help="Disable per-cluster fitting")
    parser.add_argument("--show_band", action="store_true", help="Show residual-std regression band")
    parser.add_argument("--transparent", action="store_true")
    parser.add_argument("--inv_m1", action="store_true")
    parser.add_argument("--inv_m2", action="store_true")
    parser.add_argument(
        "--spearman",
        action="store_true",
        help="Fit and compute correlation on metric ranks instead of raw metric values.",
    )
    parser.add_argument("--simplify", action="store_true", help="Use the simplified output style")
    parser.add_argument(
        "--where",
        type=str,
        default=None,
        help="Optional SQL filter/projection on merged dataframe. Table name is df.",
    )
    parser.add_argument(
        "--merge_cols",
        nargs="+",
        default=None,
        help="Columns used for merging. Default: id_col only.",
    )

    parser.add_argument("--hide_cls_legend", action="store_true", help="Hide Cluster Legend")
    parser.add_argument(
        "--cls_color_csv",
        type=str,
        default=None,
        help="Optional CSV mapping cluster names to colors and optional aliases. "
             "Expected columns: cls_col/cluster/Cluster, color/Color, "
             "and optionally alias/Alias/display_name/legend/label.",
    )

    return parser.parse_args()


if __name__ == "__main__":
    main(parse_args())