"""
Publication-ready comparison figures for BADAS model evaluation results.

Loads per-model ``metrics_by_group.csv`` files produced by the Evaluator and
renders a set of matching figures suitable for blog posts or model cards.

All display names — for both models and scenario groups — are derived
automatically from the results data; no mapping dictionaries are hard-coded.

Figures
-------
- :func:`plot_comparison_table`  — Anthropic-style metrics table with per-cell
  highlights and a rounded border on the featured (newest) model column.
- :func:`plot_f1_dumbbell`       — horizontal dumbbell chart: F1 per scenario.
- :func:`plot_kaggle`            — Kaggle AP metrics across Public/Private/Overall.
- :func:`plot_recall_fpr`        — Recall gain and FPR reduction vs baseline.
- :func:`plot_overview`          — 2×2 summary hero figure combining the above.

Convenience
-----------
- :func:`discover_models`            — auto-discover results from a directory tree.
- :func:`generate_comparison_report` — render all figures in one call.

Usage::

    from badas.validation.comparison import discover_models, generate_comparison_report

    models = discover_models("results/")
    generate_comparison_report(models, output_dir="results/blog")
"""

from __future__ import annotations

from pathlib import Path
from typing import Dict, List, Optional, Tuple

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import matplotlib.ticker
import numpy as np
import pandas as pd
from matplotlib.gridspec import GridSpec

# ---------------------------------------------------------------------------
# Type alias
# ---------------------------------------------------------------------------
# Each model is (display_name, testset_df | None, kaggle_df | None)
ModelEntry = Tuple[str, Optional[pd.DataFrame], Optional[pd.DataFrame]]

# ---------------------------------------------------------------------------
# Design tokens
# ---------------------------------------------------------------------------
_LAVEN        = "#9A8FE5"   # primary lavender
_LAVEN_LIGHT  = "#EDE9FF"   # best-cell fill
_LAVEN_DARK   = "#4A3EC0"   # best-cell text / emphasis
_SLATE        = "#94A3B8"   # baseline model colour
_BG           = "#FFFFFF"
_LINE_COLOR   = "#E8E8E8"
_TEXT_DARK    = "#111827"
_TEXT_MID     = "#6B7280"
_TEXT_LIGHT   = "#9CA3AF"

# palette[0] = baseline, palette[1] = featured (latest) model
_MODEL_COLORS = [_SLATE, _LAVEN]


# ===========================================================================
# Helpers
# ===========================================================================

def _prettify(s: str) -> str:
    """Convert snake_case / UPPER to a readable title label."""
    return s.replace("_", " ").title()


def _derive_groups(models: List[ModelEntry]) -> Tuple[List[str], Dict[str, str]]:
    """
    Read unique group names from the first available testset dataframe.

    Returns ``(group_order, group_labels)`` where any "OVERALL"-like group
    is placed last and all labels are auto-prettified from the column values.
    """
    for _, ts_df, _ in models:
        if ts_df is not None:
            return _groups_from_df(ts_df)
    # fall back to kaggle df (e.g. Public / Private / OVERALL)
    for _, _, kg_df in models:
        if kg_df is not None:
            return _groups_from_df(kg_df)
    return [], {}


def _groups_from_df(df: pd.DataFrame) -> Tuple[List[str], Dict[str, str]]:
    groups    = df["group"].astype(str).unique().tolist()
    overall   = [g for g in groups if g.upper() == "OVERALL"]
    rest      = [g for g in groups if g.upper() != "OVERALL"]
    order     = rest + overall
    labels    = {g: _prettify(g) for g in order}
    return order, labels


def _get(df: Optional[pd.DataFrame], group: str, col: str) -> Optional[float]:
    if df is None:
        return None
    r = df[df["group"].astype(str) == group]
    return float(r[col].values[0]) if len(r) and col in r.columns else None


def _pct(v: float) -> str:
    return f"{v * 100:.1f}%"


def _save_fig(fig: plt.Figure, output_path: Optional[Path | str]) -> None:
    if output_path is None:
        return
    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=180, bbox_inches="tight", facecolor=fig.get_facecolor())


# ===========================================================================
# Data loading
# ===========================================================================

def discover_models(
    results_dir: str | Path,
    display_names: Optional[Dict[str, str]] = None,
    mode: str = "sliding",
) -> List[ModelEntry]:
    """
    Auto-discover models from a results directory tree.

    Supports two layouts produced by :class:`~badas.validation.Evaluator`:

    **New layout** (``run_eval.py`` style, preferred)::

        results_dir/
          {slug}/
            testset_{mode}/metrics_by_group.csv
            kaggle_{mode}/metrics_by_group.csv

    **Legacy layout**::

        results_dir/
          {slug}_testset/metrics_by_group.csv
          {slug}_kaggle/metrics_by_group.csv

    The new layout is detected first; legacy layout is used as fallback.

    Parameters
    ----------
    results_dir:
        Root directory to search.
    display_names:
        Optional ``{slug: label}`` mapping (slug is lower-cased with
        underscores).  Slugs not present here are auto-prettified.
    mode:
        Inference mode sub-directory to load (``"sliding"`` or
        ``"single_window"``).  Only used for the new layout.

    Returns
    -------
    list of (display_name, testset_df, kaggle_df)
        Sorted alphabetically by slug.  Either DataFrame may be ``None``
        when the corresponding CSV is missing.
    """
    dn = {k.lower().replace("-", "_"): v for k, v in (display_names or {}).items()}
    results_dir = Path(results_dir)
    models: List[ModelEntry] = []
    seen: set = set()

    for p in sorted(results_dir.iterdir()):
        if not p.is_dir():
            continue

        slug = p.name.lower().replace("-", "_")

        # ── New layout: p is the model dir, containing testset_{mode}/ ──────
        ts_csv_new = p / f"testset_{mode}" / "metrics_by_group.csv"
        kg_csv_new = p / f"kaggle_{mode}"  / "metrics_by_group.csv"

        if ts_csv_new.exists() or kg_csv_new.exists():
            if slug in seen:
                continue
            seen.add(slug)
            ts_df = pd.read_csv(ts_csv_new) if ts_csv_new.exists() else None
            kg_df = pd.read_csv(kg_csv_new) if kg_csv_new.exists() else None
            name  = dn.get(slug, _prettify(p.name))
            models.append((name, ts_df, kg_df))
            continue

        # ── Legacy layout: p may be {base}_testset or {base}_kaggle ─────────
        base = p.name
        for sfx in ("_testset", "_kaggle"):
            if base.endswith(sfx):
                base = base[: -len(sfx)]
                break
        base_slug = base.lower().replace("-", "_")
        if base_slug in seen:
            continue
        seen.add(base_slug)

        ts_csv = results_dir / f"{base}_testset" / "metrics_by_group.csv"
        kg_csv = results_dir / f"{base}_kaggle"  / "metrics_by_group.csv"
        ts_df  = pd.read_csv(ts_csv) if ts_csv.exists() else None
        kg_df  = pd.read_csv(kg_csv) if kg_csv.exists() else None

        if ts_df is not None or kg_df is not None:
            name = dn.get(base_slug, _prettify(base))
            models.append((name, ts_df, kg_df))

    return models


