"""F8 - Explainability figures from the R50 checkpointed v4 run.

Sources predictions and Grad-CAM panels exclusively from the
``experiments/explainability/ensemble_refugelike_ckpt`` run (img+cd ensemble
with the refugelike R50 backbone, save_checkpoints=true). rep00 (seed=1234)
is the single rep used for the figure; matches the headline configuration in
section 3.

GradCAM machinery lives in ``v4.classes.accessory.explainability``; PAPILA
specific knowledge (disc contour rasterisation, OS→OD orientation flip) is
attached to ``ImageDataView`` in ``v4.classes.profiles.v4papila`` and consumed
here via getattr so future non-PAPILA profiles can opt in without changing
this file.

Usage
    python -m v4.figures.F8_explainability --only-fusion
    python -m v4.figures.F8_explainability --only-gradcam --run-gradcam
    python -m v4.figures.F8_explainability --run-gradcam        # both
"""
from __future__ import annotations

from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from PIL import Image
from sklearn.metrics import roc_auc_score

from v4.figures.util.loaders import REPO_ROOT


OUT_DIR = Path(__file__).parent / "output"
GRADCAM_DIR = OUT_DIR / "F8_gradcam"
V4_CKPT_RUN = REPO_ROOT / "v4" / "results" / "experiments" / "explainability" / "ensemble_refugelike_ckpt" / "rep00" / "binary"

LABEL_NAMES = {0: "Normal", 1: "Glaucoma"}
EVENT_ORDER = [
    "full_correction", "img_assist", "md_assist",
    "full_error", "img_drag", "md_drag",
    "concordant_correct", "concordant_wrong",
]
EVENT_LABELS = {
    "full_correction": "Both wrong -> fused right",
    "img_assist": "Image right, clinical wrong",
    "md_assist": "Clinical right, image wrong",
    "full_error": "Both right -> fused wrong",
    "img_drag": "Clinical right, image wrong -> fused wrong",
    "md_drag": "Image right, clinical wrong -> fused wrong",
    "concordant_correct": "All correct",
    "concordant_wrong": "All wrong",
}
EVENT_COLORS = {
    "full_correction": "#2f8f5b",
    "img_assist": "#74b66b",
    "md_assist": "#b7c85a",
    "full_error": "#b23a48",
    "img_drag": "#df7f5f",
    "md_drag": "#d2a24c",
    "concordant_correct": "#7da5c9",
    "concordant_wrong": "#9b8fc2",
    "net_positive": "#1f7a4d",
    "net_negative": "#9f2735",
}


def _prediction_frame_from_probs(
    *,
    y_true: np.ndarray,
    probs_fused: np.ndarray,
    probs_img: np.ndarray,
    probs_md: np.ndarray,
    rep: str,
    fold: str,
    source: str,
) -> pd.DataFrame:
    pred_fused = probs_fused.argmax(axis=1)
    pred_img = probs_img.argmax(axis=1)
    pred_md = probs_md.argmax(axis=1)
    df = pd.DataFrame(
        {
            "idx": np.arange(len(y_true)),
            "y_true": y_true.astype(int),
            "pred_fused": pred_fused.astype(int),
            "prob_fused_c0": probs_fused[:, 0],
            "prob_fused_c1": probs_fused[:, 1],
            "pred_img": pred_img.astype(int),
            "prob_img_c0": probs_img[:, 0],
            "prob_img_c1": probs_img[:, 1],
            "pred_md": pred_md.astype(int),
            "prob_md_c0": probs_md[:, 0],
            "prob_md_c1": probs_md[:, 1],
            "rep": rep,
            "fold": fold,
            "source": source,
        }
    )
    for head in ("fused", "img", "md"):
        pred = df[f"pred_{head}"].to_numpy()
        y = df["y_true"].to_numpy()
        df[f"tp_{head}"] = ((pred == 1) & (y == 1)).astype(int)
        df[f"fp_{head}"] = ((pred == 1) & (y == 0)).astype(int)
        df[f"tn_{head}"] = ((pred == 0) & (y == 0)).astype(int)
        df[f"fn_{head}"] = ((pred == 0) & (y == 1)).astype(int)
    return df


def _load_ckpt_predictions(split: str = "test") -> pd.DataFrame:
    """Reconstruct v4 component predictions from checkpointed fold modules.

    predictions.h5 only stores the final hb head, so the image/clinical
    component predictions needed for F8a are rebuilt from tower/head checkpoints.
    """
    if split not in {"test", "val", "both"}:
        raise ValueError(f"split must be 'test', 'val', or 'both', got {split!r}")

    import importlib
    import json
    import torch
    import torch.nn.functional as F
    from v4.classes.split_manager import SplitManager
    from v4.classes.v4_hypertower import load_data, build_towers, _make_loader
    import v4.classes.v4_hypertower as orch
    from v4.classes.heads.classifier import ClassificationHead

    summary = json.loads((V4_CKPT_RUN / "summary.json").read_text())
    cfg = summary["config"]
    orch.cfg_ref = cfg
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    data = load_data(cfg)
    label_filter = cfg.get("label_filter", None)
    df_mode = data.df.copy()
    if label_filter is not None:
        df_mode = df_mode[df_mode[data.label_col].isin(label_filter)].reset_index(drop=True)
    identity_cols = getattr(data, "identity_cols", [])
    identity_level = cfg.get("split_identity_level", 1)
    group_col = identity_cols[identity_level - 1] if identity_level and identity_cols else None
    splits = SplitManager(group_col=group_col).build_plans(
        df_mode,
        label_col=data.label_col,
        n_splits=cfg.get("folds", 5),
        seed=cfg.get("fold_seed", 100),
    )

    stage_by_name = {s["name"]: s for s in cfg["stages"]}
    nt_cfg = stage_by_name["nt"]
    hb_cfg = stage_by_name["hb"]
    rows: list[pd.DataFrame] = []

    for fold_idx, split_obj in enumerate(splits):
        ckpt_dir = V4_CKPT_RUN / "checkpoints" / f"fold{fold_idx}"
        if not ckpt_dir.exists():
            continue

        towers = build_towers(cfg["towers"], data)
        for name, tower in towers.items():
            tower.load_state_dict(torch.load(ckpt_dir / f"tower_{name}.pt", map_location="cpu"))
            tower.to(device).eval()

        nt_mod = importlib.import_module(nt_cfg["module"])
        nt = getattr(nt_mod, nt_cfg["class"])(
            [towers[n].out_dim for n in nt_cfg["inputs"]],
            **nt_cfg.get("args", {}),
        ).to(device)
        nt.load_state_dict(torch.load(ckpt_dir / "stage_nt.pt", map_location="cpu"))
        nt.eval()

        hb_mod = importlib.import_module(hb_cfg["module"])
        hb = getattr(hb_mod, hb_cfg["class"])(
            {"a": nt.out_dim, "b": nt.out_dim},
            **hb_cfg.get("args", {}),
        ).to(device)
        hb.load_state_dict(torch.load(ckpt_dir / "stage_hb.pt", map_location="cpu"))
        hb.eval()

        img_aux = ClassificationHead(towers["img"].out_dim, cfg["num_classes"]).to(device)
        cd_aux = ClassificationHead(towers["cd"].out_dim, cfg["num_classes"]).to(device)
        hb_head = ClassificationHead(hb.out_dim, cfg["num_classes"], dropout=0.3).to(device)
        img_aux.load_state_dict(torch.load(ckpt_dir / "stage_img_aux.pt", map_location="cpu"))
        cd_aux.load_state_dict(torch.load(ckpt_dir / "stage_cd_aux.pt", map_location="cpu"))
        hb_head.load_state_dict(torch.load(ckpt_dir / "stage_hb_head.pt", map_location="cpu"))
        img_aux.eval(); cd_aux.eval(); hb_head.eval()

        split_frames = []
        if split in {"val", "both"}:
            split_frames.append(("val", split_obj.val))
        if split in {"test", "both"} and split_obj.test is not None:
            split_frames.append(("test", split_obj.test))

        for split_name, split_df in split_frames:
            shell = data.build_shells(split_df, level="patient", label_filter=label_filter)
            loader = _make_loader(
                shell,
                towers,
                batch_size=cfg["training"].get("batch_size", 8),
                shuffle=False,
            )
            y_all, pf_all, pi_all, pc_all, ids_all = [], [], [], [], []
            with torch.no_grad():
                for batch in loader:
                    y = batch["label"].detach().cpu().numpy()
                    img_a = towers["img"](batch["img"]["a"].to(device))
                    img_b = towers["img"](batch["img"]["b"].to(device))
                    cd_a = towers["cd"](batch["cd"]["a"].to(device))
                    cd_b = towers["cd"](batch["cd"]["b"].to(device))

                    z_a = nt([img_a, cd_a])
                    z_b = nt([img_b, cd_b])
                    z_hb = hb({"a": z_a, "b": z_b})

                    p_fused = F.softmax(hb_head(z_hb), dim=1).cpu().numpy()
                    p_img = 0.5 * (
                        F.softmax(img_aux(img_a), dim=1).cpu().numpy()
                        + F.softmax(img_aux(img_b), dim=1).cpu().numpy()
                    )
                    p_cd = 0.5 * (
                        F.softmax(cd_aux(cd_a), dim=1).cpu().numpy()
                        + F.softmax(cd_aux(cd_b), dim=1).cpu().numpy()
                    )
                    y_all.append(y)
                    pf_all.append(p_fused)
                    pi_all.append(p_img)
                    pc_all.append(p_cd)
                    ids_all.extend(batch.get("entity_id", []))

            if y_all:
                part = _prediction_frame_from_probs(
                    y_true=np.concatenate(y_all),
                    probs_fused=np.concatenate(pf_all, axis=0),
                    probs_img=np.concatenate(pi_all, axis=0),
                    probs_md=np.concatenate(pc_all, axis=0),
                    rep="v4",
                    fold=f"fold{fold_idx}",
                    source=split_name,
                )
                part["entity_id"] = [str(e) for e in ids_all]
                rows.append(part)

        del towers, nt, hb, img_aux, cd_aux, hb_head
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    if not rows:
        raise FileNotFoundError(f"No reconstructed predictions under {V4_CKPT_RUN}")
    return pd.concat(rows, ignore_index=True)


