#!/usr/bin/env python3
"""
Post-hoc explainability for a full saved run (all folds).

Phase 1 — MD permutation feature importance (bar chart + CSV, per fold).
Phase 2 — GradCAM overlays on val/holdout patients (per fold).
Phase 3 — Fusion event analysis loaded entirely from saved prediction files
           (no re-inference needed). Runs on both train and val splits.

Usage:
    python scripts/output_analysis/explainability/explain_run.py \
        --run-dir analysis_data/v2.3_single_binary_nocrop/binary/single \
        [--checkpoint best_single.pt | best_holdout_single.pt] \
        [--split holdout]
        [--image-dir Papila/FundusImages] \
        [--clinical-dir Papila/ClinicalData] \
        [--n-permutations 30] \
        [--seed 0] \
        [--alpha 0.45]
"""
from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path
from types import SimpleNamespace

import matplotlib

matplotlib.use("Agg")
import matplotlib.cm as cm
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from PIL import Image

REPO_ROOT = Path(__file__).resolve().parents[3]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from classes.v2.data_bundle import DataBundle
from classes.v2.papila_builders import build_papila_data
from classes.v2.profiles.papila import build_papila_profile
from classes.v2.split_manager import PatientFirstSplitManager
from classes.v2.loader_factory import filter_bilateral_samples, make_loader
from classes.v2.metrics import _score_arrays
from classes.v2.models import SingleEyeHT
from classes.v2.transforms import build_eval_transform

# ---------------------------------------------------------------------------
# Label display helpers
# ---------------------------------------------------------------------------

BINARY_LABELS = {0: "Normal", 1: "Glaucoma"}
MULTICLASS_LABELS = {0: "Normal", 1: "Glaucoma", 2: "Suspect"}


def label_name(label: int, eval_mode: str) -> str:
    mapping = BINARY_LABELS if eval_mode == "binary" else MULTICLASS_LABELS
    return mapping.get(int(label), str(label))


# ---------------------------------------------------------------------------
# GradCAM
# ---------------------------------------------------------------------------


class GradCAM:
    """Minimal GradCAM using forward/backward hooks. No extra dependencies."""

    def __init__(self, target_layer: torch.nn.Module) -> None:
        self._acts: torch.Tensor | None = None
        self._grads: torch.Tensor | None = None
        self._h1 = target_layer.register_forward_hook(self._save_acts)
        self._h2 = target_layer.register_full_backward_hook(self._save_grads)

    def _save_acts(self, _m, _i, output):
        self._acts = output.detach()

    def _save_grads(self, _m, _gi, grad_output):
        self._grads = grad_output[0].detach()

    def compute(
        self,
        img: torch.Tensor,
        meta: torch.Tensor,
        model: torch.nn.Module,
        target_class: int | None = None,
    ) -> tuple[np.ndarray, int]:
        model.eval()
        with torch.enable_grad():
            out = model(img, meta)
            pred = int(out.argmax(1).item())
            tc = pred if target_class is None else target_class
            model.zero_grad()
            out[0, tc].backward()

        if self._acts is None or self._grads is None:
            raise RuntimeError("GradCAM hooks did not fire — check target_layer.")

        weights = self._grads.mean(dim=(2, 3), keepdim=True)
        cam = F.relu((weights * self._acts).sum(dim=1, keepdim=True))
        cam = F.interpolate(cam, img.shape[-2:], mode="bilinear", align_corners=False)
        cam_np = cam.squeeze().cpu().numpy()
        lo, hi = cam_np.min(), cam_np.max()
        cam_np = (cam_np - lo) / (hi - lo + 1e-8)
        return cam_np, pred

    def remove(self) -> None:
        self._h1.remove()
        self._h2.remove()


def get_gradcam_layer(model: SingleEyeHT, backbone: str) -> torch.nn.Module:
    bb = model.img_tower.backbone
    key = backbone.lower()
    if key in ("refugelike",) or "resnet" in key:
        return bb.layer4[-1]
    if "efficientnet" in key or "refuge_efficient" in key:
        return bb.features[-1]
    if "densenet" in key or key == "refuge_densenet":
        return bb.features.denseblock4
    if "mobilenet" in key:
        return bb.features[-1]
    if "vgg" in key:
        return bb.features[-1]
    raise ValueError(f"Unknown backbone for GradCAM target layer: {backbone!r}")


def overlay_gradcam(
    original_pil: "Image.Image", cam: np.ndarray, alpha: float = 0.45
) -> "Image.Image":
    cam_u8 = (cam * 255).astype(np.uint8)
    cam_resized = (
        np.array(Image.fromarray(cam_u8).resize(original_pil.size, Image.BILINEAR))
        / 255.0
    )
    colored = (cm.jet(cam_resized)[:, :, :3] * 255).astype(np.uint8)
    return Image.blend(original_pil.convert("RGB"), Image.fromarray(colored), alpha)


# ---------------------------------------------------------------------------
# Feature index map
# ---------------------------------------------------------------------------


def build_feature_index_map(data: DataBundle) -> dict[str, dict]:
    n_scalar = len(data.scalar_cols)
    cat_expanded = sum(len(m) for m in data.cat_maps.values())

    feature_map: dict[str, dict] = {}
    for i, col in enumerate(data.scalar_cols):
        missing_dim = n_scalar + cat_expanded + i
        feature_map[col] = {"value_dims": [i], "missing_dims": [missing_dim]}

    cat_offset = n_scalar
    for col in data.cat_cols:
        n_cats = len(data.cat_maps[col])
        dims = list(range(cat_offset, cat_offset + n_cats))
        feature_map[col] = {"value_dims": dims, "missing_dims": []}
        cat_offset += n_cats

    return feature_map


# ---------------------------------------------------------------------------
# Phase 1 — MD permutation importance
# ---------------------------------------------------------------------------