# ===========================================================================
# Figure 1 — Anthropic-style metrics table
# ===========================================================================

def plot_comparison_table(
    models: List[ModelEntry],
    output_path: Optional[str | Path] = None,
    group_labels: Optional[Dict[str, str]] = None,
) -> plt.Figure:
    """
    Render an Anthropic model-card–style metrics table.

    Rows = scenario groups (derived from the data), columns = one per model.
    Each cell shows F1 (top) and AP (bottom).  The featured model (last in
    ``models``) gets a rounded lavender border.  The best-value cell per
    metric row gets a lavender fill highlight — in any column, not just the
    featured one.

    Parameters
    ----------
    models:
        List of ``(display_name, testset_df, kaggle_df)`` entries.
    output_path:
        If given, the figure is saved here (PNG at 180 dpi).
    group_labels:
        Optional ``{group_key: display_label}`` override.  Missing keys are
        auto-prettified from the group name.

    Returns
    -------
    matplotlib.figure.Figure
    """
    metrics   = [("f1", "F1"), ("ap", "AP")]
    n_models  = len(models)
    n_metrics = len(metrics)
    featured  = n_models - 1

    group_order, auto_labels = _derive_groups(models)
    if group_labels:
        auto_labels.update(group_labels)
    n_groups = len(group_order)

    if n_groups == 0:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, "No testset data", ha="center", va="center")
        return fig

    # ── layout (axes-fraction coordinates) ───────────────────────────────────
    PAD_L   = 0.04
    LABEL_W = 0.30
    MODEL_W = (1.0 - LABEL_W) / n_models

    H_HEADER = 0.09
    H_ROW    = (1.0 - H_HEADER - 0.04) / n_groups
    H_METRIC = H_ROW / n_metrics
    TOP      = 1.0

    def cx(mj: int) -> float:
        return LABEL_W + (mj + 0.5) * MODEL_W

    def cell_y(gi: int, mi: int) -> float:
        row_top = TOP - H_HEADER - gi * H_ROW
        return row_top - (mi + 0.5) * H_METRIC

    # ── figure ───────────────────────────────────────────────────────────────
    fig = plt.figure(figsize=(11, 10))
    fig.patch.set_facecolor(_BG)

    model_names = " vs ".join(n for n, _, _ in models)
    fig.text(0.5, 0.970, f"Testset Evaluation  ·  {model_names}",
             ha="center", va="top", fontsize=14, fontweight="bold",
             color=_TEXT_DARK)

    ax = fig.add_axes([0.03, 0.03, 0.94, 0.87])
    ax.set_facecolor(_BG)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.axis("off")

    def hline(y: float, x0: float = 0.0, x1: float = 1.0,
              color: str = _LINE_COLOR, lw: float = 0.8) -> None:
        ax.plot([x0, x1], [y, y], color=color, lw=lw,
                transform=ax.transAxes, zorder=6, clip_on=False)

    # ── featured column border (stroke only, no fill) ─────────────────────
    fx0     = LABEL_W + featured * MODEL_W
    tbl_bot = TOP - H_HEADER - n_groups * H_ROW
    EXTEND  = 0.02
    PAD     = 0.012

    bx = fx0 + 0.006 + PAD
    by = tbl_bot - EXTEND + PAD
    bw = MODEL_W - 0.014 - 2 * PAD
    bh = (TOP - tbl_bot) + 2 * EXTEND - 2 * PAD
    ax.add_patch(mpatches.FancyBboxPatch(
        (bx, by), bw, bh,
        boxstyle=f"round,pad={PAD}",
        facecolor="none", edgecolor=_LAVEN,
        linewidth=2.0, transform=ax.transAxes, zorder=4,
        clip_on=False))

    # ── header ───────────────────────────────────────────────────────────────
    for mj, (name, _, _) in enumerate(models):
        is_feat = (mj == featured)
        ax.text(cx(mj), TOP - H_HEADER / 2, name,
                ha="center", va="center", fontsize=12,
                fontweight="bold" if is_feat else "normal",
                color=_LAVEN_DARK if is_feat else _TEXT_MID,
                transform=ax.transAxes, zorder=7)
    hline(TOP - H_HEADER, lw=1.2, color="#C8C8C8")

    # ── data rows ─────────────────────────────────────────────────────────
    for gi, group in enumerate(group_order):
        is_overall = group.upper() == "OVERALL"
        row_top = TOP - H_HEADER - gi * H_ROW

        if gi > 0:
            hline(row_top, color="#D4D4D4", lw=0.7)

        label   = auto_labels.get(group, _prettify(group))
        label_y = (cell_y(gi, 0) + cell_y(gi, n_metrics - 1)) / 2
        ax.text(PAD_L, label_y, label,
                ha="left", va="center",
                fontsize=13 if is_overall else 12,
                fontweight="bold" if is_overall else "normal",
                color=_LAVEN_DARK if is_overall else _TEXT_DARK,
                transform=ax.transAxes, zorder=7)

        for mi, (col, metric_label) in enumerate(metrics):
            y = cell_y(gi, mi)

            ax.text(LABEL_W - 0.012, y, metric_label,
                    ha="right", va="center", fontsize=8.5,
                    color=_TEXT_LIGHT, transform=ax.transAxes, zorder=7)

            vals  = [_get(df, group, col) for _, df, _ in models]
            valid = [(i, v) for i, v in enumerate(vals) if v is not None]
            best_i = max(valid, key=lambda t: t[1])[0] if valid else None

            if best_i is not None:
                cm = 0.005
                ax.add_patch(mpatches.FancyBboxPatch(
                    (LABEL_W + best_i * MODEL_W + cm, y - H_METRIC / 2 + cm),
                    MODEL_W - 2 * cm, H_METRIC - 2 * cm,
                    boxstyle="round,pad=0.005",
                    facecolor=_LAVEN_LIGHT, edgecolor="none",
                    transform=ax.transAxes, zorder=2))

            for mj, v in enumerate(vals):
                if v is None:
                    ax.text(cx(mj), y, "—", ha="center", va="center",
                            fontsize=12, color=_TEXT_LIGHT,
                            transform=ax.transAxes, zorder=7)
                    continue
                is_feat = (mj == featured)
                is_best = (mj == best_i)
                if is_best and is_feat:
                    color = _LAVEN_DARK
                elif is_best:
                    color = _TEXT_DARK
                else:
                    color = _TEXT_MID
                ax.text(cx(mj), y, _pct(v),
                        ha="center", va="center",
                        fontsize=14 if is_best else 13,
                        fontweight="bold" if is_best else "normal",
                        color=color, transform=ax.transAxes, zorder=7)

    _save_fig(fig, output_path)
    return fig


