#!/usr/bin/env python
"""Compare A1 fit results across multiple model slugs.

Reads ``core_roi_layer_summary.csv`` and ``core_roi_best_layer_summary.csv``
from each ``--fit-dir`` and renders side-by-side plots tagged by model slug.

Outputs go to a dedicated comparison directory so the per-model visualize
step (run_a1_visualize.py) is never overwritten.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt  # noqa: E402
import pandas as pd  # noqa: E402
import seaborn as sns  # noqa: E402


METRIC_ALIAS_MAP: dict[str, str] = {
    "2v2": "mean_2v2_accuracy",
    "2v2_accuracy": "mean_2v2_accuracy",
    "two_v_two_accuracy": "mean_2v2_accuracy",
}


def _resolve_metric_column(metric: str) -> str:
    token = str(metric).strip().lower()
    return METRIC_ALIAS_MAP.get(token, str(metric))


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Compare A1 fits across models")
    parser.add_argument(
        "--fit-dir",
        action="append",
        required=True,
        help=(
            "Path to a fit_results/<model-slug> directory. Repeat for each model. "
            "Optionally prefix with a label, e.g. 'qwen3=outputs/.../fit_results/Qwen_Qwen3-0.6B'."
        ),
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        required=True,
        help="Where to write comparison plots and merged CSVs",
    )
    parser.add_argument(
        "--metric",
        type=str,
        default="mean_corr",
        choices=[
            "mean_corr",
            "mean_r2",
            "mean_2v2_accuracy",
            "2v2",
            "2v2_accuracy",
            "two_v_two_accuracy",
        ],
    )
    parser.add_argument(
        "--title-suffix",
        type=str,
        default="",
        help="Optional title suffix to add to comparison plots",
    )
    return parser


def _split_label(entry: str) -> tuple[str, Path]:
    if "=" in entry:
        label, path = entry.split("=", 1)
        return label.strip() or Path(path).name, Path(path).expanduser().resolve()
    p = Path(entry).expanduser().resolve()
    return p.name, p


def _load_fit_summaries(
    entries: list[str],
) -> tuple[pd.DataFrame, pd.DataFrame]:
    layer_frames: list[pd.DataFrame] = []
    best_frames: list[pd.DataFrame] = []
    for entry in entries:
        label, fit_dir = _split_label(entry)
        layer_path = fit_dir / "core_roi_layer_summary.csv"
        best_path = fit_dir / "core_roi_best_layer_summary.csv"
        missing = [p for p in (layer_path, best_path) if not p.exists()]
        if missing:
            raise FileNotFoundError(
                "Missing fit summaries for "
                f"{label} ({fit_dir}): {[str(p) for p in missing]}"
            )

        layer_df = pd.read_csv(layer_path)
        layer_df["model_label"] = label
        layer_df["fit_dir"] = str(fit_dir)
        layer_frames.append(layer_df)

        best_df = pd.read_csv(best_path)
        best_df["model_label"] = label
        best_df["fit_dir"] = str(fit_dir)
        best_frames.append(best_df)

    return (
        pd.concat(layer_frames, ignore_index=True),
        pd.concat(best_frames, ignore_index=True),
    )


def _plot_layer_curve(
    layer_df: pd.DataFrame,
    metric: str,
    output_path: Path,
    title_suffix: str,
) -> None:
    if layer_df.empty:
        return
    curve_df = (
        layer_df.groupby(["model_label", "protocol", "layer_idx"], as_index=False)[metric]
        .mean()
        .sort_values(["model_label", "protocol", "layer_idx"])
    )

    protocols = sorted(curve_df["protocol"].unique())
    fig, axes = plt.subplots(
        1, len(protocols), figsize=(6 * len(protocols), 4.5), sharey=True, squeeze=False
    )
    for ax, protocol in zip(axes[0], protocols):
        sub = curve_df[curve_df["protocol"] == protocol]
        sns.lineplot(data=sub, x="layer_idx", y=metric, hue="model_label", marker="o", ax=ax)
        ax.set_title(f"Protocol {protocol}")
        ax.set_xlabel("Layer")
        ax.set_ylabel(metric)
    suffix = f" — {title_suffix}" if title_suffix else ""
    fig.suptitle(f"Avg target-mask {metric} by layer across models{suffix}")
    fig.tight_layout()
    fig.savefig(output_path, dpi=180)
    plt.close(fig)


def _plot_best_layer_bar(
    best_df: pd.DataFrame,
    metric: str,
    output_path: Path,
    title_suffix: str,
) -> None:
    if best_df.empty:
        return
    chart_df = best_df.copy()
    chart_df["roi_label"] = chart_df["roi_name"].astype(str)

    protocols = sorted(chart_df["protocol"].unique())
    fig, axes = plt.subplots(
        len(protocols), 1, figsize=(12, 4.5 * len(protocols)), squeeze=False
    )
    for ax, protocol in zip(axes[:, 0], protocols):
        sub = chart_df[chart_df["protocol"] == protocol]
        sns.barplot(data=sub, x="roi_label", y=metric, hue="model_label", ax=ax)
        ax.set_title(f"Best-layer {metric} per ROI — protocol {protocol}")
        ax.set_xlabel("ROI")
        ax.set_ylabel(metric)
        ax.tick_params(axis="x", rotation=20)
    suffix = f" — {title_suffix}" if title_suffix else ""
    fig.suptitle(f"Best layer {metric} per target mask across models{suffix}")
    fig.tight_layout()
    fig.savefig(output_path, dpi=180)
    plt.close(fig)


def _plot_overall_summary(
    best_df: pd.DataFrame,
    metric: str,
    output_path: Path,
    title_suffix: str,
) -> None:
    if best_df.empty:
        return
    summary_df = (
        best_df.groupby(["model_label", "protocol"], as_index=False)[metric]
        .mean()
        .sort_values(["protocol", "model_label"])
    )
    fig, ax = plt.subplots(figsize=(8, 4.8))
    sns.barplot(data=summary_df, x="protocol", y=metric, hue="model_label", ax=ax)
    ax.set_title(
        f"Mean best-layer {metric} per protocol{(' — ' + title_suffix) if title_suffix else ''}"
    )
    ax.set_ylabel(metric)
    ax.set_xlabel("Protocol")
    fig.tight_layout()
    fig.savefig(output_path, dpi=180)
    plt.close(fig)


def main() -> None:
    args = _build_parser().parse_args()
    output_dir = Path(args.output_dir).expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    layer_df, best_df = _load_fit_summaries(entries=list(args.fit_dir))
    metric = _resolve_metric_column(args.metric)
    if metric not in layer_df.columns:
        raise ValueError(
            f"Metric column '{metric}' missing from layer summary. "
            f"Available: {sorted(layer_df.columns.tolist())}"
        )
    if metric not in best_df.columns:
        raise ValueError(
            f"Metric column '{metric}' missing from best summary. "
            f"Available: {sorted(best_df.columns.tolist())}"
        )

    layer_df.to_csv(output_dir / "merged_core_roi_layer_summary.csv", index=False)
    best_df.to_csv(output_dir / "merged_core_roi_best_layer_summary.csv", index=False)

    _plot_layer_curve(
        layer_df=layer_df,
        metric=metric,
        output_path=output_dir / f"compare_layer_curve_{metric}.png",
        title_suffix=args.title_suffix,
    )
    _plot_best_layer_bar(
        best_df=best_df,
        metric=metric,
        output_path=output_dir / f"compare_best_layer_bar_{metric}.png",
        title_suffix=args.title_suffix,
    )
    _plot_overall_summary(
        best_df=best_df,
        metric=metric,
        output_path=output_dir / f"compare_overall_{metric}.png",
        title_suffix=args.title_suffix,
    )

    print("=" * 72)
    print("A1 model comparison complete")
    print(f"Models compared : {sorted(set(layer_df['model_label']))}")
    print(f"Metric          : {metric}")
    print(f"Output directory: {output_dir}")
    print("=" * 72)


if __name__ == "__main__":
    main()