def run_permutation_importance(
    model: SingleEyeHT,
    loader,
    data: DataBundle,
    num_classes: int,
    device: torch.device,
    n_permutations: int,
    seed: int,
    out_dir: Path,
) -> None:
    print("\n[Phase 1] MD permutation importance ...", flush=True)

    img1_feats_list, img2_feats_list = [], []
    md1_list, md2_list, label_list = [], [], []
    model.eval()
    with torch.no_grad():
        for batch in loader:
            img1 = batch["image_1"].to(device)
            img2 = batch["image_2"].to(device)
            md1 = batch["matrix_1"].to(device)
            md2 = batch["matrix_2"].to(device)
            labels = batch["label_1"]
            img1_feats_list.append(model.img_tower(img1))
            img2_feats_list.append(model.img_tower(img2))
            md1_list.append(md1)
            md2_list.append(md2)
            if isinstance(labels, torch.Tensor):
                label_list.append(labels)
            else:
                label_list.append(torch.tensor(labels, dtype=torch.long))

    img1_feats = torch.cat(img1_feats_list)
    img2_feats = torch.cat(img2_feats_list)
    md1_tensor = torch.cat(md1_list)
    md2_tensor = torch.cat(md2_list)
    y_true = torch.cat(label_list).numpy()
    N = len(y_true)

    if N == 0:
        print("  [Phase 1] No samples — skipping.", flush=True)
        return

    with torch.no_grad():
        md1_feats = model.md_tower(md1_tensor)
        md2_feats = model.md_tower(md2_tensor)
        fused1, _, _ = model.bridge(img1_feats, md1_feats)
        fused2, _, _ = model.bridge(img2_feats, md2_feats)
        probs_baseline = (
            0.5 * (torch.softmax(fused1, dim=1) + torch.softmax(fused2, dim=1))
        ).cpu().numpy()
    _, baseline_auc, _ = _score_arrays(y_true, probs_baseline, num_classes)
    print(f"  Baseline AUC: {baseline_auc:.4f}  (N={N})", flush=True)

    feat_map = build_feature_index_map(data)
    rng = np.random.default_rng(seed)

    results = []
    for feat_name, dims in feat_map.items():
        all_dims = dims["value_dims"] + dims["missing_dims"]
        drops = []
        for _ in range(n_permutations):
            perm1 = md1_tensor.clone()
            perm2 = md2_tensor.clone()
            perm_idx = torch.from_numpy(rng.permutation(N)).to(device)
            perm1[:, all_dims] = perm1[perm_idx][:, all_dims]
            perm2[:, all_dims] = perm2[perm_idx][:, all_dims]
            with torch.no_grad():
                md1_p = model.md_tower(perm1)
                md2_p = model.md_tower(perm2)
                fused1_p, _, _ = model.bridge(img1_feats, md1_p)
                fused2_p, _, _ = model.bridge(img2_feats, md2_p)
                probs_p = (
                    0.5 * (torch.softmax(fused1_p, dim=1) + torch.softmax(fused2_p, dim=1))
                ).cpu().numpy()
            _, auc_p, _ = _score_arrays(y_true, probs_p, num_classes)
            drops.append(baseline_auc - auc_p)

        mean_drop = float(np.mean(drops))
        std_drop = float(np.std(drops))
        results.append({"feature": feat_name, "importance": mean_drop, "std": std_drop})
        print(f"  {feat_name:30s}  Δ AUC = {mean_drop:+.4f} ± {std_drop:.4f}", flush=True)

    results.sort(key=lambda r: r["importance"], reverse=True)

    # total MD ablation
    print("  Running total MD ablation ...", flush=True)
    total_drops = []
    for _ in range(n_permutations):
        perm_idx = torch.from_numpy(rng.permutation(N)).to(device)
        perm1_all = md1_tensor[perm_idx]
        perm2_all = md2_tensor[perm_idx]
        with torch.no_grad():
            md1_all = model.md_tower(perm1_all)
            md2_all = model.md_tower(perm2_all)
            f1, _, _ = model.bridge(img1_feats, md1_all)
            f2, _, _ = model.bridge(img2_feats, md2_all)
            probs_all = (
                0.5 * (torch.softmax(f1, dim=1) + torch.softmax(f2, dim=1))
            ).cpu().numpy()
        _, auc_all, _ = _score_arrays(y_true, probs_all, num_classes)
        total_drops.append(baseline_auc - auc_all)
    total_mean = float(np.mean(total_drops))
    total_std = float(np.std(total_drops))
    print(f"  Total MD ablation  Δ AUC = {total_mean:+.4f} ± {total_std:.4f}", flush=True)

    # Gaussian noise ablation
    print("  Running Gaussian noise ablation ...", flush=True)
    noise_drops = []
    for _ in range(n_permutations):
        noise1 = torch.randn_like(md1_tensor)
        noise2 = torch.randn_like(md2_tensor)
        with torch.no_grad():
            md1_noise = model.md_tower(noise1)
            md2_noise = model.md_tower(noise2)
            f1, _, _ = model.bridge(img1_feats, md1_noise)
            f2, _, _ = model.bridge(img2_feats, md2_noise)
            probs_noise = (
                0.5 * (torch.softmax(f1, dim=1) + torch.softmax(f2, dim=1))
            ).cpu().numpy()
        _, auc_noise, _ = _score_arrays(y_true, probs_noise, num_classes)
        noise_drops.append(baseline_auc - auc_noise)
    noise_mean = float(np.mean(noise_drops))
    noise_std = float(np.std(noise_drops))
    print(f"  Gaussian noise ablation  Δ AUC = {noise_mean:+.4f} ± {noise_std:.4f}", flush=True)
    print(
        f"  [interpretation] permutation Δ={total_mean:+.4f}  noise Δ={noise_mean:+.4f}  "
        f"informational gain = {total_mean - noise_mean:+.4f}",
        flush=True,
    )

    csv_path = out_dir / "md_permutation_importance.csv"
    with csv_path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["feature", "importance", "std"])
        writer.writeheader()
        writer.writerows(results)
        writer.writerow({"feature": "TOTAL_MD_ABLATION", "importance": total_mean, "std": total_std})
        writer.writerow({"feature": "GAUSSIAN_NOISE_ABLATION", "importance": noise_mean, "std": noise_std})

    names = [r["feature"] for r in results]
    imps = [r["importance"] for r in results]
    stds = [r["std"] for r in results]
    colors = ["#e05c5c" if v >= 0 else "#5c9ee0" for v in imps]

    fig, ax = plt.subplots(figsize=(9, max(4, (len(names) + 3) * 0.45)))
    y_pos = np.arange(len(names))
    ax.barh(y_pos, imps, xerr=stds, color=colors, ecolor="grey", capsize=3, height=0.6)
    ax.axhline(len(names) - 0.25, color="grey", linewidth=0.6, linestyle="--")
    ax.barh(len(names) + 0.5, total_mean, xerr=total_std,
            color="#c45ce0" if total_mean >= 0 else "#5c9ee0",
            ecolor="grey", capsize=3, height=0.6)
    ax.barh(len(names) + 1.5, noise_mean, xerr=noise_std,
            color="#e08c2a" if noise_mean >= 0 else "#5c9ee0",
            ecolor="grey", capsize=3, height=0.6)
    ax.set_yticks(list(y_pos) + [len(names) + 0.5, len(names) + 1.5])
    ax.set_yticklabels(names + ["ALL MD (permute)", "ALL MD (noise)"], fontsize=9)
    ax.invert_yaxis()
    ax.axvline(0, color="black", linewidth=0.8)
    ax.set_xlabel("Mean AUC drop (baseline − permuted)", fontsize=10)
    ax.set_title(
        f"MD Tower — Permutation Feature Importance\n"
        f"baseline AUC={baseline_auc:.4f}  N={N}  repeats={n_permutations}",
        fontsize=11,
    )
    fig.tight_layout()
    fig.savefig(out_dir / "md_permutation_importance.png", dpi=150)
    plt.close(fig)
    print(f"  Saved → {out_dir / 'md_permutation_importance.png'}", flush=True)