# ===========================================================================
# Figure 2 — Dumbbell: F1 by group
# ===========================================================================

def plot_f1_dumbbell(
    models: List[ModelEntry],
    output_path: Optional[str | Path] = None,
    group_labels: Optional[Dict[str, str]] = None,
) -> plt.Figure:
    """
    Horizontal dumbbell chart showing F1 for each scenario group.

    Parameters
    ----------
    models:
        List of ``(display_name, testset_df, kaggle_df)`` entries.
    output_path:
        If given, the figure is saved here (PNG at 180 dpi).
    group_labels:
        Optional ``{group_key: display_label}`` override.

    Returns
    -------
    matplotlib.figure.Figure
    """
    group_order, auto_labels = _derive_groups(models)
    if group_labels:
        auto_labels.update(group_labels)

    fig, ax = plt.subplots(figsize=(10, 5.5))
    fig.patch.set_facecolor(_BG)
    ax.set_facecolor(_BG)

    groups   = group_order[::-1]   # bottom-to-top: OVERALL at top
    labels   = [auto_labels.get(g, _prettify(g)) for g in groups]
    y_pos    = np.arange(len(groups))
    n_models = len(models)

    all_vals: List[float] = []
    for gi, group in enumerate(groups):
        vals = [_get(df, group, "f1") for _, df, _ in models]
        all_vals.extend(v for v in vals if v is not None)
        is_overall = group.upper() == "OVERALL"

        v0 = vals[0] if len(vals) > 0 else None
        v1 = vals[1] if len(vals) > 1 else None
        if v0 is not None and v1 is not None:
            ax.plot([v0, v1], [gi, gi], color=_LINE_COLOR,
                    lw=2.5 if is_overall else 1.8, zorder=1,
                    solid_capstyle="round")

        for mj, v in enumerate(vals):
            if v is None:
                continue
            color = _MODEL_COLORS[mj % len(_MODEL_COLORS)]
            ax.scatter(v, gi,
                       s=140 if is_overall else 100,
                       color=color, zorder=3,
                       edgecolors="white", linewidths=1.5)
            offset = 0.007 if mj == n_models - 1 else -0.007
            ha     = "left"  if mj == n_models - 1 else "right"
            ax.text(v + offset, gi, f"{v*100:.1f}",
                    va="center", ha=ha, fontsize=8.5,
                    fontweight="bold" if is_overall else "normal",
                    color=_TEXT_DARK)

    overall_idx = next((i for i, g in enumerate(groups) if g.upper() == "OVERALL"), None)
    if overall_idx is not None:
        ax.axhline(overall_idx - 0.5, color=_LINE_COLOR, lw=1.0, ls="--")

    x_min = max(0.60, min(all_vals) - 0.04) if all_vals else 0.60
    x_max = min(1.02, max(all_vals) + 0.06) if all_vals else 1.02
    ax.set_xlim(x_min, x_max)
    ax.set_ylim(-0.6, len(groups) - 0.4)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=10)
    ax.set_xlabel("F1 Score", fontsize=10, color=_TEXT_MID)
    ax.xaxis.set_major_formatter(
        matplotlib.ticker.FuncFormatter(lambda v, _: f"{v:.2f}"))
    ax.tick_params(axis="x", labelsize=8.5, colors=_TEXT_MID)
    ax.tick_params(axis="y", labelsize=10,  colors=_TEXT_DARK)
    ax.spines[["left", "top", "right"]].set_visible(False)
    ax.spines["bottom"].set_color(_LINE_COLOR)
    ax.grid(axis="x", color=_LINE_COLOR, lw=0.6, zorder=0)

    handles = [mpatches.Patch(color=_MODEL_COLORS[i % len(_MODEL_COLORS)], label=n)
               for i, (n, _, _) in enumerate(models)]
    ax.legend(handles=handles, fontsize=9, framealpha=0, loc="lower right")
    ax.set_title("F1 Score by Scenario Group", fontsize=13,
                 fontweight="bold", color=_TEXT_DARK, pad=12)

    fig.tight_layout()
    _save_fig(fig, output_path)
    return fig


# ===========================================================================
# Figure 3 — Kaggle results
# ===========================================================================