def _classify_events(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    y = df["y_true"].to_numpy()
    fused_ok = df["pred_fused"].to_numpy() == y
    img_ok = df["pred_img"].to_numpy() == y
    md_ok = df["pred_md"].to_numpy() == y

    def classify(fo: bool, io: bool, mo: bool) -> str:
        if fo and io and mo:
            return "concordant_correct"
        if not fo and not io and not mo:
            return "concordant_wrong"
        if fo and not io and not mo:
            return "full_correction"
        if fo and io and not mo:
            return "img_assist"
        if fo and not io and mo:
            return "md_assist"
        if not fo and io and mo:
            return "full_error"
        if not fo and not io and mo:
            return "img_drag"
        if not fo and io and not mo:
            return "md_drag"
        return "other"

    df["event_type"] = [classify(fo, io, mo) for fo, io, mo in zip(fused_ok, img_ok, md_ok)]
    df["fused_ok"] = fused_ok
    df["img_ok"] = img_ok
    df["md_ok"] = md_ok
    df["tower_state"] = np.select(
        [
            (~img_ok) & (~md_ok),
            img_ok & (~md_ok),
            (~img_ok) & md_ok,
            img_ok & md_ok,
        ],
        [
            "Both networks wrong",
            "Image only correct",
            "Clinical only correct",
            "Both networks correct",
        ],
        default="Other",
    )
    df["tower_mean_c1"] = 0.5 * (df["prob_img_c1"] + df["prob_md_c1"])
    df["fusion_delta_c1"] = df["prob_fused_c1"] - df["tower_mean_c1"]
    df["fusion_margin"] = np.where(
        df["y_true"] == 1,
        df["prob_fused_c1"] - df["tower_mean_c1"],
        (1.0 - df["prob_fused_c1"]) - (1.0 - df["tower_mean_c1"]),
    )
    df["tower_gap_abs"] = (df["prob_img_c1"] - df["prob_md_c1"]).abs()
    return df


def make_fusion_event_panel(split: str = "test") -> None:
    df = _classify_events(_load_ckpt_predictions(split=split))
    auc = roc_auc_score(df["y_true"], df["prob_fused_c1"])
    fold_groups = list(df.groupby(["rep", "fold"]))
    counts = df["event_type"].value_counts().reindex(EVENT_ORDER, fill_value=0)
    positive_keys = EVENT_ORDER[:3]
    negative_keys = EVENT_ORDER[3:6]
    total_positive = int(counts[positive_keys].sum())
    total_negative = int(counts[negative_keys].sum())

    fig = plt.figure(figsize=(15, 8.4))
    gs = fig.add_gridspec(2, 2, height_ratios=[0.9, 1.45], width_ratios=[1.05, 1.0],
                          wspace=0.32, hspace=0.38)

    ax = fig.add_subplot(gs[0, 0])
    display_counts = counts.to_dict()
    display_counts["net_positive"] = total_positive
    display_counts["net_negative"] = total_negative
    display_labels = {
        **EVENT_LABELS,
        "net_positive": "Net positive",
        "net_negative": "Net negative",
    }
    display_order = [
        "full_correction", "img_assist", "md_assist",
        "full_error", "img_drag", "md_drag",
        "net_positive", "net_negative",
    ]
    bars = [k for k in display_order if display_counts.get(k, 0) > 0]
    y = np.arange(len(bars))
    ax.barh(y, [display_counts[k] for k in bars],
            color=[EVENT_COLORS[k] for k in bars], height=0.68)
    ax.set_yticks(y)
    ax.set_yticklabels([display_labels[k] for k in bars], fontsize=8)
    ax.invert_yaxis()
    ax.set_xlabel("Count")
    ax.set_title("Fusion event taxonomy", fontsize=10, fontweight="bold")
    max_count = max(display_counts[k] for k in bars) if bars else 0
    ax.set_xlim(0, max_count * 1.10 + 1)
    for yi, k in enumerate(bars):
        ax.text(display_counts[k] + 0.8, yi, str(int(display_counts[k])),
                va="center", fontsize=8)

    ax = fig.add_subplot(gs[0, 1])
    per_fold = pd.DataFrame(
        {
            "fold": [str(f) for (_, f), _ in fold_groups],
            "positive": [sum((g["event_type"] == k).sum() for k in positive_keys) for _, g in fold_groups],
            "negative": [sum((g["event_type"] == k).sum() for k in negative_keys) for _, g in fold_groups],
        }
    )
    x = np.arange(len(per_fold))
    ax.bar(x, per_fold["positive"], color="#2f8f5b", width=0.72, label="positive correction")
    ax.bar(x, -per_fold["negative"], color="#b23a48", width=0.72, label="negative correction")
    ax.plot(x, per_fold["positive"] - per_fold["negative"], color="#222", lw=1.2,
            marker="o", ms=3, label="net")
    ax.axhline(0, color="#222", lw=0.8)
    ax.set_xticks(x)
    ax.set_xticklabels(per_fold["fold"], rotation=45, ha="right", fontsize=7)
    ax.set_ylabel("Events per fold")
    ax.set_title("Per-fold correction balance", fontsize=10, fontweight="bold")
    ax.legend(fontsize=7, loc="upper right", frameon=False)

    ax = fig.add_subplot(gs[1, :])
    disagree = df[df["event_type"].isin(positive_keys + negative_keys)].copy()
    concordant = df[df["event_type"].isin(["concordant_correct", "concordant_wrong"])].copy()
    n_bins = 28
    agreement_grid = np.zeros((n_bins, n_bins), dtype=float)
    pos_grid = np.zeros((n_bins, n_bins), dtype=float)
    neg_grid = np.zeros((n_bins, n_bins), dtype=float)
    for _, row in concordant.iterrows():
        xi = min(n_bins - 1, max(0, int(row["prob_img_c1"] * n_bins)))
        yi = min(n_bins - 1, max(0, int(row["prob_md_c1"] * n_bins)))
        agreement_grid[yi, xi] += 1
    for _, row in disagree.iterrows():
        xi = min(n_bins - 1, max(0, int(row["prob_img_c1"] * n_bins)))
        yi = min(n_bins - 1, max(0, int(row["prob_md_c1"] * n_bins)))
        if row["event_type"] in positive_keys:
            pos_grid[yi, xi] += 1
        else:
            neg_grid[yi, xi] += 1

    grey_rgba = np.zeros((n_bins, n_bins, 4), dtype=float)
    grey_rgba[:, :, :3] = 0.52
    if agreement_grid.max() > 0:
        grey_rgba[:, :, 3] = 0.06 + 0.22 * np.sqrt(agreement_grid / agreement_grid.max())
    grey_rgba[agreement_grid == 0, 3] = 0.0
    ax.imshow(grey_rgba, extent=[0, 1, 0, 1], origin="lower", aspect="auto",
              interpolation="nearest")

    support = pos_grid + neg_grid
    dominance = np.divide(pos_grid - neg_grid, support,
                          out=np.zeros_like(pos_grid), where=support > 0)
    cmap = plt.get_cmap("RdYlGn")
    rgba = cmap((dominance + 1.0) / 2.0)
    if support.max() > 0:
        rgba[:, :, 3] = 0.08 + 0.46 * np.sqrt(support / support.max())
    rgba[support == 0, 3] = 0.0
    ax.imshow(rgba, extent=[0, 1, 0, 1], origin="lower", aspect="auto",
              interpolation="nearest")

    point_handles = []
    for correct, label, color in [
        (True, "fused correct", "#2f8f5b"),
        (False, "fused wrong", "#b23a48"),
    ]:
        sub = df[df["fused_ok"] == correct]
        h = ax.scatter(
            sub["prob_img_c1"], sub["prob_md_c1"], s=18, alpha=0.45,
            color=color, edgecolors="white", linewidths=0.18,
            label=label,
        )
        point_handles.append(h)
    ax.axvline(0.5, color="#222", lw=0.9, ls="--", alpha=0.75)
    ax.axhline(0.5, color="#222", lw=0.9, ls="--", alpha=0.75)
    ax.plot([0, 1], [0, 1], color="#222", lw=0.7, ls=":", alpha=0.55)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_xlabel("Image network P(glaucoma)")
    ax.set_ylabel("Clinical network P(glaucoma)")
    ax.set_title(
        "Network confidence space - shaded by dominant empirical fusion outcome "
        "(grey=agreement, green=positive correction, red=negative correction)",
        fontsize=10, fontweight="bold",
    )
    ax.text(0.02, 0.96, "clinical says glaucoma", transform=ax.transAxes,
            ha="left", va="top", fontsize=8, color="#333")
    ax.text(0.98, 0.04, "image says glaucoma", transform=ax.transAxes,
            ha="right", va="bottom", fontsize=8, color="#333")
    shade_handles = [
        mpatches.Patch(facecolor="#858585", alpha=0.28, label="agreement density"),
        mpatches.Patch(facecolor="#2f8f5b", alpha=0.34, label="positive correction shade"),
        mpatches.Patch(facecolor="#b23a48", alpha=0.34, label="negative correction shade"),
    ]
    ax.legend(handles=point_handles + shade_handles, ncol=5, fontsize=7,
              loc="upper center", bbox_to_anchor=(0.5, -0.14), frameon=False)

    out = OUT_DIR / "S8a_comparison_panel.png"
    fig.savefig(out, dpi=180, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved: {out}")


def _drop_structural_features(
    names: list[str],
    groups: list[list[int]],
    drop: set[str],
) -> tuple[list[str], list[list[int]]]:
    """Return (names, groups) with entries in ``drop`` removed.

    Encoded-vector indices are left intact - the model still sees them - but
    we simply do not permute/report those dims. Used to hide index-level
    columns like eyeID that appear in the clinical feature list only because
    they double as an index level.
    """
    kept = [(n, g) for n, g in zip(names, groups) if n not in drop]
    if not kept:
        return list(names), list(groups)
    kept_names, kept_groups = zip(*kept)
    return list(kept_names), [list(g) for g in kept_groups]


def make_clinical_importance(n_permutations: int = 30, seed: int = 0) -> None:
    """S8e clinical permutation importance via the cd-tower → cd_aux head.

    Isolates the clinical-only prediction path at the R50 ckpt run, then
    column-shuffles the encoded clinical vector to measure per-feature AUC
    drop. Per-original-column grouping comes from
    ``ClinicalDataView.feature_groups`` (one-hot encoded dims for a
    categorical column shuffle together).

    Aggregates per-fold mean drops, then averages across folds.
    """
    import json
    import torch
    import torch.nn.functional as TF_
    from tqdm import tqdm
    from v4.classes.accessory.explainability import permutation_importance
    from v4.classes.split_manager import SplitManager
    from v4.classes.v4_hypertower import load_data, _make_loader
    import v4.classes.v4_hypertower as orch

    summary_path = V4_CKPT_RUN / "summary.json"
    if not summary_path.exists():
        print(f"[F8] make_clinical_importance: no summary.json under {V4_CKPT_RUN}")
        return

    cfg = json.loads(summary_path.read_text())["config"]
    orch.cfg_ref = cfg
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    data = load_data(cfg)
    clinical_view = data.matrix
    feature_names = getattr(clinical_view, "feature_names", None)
    feature_groups = getattr(clinical_view, "feature_groups", None)
    if feature_names is None or feature_groups is None:
        print(
            "[F8] make_clinical_importance: profile lacks feature_names / "
            "feature_groups; skipping S8e."
        )
        return
    # eyeID is structural (used as the OD/OS index level in DataViews), not a
    # learned variable. Strip it from the reported permutation analysis.
    feature_names, feature_groups = _drop_structural_features(
        feature_names, feature_groups, drop={"eyeID"},
    )

    label_filter = cfg.get("label_filter", None)
    df_mode = data.df.copy()
    if label_filter is not None:
        df_mode = df_mode[df_mode[data.label_col].isin(label_filter)].reset_index(drop=True)
    identity_cols = getattr(data, "identity_cols", [])
    identity_level = cfg.get("split_identity_level", 1)
    group_col = identity_cols[identity_level - 1] if identity_level and identity_cols else None
    splits = SplitManager(group_col=group_col).build_plans(
        df_mode,
        label_col=data.label_col,
        n_splits=cfg.get("folds", 5),
        seed=cfg.get("fold_seed", 100),
    )

    all_results = []
    for fold_idx in tqdm(range(cfg.get("folds", 5)), desc="Importance folds", unit="fold"):
        ckpt_dir = V4_CKPT_RUN / "checkpoints" / f"fold{fold_idx}"
        if not ckpt_dir.exists():
            continue
        towers, nt, hb, img_aux, cd_aux, hb_head = _build_v4_fold_modules(
            fold_idx, cfg, data, device,
        )

        split_obj = splits[fold_idx]
        if split_obj.test is None:
            del towers, nt, hb, img_aux, cd_aux, hb_head
            continue
        shell = data.build_shells(split_obj.test, level="patient", label_filter=label_filter)
        loader = _make_loader(
            shell,
            towers,
            batch_size=cfg["training"].get("batch_size", 8),
            shuffle=False,
        )

        X_list, y_list = [], []
        with torch.no_grad():
            for batch in loader:
                X_list.append(batch["cd"]["a"].cpu().numpy())   # OD-side clinical
                y_list.append(batch["label"].cpu().numpy())
        if not X_list:
            del towers, nt, hb, img_aux, cd_aux, hb_head
            continue
        X = np.concatenate(X_list).astype(np.float32)
        y = np.concatenate(y_list).astype(int)
        if len(np.unique(y)) < 2:
            del towers, nt, hb, img_aux, cd_aux, hb_head
            continue

        def score_fn(X_in: np.ndarray) -> float:
            x_t = torch.from_numpy(X_in).to(device)
            with torch.no_grad():
                z = towers["cd"](x_t)
                logits = cd_aux(z)
                probs = TF_.softmax(logits, dim=1).cpu().numpy()
            try:
                return roc_auc_score(y, probs[:, 1])
            except Exception:
                return float("nan")

        result = permutation_importance(
            score_fn=score_fn,
            X=X,
            groups=feature_groups,
            n_permutations=n_permutations,
            seed=seed + fold_idx,
            feature_names=feature_names,
        )
        all_results.append(result)

        del towers, nt, hb, img_aux, cd_aux, hb_head
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    if not all_results:
        print("[F8] make_clinical_importance: no folds with usable data; skipping S8e.")
        return

    fold_drops = np.stack([r["mean_drop"] for r in all_results])
    mean_across_folds = fold_drops.mean(axis=0)
    std_across_folds = fold_drops.std(axis=0)
    baseline_mean = float(np.mean([r["baseline"] for r in all_results]))
    names = all_results[0]["feature_names"]

    order = np.argsort(mean_across_folds)[::-1]
    sorted_names = [names[i] for i in order]
    sorted_means = mean_across_folds[order]
    sorted_stds = std_across_folds[order]

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    fig, ax = plt.subplots(figsize=(7.5, max(3.5, 0.32 * len(sorted_names))))
    y_pos = np.arange(len(sorted_names))
    ax.barh(
        y_pos, sorted_means, xerr=sorted_stds,
        color="#3B6FB5", edgecolor="black", height=0.7,
        error_kw=dict(ecolor="#444", lw=0.8, capsize=2),
    )
    ax.set_yticks(y_pos)
    ax.set_yticklabels(sorted_names, fontsize=9)
    ax.invert_yaxis()
    ax.set_xlabel("Mean AUC drop on shuffling (averaged across folds)", fontsize=10)
    ax.axvline(0, color="black", linewidth=0.7)
    ax.set_title(
        f"S8e — Clinical permutation importance (cd-only head, R50 ckpt run)\n"
        f"baseline AUC = {baseline_mean:.3f}; n_permutations = {n_permutations}",
        fontsize=10, fontweight="bold",
    )
    ax.grid(axis="x", alpha=0.3, linestyle="--")
    fig.tight_layout()

    out_png = OUT_DIR / "S8e_clinical_importance.png"
    out_csv = OUT_DIR / "S8e_clinical_importance.csv"
    fig.savefig(out_png, dpi=180, bbox_inches="tight")
    plt.close(fig)

    pd.DataFrame({
        "feature": sorted_names,
        "mean_drop": sorted_means,
        "std_drop": sorted_stds,
        "baseline_auc_mean": baseline_mean,
    }).to_csv(out_csv, index=False)

    print(f"Saved: {out_png}")
    print(f"Saved: {out_csv}")


def make_fused_clinical_importance(n_permutations: int = 30, seed: int = 0) -> None:
    """S8e-b clinical permutation importance measured at the fused L2 head.

    Answers a different question from the L1 (cd-only) importance: given the
    image tower is already contributing, which clinical features still change
    the patient-level fused prediction? For each permutation, the clinical
    vector is re-embedded through the cd tower, fused with cached image
    embeddings through the L1 (nt) bridge per eye, aggregated through the L2
    (hb) bridge, and scored against the patient-level label at hb_head.

    OD and OS clinical vectors are permuted together (same source-patient
    replaces both eyes) so patient-level pairing is preserved.

    Aggregation: all 10 reps of ``ensemble_refugelike_ckpt`` are looped; each
    rep produces a per-fold mean AUC drop, averaged within-rep to a rep-level
    drop; final bars show mean ± SD **across reps** (n=10). Baseline is the
    rep-mean of within-rep fold-mean baseline AUCs. This matches the
    rep-mean-of-fold-means aggregation used elsewhere in the manuscript.
    """
    import json
    import torch
    import torch.nn.functional as TF_
    from tqdm import tqdm
    from v4.classes.accessory.explainability import permutation_importance
    from v4.classes.split_manager import SplitManager
    from v4.classes.v4_hypertower import load_data, _make_loader
    import v4.classes.v4_hypertower as orch

    rep_base = V4_CKPT_RUN.parent.parent  # .../ensemble_refugelike_ckpt
    rep_dirs = sorted(
        p for p in rep_base.glob("rep*/binary")
        if (p / "summary.json").exists() and (p / "checkpoints").exists()
    )
    if not rep_dirs:
        print(f"[F8] make_fused_clinical_importance: no reps found under {rep_base}")
        return
    print(f"[F8] Fused clinical importance across {len(rep_dirs)} reps")

    cfg = json.loads((rep_dirs[0] / "summary.json").read_text())["config"]
    orch.cfg_ref = cfg
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    data = load_data(cfg)
    clinical_view = data.matrix
    feature_names = getattr(clinical_view, "feature_names", None)
    feature_groups = getattr(clinical_view, "feature_groups", None)
    if feature_names is None or feature_groups is None:
        print("[F8] make_fused_clinical_importance: profile lacks feature_names / "
              "feature_groups; skipping.")
        return
    # eyeID is structural (index level for OD/OS in DataViews), not a learned
    # variable. Skip it in the reported analysis.
    feature_names, feature_groups = _drop_structural_features(
        feature_names, feature_groups, drop={"eyeID"},
    )

    label_filter = cfg.get("label_filter", None)
    df_mode = data.df.copy()
    if label_filter is not None:
        df_mode = df_mode[df_mode[data.label_col].isin(label_filter)].reset_index(drop=True)
    identity_cols = getattr(data, "identity_cols", [])
    identity_level = cfg.get("split_identity_level", 1)
    group_col = identity_cols[identity_level - 1] if identity_level and identity_cols else None
    splits = SplitManager(group_col=group_col).build_plans(
        df_mode,
        label_col=data.label_col,
        n_splits=cfg.get("folds", 5),
        seed=cfg.get("fold_seed", 100),
    )

    rep_mean_drops = []      # list of (F,) arrays: within-rep fold-mean drops
    rep_mean_baselines = []  # list of floats: within-rep fold-mean baseline AUCs
    reference_names = None

    for rep_idx, run_dir in enumerate(rep_dirs):
        rep_cfg = json.loads((run_dir / "summary.json").read_text())["config"]
        # matched-seed reps share the same fold-splitting seed / label_filter;
        # verify config alignment before reusing splits computed from cfg above.
        if rep_cfg.get("fold_seed") != cfg.get("fold_seed"):
            print(f"[F8]   {run_dir.parent.name}: fold_seed mismatch, rebuilding splits.")
            splits_here = SplitManager(group_col=group_col).build_plans(
                df_mode, label_col=data.label_col,
                n_splits=rep_cfg.get("folds", 5), seed=rep_cfg.get("fold_seed", 100),
            )
        else:
            splits_here = splits

        fold_results = []
        for fold_idx in tqdm(
            range(rep_cfg.get("folds", 5)),
            desc=f"rep {rep_idx:02d}/{len(rep_dirs)-1}",
            unit="fold",
            leave=False,
        ):
            ckpt_dir = run_dir / "checkpoints" / f"fold{fold_idx}"
            if not ckpt_dir.exists():
                continue
            towers, nt, hb, img_aux, cd_aux, hb_head = _build_v4_fold_modules(
                fold_idx, rep_cfg, data, device, run_dir=run_dir,
            )

            split_obj = splits_here[fold_idx]
            if split_obj.test is None:
                del towers, nt, hb, img_aux, cd_aux, hb_head
                continue
            shell = data.build_shells(split_obj.test, level="patient", label_filter=label_filter)
            loader = _make_loader(
                shell,
                towers,
                batch_size=rep_cfg["training"].get("batch_size", 8),
                shuffle=False,
            )

            cd_a_list, cd_b_list = [], []
            z_img_a_list, z_img_b_list = [], []
            y_list = []
            with torch.no_grad():
                for batch in loader:
                    cd_a_list.append(batch["cd"]["a"].cpu().numpy())
                    cd_b_list.append(batch["cd"]["b"].cpu().numpy())
                    z_img_a = towers["img"](batch["img"]["a"].to(device))
                    z_img_b = towers["img"](batch["img"]["b"].to(device))
                    z_img_a_list.append(z_img_a.cpu().numpy())
                    z_img_b_list.append(z_img_b.cpu().numpy())
                    y_list.append(batch["label"].cpu().numpy())
            if not cd_a_list:
                del towers, nt, hb, img_aux, cd_aux, hb_head
                continue
            cd_a  = np.concatenate(cd_a_list).astype(np.float32)
            cd_b  = np.concatenate(cd_b_list).astype(np.float32)
            z_img_a_cached = torch.from_numpy(np.concatenate(z_img_a_list)).to(device)
            z_img_b_cached = torch.from_numpy(np.concatenate(z_img_b_list)).to(device)
            y     = np.concatenate(y_list).astype(int)
            if len(np.unique(y)) < 2:
                del towers, nt, hb, img_aux, cd_aux, hb_head
                continue

            F = cd_a.shape[1]
            X_wide = np.concatenate([cd_a, cd_b], axis=1)  # (N, 2F)
            wide_groups = [[c for c in g] + [c + F for c in g] for g in feature_groups]

            def score_fn(X_in: np.ndarray) -> float:
                cd_a_in = torch.from_numpy(X_in[:, :F]).to(device)
                cd_b_in = torch.from_numpy(X_in[:, F:]).to(device)
                with torch.no_grad():
                    z_cd_a = towers["cd"](cd_a_in)
                    z_cd_b = towers["cd"](cd_b_in)
                    z_nt_a = nt([z_img_a_cached, z_cd_a])
                    z_nt_b = nt([z_img_b_cached, z_cd_b])
                    z_hb   = hb({"a": z_nt_a, "b": z_nt_b})
                    logits = hb_head(z_hb)
                    probs  = TF_.softmax(logits, dim=1).cpu().numpy()
                try:
                    return roc_auc_score(y, probs[:, 1])
                except Exception:
                    return float("nan")

            result = permutation_importance(
                score_fn=score_fn,
                X=X_wide,
                groups=wide_groups,
                n_permutations=n_permutations,
                seed=seed + rep_idx * 100 + fold_idx,
                feature_names=feature_names,
            )
            fold_results.append(result)

            del towers, nt, hb, img_aux, cd_aux, hb_head
            del z_img_a_cached, z_img_b_cached
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        if not fold_results:
            print(f"[F8]   {run_dir.parent.name}: no folds produced usable data; skipping rep.")
            continue

        rep_drops_matrix = np.stack([r["mean_drop"] for r in fold_results])
        rep_mean_drops.append(rep_drops_matrix.mean(axis=0))
        rep_mean_baselines.append(float(np.mean([r["baseline"] for r in fold_results])))
        if reference_names is None:
            reference_names = fold_results[0]["feature_names"]
        print(f"[F8]   {run_dir.parent.name}: baseline={rep_mean_baselines[-1]:.4f} "
              f"(n_folds={len(fold_results)})")

    if not rep_mean_drops:
        print("[F8] make_fused_clinical_importance: no reps produced usable data; skipping.")
        return

    rep_stack        = np.stack(rep_mean_drops)      # (n_reps, F)
    mean_across_reps = rep_stack.mean(axis=0)
    std_across_reps  = rep_stack.std(axis=0)
    baseline_mean    = float(np.mean(rep_mean_baselines))
    baseline_std     = float(np.std(rep_mean_baselines))
    n_reps           = rep_stack.shape[0]
    names            = reference_names

    order = np.argsort(mean_across_reps)[::-1]
    sorted_names = [names[i] for i in order]
    sorted_means = mean_across_reps[order]
    sorted_stds  = std_across_reps[order]

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    fig, ax = plt.subplots(figsize=(7.5, max(3.5, 0.32 * len(sorted_names))))
    y_pos = np.arange(len(sorted_names))
    ax.barh(
        y_pos, sorted_means, xerr=sorted_stds,
        color="#c44e52", edgecolor="black", height=0.7,
        error_kw=dict(ecolor="#444", lw=0.8, capsize=2),
    )
    ax.set_yticks(y_pos)
    ax.set_yticklabels(sorted_names, fontsize=9)
    ax.invert_yaxis()
    ax.set_xlabel(f"Mean fused-head AUC drop on clinical shuffling (mean +/- SD across {n_reps} reps)", fontsize=10)
    ax.axvline(0, color="black", linewidth=0.7)
    ax.set_title(
        "Fused-head clinical permutation importance",
        fontsize=10, fontweight="bold",
    )
    ax.grid(axis="x", alpha=0.3, linestyle="--")
    fig.tight_layout()

    out_png = OUT_DIR / "S8e_fused_clinical_importance.png"
    out_csv = OUT_DIR / "S8e_fused_clinical_importance.csv"
    fig.savefig(out_png, dpi=180, bbox_inches="tight")
    plt.close(fig)

    pd.DataFrame({
        "feature":                 sorted_names,
        "mean_drop_across_reps":   sorted_means,
        "std_drop_across_reps":    sorted_stds,
        "baseline_auc_mean":       baseline_mean,
        "baseline_auc_std":        baseline_std,
        "n_reps":                  n_reps,
    }).to_csv(out_csv, index=False)

    print(f"Saved: {out_png}")
    print(f"Saved: {out_csv}")


def _build_v4_fold_modules(fold_idx: int, cfg: dict, data, device, run_dir: Path | None = None):
    import importlib
    import torch
    from v4.classes.v4_hypertower import build_towers
    from v4.classes.heads.classifier import ClassificationHead

    stage_by_name = {s["name"]: s for s in cfg["stages"]}
    nt_cfg = stage_by_name["nt"]
    hb_cfg = stage_by_name["hb"]
    base = run_dir if run_dir is not None else V4_CKPT_RUN
    ckpt_dir = base / "checkpoints" / f"fold{fold_idx}"
    if not ckpt_dir.exists():
        raise FileNotFoundError(f"No v4 checkpoint directory: {ckpt_dir}")

    towers = build_towers(cfg["towers"], data)
    for name, tower in towers.items():
        tower.load_state_dict(torch.load(ckpt_dir / f"tower_{name}.pt", map_location="cpu"))
        tower.to(device).eval()

    nt_mod = importlib.import_module(nt_cfg["module"])
    nt = getattr(nt_mod, nt_cfg["class"])(
        [towers[n].out_dim for n in nt_cfg["inputs"]],
        **nt_cfg.get("args", {}),
    ).to(device)
    nt.load_state_dict(torch.load(ckpt_dir / "stage_nt.pt", map_location="cpu"))
    nt.eval()

    hb_mod = importlib.import_module(hb_cfg["module"])
    hb = getattr(hb_mod, hb_cfg["class"])(
        {"a": nt.out_dim, "b": nt.out_dim},
        **hb_cfg.get("args", {}),
    ).to(device)
    hb.load_state_dict(torch.load(ckpt_dir / "stage_hb.pt", map_location="cpu"))
    hb.eval()

    img_aux = ClassificationHead(towers["img"].out_dim, cfg["num_classes"]).to(device)
    cd_aux = ClassificationHead(towers["cd"].out_dim, cfg["num_classes"]).to(device)
    hb_head = ClassificationHead(hb.out_dim, cfg["num_classes"], dropout=0.3).to(device)
    img_aux.load_state_dict(torch.load(ckpt_dir / "stage_img_aux.pt", map_location="cpu"))
    cd_aux.load_state_dict(torch.load(ckpt_dir / "stage_cd_aux.pt", map_location="cpu"))
    hb_head.load_state_dict(torch.load(ckpt_dir / "stage_hb_head.pt", map_location="cpu"))
    img_aux.eval(); cd_aux.eval(); hb_head.eval()
    return towers, nt, hb, img_aux, cd_aux, hb_head


def _orient_eye_array(arr: np.ndarray, eye: str) -> np.ndarray:
    """Mirror OS so temporal/nasal anatomy is aligned to OD-style orientation."""
    return np.fliplr(arr) if eye.upper() == "OS" else arr


def _disc_centred_patch_array(
    arr: np.ndarray,
    disc_mask: np.ndarray,
    *,
    span: int = 5,
    out: int = 96,
) -> tuple[np.ndarray | None, float | None]:
    if disc_mask is None or disc_mask.sum() == 0:
        return None, None
    ys, xs = np.where(disc_mask)
    cy, cx = ys.mean(), xs.mean()
    disc_r = float(np.sqrt(disc_mask.sum() / np.pi))
    half = max(1, int(round(span * disc_r / 2)))
    h, w = disc_mask.shape
    y0, y1 = int(round(cy)) - half, int(round(cy)) + half
    x0, x1 = int(round(cx)) - half, int(round(cx)) + half
    pad_spec = ((max(0, -y0), max(0, y1 - h)), (max(0, -x0), max(0, x1 - w)))
    if arr.ndim == 3:
        pad_spec = (*pad_spec, (0, 0))
    arr_pad = np.pad(arr, pad_spec, constant_values=0)
    patch = arr_pad[y0 + pad_spec[0][0]: y1 + pad_spec[0][0],
                    x0 + pad_spec[1][0]: x1 + pad_spec[1][0]]
    if arr.ndim == 2:
        pil = Image.fromarray((np.clip(patch, 0, 1) * 255).astype(np.uint8))
        patch_out = np.array(pil.resize((out, out), Image.BILINEAR)) / 255.0
    else:
        patch_out = np.array(Image.fromarray(patch.astype(np.uint8)).resize((out, out), Image.BILINEAR))
    return patch_out.astype(np.float32), out * disc_r / (2 * half)


QUAD_ORDER = ("ST", "SN", "IT", "IN")  # superotemporal, superonasal, inferotemporal, inferonasal


def _quadrant_fractions(
    cam: np.ndarray,
    disc_mask: np.ndarray,
    *,
    peri_inner: float = 1.0,
    peri_outer: float = 2.0,
) -> tuple[dict[str, float], dict[str, float], dict[str, float]] | tuple[None, None, None]:
    """Per-quadrant Grad-CAM fractions in OD-oriented coordinates, for three
    region scopes.

    Quadrant boundaries are the disc-mask centroid (cx, cy). In the OD-oriented
    frame nasal is left (x < cx) and temporal is right (x > cx); superior is
    top (y < cy) and inferior is bottom (y > cy):

        ST = x > cx,  y < cy
        SN = x < cx,  y < cy
        IT = x > cx,  y > cy
        IN = x < cx,  y > cy

    Three region scopes are returned:
        disc_q : fractions of CAM intensity that fall inside the GT disc mask
        peri_q : fractions inside a peri-disc annulus of disc-radius units
                 (peri_inner to peri_outer, default 1x-2x), excluding the disc
        full_q : fractions over the entire image
    Each dict sums to 1 (within floating-point error). Returns (None, None, None)
    if the disc mask is empty.
    """
    if disc_mask is None or disc_mask.sum() == 0:
        return None, None, None
    ys, xs = np.where(disc_mask)
    cy = float(ys.mean()); cx = float(xs.mean())
    disc_r = float(np.sqrt(disc_mask.sum() / np.pi))
    h, w = cam.shape
    yy, xx = np.mgrid[0:h, 0:w]
    dist = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2)
    peri_mask = (dist >= peri_inner * disc_r) & (dist <= peri_outer * disc_r) & ~disc_mask
    quads = {
        "ST": (xx > cx) & (yy < cy),
        "SN": (xx < cx) & (yy < cy),
        "IT": (xx > cx) & (yy > cy),
        "IN": (xx < cx) & (yy > cy),
    }
    disc_total = float(cam[disc_mask].sum()) + 1e-8
    peri_total = float(cam[peri_mask].sum()) + 1e-8
    full_total = float(cam.sum()) + 1e-8
    disc_q = {k: float(cam[disc_mask & q].sum()) / disc_total for k, q in quads.items()}
    peri_q = {k: float(cam[peri_mask & q].sum()) / peri_total for k, q in quads.items()}
    full_q = {k: float(cam[q].sum())             / full_total for k, q in quads.items()}
    return disc_q, peri_q, full_q


def _annotate_nasal_temporal(ax, *, fontsize: int = 9, color: str = "white",
                             pad: float = 2.5) -> None:
    """Label the disc-side (nasal) and macula-side (temporal) edges of an
    OD-oriented fundus axis. After OS is mirrored to OD orientation, the disc
    sits on the LEFT (nasal) side and the macula on the RIGHT (temporal) side.
    """
    ax.text(0.015, 0.5, "N", transform=ax.transAxes,
            ha="left", va="center", fontsize=fontsize, fontweight="bold",
            color=color,
            bbox=dict(facecolor="black", alpha=0.55, edgecolor="none", pad=pad))
    ax.text(0.985, 0.5, "T", transform=ax.transAxes,
            ha="right", va="center", fontsize=fontsize, fontweight="bold",
            color=color,
            bbox=dict(facecolor="black", alpha=0.55, edgecolor="none", pad=pad))


def _make_oriented_disc_detail(mean_patches: dict, examples: dict, out_path: Path) -> None:
    """2x2 grid of mean Grad-CAM patches centred on the optic disc.

    Rows = ground-truth class (Normal / Glaucoma)
    Cols = prediction outcome relative to truth (Correct / Incorrect)

    Each cell shows the per-pixel mean Grad-CAM intensity in a disc-centred,
    OD-oriented patch (OS mirrored). A dashed circle traces the mean disc
    boundary. Nasal (N) and Temporal (T) edges are labelled.

    Sample overlays are intentionally NOT shown here; see overlay_grid_*.png
    for browseable per-image overlays and hand-pick from those for any
    figure that wants concrete examples.
    """
    from matplotlib.patches import Circle

    classes = ["Normal", "Glaucoma"]
    splits  = ["correct", "incorrect"]

    fig, axes = plt.subplots(2, 2, figsize=(8, 7.8), constrained_layout=True)

    # Column headers
    for ci, split in enumerate(splits):
        axes[0, ci].set_title(
            f"{split.capitalize()} predictions",
            fontsize=11, fontweight="bold", pad=10, color="#222",
        )

    for ri, cls in enumerate(classes):
        # Row label drawn outside the leftmost cell
        axes[ri, 0].text(
            -0.13, 0.5, cls, transform=axes[ri, 0].transAxes,
            ha="right", va="center", fontsize=12, fontweight="bold",
            color="#c44e52" if cls == "Glaucoma" else "#3B6FB5", rotation=90,
        )
        for ci, split in enumerate(splits):
            ax  = axes[ri, ci]
            ax.set_facecolor("#202020")
            key = (cls, split)
            if key in mean_patches:
                patch, radius, count, *rest = mean_patches[key]
                disc_frac = rest[0] if rest else None
                ax.imshow(patch, cmap="jet", vmin=0, vmax=1)
                ax.add_patch(Circle((48, 48), radius, fill=False,
                                    edgecolor="white", linewidth=1.8, linestyle="--"))
                _annotate_nasal_temporal(ax)
                cap = f"n = {count}"
                if disc_frac is not None and not np.isnan(disc_frac):
                    cap += f"   ·   disc-frac = {disc_frac:.2f}"
                ax.text(0.5, -0.06, cap, transform=ax.transAxes,
                        ha="center", va="top", fontsize=9, color="#222")
            else:
                ax.text(0.5, 0.5, "no data", transform=ax.transAxes,
                        ha="center", va="center", color="white", fontsize=9)
            ax.set_xticks([]); ax.set_yticks([])

    fig.savefig(out_path, dpi=180, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved: {out_path}")


def _patient_id_from_entity_id(entity_id) -> int:
    if isinstance(entity_id, (tuple, list)):
        return int(entity_id[0])
    return int(entity_id)


def _v4_gradcam_target_layer(img_tower):
    blocks = getattr(img_tower, "_blocks", None)
    if blocks:
        return blocks[-1]
    backbone = getattr(img_tower, "backbone", None)
    if hasattr(backbone, "features"):
        return backbone.features[-1]
    children = list(backbone.children()) if backbone is not None else []
    if children:
        return children[-1]
    raise RuntimeError(f"Could not infer GradCAM target layer for {type(img_tower).__name__}")


def _maybe(profile_view, method: str, *args, **kwargs):
    """Call profile_view.method(*args, **kwargs) if it exists, else return None."""
    fn = getattr(profile_view, method, None)
    return fn(*args, **kwargs) if fn is not None else None


def _orient_via_profile(profile_view, image_or_array, side: str):
    """Use the profile's orient_for_display if it provides one; else passthrough."""
    return _maybe(profile_view, "orient_for_display", image_or_array, side) \
        if hasattr(profile_view, "orient_for_display") else image_or_array


def _compute_image_aux_cam(
    gcam,
    batch: dict,
    eye: str,
    towers: dict,
    img_aux,
    device,
    target_class: int | None,
) -> tuple[np.ndarray, int]:
    """Run image-tower GradCAM against the img_aux head for one (batch, eye) pair."""
    for mod in [towers["img"], img_aux]:
        mod.zero_grad(set_to_none=True)

    side = "a" if eye.upper() == "OD" else "b"
    img_t = batch["img"][side].to(device)

    return gcam.compute(
        forward_fn=lambda: img_aux(towers["img"](img_t)),
        output_shape=img_t.shape[-2:],
        target_class=target_class,
    )


def _make_oriented_gradcam(n_grid: int = 16, alpha: float = 0.45,
                            target_class: int | None = None) -> None:
    import json
    import torch
    from tqdm import tqdm
    from v4.classes.accessory.explainability import GradCAM, overlay_gradcam
    from v4.classes.split_manager import SplitManager
    from v4.classes.v4_hypertower import load_data, _make_loader
    import v4.classes.v4_hypertower as orch

    summary = json.loads((V4_CKPT_RUN / "summary.json").read_text())
    cfg = summary["config"]
    orch.cfg_ref = cfg
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"GradCAM device: {device}")

    data = load_data(cfg)
    image_view = data.image  # profile data view — owns dataset-specific knowledge
    label_filter = cfg.get("label_filter", None)
    df_mode = data.df.copy()
    if label_filter is not None:
        df_mode = df_mode[df_mode[data.label_col].isin(label_filter)].reset_index(drop=True)
    identity_cols = getattr(data, "identity_cols", [])
    identity_level = cfg.get("split_identity_level", 1)
    group_col = identity_cols[identity_level - 1] if identity_level and identity_cols else None
    splits = SplitManager(group_col=group_col).build_plans(
        df_mode,
        label_col=data.label_col,
        n_splits=cfg.get("folds", 5),
        seed=cfg.get("fold_seed", 100),
    )

    cam_sum = {0: None, 1: None}
    cam_count = {0: 0, 1: 0}
    overlay_items = {0: [], 1: []}
    disc_patch_sum: dict[tuple[str, str], np.ndarray] = {}
    disc_patch_count: dict[tuple[str, str], int] = {}
    disc_radius_sum: dict[tuple[str, str], float] = {}
    disc_frac_sum: dict[tuple[str, str], float] = {}
    # Per-eye quadrant fractions for three region scopes; aggregated per cell.
    quad_disc_list: dict[tuple[str, str], list[dict[str, float]]] = {}
    quad_peri_list: dict[tuple[str, str], list[dict[str, float]]] = {}
    quad_full_list: dict[tuple[str, str], list[dict[str, float]]] = {}
    examples: dict[tuple[str, str], tuple[np.ndarray, int, str, int]] = {}

    fold_range = range(cfg.get("folds", 5))
    for fold_idx in tqdm(fold_range, desc="GradCAM folds", unit="fold"):
        ckpt_dir = V4_CKPT_RUN / "checkpoints" / f"fold{fold_idx}"
        if not ckpt_dir.exists():
            continue
        towers, nt, hb, img_aux, _cd_aux, hb_head = _build_v4_fold_modules(
            fold_idx, cfg, data, device,
        )
        target_layer = _v4_gradcam_target_layer(towers["img"])
        gcam = GradCAM(target_layer)

        split_obj = splits[fold_idx]
        if split_obj.test is None:
            continue
        shell = data.build_shells(split_obj.test, level="patient", label_filter=label_filter)
        loader = _make_loader(shell, towers, batch_size=1, shuffle=False)

        for batch in tqdm(loader, desc=f"fold{fold_idx}", leave=False, unit="pt"):
            pid = _patient_id_from_entity_id(batch["entity_id"][0])
            label = int(batch["label"][0].item())
            for eye in ("OD", "OS"):
                if not image_view.get_image_path(pid, eye).exists():
                    continue
                pil_eval = image_view.eval_image_pil(pid, eye)
                roi_mask = _maybe(image_view, "build_roi_mask",
                                  pid, eye, target_h=pil_eval.size[1], target_w=pil_eval.size[0])

                cam_np, pred = _compute_image_aux_cam(
                    gcam, batch, eye, towers, img_aux, device, target_class,
                )

                # Orient for OD-style display (profile decides; passthrough if it doesn't define)
                cam_np = _orient_via_profile(image_view, cam_np, eye)
                pil_oriented = _orient_via_profile(image_view, pil_eval, eye)
                if roi_mask is not None:
                    roi_mask = _orient_via_profile(image_view, roi_mask, eye)

                if cam_sum[label] is None:
                    cam_sum[label] = cam_np.copy()
                else:
                    cam_sum[label] += cam_np
                cam_count[label] += 1

                if len(overlay_items[label]) < n_grid:
                    overlay_items[label].append(
                        (overlay_gradcam(pil_oriented, cam_np, alpha), pid, eye, pred)
                    )

                patch, roi_r_out = _disc_centred_patch_array(cam_np, roi_mask)
                if patch is not None:
                    key = (LABEL_NAMES[label], "correct" if pred == label else "incorrect")
                    disc_patch_sum[key] = disc_patch_sum.get(key, 0) + patch
                    disc_patch_count[key] = disc_patch_count.get(key, 0) + 1
                    disc_radius_sum[key] = disc_radius_sum.get(key, 0.0) + float(roi_r_out)
                    disc_frac_sum[key] = disc_frac_sum.get(key, 0.0) + float(
                        cam_np[roi_mask].sum() / (cam_np.sum() + 1e-8)
                    )
                    disc_q, peri_q, full_q = _quadrant_fractions(cam_np, roi_mask)
                    if disc_q is not None:
                        quad_disc_list.setdefault(key, []).append(disc_q)
                        quad_peri_list.setdefault(key, []).append(peri_q)
                        quad_full_list.setdefault(key, []).append(full_q)
                    if key not in examples:
                        ov = overlay_gradcam(pil_oriented, cam_np, alpha)
                        ov_small = np.array(ov.resize(cam_np.shape[::-1], Image.BILINEAR))
                        ex_patch, _ = _disc_centred_patch_array(ov_small, roi_mask)
                        if ex_patch is not None:
                            examples[key] = (ex_patch, pid, eye, pred)

                del cam_np

        gcam.remove()
        del towers, nt, hb, img_aux, _cd_aux, hb_head
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    GRADCAM_DIR.mkdir(parents=True, exist_ok=True)

    for cls in (0, 1):
        if cam_count[cls] == 0:
            continue
        mean_cam = cam_sum[cls] / cam_count[cls]
        mean_cam = (mean_cam - mean_cam.min()) / (mean_cam.max() - mean_cam.min() + 1e-8)
        fig, ax = plt.subplots(figsize=(5, 5))
        im = ax.imshow(mean_cam, cmap="jet", vmin=0, vmax=1)
        ax.axis("off")
        ax.set_title(f"Mean v4 image-aux GradCAM - {LABEL_NAMES[cls]} (oriented, n={cam_count[cls]})",
                     fontsize=10, fontweight="bold")
        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        fig.savefig(GRADCAM_DIR / f"mean_cam_{LABEL_NAMES[cls].lower()}.png",
                    dpi=180, bbox_inches="tight")
        plt.close(fig)

    if cam_count[0] and cam_count[1]:
        fig, axes = plt.subplots(1, 2, figsize=(11, 5.5), constrained_layout=True)
        fig.patch.set_facecolor("#f6f6f6")
        fig.suptitle(
            "Mean image-tower Grad-CAM by ground-truth class (OD-oriented)",
            fontsize=12, fontweight="bold",
        )
        cls_colors = {0: "#3B6FB5", 1: "#c44e52"}
        for ax, cls in zip(axes, (0, 1)):
            mean_cam = cam_sum[cls] / cam_count[cls]
            mean_cam = (mean_cam - mean_cam.min()) / (mean_cam.max() - mean_cam.min() + 1e-8)
            im = ax.imshow(mean_cam, cmap="jet", vmin=0, vmax=1)
            _annotate_nasal_temporal(ax)
            ax.set_title(
                f"{LABEL_NAMES[cls]}    n = {cam_count[cls]}",
                fontsize=12, fontweight="bold", color=cls_colors[cls],
            )
            ax.set_xticks([]); ax.set_yticks([])
            fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        fig.savefig(GRADCAM_DIR / "mean_cam_comparison.png", dpi=180, bbox_inches="tight")
        plt.close(fig)

    for cls in (0, 1):
        items = overlay_items[cls]
        if not items:
            continue
        items.sort(key=lambda x: x[3] == cls)
        ncols = 4
        nrows = int(np.ceil(len(items) / ncols))
        fig, axes = plt.subplots(nrows, ncols, figsize=(ncols * 3.1, nrows * 3.1))
        axes = np.array(axes).reshape(-1)
        fig.suptitle(
            f"Per-eye Grad-CAM samples for browsing — ground truth = {LABEL_NAMES[cls]}\n"
            "(OD-oriented; green title = model agreed, red = disagreed)",
            fontsize=11, fontweight="bold",
        )
        for i, ax in enumerate(axes):
            if i < len(items):
                ov, pid, eye, pred = items[i]
                ax.imshow(ov)
                _annotate_nasal_temporal(ax, fontsize=7, pad=1.5)
                color = "#2f7d46" if pred == cls else "#b23a48"
                ax.set_title(f"RET{pid:03d}{eye} -> {LABEL_NAMES[pred]}",
                             fontsize=7.5, color=color)
            ax.axis("off")
        fig.tight_layout()
        fig.savefig(GRADCAM_DIR / f"overlay_grid_{LABEL_NAMES[cls].lower()}.png",
                    dpi=160, bbox_inches="tight")
        plt.close(fig)

    mean_patches = {
        key: (
            disc_patch_sum[key] / disc_patch_count[key],
            disc_radius_sum[key] / disc_patch_count[key],
            disc_patch_count[key],
            disc_frac_sum.get(key, float("nan")) / disc_patch_count[key],
        )
        for key in disc_patch_sum
    }
    if mean_patches:
        _make_oriented_disc_detail(mean_patches, examples, GRADCAM_DIR / "disc_attention_detail.png")

    # Per-quadrant CAM fractions (within-disc and full-image scopes).
    # One CSV row per (class, outcome, scope, quadrant) cell, with mean and SD
    # computed over the per-eye fractions in that cell.
    rows = []
    for key in sorted(quad_disc_list.keys()):
        cls_name, outcome = key
        n_eyes = len(quad_disc_list[key])
        for scope_label, scope_list in (("disc", quad_disc_list[key]),
                                        ("peri", quad_peri_list[key]),
                                        ("full", quad_full_list[key])):
            for q in QUAD_ORDER:
                vals = np.array([d[q] for d in scope_list], dtype=np.float64)
                rows.append({
                    "class":    cls_name,
                    "outcome":  outcome,
                    "scope":    scope_label,
                    "quadrant": q,
                    "n_eyes":   n_eyes,
                    "mean":     float(vals.mean()),
                    "sd":       float(vals.std(ddof=1)) if n_eyes > 1 else float("nan"),
                })
    if rows:
        out_csv = OUT_DIR / "F8_quadrant_fractions.csv"
        pd.DataFrame(rows).to_csv(out_csv, index=False)
        print(f"saved quadrant fractions: {out_csv}")


def make_oriented_gradcam(n_grid: int = 16, alpha: float = 0.45,
                          target_class: int | None = None) -> None:
    _make_oriented_gradcam(n_grid=n_grid, alpha=alpha, target_class=target_class)


def main(*, run_gradcam: bool = False, n_permutations: int = 30,
         fusion_split: str = "test") -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    make_fusion_event_panel(split=fusion_split)
    make_clinical_importance(n_permutations=n_permutations)
    make_fused_clinical_importance(n_permutations=n_permutations)
    if run_gradcam:
        make_oriented_gradcam()
    else:
        print("Skipped GradCAM. Re-run with --run-gradcam to generate oriented heatmaps.")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--only-fusion", action="store_true",
                        help="Generate only the fusion event panel; skip clinical and GradCAM figures.")
    parser.add_argument("--only-gradcam", action="store_true",
                        help="Generate only oriented GradCAM outputs; skip fusion and clinical figures.")
    parser.add_argument("--run-gradcam", action="store_true",
                        help="Generate GPU-heavy oriented GradCAM outputs.")
    parser.add_argument("--n-grid", type=int, default=16,
                        help="Max overlays per class for GradCAM grids.")
    parser.add_argument("--alpha", type=float, default=0.45,
                        help="GradCAM overlay opacity.")
    parser.add_argument("--target-class", type=int, default=None,
                        help="GradCAM target class; default uses predicted class.")
    parser.add_argument("--n-permutations", type=int, default=30,
                        help="Clinical permutation repeats per feature.")
    parser.add_argument("--fusion-split", choices=["test", "val", "both"], default="test",
                        help="Prediction split for the fusion event panel.")
    args = parser.parse_args()
    if args.only_fusion:
        OUT_DIR.mkdir(parents=True, exist_ok=True)
        make_fusion_event_panel(split=args.fusion_split)
    elif args.only_gradcam:
        make_oriented_gradcam(
            n_grid=args.n_grid,
            alpha=args.alpha,
            target_class=args.target_class,
        )
    else:
        main(
            run_gradcam=args.run_gradcam,
            n_permutations=args.n_permutations,
            fusion_split=args.fusion_split,
        )