# ---------------------------------------------------------------------------
# Phase 2 — GradCAM overlays
# ---------------------------------------------------------------------------


def run_gradcam(
    model: SingleEyeHT,
    loader,
    data: DataBundle,
    eval_df,
    eval_mode: str,
    backbone: str,
    device: torch.device,
    alpha: float,
    out_dir: Path,
) -> None:
    print("\n[Phase 2] GradCAM overlays ...", flush=True)
    gradcam_dir = out_dir / "gradcam"
    gradcam_dir.mkdir(exist_ok=True)

    manifest_path = gradcam_dir / "gradcam_manifest.csv"
    skip_overlays = manifest_path.exists()
    overlay_grid_items = []

    if skip_overlays:
        print("  Overlays already exist — skipping computation, loading from disk.", flush=True)
        manifest_df = pd.read_csv(manifest_path)
        for _, row in manifest_df.iterrows():
            pid = str(row["pid"])
            od_path = gradcam_dir / f"gradcam_od_{pid}.png"
            os_path = gradcam_dir / f"gradcam_os_{pid}.png"
            od_ov = Image.open(od_path).convert("RGB") if od_path.exists() else None
            os_ov = Image.open(os_path).convert("RGB") if os_path.exists() else None
            overlay_grid_items.append((od_ov, os_ov, str(row["short_lbl"]), bool(row["correct"])))
    else:
        target_layer = get_gradcam_layer(model, backbone)
        gcam = GradCAM(target_layer)
        manifest_rows = []

        model.eval()
        for batch in loader:
            img_od = batch["image_1"].to(device)
            img_os = batch["image_2"].to(device)
            meta_od = batch["matrix_1"].to(device)
            meta_os = batch["matrix_2"].to(device)
            lbl_raw = batch["label_1"][0]
            label = int(lbl_raw.item() if isinstance(lbl_raw, torch.Tensor) else lbl_raw)
            pid = batch["id_1"][0]

            cam_od, pred = gcam.compute(img_od, meta_od, model)
            cam_os, _ = gcam.compute(img_os, meta_os, model)

            with torch.no_grad():
                out_od = model(img_od, meta_od)
                conf = float(torch.softmax(out_od, dim=1)[0, pred].item())

            row_od = eval_df[
                (eval_df["Patient ID"] == int(pid)) & (eval_df["eyeID"] == "OD")
            ]
            row_os = eval_df[
                (eval_df["Patient ID"] == int(pid)) & (eval_df["eyeID"] == "OS")
            ]
            orig_od = (
                Image.open(data.get_image_path(row_od.iloc[0])).convert("RGB")
                if len(row_od) else None
            )
            orig_os = (
                Image.open(data.get_image_path(row_os.iloc[0])).convert("RGB")
                if len(row_os) else None
            )

            true_name = label_name(label, eval_mode)
            pred_name = label_name(pred, eval_mode)
            correct = label == pred
            title = (
                f"Patient {pid}  |  True: {true_name}  |  Pred: {pred_name}  "
                f"|  conf={conf:.2f}  {'✓' if correct else '✗'}"
            )

            fig, axes = plt.subplots(2, 2, figsize=(10, 9))
            fig.suptitle(title, fontsize=11, fontweight="bold", color="green" if correct else "red")

            if orig_od is not None:
                axes[0, 0].imshow(orig_od)
                axes[0, 0].set_title("OD — original", fontsize=9)
                axes[0, 1].imshow(overlay_gradcam(orig_od, cam_od, alpha))
                axes[0, 1].set_title("OD — GradCAM", fontsize=9)
            else:
                axes[0, 0].set_title("OD — (missing)", fontsize=9)
            axes[0, 0].axis("off")
            axes[0, 1].axis("off")

            if orig_os is not None:
                axes[1, 0].imshow(orig_os)
                axes[1, 0].set_title("OS — original", fontsize=9)
                axes[1, 1].imshow(overlay_gradcam(orig_os, cam_os, alpha))
                axes[1, 1].set_title("OS — GradCAM", fontsize=9)
            else:
                axes[1, 0].set_title("OS — (missing)", fontsize=9)
            axes[1, 0].axis("off")
            axes[1, 1].axis("off")

            fig.tight_layout()
            out_path = gradcam_dir / f"patient_{pid}_OD_OS.png"
            fig.savefig(out_path, dpi=120)
            plt.close(fig)
            print(f"  Patient {pid}: {true_name} → {pred_name} ({conf:.2f})  → {out_path.name}", flush=True)

            od_overlay = overlay_gradcam(orig_od, cam_od, alpha) if orig_od else None
            os_overlay = overlay_gradcam(orig_os, cam_os, alpha) if orig_os else None
            if od_overlay is not None:
                od_overlay.save(gradcam_dir / f"gradcam_od_{pid}.png")
            if os_overlay is not None:
                os_overlay.save(gradcam_dir / f"gradcam_os_{pid}.png")
            short_lbl = f"P{pid}  {true_name[:3]}→{pred_name[:3]} {'✓' if correct else '✗'}"
            overlay_grid_items.append((od_overlay, os_overlay, short_lbl, correct))
            manifest_rows.append({"pid": pid, "short_lbl": short_lbl, "correct": correct})

        gcam.remove()
        pd.DataFrame(manifest_rows).to_csv(manifest_path, index=False)

    # Always regenerate the summary grid
    n = len(overlay_grid_items)
    if n == 0:
        print("  [Phase 2] No patients to visualise.", flush=True)
        return

    fig, axes = plt.subplots(n, 2, figsize=(8, n * 3.2 + 1.5))
    if n == 1:
        axes = axes[np.newaxis, :]
    fig.suptitle("GradCAM Summary Grid — all holdout patients", fontsize=12)

    for i, (od_ov, os_ov, lbl, correct) in enumerate(overlay_grid_items):
        color = "green" if correct else "red"
        for j in range(2):
            axes[i, j].axis("off")
        if od_ov is not None:
            axes[i, 0].imshow(od_ov)
        axes[i, 0].set_title(f"{lbl}\nOD", fontsize=7, color=color)
        if os_ov is not None:
            axes[i, 1].imshow(os_ov)
        axes[i, 1].set_title(f"{lbl}\nOS", fontsize=7, color=color)

    fig.tight_layout(rect=[0, 0, 1, 0.97])
    grid_path = out_dir / "gradcam_summary_grid.png"
    fig.savefig(grid_path, dpi=120)
    plt.close(fig)
    print(f"  Summary grid → {grid_path}", flush=True)


# ---------------------------------------------------------------------------
# Phase 3 — Fusion event analysis (disk-based, no re-inference)
# ---------------------------------------------------------------------------