def plot_kaggle(
    models: List[ModelEntry],
    output_path: Optional[str | Path] = None,
) -> plt.Figure:
    """
    Grouped bar chart for Kaggle AP metrics.

    Splits (Public / Private / Overall) are derived from the kaggle dataframe.

    Parameters
    ----------
    models:
        List of ``(display_name, testset_df, kaggle_df)`` entries.
    output_path:
        If given, the figure is saved here (PNG at 180 dpi).

    Returns
    -------
    matplotlib.figure.Figure
    """
    # Derive splits from the kaggle dataframe
    kg_order, kg_labels = [], {}
    for _, _, kg_df in models:
        if kg_df is not None:
            kg_order, kg_labels = _groups_from_df(kg_df)
            break

    metrics = [
        ("mean_ap", "mAP"),
        ("ap_0.5s", "AP@0.5s"),
        ("ap_1.0s", "AP@1.0s"),
        ("ap_1.5s", "AP@1.5s"),
    ]
    n_models  = len(models)
    n_splits  = len(kg_order)
    n_metrics = len(metrics)

    if n_splits == 0:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, "No Kaggle data", ha="center", va="center")
        return fig

    data = np.full((n_models, n_splits, n_metrics), np.nan)
    for mi, (_, _, kg_df) in enumerate(models):
        for si, split_key in enumerate(kg_order):
            for meti, (col, _) in enumerate(metrics):
                v = _get(kg_df, split_key, col)
                if v is not None:
                    data[mi, si, meti] = v

    fig, axes = plt.subplots(1, n_splits, figsize=(13, 5.5), sharey=True)
    fig.patch.set_facecolor(_BG)
    if n_splits == 1:
        axes = [axes]

    bar_w    = 0.65 / n_models
    x        = np.arange(n_metrics)
    all_vals = data[~np.isnan(data)]
    y_min    = max(0.84, np.nanmin(all_vals) - 0.02) if len(all_vals) else 0.84
    y_max    = min(1.01, np.nanmax(all_vals) + 0.025) if len(all_vals) else 1.00

    for si, (ax, split_key) in enumerate(zip(axes, kg_order)):
        ax.set_facecolor(_BG)
        is_overall = split_key.upper() == "OVERALL"
        split_label = kg_labels.get(split_key, _prettify(split_key))

        for mj, ((name, _, _), color) in enumerate(
                zip(models, _MODEL_COLORS * ((n_models // len(_MODEL_COLORS)) + 1))):
            offset = (mj - (n_models - 1) / 2) * bar_w
            vals   = data[mj, si]
            bars   = ax.bar(x + offset, vals, width=bar_w * 0.88,
                            color=color, alpha=0.90, label=name,
                            edgecolor="white", linewidth=0.5)
            for bar, v in zip(bars, vals):
                if not np.isnan(v):
                    ax.text(bar.get_x() + bar.get_width() / 2,
                            bar.get_height() + 0.002,
                            f"{v:.3f}", ha="center", va="bottom",
                            fontsize=7.5, color=_TEXT_DARK,
                            fontweight="bold" if is_overall else "normal")

        ax.set_xticks(x)
        ax.set_xticklabels([lbl for _, lbl in metrics], fontsize=9)
        ax.set_ylim(y_min, y_max)
        ax.yaxis.set_major_formatter(matplotlib.ticker.FormatStrFormatter("%.3f"))
        ax.tick_params(axis="y", labelsize=8)
        ax.spines[["top", "right"]].set_visible(False)
        ax.spines[["bottom", "left"]].set_color(_LINE_COLOR)
        ax.grid(axis="y", color=_LINE_COLOR, lw=0.6, zorder=0)
        ax.set_title(split_label, fontsize=11,
                     fontweight="bold" if is_overall else "normal",
                     color=_LAVEN_DARK if is_overall else _TEXT_DARK, pad=8)
        if si == 0:
            ax.set_ylabel("Average Precision", fontsize=9, color=_TEXT_MID)
            ax.legend(fontsize=8.5, framealpha=0, loc="lower right")

    fig.suptitle("Kaggle Competition Results", fontsize=13,
                 fontweight="bold", color=_TEXT_DARK, y=1.02)
    fig.tight_layout(w_pad=1.5)
    _save_fig(fig, output_path)
    return fig


# ===========================================================================
# Figure 4 — Recall gain vs FPR reduction
# ===========================================================================

def plot_recall_fpr(
    models: List[ModelEntry],
    output_path: Optional[str | Path] = None,
    group_labels: Optional[Dict[str, str]] = None,
) -> plt.Figure:
    """
    Side-by-side delta bars: recall improvement and FPR reduction vs baseline.

    ``models[0]`` is the baseline; ``models[1]`` is the new model.

    Parameters
    ----------
    models:
        List of ``(display_name, testset_df, kaggle_df)`` entries.
    output_path:
        If given, the figure is saved here (PNG at 180 dpi).
    group_labels:
        Optional ``{group_key: display_label}`` override.

    Returns
    -------
    matplotlib.figure.Figure

    Raises
    ------
    ValueError
        If fewer than two models are provided.
    """
    if len(models) < 2:
        raise ValueError("plot_recall_fpr requires at least 2 models.")

    baseline_name, baseline_df, _ = models[0]
    new_name,      new_df,      _ = models[1]

    group_order, auto_labels = _derive_groups(models)
    if group_labels:
        auto_labels.update(group_labels)

    labels = [auto_labels.get(g, _prettify(g)) for g in group_order]
    y_pos  = np.arange(len(group_order))

    recall_base = np.array([_get(baseline_df, g, "recall") or 0 for g in group_order])
    recall_new  = np.array([_get(new_df,      g, "recall") or 0 for g in group_order])
    fpr_base    = np.array([_get(baseline_df, g, "fpr")    or 0 for g in group_order])
    fpr_new     = np.array([_get(new_df,      g, "fpr")    or 0 for g in group_order])
    d_recall = recall_new - recall_base
    d_fpr    = fpr_base   - fpr_new

    fig, (ax_r, ax_f) = plt.subplots(1, 2, figsize=(13, 5.5))
    fig.patch.set_facecolor(_BG)

    def _delta_bars(ax: plt.Axes, deltas: np.ndarray,
                    title: str, xlabel: str) -> None:
        ax.set_facecolor(_BG)
        colors = [
            _LAVEN if d > 0.001 else (_LINE_COLOR if abs(d) <= 0.001 else "#EF4444")
            for d in deltas
        ]
        bars = ax.barh(y_pos, deltas * 100, color=colors, alpha=0.90,
                       edgecolor="white", height=0.55, zorder=2)
        for bar, d, gi in zip(bars, deltas, range(len(group_order))):
            w  = bar.get_width()
            ha = "left" if w >= 0 else "right"
            offset = 0.15 if w >= 0 else -0.15
            ax.text(w + offset, bar.get_y() + bar.get_height() / 2,
                    f"{d*100:+.1f}pp", va="center", ha=ha, fontsize=9,
                    fontweight="bold" if group_order[gi].upper() == "OVERALL" else "normal",
                    color=_TEXT_DARK)

        ax.axvline(0, color=_TEXT_MID, lw=1.0, zorder=3)
        overall_idx = next(
            (i for i, g in enumerate(group_order) if g.upper() == "OVERALL"), None)
        if overall_idx is not None:
            ax.axhline(overall_idx - 0.5, color=_LINE_COLOR, lw=1.0, ls="--")
        ax.set_yticks(y_pos)
        ax.set_yticklabels(labels, fontsize=10)
        ax.set_xlabel(xlabel, fontsize=9, color=_TEXT_MID)
        ax.spines[["top", "right", "left"]].set_visible(False)
        ax.spines["bottom"].set_color(_LINE_COLOR)
        ax.grid(axis="x", color=_LINE_COLOR, lw=0.6, zorder=0)
        ax.tick_params(axis="x", labelsize=8.5, colors=_TEXT_MID)
        ax.tick_params(axis="y", colors=_TEXT_DARK)
        ax.set_title(title, fontsize=12, fontweight="bold",
                     color=_TEXT_DARK, pad=10)

    _delta_bars(ax_r, d_recall, "Recall Improvement",
                "Δ Recall (percentage points)")
    _delta_bars(ax_f, d_fpr,   "False Positive Rate Reduction",
                "Δ FPR (pp, positive = fewer false alarms)")

    fig.suptitle(f"Where does the new model win?  ·  {new_name}  vs  {baseline_name}",
                 fontsize=13, fontweight="bold", color=_TEXT_DARK, y=1.02)
    fig.tight_layout(w_pad=2.0)
    _save_fig(fig, output_path)
    return fig


# ===========================================================================
# Figure 5 — 2×2 overview (blog hero)
# ===========================================================================

def plot_overview(
    models: List[ModelEntry],
    output_path: Optional[str | Path] = None,
    group_labels: Optional[Dict[str, str]] = None,
) -> plt.Figure:
    """
    2×2 summary hero figure: F1 dumbbell | Kaggle mAP | Recall Δ | FPR Δ.

    ``models[0]`` is the baseline; ``models[1]`` is the new model.

    Parameters
    ----------
    models:
        List of ``(display_name, testset_df, kaggle_df)`` entries.
    output_path:
        If given, the figure is saved here (PNG at 180 dpi).
    group_labels:
        Optional ``{group_key: display_label}`` override.

    Returns
    -------
    matplotlib.figure.Figure

    Raises
    ------
    ValueError
        If fewer than two models are provided.
    """
    if len(models) < 2:
        raise ValueError("plot_overview requires at least 2 models.")

    baseline_name, baseline_df, baseline_kg = models[0]
    new_name,      new_df,      new_kg      = models[1]
    n_models = len(models)

    group_order, auto_labels = _derive_groups(models)
    if group_labels:
        auto_labels.update(group_labels)
    glabels  = [auto_labels.get(g, _prettify(g)) for g in group_order]
    y_pos    = np.arange(len(group_order))

    # kaggle splits
    kg_order, kg_labels = [], {}
    for _, _, kg_df in models:
        if kg_df is not None:
            kg_order, kg_labels = _groups_from_df(kg_df)
            break

    fig = plt.figure(figsize=(16, 11))
    fig.patch.set_facecolor(_BG)
    gs = GridSpec(2, 2, figure=fig, hspace=0.42, wspace=0.32,
                  left=0.07, right=0.97, top=0.91, bottom=0.07)

    # ── Panel A: F1 dumbbell ─────────────────────────────────────────────────
    ax_a = fig.add_subplot(gs[0, 0])
    ax_a.set_facecolor(_BG)
    g_rev = group_order[::-1]
    l_rev = [auto_labels.get(g, _prettify(g)) for g in g_rev]
    all_f1: List[float] = []
    for gi, group in enumerate(g_rev):
        vals = [_get(df, group, "f1") for _, df, _ in models]
        all_f1.extend(v for v in vals if v is not None)
        is_overall = group.upper() == "OVERALL"
        v0 = vals[0] if len(vals) > 0 else None
        v1 = vals[1] if len(vals) > 1 else None
        if v0 is not None and v1 is not None:
            ax_a.plot([v0, v1], [gi, gi], color=_LINE_COLOR,
                      lw=2.0 if is_overall else 1.5, zorder=1)
        for mj, v in enumerate(vals):
            if v is None:
                continue
            ax_a.scatter(v, gi,
                         s=120 if is_overall else 80,
                         color=_MODEL_COLORS[mj % len(_MODEL_COLORS)],
                         zorder=3, edgecolors="white", linewidths=1.5)
    xmin = max(0.60, min(all_f1) - 0.04) if all_f1 else 0.60
    xmax = min(1.02, max(all_f1) + 0.07) if all_f1 else 1.02
    ax_a.set_xlim(xmin, xmax)
    ax_a.set_ylim(-0.6, len(group_order) - 0.4)
    ax_a.set_yticks(np.arange(len(group_order)))
    ax_a.set_yticklabels(l_rev, fontsize=9)
    overall_rev_idx = next(
        (i for i, g in enumerate(g_rev) if g.upper() == "OVERALL"), None)
    if overall_rev_idx is not None:
        ax_a.axhline(overall_rev_idx - 0.5, color=_LINE_COLOR, lw=1.0, ls="--")
    ax_a.xaxis.set_major_formatter(matplotlib.ticker.FormatStrFormatter("%.2f"))
    ax_a.tick_params(axis="x", labelsize=8, colors=_TEXT_MID)
    ax_a.spines[["left", "top", "right"]].set_visible(False)
    ax_a.spines["bottom"].set_color(_LINE_COLOR)
    ax_a.grid(axis="x", color=_LINE_COLOR, lw=0.5, zorder=0)
    ax_a.set_title("F1 Score by Group", fontsize=11,
                   fontweight="bold", color=_TEXT_DARK, pad=8)
    handles_a = [mpatches.Patch(color=_MODEL_COLORS[i % len(_MODEL_COLORS)], label=n)
                 for i, (n, _, _) in enumerate(models)]
    ax_a.legend(handles=handles_a, fontsize=8, framealpha=0, loc="lower right")

    # ── Panel B: Kaggle mAP ──────────────────────────────────────────────────
    ax_b = fig.add_subplot(gs[0, 1])
    ax_b.set_facecolor(_BG)
    x     = np.arange(len(kg_order))
    bar_w = 0.32
    kg_vals: List[float] = []
    for mj, (name, _, kg_df) in enumerate(models):
        vals = [_get(kg_df, sk, "mean_ap") or 0 for sk in kg_order]
        kg_vals.extend(vals)
        offset = (mj - (n_models - 1) / 2) * bar_w
        bars = ax_b.bar(x + offset, vals, width=bar_w * 0.88,
                        color=_MODEL_COLORS[mj % len(_MODEL_COLORS)], alpha=0.90,
                        label=name, edgecolor="white")
        for bar, v in zip(bars, vals):
            ax_b.text(bar.get_x() + bar.get_width() / 2,
                      bar.get_height() + 0.001, f"{v:.3f}",
                      ha="center", va="bottom", fontsize=8, color=_TEXT_DARK)
    y_min_b = max(0.89, min(kg_vals) - 0.015) if kg_vals else 0.89
    y_max_b = min(1.01, max(kg_vals) + 0.02)  if kg_vals else 1.00
    ax_b.set_ylim(y_min_b, y_max_b)
    ax_b.set_xticks(x)
    ax_b.set_xticklabels([kg_labels.get(k, _prettify(k)) for k in kg_order], fontsize=9)
    ax_b.yaxis.set_major_formatter(matplotlib.ticker.FormatStrFormatter("%.3f"))
    ax_b.tick_params(axis="y", labelsize=8)
    ax_b.spines[["top", "right"]].set_visible(False)
    ax_b.spines[["bottom", "left"]].set_color(_LINE_COLOR)
    ax_b.grid(axis="y", color=_LINE_COLOR, lw=0.5, zorder=0)
    ax_b.set_title("Kaggle mAP", fontsize=11,
                   fontweight="bold", color=_TEXT_DARK, pad=8)
    ax_b.legend(fontsize=8, framealpha=0, loc="lower right")

    # ── Panel C: Recall Δ ────────────────────────────────────────────────────
    ax_c = fig.add_subplot(gs[1, 0])
    ax_c.set_facecolor(_BG)
    recall_base = np.array([_get(baseline_df, g, "recall") or 0 for g in group_order])
    recall_new  = np.array([_get(new_df,      g, "recall") or 0 for g in group_order])
    d_recall    = recall_new - recall_base
    colors_r    = [_LAVEN if d > 0.001 else (_LINE_COLOR if abs(d) <= 0.001 else "#EF4444")
                   for d in d_recall]
    ax_c.barh(y_pos, d_recall * 100, color=colors_r, alpha=0.90,
              edgecolor="white", height=0.55, zorder=2)
    for gi, d in enumerate(d_recall):
        offset = 0.12 if d >= 0 else -0.12
        ax_c.text(d * 100 + offset, gi, f"{d*100:+.1f}pp",
                  va="center", ha="left" if d >= 0 else "right", fontsize=9,
                  fontweight="bold" if group_order[gi].upper() == "OVERALL" else "normal",
                  color=_TEXT_DARK)
    ax_c.axvline(0, color=_TEXT_MID, lw=1.0, zorder=3)
    if overall_idx := next((i for i, g in enumerate(group_order) if g.upper() == "OVERALL"), None):
        ax_c.axhline(overall_idx - 0.5, color=_LINE_COLOR, lw=1.0, ls="--")
    ax_c.set_yticks(y_pos)
    ax_c.set_yticklabels(glabels, fontsize=9)
    ax_c.spines[["top", "right", "left"]].set_visible(False)
    ax_c.spines["bottom"].set_color(_LINE_COLOR)
    ax_c.grid(axis="x", color=_LINE_COLOR, lw=0.5, zorder=0)
    ax_c.tick_params(axis="x", labelsize=8, colors=_TEXT_MID)
    ax_c.set_xlabel("Δ Recall (pp)", fontsize=9, color=_TEXT_MID)
    ax_c.set_title("Recall Improvement", fontsize=11,
                   fontweight="bold", color=_TEXT_DARK, pad=8)

    # ── Panel D: FPR Δ ──────────────────────────────────────────────────────
    ax_d = fig.add_subplot(gs[1, 1])
    ax_d.set_facecolor(_BG)
    fpr_base = np.array([_get(baseline_df, g, "fpr") or 0 for g in group_order])
    fpr_new  = np.array([_get(new_df,      g, "fpr") or 0 for g in group_order])
    d_fpr    = fpr_base - fpr_new
    colors_f = [_LAVEN if d > 0.001 else (_LINE_COLOR if abs(d) <= 0.001 else "#EF4444")
                for d in d_fpr]
    ax_d.barh(y_pos, d_fpr * 100, color=colors_f, alpha=0.90,
              edgecolor="white", height=0.55, zorder=2)
    for gi, d in enumerate(d_fpr):
        offset = 0.08 if d >= 0 else -0.08
        ax_d.text(d * 100 + offset, gi, f"{d*100:+.1f}pp",
                  va="center", ha="left" if d >= 0 else "right", fontsize=9,
                  fontweight="bold" if group_order[gi].upper() == "OVERALL" else "normal",
                  color=_TEXT_DARK)
    ax_d.axvline(0, color=_TEXT_MID, lw=1.0, zorder=3)
    if overall_idx is not None:
        ax_d.axhline(overall_idx - 0.5, color=_LINE_COLOR, lw=1.0, ls="--")
    ax_d.set_yticks(y_pos)
    ax_d.set_yticklabels(glabels, fontsize=9)
    ax_d.spines[["top", "right", "left"]].set_visible(False)
    ax_d.spines["bottom"].set_color(_LINE_COLOR)
    ax_d.grid(axis="x", color=_LINE_COLOR, lw=0.5, zorder=0)
    ax_d.tick_params(axis="x", labelsize=8, colors=_TEXT_MID)
    ax_d.set_xlabel("Δ FPR (pp, positive = fewer false alarms)", fontsize=9,
                    color=_TEXT_MID)
    ax_d.set_title("False Positive Rate Reduction", fontsize=11,
                   fontweight="bold", color=_TEXT_DARK, pad=8)

    fig.suptitle(f"Model Comparison  ·  {new_name}  vs  {baseline_name}",
                 fontsize=14, fontweight="bold", color=_TEXT_DARK, y=0.97)
    _save_fig(fig, output_path)
    return fig


# ===========================================================================
# Figure 6 — AUC-ROC by scenario group (all models)
# ===========================================================================

_DARK_BG     = "#1e1e2e"
_DARK_GRID   = "#333355"
_DARK_BORDER = "#555555"
_DARK_TEXT   = "#cccccc"
_DARK_TITLE  = "#ffffff"
_DARK_NOTE   = "#888888"

_BAR_PALETTE = [
    "#888888",  # grey    (baseline / BADAS-1.0)
    "#4C72B0",  # blue    (BADAS-1.5)
    "#55A868",  # green   (Cosmos-BADAS)
    "#C44E52",  # red     (Cosmos Vanilla)
    "#DD8452",  # orange  (Qwen3-VL)
    "#8172B2",  # purple
    "#937860",  # brown
]


def plot_auc_by_group(
    models: List[ModelEntry],
    output_path: Optional[str | Path] = None,
    group_labels: Optional[Dict[str, str]] = None,
    colors: Optional[List[str]] = None,
) -> plt.Figure:
    """
    Grouped bar chart of AUC-ROC across scenario groups for all models.

    Uses the testset DataFrame from each :data:`ModelEntry`.

    Parameters
    ----------
    models:
        List of ``(display_name, testset_df, kaggle_df)`` entries.
    output_path:
        If given, the figure is saved here (PNG at 180 dpi).
    group_labels:
        Optional ``{group_key: display_label}`` override.
    colors:
        Optional list of bar colours (one per model).  Defaults to the
        built-in :data:`_BAR_PALETTE`.

    Returns
    -------
    matplotlib.figure.Figure
    """
    group_order, auto_labels = _derive_groups(models)
    if group_labels:
        auto_labels.update(group_labels)
    # Exclude OVERALL from bar chart — too compressed to be useful
    plot_groups = [g for g in group_order if g.upper() != "OVERALL"]
    x_labels    = [auto_labels.get(g, _prettify(g)) for g in plot_groups]

    palette     = colors or _BAR_PALETTE
    n_models    = len(models)
    x           = np.arange(len(plot_groups))
    bar_w       = 0.72 / max(n_models, 1)
    offsets     = np.linspace(-(n_models - 1) / 2, (n_models - 1) / 2, n_models) * bar_w

    fig, ax = plt.subplots(figsize=(10, 5))
    fig.patch.set_facecolor(_DARK_BG)
    ax.set_facecolor(_DARK_BG)

    for i, (name, ts_df, _) in enumerate(models):
        vals = [_get(ts_df, g, "auc_roc") or 0.0 for g in plot_groups]
        ax.bar(x + offsets[i], vals, bar_w * 0.9,
               label=name, color=palette[i % len(palette)], zorder=3)

    ax.set_xticks(x)
    ax.set_xticklabels(x_labels, color=_DARK_TEXT, fontsize=11)
    ax.set_ylabel("AUC-ROC", color=_DARK_TEXT, fontsize=11)
    ax.set_title("Collision Anticipation — AUC-ROC by Scenario Group",
                 color=_DARK_TITLE, fontsize=13, pad=10)
    ax.set_ylim(0.5, 1.05)
    ax.tick_params(colors=_DARK_TEXT)
    ax.spines[:].set_color(_DARK_BORDER)
    ax.grid(axis="y", color=_DARK_GRID, zorder=0)
    ax.axhline(0.5, color=_DARK_BORDER, linewidth=0.8, linestyle="--", zorder=2)
    ax.legend(loc="lower right", framealpha=0.3, labelcolor=_DARK_TEXT,
              facecolor="#2a2a3e", edgecolor=_DARK_BORDER, fontsize=9)

    plt.tight_layout()
    _save_fig(fig, output_path)
    return fig


# ===========================================================================
# Figure 7 — F1 + MTTA side-by-side (2+ models, any group set)
# ===========================================================================

def plot_f1_mtta(
    models: List[ModelEntry],
    output_path: Optional[str | Path] = None,
    group_labels: Optional[Dict[str, str]] = None,
    colors: Optional[List[str]] = None,
) -> plt.Figure:
    """
    Side-by-side bar chart of F1 and MTTA across scenario groups.

    F1 and MTTA are read from the testset DataFrame (columns ``f1`` and
    ``mtta``), which are computed at the evaluation threshold (default 0.75)
    using sustained-alert logic.  Missing MTTA values are shown as N/A.

    Parameters
    ----------
    models:
        List of ``(display_name, testset_df, kaggle_df)`` entries.
    output_path:
        If given, the figure is saved here (PNG at 180 dpi).
    group_labels:
        Optional ``{group_key: display_label}`` override.
    colors:
        Optional list of bar colours (one per model).

    Returns
    -------
    matplotlib.figure.Figure
    """
    group_order, auto_labels = _derive_groups(models)
    if group_labels:
        auto_labels.update(group_labels)
    plot_groups = [g for g in group_order if g.upper() != "OVERALL"]
    x_labels    = [auto_labels.get(g, _prettify(g)) for g in plot_groups]

    palette  = colors or _BAR_PALETTE
    n_models = len(models)
    x        = np.arange(len(plot_groups))
    bar_w    = 0.72 / max(n_models, 1)
    offsets  = np.linspace(-(n_models - 1) / 2, (n_models - 1) / 2, n_models) * bar_w

    fig, axes = plt.subplots(1, 2, figsize=(13, 5))
    fig.patch.set_facecolor(_DARK_BG)
    for ax in axes:
        ax.set_facecolor(_DARK_BG)

    for (col, title, ylabel, ylim), ax in zip(
        [
            ("f1",   "F1 Score",           "F1  (threshold = 0.75)", (0.0, 1.1)),
            ("mtta", "Mean Time-To-Alert",  "Seconds before event",   (0.0, 3.2)),
        ],
        axes,
    ):
        for i, (name, ts_df, _) in enumerate(models):
            vals  = [_get(ts_df, g, col) for g in plot_groups]
            finite = [v if v is not None and not np.isnan(v) else 0.0 for v in vals]
            bars  = ax.bar(x + offsets[i], finite, bar_w * 0.9,
                           label=name, color=palette[i % len(palette)], zorder=3)
            # value labels
            for bar, v in zip(bars, vals):
                if v is None or np.isnan(v):
                    ax.text(bar.get_x() + bar.get_width() / 2, 0.05,
                            "N/A", ha="center", va="bottom",
                            fontsize=6, color="#ff8888", rotation=90)
                elif v > 0:
                    ax.text(bar.get_x() + bar.get_width() / 2,
                            bar.get_height() + 0.02,
                            f"{v:.2f}", ha="center", va="bottom",
                            fontsize=7, color=_DARK_TEXT)

        ax.set_xticks(x)
        ax.set_xticklabels(x_labels, color=_DARK_TEXT, fontsize=10)
        ax.set_ylabel(ylabel, color=_DARK_TEXT, fontsize=10)
        ax.set_title(title, color=_DARK_TITLE, fontsize=12, pad=10)
        ax.set_ylim(*ylim)
        ax.tick_params(colors=_DARK_TEXT)
        ax.spines[:].set_color(_DARK_BORDER)
        ax.grid(axis="y", color=_DARK_GRID, zorder=0)
        ax.legend(loc="lower right" if col == "f1" else "upper left",
                  framealpha=0.3, labelcolor=_DARK_TEXT,
                  facecolor="#2a2a3e", edgecolor=_DARK_BORDER, fontsize=8)

    names = " vs ".join(n for n, _, _ in models)
    fig.suptitle(f"Deep Dive — {names}", color=_DARK_TITLE, fontsize=13, y=1.01)
    fig.text(
        0.5, -0.02,
        "Metrics at threshold = 0.75 with sustained-alert logic "
        "(up to 8 sub-threshold frames tolerated).  N/A = no alert before event.",
        ha="center", fontsize=8, color=_DARK_NOTE,
    )
    plt.tight_layout()
    _save_fig(fig, output_path)
    return fig


# ===========================================================================
# Convenience: generate all figures at once
# ===========================================================================

def generate_comparison_report(
    models: List[ModelEntry],
    output_dir: str | Path,
    group_labels: Optional[Dict[str, str]] = None,
) -> Dict[str, Path]:
    """
    Generate all comparison figures and save them to ``output_dir``.

    Parameters
    ----------
    models:
        List of ``(display_name, testset_df, kaggle_df)`` entries, typically
        from :func:`discover_models`.
    output_dir:
        Directory where figures are written.  Created if it does not exist.
    group_labels:
        Optional ``{group_key: display_label}`` override passed to all figures.

    Returns
    -------
    dict mapping figure name → saved path
        Keys: ``"auc_by_group"``, ``"f1_mtta"``, ``"table"``,
        ``"f1_dumbbell"``, ``"kaggle"``, ``"recall_fpr"``, ``"overview"``.
        A key is omitted when data are insufficient to render that figure.
    """
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)

    has_testset = any(m[1] is not None for m in models)
    has_kaggle  = any(m[2] is not None for m in models)
    saved: Dict[str, Path] = {}

    if has_testset:
        p = out / "auc_by_group.png"
        plot_auc_by_group(models, output_path=p, group_labels=group_labels)
        plt.close("all")
        saved["auc_by_group"] = p

        p = out / "f1_mtta.png"
        plot_f1_mtta(models, output_path=p, group_labels=group_labels)
        plt.close("all")
        saved["f1_mtta"] = p

        p = out / "table.png"
        plot_comparison_table(models, output_path=p, group_labels=group_labels)
        plt.close("all")
        saved["table"] = p

        p = out / "f1_dumbbell.png"
        plot_f1_dumbbell(models, output_path=p, group_labels=group_labels)
        plt.close("all")
        saved["f1_dumbbell"] = p

        if len(models) >= 2:
            p = out / "recall_fpr.png"
            plot_recall_fpr(models, output_path=p, group_labels=group_labels)
            plt.close("all")
            saved["recall_fpr"] = p

    if has_kaggle:
        p = out / "kaggle.png"
        plot_kaggle(models, output_path=p)
        plt.close("all")
        saved["kaggle"] = p

    if has_testset and has_kaggle and len(models) >= 2:
        p = out / "overview.png"
        plot_overview(models, output_path=p, group_labels=group_labels)
        plt.close("all")
        saved["overview"] = p

    return saved
