#!/usr/bin/env python3
"""
Post-hoc explainability for a single saved fold.

Phase 1 — MD permutation feature importance (bar chart + CSV).
Phase 2 — GradCAM overlays on all holdout (or val) patients.

Usage:
    python scripts/output_analysis/explainability/explain_fold.py \
        --fold-dir analysis_data/.../binary/single/fold0 \
        [--checkpoint best_single.pt | best_holdout_single.pt] \
        [--split holdout]              # falls back to val if no holdout
        [--image-dir Papila/FundusImages] \
        [--clinical-dir Papila/ClinicalData] \
        [--n-permutations 30] \
        [--seed 0] \
        [--alpha 0.45]
"""
from __future__ import annotations

import argparse
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 torch
import torch.nn.functional as F
from PIL import Image
from sklearn.metrics import roc_auc_score

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]:
        """Return (cam [H,W] in [0,1], predicted_class_index)."""
        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)  # [1,C,1,1]
        cam = F.relu((weights * self._acts).sum(dim=1, keepdim=True))  # [1,1,h,w]
        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:
    """Return the final spatial feature map layer for GradCAM."""
    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:
    """Blend a jet-coloured GradCAM map onto the original 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]:
    """
    Return a mapping  feature_name → {"value_dims": [...], "missing_dims": [...]}
    that covers every input dimension of the MD tower vector.

    Layout (from DataBundle.vectorize_row):
        [scalar_0..scalar_n-1 | cat_onehot | scalar_missing_0..scalar_missing_n-1]
    """
    n_scalar = len(data.scalar_cols)
    cat_expanded = sum(len(m) for m in data.cat_maps.values())

    feature_map: dict[str, dict] = {}
    idx = 0

    # Scalar features: value_dim + corresponding missing flag
    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]}
        idx += 1

    # Categorical features: permute the entire one-hot block
    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)

    # ---- cache bilateral image embeddings + metadata tensors + labels ----
    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)  # [N, img_dim]
    img2_feats = torch.cat(img2_feats_list)  # [N, img_dim]
    md1_tensor = torch.cat(md1_list)  # [N, feature_dim]
    md2_tensor = torch.cat(md2_list)  # [N, feature_dim]
    y_true = torch.cat(label_list).numpy()
    N = len(y_true)

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

    # ---- baseline AUC (patient-level: average OD/OS fused probabilities) ----
    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)

    # ---- feature index map ----
    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)
            # Apply the same donor patient permutation to both eyes to preserve
            # within-patient coherence while breaking feature-label association.
            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 (all features permuted simultaneously) ----
    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 (tests architectural vs informational benefit) ----
    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,
    )

    # ---- save CSV ----
    import csv

    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})

    # ---- bar chart ----
    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="--")
    # total ablation
    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,
    )
    # gaussian noise ablation
    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)

    target_layer = get_gradcam_layer(model, backbone)
    gcam = GradCAM(target_layer)

    num_classes = model.bridge.classifier_fused[-1].out_features

    overlay_grid_items: list[
        tuple[Image.Image | None, Image.Image | None, str, bool]
    ] = []

    model.eval()
    for batch in loader:
        img_od = batch["image_1"].to(device)  # [1, 3, H, W]
        img_os = batch["image_2"].to(device)  # [1, 3, H, W]
        meta_od = batch["matrix_1"].to(device)  # [1, feature_dim]
        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]

        # GradCAM for each eye (OD drives the prediction label)
        cam_od, pred = gcam.compute(img_od, meta_od, model)
        cam_os, _ = gcam.compute(img_os, meta_os, model)

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

        # Load original (un-normalised) images from disk
        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 '✗'}"
        )

        # ---- per-patient 2×2 figure (OD raw | OD overlay / OS raw | OS overlay) ----
        fig, axes = plt.subplots(2, 2, figsize=(10, 9))
        fig.suptitle(
            title, fontsize=11, fontweight="bold", color="green" if correct else "red"
        )

        # Row 0: OD
        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")

        # Row 1: OS
        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,
        )

        # Accumulate for summary grid
        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
        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))

    gcam.remove()

    # ---- summary grid: N_patients rows × 2 cols (OD overlay | OS overlay) ----
    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 + 0.8))
    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()
    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)


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


def parse_args():
    ap = argparse.ArgumentParser(
        description="Post-hoc explainability for a saved fold."
    )
    ap.add_argument(
        "--fold-dir",
        type=Path,
        required=True,
        help="Path to fold directory, e.g. analysis_data/.../binary/single/fold0",
    )
    ap.add_argument(
        "--checkpoint",
        default="best_single.pt",
        help="Checkpoint filename inside fold_dir (default: best_single.pt; "
        "use best_holdout_single.pt for holdout-selected model)",
    )
    ap.add_argument(
        "--split",
        choices=["holdout", "val"],
        default="holdout",
        help="Which patient set to analyse (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,
        help="Repetitions per feature for permutation importance (default: 30)",
    )
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument(
        "--alpha",
        type=float,
        default=0.45,
        help="GradCAM overlay opacity (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()


# ---------------------------------------------------------------------------
# Phase 3 — Fusion event analysis
# ---------------------------------------------------------------------------


def run_fusion_event_analysis(
    model: SingleEyeHT,
    loader,
    device: torch.device,
    out_dir: Path,
) -> None:
    print("\n[Phase 3] Fusion event analysis ...", flush=True)

    from classes.v2.models import collect_probs_single_components

    y_true, pf, pi, pm = collect_probs_single_components(
        model, loader, device, aggregate_patient=True
    )
    N = len(y_true)
    if N == 0:
        print("  [Phase 3] No samples — skipping.", flush=True)
        return

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

    corrections = (pred_f == y_true) & (pred_i != y_true) & (pred_m != y_true)
    errors      = (pred_f != y_true) & (pred_i == y_true) & (pred_m == y_true)
    n_corr = corrections.sum()
    n_err  = errors.sum()
    both_wrong   = ((pred_i != y_true) & (pred_m != y_true)).sum()
    both_correct = ((pred_i == y_true) & (pred_m == y_true)).sum()

    print(f"  N={N}  corrections={n_corr}  errors={n_err}  ratio={n_corr}/{n_err}", flush=True)
    print(f"  correction rate: {n_corr}/{both_wrong} = {n_corr/max(both_wrong,1):.2%} of both-wrong cases", flush=True)
    print(f"  error rate:      {n_err}/{both_correct} = {n_err/max(both_correct,1):.2%} of both-correct cases", flush=True)

    # ---- cache intermediate hm/hi vectors for all patients ----
    model.eval()
    hm1_list, hm2_list, hi1_list, hi2_list = [], [], [], []
    with torch.no_grad():
        for batch in loader:
            x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
            x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
            if not (torch.is_tensor(x1) and torch.is_tensor(m1)):
                continue
            hi1 = model.bridge.ln_img(model.bridge.W_img(model.img_tower(x1.to(device))))
            hi2 = model.bridge.ln_img(model.bridge.W_img(model.img_tower(x2.to(device))))
            hm1 = model.bridge.ln_md(model.bridge.W_md(model.md_tower(m1.to(device))))
            hm2 = model.bridge.ln_md(model.bridge.W_md(model.md_tower(m2.to(device))))
            hi1_list.append(hi1.cpu()); hi2_list.append(hi2.cpu())
            hm1_list.append(hm1.cpu()); hm2_list.append(hm2.cpu())

    hi1 = torch.cat(hi1_list)  # [N, fusion_dim]
    hi2 = torch.cat(hi2_list)
    hm1 = torch.cat(hm1_list)  # [N, fusion_dim]
    hm2 = torch.cat(hm2_list)

    hm1_mean = hm1.mean(dim=0, keepdim=True)
    hm2_mean = hm2.mean(dim=0, keepdim=True)

    # ---- for each patient: compare logit[true_class] with real hm vs mean hm ----
    gains = []
    with torch.no_grad():
        for idx in range(N):
            true_cls = int(y_true[idx])
            # patient-level average of OD/OS fused vectors (SE skipped: hard to replicate outside forward)
            fused_real = (hi1[idx:idx+1] * hm1[idx:idx+1] + hi2[idx:idx+1] * hm2[idx:idx+1]) * 0.5
            fused_mean = (hi1[idx:idx+1] * hm1_mean + hi2[idx:idx+1] * hm2_mean) * 0.5
            logit_real = model.bridge.classifier_fused(fused_real.to(device))
            logit_mean = model.bridge.classifier_fused(fused_mean.to(device))
            gain = (logit_real[0, true_cls] - logit_mean[0, true_cls]).item()
            gains.append(gain)

    gains = np.array(gains)

    if n_corr > 0:
        corr_gains = gains[corrections]
        helped = (corr_gains > 0).sum()
        print(f"\n  Fusion corrections — MD gate gain vs mean gate:", flush=True)
        print(f"    mean gain = {corr_gains.mean():+.4f}  median = {np.median(corr_gains):+.4f}", flush=True)
        print(f"    real MD helped {helped}/{n_corr} correction patients ({helped/n_corr:.0%})", flush=True)

    if n_err > 0:
        err_gains = gains[errors]
        print(f"\n  Fusion errors — MD gate gain vs mean gate:", flush=True)
        print(f"    mean gain = {err_gains.mean():+.4f}  median = {np.median(err_gains):+.4f}", flush=True)

    # ---- save CSV ----
    import csv
    rows = []
    for idx in range(N):
        rows.append({
            "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(pf[idx].max()),
            "conf_img": float(pi[idx].max()),
            "conf_md": float(pm[idx].max()),
            "is_correction": bool(corrections[idx]),
            "is_error": bool(errors[idx]),
            "md_gate_gain": float(gains[idx]),
        })
    csv_path = out_dir / "fusion_events.csv"
    with csv_path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)
    print(f"  Saved → {csv_path}", flush=True)

    # ---- chart ----
    fig, axes = plt.subplots(1, 2, figsize=(11, 4))

    categories = ["corrections\n(both wrong→fused right)", "errors\n(both right→fused wrong)"]
    counts = [int(n_corr), int(n_err)]
    axes[0].bar(categories, counts, color=["#e05c5c", "#5c9ee0"], width=0.5)
    axes[0].set_ylabel("Count")
    axes[0].set_title(f"Fusion Events  (N={N})")
    for i, v in enumerate(counts):
        axes[0].text(i, v + 0.1, str(v), ha="center", fontsize=11)

    if n_corr > 0:
        axes[1].hist(gains[corrections], bins=10, alpha=0.7, color="#e05c5c", label=f"corrections (n={n_corr})")
    if n_err > 0:
        axes[1].hist(gains[errors], bins=10, alpha=0.7, color="#5c9ee0", label=f"errors (n={n_err})")
    axes[1].axvline(0, color="black", linewidth=0.8)
    axes[1].set_xlabel("MD gate gain vs mean gate\n(logit[true class]: real − mean)")
    axes[1].set_title("Does real MD help the fused prediction?")
    axes[1].legend(fontsize=9)

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


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

    ckpt_path = fold_dir / args.checkpoint
    if not ckpt_path.exists():
        sys.exit(
            f"[ERROR] Checkpoint not found: {ckpt_path}\n"
            f"  Run training with --save-checkpoints (now the default) to produce checkpoints."
        )

    # ---- read config from summary.json in parent (tower-mode) dir ----
    summary_path = fold_dir.parent / "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")
    fold_idx = int(fold_dir.name.replace("fold", ""))
    print(
        f"[explain_fold] fold={fold_idx}  backbone={backbone}  eval_mode={eval_mode}  tower_mode={tower_mode}"
    )

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

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[explain_fold] device={device}  checkpoint={args.checkpoint}")

    # ---- build DataBundle ----
    print("[explain_fold] 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 the exact same split ----
    print("[explain_fold] Reconstructing split ...", flush=True)
    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)
    if fold_idx >= len(plans):
        sys.exit(f"[ERROR] fold_idx={fold_idx} but only {len(plans)} plans built.")
    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 available; falling back to val.", flush=True)
        eval_df = split.val
        split_name = "val"
    print(
        f"  Using {split_name} set: {eval_df['Patient ID'].nunique()} patients",
        flush=True,
    )

    # ---- build loader ----
    profile_patient = build_papila_profile(
        patient_col="Patient ID", label_col=args.label_col, sample_mode="patient"
    )
    samples = filter_bilateral_samples(
        profile_patient.build_samples(df=eval_df, clinical=data)
    )
    if not samples:
        sys.exit("[ERROR] No bilateral samples found in the eval set.")
    loader = make_loader(
        samples,
        profile_patient.slot_descriptors(),
        image_transform=build_eval_transform(backbone),
        image_preprocessor=None,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=0,
    )

    # ---- load model ----
    print(f"[explain_fold] 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()

    # ---- output directory ----
    out_dir = fold_dir / "explainability"
    out_dir.mkdir(exist_ok=True)
    print(f"[explain_fold] Output → {out_dir}", flush=True)

    # ---- Phase 1 ----
    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,
        )

    # ---- Phase 2 ----
    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,
        )

    # ---- Phase 3 ----
    if not args.no_phase3:
        run_fusion_event_analysis(
            model=model,
            loader=loader,
            device=device,
            out_dir=out_dir,
        )

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


if __name__ == "__main__":
    main()