_EVENT_LABELS = [
    "full correction\n(both wrong→fused right)",
    "img assist\n(img wrong, md right→right)",
    "md assist\n(md wrong, img right→right)",
    "full error\n(both right→fused wrong)",
    "img drag\n(img wrong, md right→wrong)",
    "md drag\n(md wrong, img right→wrong)",
]
_EVENT_KEYS   = ["full_correction", "img_assist", "md_assist",
                 "full_error", "img_drag", "md_drag"]
_EVENT_COLORS = ["#2ca02c", "#98df8a", "#b5d46e", "#d62728", "#ff9896", "#ffbf9b"]


def _fusion_event_stats(
    y_true: np.ndarray,
    pf: np.ndarray,
    pi: np.ndarray,
    pm: np.ndarray,
    split_name: str,
    out_dir: Path,
    component_labels: tuple[str, str] = ("img", "md"),
) -> dict:
    """Compute, save, and plot fusion events for one split. Returns summary dict."""
    N = len(y_true)
    if N == 0:
        print(f"  [{split_name}] No samples — skipping.", flush=True)
        return {}

    a, b = component_labels
    event_labels = [
        "full correction\n(both wrong→fused right)",
        f"{a} assist\n({a} wrong, {b} right→right)",
        f"{b} assist\n({b} wrong, {a} right→right)",
        "full error\n(both right→fused wrong)",
        f"{a} drag\n({a} wrong, {b} right→wrong)",
        f"{b} drag\n({b} wrong, {a} right→wrong)",
    ]

    pred_f = pf.argmax(axis=1)
    pred_i = pi.argmax(axis=1)
    pred_m = pm.argmax(axis=1)

    conf_f = np.take_along_axis(pf, pred_f[:, None], axis=1).squeeze(1)
    conf_i = np.take_along_axis(pi, pred_i[:, None], axis=1).squeeze(1)
    conf_m = np.take_along_axis(pm, pred_m[:, None], axis=1).squeeze(1)
    conf_delta = conf_f - 0.5 * (conf_i + conf_m)

    f_ok = pred_f == y_true
    i_ok = pred_i == y_true
    m_ok = pred_m == y_true

    full_correction = f_ok  & ~i_ok & ~m_ok
    img_assist      = f_ok  & ~i_ok &  m_ok
    md_assist       = f_ok  &  i_ok & ~m_ok
    full_error      = ~f_ok &  i_ok &  m_ok
    img_drag        = ~f_ok & ~i_ok &  m_ok
    md_drag         = ~f_ok &  i_ok & ~m_ok
    concordant_ok   =  f_ok &  i_ok &  m_ok
    concordant_bad  = ~f_ok & ~i_ok & ~m_ok

    event_masks = [full_correction, img_assist, md_assist, full_error, img_drag, md_drag]
    counts = [int(m.sum()) for m in event_masks]

    print(f"\n  [{split_name}] N={N}", flush=True)
    for label, count in zip(event_labels, counts):
        print(f"    {label.replace(chr(10), ' '):55s}: {count}", flush=True)
    n_corr, n_err = counts[0], counts[3]
    print(f"    full correction/error ratio: {n_corr}/{n_err}", flush=True)
    print(f"    conf_delta  mean={conf_delta.mean():+.4f}  median={np.median(conf_delta):+.4f}",
          flush=True)

    # CSV
    event_type = np.where(concordant_ok, "concordant_correct",
                 np.where(concordant_bad, "concordant_wrong", "other")).astype(object)
    for mask, key in zip(event_masks, _EVENT_KEYS):
        event_type[mask] = key

    rows = [
        {
            "patient_idx": idx,
            "y_true":      int(y_true[idx]),
            "pred_fused":  int(pred_f[idx]),
            "pred_img":    int(pred_i[idx]),
            "pred_md":     int(pred_m[idx]),
            "conf_fused":  float(conf_f[idx]),
            "conf_img":    float(conf_i[idx]),
            "conf_md":     float(conf_m[idx]),
            "conf_delta":  float(conf_delta[idx]),
            "event_type":  event_type[idx],
        }
        for idx in range(N)
    ]
    csv_path = out_dir / f"fusion_events_{split_name}.csv"
    with csv_path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)

    # Plot
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    fig.suptitle(f"Fusion Events — {split_name}  (N={N})", fontsize=11)

    for bar_x, bar_counts, bar_colors in (
        (0, counts[:3], _EVENT_COLORS[:3]),
        (1, counts[3:], _EVENT_COLORS[3:]),
    ):
        bot = 0
        for c, col in zip(bar_counts, bar_colors):
            axes[0].bar(bar_x, c, bottom=bot, color=col, width=0.5)
            if c > 0:
                axes[0].text(bar_x, bot + c / 2, str(c), ha="center", va="center",
                             fontsize=9, fontweight="bold")
            bot += c
    axes[0].set_xticks([0, 1])
    axes[0].set_xticklabels(["Positive\nevents", "Negative\nevents"])
    axes[0].set_ylabel("Count")
    patches = [mpatches.Patch(color=c, label=l.replace("\n", " "))
               for c, l in zip(_EVENT_COLORS, event_labels)]
    axes[0].legend(handles=patches, fontsize=6, loc="upper right")

    box_data   = [conf_delta[m] for m in event_masks if m.sum() > 0]
    box_labels = [l.split("\n")[0] for m, l in zip(event_masks, event_labels) if m.sum() > 0]
    box_cols   = [c for m, c in zip(event_masks, _EVENT_COLORS) if m.sum() > 0]
    if box_data:
        bp = axes[1].boxplot(box_data, patch_artist=True, widths=0.5)
        for patch, color in zip(bp["boxes"], box_cols):
            patch.set_facecolor(color)
        axes[1].set_xticks(range(1, len(box_labels) + 1))
        axes[1].set_xticklabels(box_labels, rotation=35, ha="right", fontsize=7)
    axes[1].axhline(0, color="black", linewidth=0.8, linestyle="--")
    axes[1].set_ylabel(f"conf_delta\n(fused − avg({a}, {b}))")
    axes[1].set_title("Confidence delta by event type")

    for mask, color, label in zip(event_masks, _EVENT_COLORS, event_labels):
        if mask.sum() > 0:
            axes[2].scatter(conf_i[mask], conf_m[mask], c=color,
                            label=label.split("\n")[0], alpha=0.85, s=45, edgecolors="none")
    if concordant_ok.sum() > 0:
        axes[2].scatter(conf_i[concordant_ok], conf_m[concordant_ok],
                        c="lightgrey", alpha=0.4, s=20, edgecolors="none", label="concordant correct")
    if concordant_bad.sum() > 0:
        axes[2].scatter(conf_i[concordant_bad], conf_m[concordant_bad],
                        c="darkgrey", alpha=0.4, s=20, edgecolors="none", label="concordant wrong")
    axes[2].plot([0, 1], [0, 1], "k--", linewidth=0.5, alpha=0.4)
    axes[2].set_xlabel(f"conf_{a}")
    axes[2].set_ylabel(f"conf_{b}")
    axes[2].set_title("Tower confidence space\ncoloured by fusion event")
    axes[2].legend(fontsize=6, loc="lower right")

    fig.tight_layout()
    fig.savefig(out_dir / f"fusion_events_{split_name}.png", dpi=150)
    plt.close(fig)
    print(f"  Saved → {out_dir / f'fusion_events_{split_name}.png'}", flush=True)

    return {
        "split": split_name,
        "N": N,
        **{key: cnt for key, cnt in zip(_EVENT_KEYS, counts)},
        "conf_delta_mean": float(conf_delta.mean()),
        "conf_delta_median": float(np.median(conf_delta)),
    }


def _load_best_epoch_idx(fold_dir: Path, n_epochs: int) -> int:
    """Return 0-based array index of the best single checkpoint.

    Val/train prediction arrays skip md_warmup epochs but include all other
    phases (tower_warmup, fused_warmup, main).  The array index is therefore
    the row position within epoch_log *after* filtering out md_warmup rows.
    """
    log_path = fold_dir / "epoch_log.csv"
    if not log_path.exists():
        return n_epochs - 1
    try:
        log = pd.read_csv(log_path)
        if "is_best_single" not in log.columns or "phase_single" not in log.columns:
            return n_epochs - 1
        # Keep only phases that produce predictions (everything except md_warmup)
        pred_log = log[log["phase_single"] != "md_warmup"].reset_index(drop=True)
        best_rows = pred_log[pred_log["is_best_single"] == True]
        if not best_rows.empty:
            return int(best_rows.index[-1])
    except Exception:
        pass
    return n_epochs - 1


def _aggregate_by_patient(
    patient_ids: np.ndarray,
    y_true: np.ndarray,
    *prob_arrays: np.ndarray,
) -> tuple:
    """Average per-eye probs to patient level. Returns (y, *averaged_probs)."""
    unique_pids = np.unique(patient_ids)
    y_bilat = np.array([y_true[patient_ids == pid][0] for pid in unique_pids])
    averaged = tuple(
        np.array([arr[patient_ids == pid].mean(axis=0) for pid in unique_pids])
        for arr in prob_arrays
    )
    return (y_bilat,) + averaged


def run_fusion_event_analysis(fold_dir: Path, out_dir: Path) -> list[dict]:
    """
    Load pre-saved predictions from disk and run fusion event analysis
    for both the train and val splits at the best checkpoint epoch.
    Returns list of summary dicts (one per split).
    """
    print("\n[Phase 3] Fusion event analysis (from saved predictions) ...", flush=True)
    summaries = []

    # ---- val: separate OD/OS per-epoch files, already patient-level ----
    od_f = fold_dir / "val_probs_fused_od_epochs.npy"
    os_f = fold_dir / "val_probs_fused_os_epochs.npy"
    od_i = fold_dir / "val_probs_img_od_epochs.npy"
    os_i = fold_dir / "val_probs_img_os_epochs.npy"
    od_m = fold_dir / "val_probs_md_od_epochs.npy"
    os_m = fold_dir / "val_probs_md_os_epochs.npy"
    y_f  = fold_dir / "val_y_true_epochs.npy"

    tower_mode = fold_dir.parent.name  # "single", "ensemble", "bilateral", …
    is_ensemble_like = tower_mode in ("ensemble", "bilateral")

    if all(p.exists() for p in [od_f, os_f, od_i, os_i, od_m, os_m, y_f]):
        n_epochs = np.load(od_f).shape[0]
        epoch_idx = _load_best_epoch_idx(fold_dir, n_epochs)
        print(f"  Val best epoch index: {epoch_idx}", flush=True)
        y_val   = np.load(y_f)[epoch_idx]
        pf_od   = np.load(od_f)[epoch_idx]
        pf_os   = np.load(os_f)[epoch_idx]
        pi_od   = np.load(od_i)[epoch_idx]
        pi_os   = np.load(os_i)[epoch_idx]
        pm_od   = np.load(od_m)[epoch_idx]
        pm_os   = np.load(os_m)[epoch_idx]

        # Per-eye bridge (ensemble/bilateral only — redundant in single mode)
        if is_ensemble_like:
            summaries.append(_fusion_event_stats(y_val, pf_od, pi_od, pm_od, "val_OD", out_dir))
            summaries.append(_fusion_event_stats(y_val, pf_os, pi_os, pm_os, "val_OS", out_dir))

        # Ensemble: OD+OS averaged (bilateral patient-level)
        pf_ens = 0.5 * (pf_od + pf_os)
        pi_ens = 0.5 * (pi_od + pi_os)
        pm_ens = 0.5 * (pm_od + pm_os)
        summaries.append(_fusion_event_stats(y_val, pf_ens, pi_ens, pm_ens, "val", out_dir))

        # Fused head (if available): learned bilateral combination vs per-eye bridge outputs
        fused_head_f = fold_dir / "probs_fused_head.npy"
        if fused_head_f.exists():
            pf_head = np.load(fused_head_f)
            summaries.append(_fusion_event_stats(
                y_val, pf_head, pf_od, pf_os, "val_fused_head", out_dir,
                component_labels=("OD", "OS")))
    else:
        print("  Val epoch files not found — skipping val.", flush=True)

    # ---- train: eye-level per-epoch files, aggregate to patient level ----
    tr_pf_f = fold_dir / "train_probs_fused.npy"
    tr_pi_f = fold_dir / "train_probs_img.npy"
    tr_pm_f = fold_dir / "train_probs_md.npy"
    tr_y_f  = fold_dir / "train_y_true.npy"
    tr_id_f = fold_dir / "train_patient_ids.npy"

    if all(p.exists() for p in [tr_pf_f, tr_pi_f, tr_pm_f, tr_y_f, tr_id_f]):
        n_epochs = np.load(tr_pf_f).shape[0]
        epoch_idx = _load_best_epoch_idx(fold_dir, n_epochs)
        print(f"  Train best epoch index: {epoch_idx}", flush=True)
        pf_eyes = np.load(tr_pf_f)[epoch_idx]
        pi_eyes = np.load(tr_pi_f)[epoch_idx]
        pm_eyes = np.load(tr_pm_f)[epoch_idx]
        y_eyes  = np.load(tr_y_f)
        pids    = np.load(tr_id_f, allow_pickle=True).astype(str)
        y_tr, pf_tr, pi_tr, pm_tr = _aggregate_by_patient(pids, y_eyes,
                                                           pf_eyes, pi_eyes, pm_eyes)
        summaries.append(_fusion_event_stats(y_tr, pf_tr, pi_tr, pm_tr, "train", out_dir))
    else:
        print("  Train epoch files not found — skipping train.", flush=True)

    return summaries


# ---------------------------------------------------------------------------
# Cross-fold MD importance summary plot
# ---------------------------------------------------------------------------


def _plot_run_md_importance_summary(run_dir: Path, folds: list) -> None:
    """Aggregate per-fold MD permutation importance CSVs into a run-level summary plot."""
    all_dfs = []
    for fold_idx, fold_dir, _ in folds:
        csv_path = fold_dir / "explainability" / "md_permutation_importance.csv"
        if csv_path.exists():
            df = pd.read_csv(csv_path)
            df["fold"] = fold_idx
            all_dfs.append(df)

    if not all_dfs:
        print("  [MD summary] No per-fold importance CSVs found — skipping.", flush=True)
        return

    combined = pd.concat(all_dfs, ignore_index=True)
    _SPECIAL = {"TOTAL_MD_ABLATION", "GAUSSIAN_NOISE_ABLATION"}
    feature_rows = combined[~combined["feature"].isin(_SPECIAL)]
    special_rows = combined[combined["feature"].isin(_SPECIAL)]

    agg = (
        feature_rows.groupby("feature")["importance"]
        .agg(["mean", "std"])
        .reset_index()
        .rename(columns={"mean": "mean_importance", "std": "std_importance"})
        .sort_values("mean_importance", ascending=False)
        .reset_index(drop=True)
    )
    special_agg = (
        special_rows.groupby("feature")["importance"]
        .agg(["mean", "std"])
        .reset_index()
    )

    names = agg["feature"].tolist()
    imps  = agg["mean_importance"].tolist()
    stds  = agg["std_importance"].fillna(0).tolist()
    colors = ["#e05c5c" if v >= 0 else "#5c9ee0" for v in imps]

    fig, ax = plt.subplots(figsize=(9, max(4, (len(names) + 3) * 0.45)))
    y_pos = np.arange(len(names))
    ax.barh(y_pos, imps, xerr=stds, color=colors, ecolor="grey", capsize=3, height=0.6)
    ax.axhline(len(names) - 0.25, color="grey", linewidth=0.6, linestyle="--")

    special_label_map = {
        "TOTAL_MD_ABLATION":    "ALL MD (permute)",
        "GAUSSIAN_NOISE_ABLATION": "ALL MD (noise)",
    }
    special_colors = {
        "TOTAL_MD_ABLATION":    "#c45ce0",
        "GAUSSIAN_NOISE_ABLATION": "#e08c2a",
    }
    extra_ytick_pos = []
    extra_ytick_labels = []
    for i, feat in enumerate(["TOTAL_MD_ABLATION", "GAUSSIAN_NOISE_ABLATION"]):
        row = special_agg[special_agg["feature"] == feat]
        if row.empty:
            continue
        offset = len(names) + 0.5 + i
        val, err = float(row["mean"].iloc[0]), float(row["std"].iloc[0])
        ax.barh(offset, val, xerr=err,
                color=special_colors[feat] if val >= 0 else "#5c9ee0",
                ecolor="grey", capsize=3, height=0.6)
        extra_ytick_pos.append(offset)
        extra_ytick_labels.append(special_label_map[feat])

    ax.set_yticks(list(y_pos) + extra_ytick_pos)
    ax.set_yticklabels(names + extra_ytick_labels, fontsize=9)
    ax.invert_yaxis()
    ax.axvline(0, color="black", linewidth=0.8)
    ax.set_xlabel("Mean AUC drop (baseline − permuted)", fontsize=10)
    ax.set_title(
        f"MD Tower — Permutation Feature Importance ({len(all_dfs)}-fold summary)\n"
        f"error bars = std across folds",
        fontsize=11,
    )
    fig.tight_layout()
    out_path = run_dir / "explainability_md_importance_summary.png"
    fig.savefig(out_path, dpi=150)
    plt.close(fig)
    print(f"  MD importance summary → {out_path}", flush=True)

    agg.to_csv(run_dir / "explainability_md_importance_summary.csv", index=False)


# ---------------------------------------------------------------------------
# Cross-fold fusion summary plot
# ---------------------------------------------------------------------------


def _plot_cross_fold_fusion_summary(df_sum: pd.DataFrame, run_dir: Path) -> None:
    """Generate aggregate fusion event plots from the cross-fold summary DataFrame.

    Loads per-fold fusion_events_*.csv files to get sample-level data for the
    conf_delta and tower-confidence-space panels.
    """
    for split_name, grp in df_sum.groupby("split"):
        grp = grp.sort_values("fold").reset_index(drop=True)
        n_folds = len(grp)
        fold_ids = grp["fold"].values

        # For the fused_head split the two comparators are OD/OS bridges, not img/md towers
        comp_a, comp_b = ("OD", "OS") if "fused_head" in split_name else ("img", "md")
        summary_event_labels = [
            "full correction\n(both wrong→fused right)",
            f"{comp_a} assist\n({comp_a} wrong, {comp_b} right→right)",
            f"{comp_b} assist\n({comp_b} wrong, {comp_a} right→right)",
            "full error\n(both right→fused wrong)",
            f"{comp_a} drag\n({comp_a} wrong, {comp_b} right→wrong)",
            f"{comp_b} drag\n({comp_b} wrong, {comp_a} right→wrong)",
        ]

        # Load all per-fold CSVs for this split to get sample-level data
        sample_dfs = []
        for fold_idx in fold_ids:
            csv_path = (run_dir / f"fold{fold_idx}" / "explainability"
                        / f"fusion_events_{split_name}.csv")
            if csv_path.exists():
                sample_dfs.append(pd.read_csv(csv_path))
        sample_df = pd.concat(sample_dfs, ignore_index=True) if sample_dfs else pd.DataFrame()

        fig, axes = plt.subplots(1, 5, figsize=(25, 4))
        fig.suptitle(f"Fusion Events — {split_name}  (all {n_folds} folds combined)", fontsize=11)

        # Panel 1: stacked bar of totals (positive vs negative)
        for bar_x, keys, colors in (
            (0, _EVENT_KEYS[:3], _EVENT_COLORS[:3]),
            (1, _EVENT_KEYS[3:], _EVENT_COLORS[3:]),
        ):
            bot = 0
            for key, col in zip(keys, colors):
                c = int(grp[key].sum())
                axes[0].bar(bar_x, c, bottom=bot, color=col, width=0.5)
                if c > 0:
                    axes[0].text(bar_x, bot + c / 2, str(c), ha="center", va="center",
                                 fontsize=9, fontweight="bold")
                bot += c
        axes[0].set_xticks([0, 1])
        axes[0].set_xticklabels(["Positive\nevents", "Negative\nevents"])
        axes[0].set_ylabel("Count (all folds)")
        patches = [mpatches.Patch(color=c, label=l.replace("\n", " "))
                   for c, l in zip(_EVENT_COLORS, summary_event_labels)]
        axes[0].legend(handles=patches, fontsize=6, loc="upper right")

        # Panel 2: per-fold stacked bar (fold variance)
        x = np.arange(n_folds)
        pos_bot = np.zeros(n_folds)
        neg_bot = np.zeros(n_folds)
        for key, color in zip(_EVENT_KEYS[:3], _EVENT_COLORS[:3]):
            vals = grp[key].values.astype(float)
            axes[1].bar(x, vals, bottom=pos_bot, color=color, width=0.6)
            pos_bot += vals
        for key, color in zip(_EVENT_KEYS[3:], _EVENT_COLORS[3:]):
            vals = grp[key].values.astype(float)
            axes[1].bar(x + 0.65, vals, bottom=neg_bot, color=color, width=0.6)
            neg_bot += vals
        axes[1].set_xticks(x + 0.325)
        axes[1].set_xticklabels([f"fold {f}" for f in fold_ids], fontsize=8)
        axes[1].set_ylabel("Count")
        axes[1].set_title("Per-fold breakdown\n(left=positive, right=negative)")

        # Panel 3: conf_delta_mean per fold (bar + mean line)
        cd = grp["conf_delta_mean"].values
        bar_colors = ["#2ca02c" if v >= 0 else "#d62728" for v in cd]
        axes[2].bar(x, cd, color=bar_colors, width=0.6, alpha=0.8)
        axes[2].axhline(cd.mean(), color="black", linewidth=1.2, linestyle="--",
                        label=f"mean={cd.mean():+.4f}")
        axes[2].axhline(0, color="grey", linewidth=0.7)
        axes[2].set_xticks(x)
        axes[2].set_xticklabels([f"fold {f}" for f in fold_ids], fontsize=8)
        axes[2].set_ylabel(f"conf_delta mean\n(fused − avg({comp_a}, {comp_b}))")
        axes[2].set_title("Confidence delta per fold")
        axes[2].legend(fontsize=8)

        # Panel 4: conf_delta boxplot by event type (all folds combined)
        if not sample_df.empty and "event_type" in sample_df.columns:
            key_order = [k for k in _EVENT_KEYS if k in sample_df["event_type"].values]
            box_data   = [sample_df.loc[sample_df["event_type"] == k, "conf_delta"].values
                          for k in key_order]
            box_labels = [l.split("\n")[0]
                          for k, l in zip(_EVENT_KEYS, summary_event_labels) if k in key_order]
            box_cols   = [c for k, c in zip(_EVENT_KEYS, _EVENT_COLORS) if k in key_order]
            if box_data:
                bp = axes[3].boxplot(box_data, patch_artist=True, widths=0.5)
                for patch, color in zip(bp["boxes"], box_cols):
                    patch.set_facecolor(color)
                axes[3].set_xticks(range(1, len(box_labels) + 1))
                axes[3].set_xticklabels(box_labels, rotation=35, ha="right", fontsize=7)
        axes[3].axhline(0, color="black", linewidth=0.8, linestyle="--")
        axes[3].set_ylabel(f"conf_delta\n(fused − avg({comp_a}, {comp_b}))")
        axes[3].set_title("Confidence delta by event type\n(all folds)")

        # Panel 5: tower confidence space scatter (all folds combined)
        if not sample_df.empty:
            event_color_map = dict(zip(_EVENT_KEYS, _EVENT_COLORS))
            for key, color in zip(_EVENT_KEYS, _EVENT_COLORS):
                sub = sample_df[sample_df["event_type"] == key]
                if len(sub):
                    label = next(l.split("\n")[0] for k, l in zip(_EVENT_KEYS, summary_event_labels)
                                 if k == key)
                    axes[4].scatter(sub["conf_img"], sub["conf_md"], c=color,
                                    label=label, alpha=0.7, s=30, edgecolors="none")
            for conc_key, conc_color, conc_label in [
                ("concordant_correct", "lightgrey", "concordant correct"),
                ("concordant_wrong",   "darkgrey",  "concordant wrong"),
            ]:
                sub = sample_df[sample_df["event_type"] == conc_key]
                if len(sub):
                    axes[4].scatter(sub["conf_img"], sub["conf_md"], c=conc_color,
                                    alpha=0.3, s=15, edgecolors="none", label=conc_label)
            axes[4].plot([0, 1], [0, 1], "k--", linewidth=0.5, alpha=0.4)
            axes[4].set_xlabel(f"conf_{comp_a}")
            axes[4].set_ylabel(f"conf_{comp_b}")
            axes[4].legend(fontsize=6, loc="lower right")
        axes[4].set_title("Tower confidence space\n(all folds)")

        fig.tight_layout()
        out_path = run_dir / f"explainability_fusion_summary_{split_name}.png"
        fig.savefig(out_path, dpi=150)
        plt.close(fig)
        print(f"  Summary plot → {out_path}", flush=True)


# ---------------------------------------------------------------------------
# Fold discovery
# ---------------------------------------------------------------------------


def find_folds(run_dir: Path, checkpoint: str) -> list[tuple[int, Path, Path]]:
    """Return sorted list of (fold_idx, fold_dir, ckpt_path) for existing folds."""
    folds = []
    for fold_dir in sorted(run_dir.glob("fold*")):
        if not fold_dir.is_dir():
            continue
        try:
            fold_idx = int(fold_dir.name.replace("fold", ""))
        except ValueError:
            continue
        ckpt = fold_dir / checkpoint
        if not ckpt.exists():
            print(f"  [WARN] Checkpoint not found: {ckpt} — skipping fold {fold_idx}",
                  flush=True)
            continue
        folds.append((fold_idx, fold_dir, ckpt))
    return folds


# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------


def parse_args():
    ap = argparse.ArgumentParser(
        description="Post-hoc explainability for all folds of a saved run."
    )
    ap.add_argument(
        "--run-dir",
        type=Path,
        required=True,
        help="Path to the tower-mode run directory, e.g. "
             "analysis_data/v2.3_single_binary_nocrop/binary/single",
    )
    ap.add_argument(
        "--checkpoint",
        default="best_single.pt",
        help="Checkpoint filename inside each fold_dir (default: best_single.pt)",
    )
    ap.add_argument(
        "--split",
        choices=["holdout", "val"],
        default="holdout",
        help="Which patient set to use for Phases 1 & 2 (default: holdout, falls back to val)",
    )
    ap.add_argument("--image-dir", default="Papila/FundusImages")
    ap.add_argument("--clinical-dir", default="Papila/ClinicalData")
    ap.add_argument("--label-col", default="Diagnosis")
    ap.add_argument("--cat-cols", nargs="*", default=["Gender", "Phakic/Pseudophakic"])
    ap.add_argument("--fold-seed", type=int, default=42)
    ap.add_argument("--holdout-seed", type=int, default=123)
    ap.add_argument("--holdout-per-class", type=int, default=5)
    ap.add_argument("--n-splits", type=int, default=5)
    ap.add_argument("--n-permutations", type=int, default=30)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--alpha", type=float, default=0.45)
    ap.add_argument("--batch-size", type=int, default=1)
    ap.add_argument("--no-phase1", action="store_true", help="Skip MD importance")
    ap.add_argument("--no-phase2", action="store_true", help="Skip GradCAM")
    ap.add_argument("--no-phase3", action="store_true", help="Skip fusion event analysis")
    return ap.parse_args()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main():
    args = parse_args()
    run_dir = args.run_dir.resolve()
    if not run_dir.is_dir():
        sys.exit(f"[ERROR] run_dir does not exist: {run_dir}")

    summary_path = run_dir / "summary.json"
    if not summary_path.exists():
        sys.exit(f"[ERROR] summary.json not found: {summary_path}")
    summary = json.loads(summary_path.read_text())
    backbone   = summary["backbone"]
    eval_mode  = summary["eval_mode"]
    tower_mode = summary.get("tower_mode", "single")
    print(f"[explain_run] backbone={backbone}  eval_mode={eval_mode}  tower_mode={tower_mode}")

    if tower_mode not in ("single", "ensemble"):
        sys.exit(
            f"[ERROR] explain_run supports single/ensemble tower modes, got: {tower_mode!r}"
        )

    folds = find_folds(run_dir, args.checkpoint)
    if not folds:
        sys.exit(f"[ERROR] No fold directories with checkpoint '{args.checkpoint}' found in {run_dir}")
    print(f"[explain_run] Found {len(folds)} fold(s): {[f[0] for f in folds]}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[explain_run] device={device}", flush=True)

    # Build data bundle once (shared across folds)
    print("[explain_run] Loading clinical data ...", flush=True)
    data = build_papila_data(
        image_dir=args.image_dir,
        clinical_dir=args.clinical_dir,
        label_col=args.label_col,
        cat_cols=args.cat_cols,
        n_splits=args.n_splits,
        random_seed=args.fold_seed,
    )
    df_mode = data.df.copy()
    if eval_mode == "binary":
        df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True)
    num_classes = 2 if eval_mode == "binary" else int(df_mode[args.label_col].nunique())

    # Reconstruct all splits once
    splitter = PatientFirstSplitManager(
        patient_col="Patient ID", label_col=args.label_col
    )
    split_args = SimpleNamespace(
        eval_mode=eval_mode,
        holdout_per_class=args.holdout_per_class,
        holdout_seed=args.holdout_seed,
        n_splits=args.n_splits,
        fold_seed=args.fold_seed,
    )
    clinical_ns = SimpleNamespace(df=df_mode, label_col=args.label_col)
    plans = splitter.build_plans(clinical=clinical_ns, args=split_args, profile=None)

    profile_patient = build_papila_profile(
        patient_col="Patient ID", label_col=args.label_col, sample_mode="patient"
    )
    eval_transform = build_eval_transform(backbone)

    all_phase3_summaries = []

    for fold_idx, fold_dir, ckpt_path in folds:
        print(f"\n{'='*60}", flush=True)
        print(f"[explain_run] === Fold {fold_idx} ===", flush=True)

        out_dir = fold_dir / "explainability"
        out_dir.mkdir(exist_ok=True)

        # Phase 3 needs no model — run first so the model load can be skipped
        # if only phase3 is requested
        if not args.no_phase3:
            p3_summaries = run_fusion_event_analysis(fold_dir=fold_dir, out_dir=out_dir)
            for s in p3_summaries:
                s["fold"] = fold_idx
            all_phase3_summaries.extend(p3_summaries)

        if args.no_phase1 and args.no_phase2:
            continue

        # ---- Phases 1 & 2 require model + loader ----
        if fold_idx >= len(plans):
            print(f"  [WARN] fold_idx={fold_idx} >= n_plans={len(plans)} — skipping Phases 1/2",
                  flush=True)
            continue
        split = plans[fold_idx]

        if args.split == "holdout" and split.holdout is not None and not split.holdout.empty:
            eval_df = split.holdout
            split_name = "holdout"
        else:
            if args.split == "holdout":
                print("  [WARN] No holdout set; falling back to val.", flush=True)
            eval_df = split.val
            split_name = "val"
        print(f"  Phases 1/2 using {split_name}: {eval_df['Patient ID'].nunique()} patients",
              flush=True)

        samples = filter_bilateral_samples(
            profile_patient.build_samples(df=eval_df, clinical=data)
        )
        if not samples:
            print(f"  [WARN] No bilateral samples for fold {fold_idx} — skipping Phases 1/2",
                  flush=True)
            continue

        loader = make_loader(
            samples,
            profile_patient.slot_descriptors(),
            image_transform=eval_transform,
            image_preprocessor=None,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=0,
        )

        print(f"  Loading model from {ckpt_path} ...", flush=True)
        model = SingleEyeHT(
            backbone=backbone,
            freeze_ratio=0.0,
            augment=False,
            clinical_data=data,
            num_classes=num_classes,
        ).to(device)
        state = torch.load(ckpt_path, map_location=device)
        model.load_state_dict(state)
        model.eval()

        if not args.no_phase1:
            run_permutation_importance(
                model=model,
                loader=loader,
                data=data,
                num_classes=num_classes,
                device=device,
                n_permutations=args.n_permutations,
                seed=args.seed,
                out_dir=out_dir,
            )

        if not args.no_phase2:
            run_gradcam(
                model=model,
                loader=loader,
                data=data,
                eval_df=eval_df,
                eval_mode=eval_mode,
                backbone=backbone,
                device=device,
                alpha=args.alpha,
                out_dir=out_dir,
            )

        # free GPU memory between folds
        del model
        torch.cuda.empty_cache()

    # ---- Cross-fold MD importance summary ----
    if not args.no_phase1:
        print(f"\n{'='*60}", flush=True)
        print("[explain_run] === Cross-fold MD importance summary ===", flush=True)
        _plot_run_md_importance_summary(run_dir, folds)

    # ---- Cross-fold Phase 3 summary ----
    if all_phase3_summaries and not args.no_phase3:
        print(f"\n{'='*60}", flush=True)
        print("[explain_run] === Cross-fold fusion event summary ===", flush=True)
        df_sum = pd.DataFrame(all_phase3_summaries)
        for split_name, grp in df_sum.groupby("split"):
            print(f"\n  [{split_name}]", flush=True)
            for col in _EVENT_KEYS:
                vals = grp[col].values
                print(f"    {col:20s}  mean={vals.mean():.1f}  total={vals.sum()}", flush=True)
            cd = grp["conf_delta_mean"].values
            print(f"    conf_delta_mean   mean={cd.mean():+.4f}  std={cd.std():.4f}", flush=True)
        summary_csv = run_dir / "explainability_fusion_summary.csv"
        df_sum.to_csv(summary_csv, index=False)
        print(f"\n  Cross-fold summary → {summary_csv}", flush=True)
        _plot_cross_fold_fusion_summary(df_sum, run_dir)

    print("\n[explain_run] Done.", flush=True)


if __name__ == "__main__":
    main()