hypertower/v4/figures/V2M_F3_hyperfeature_ablation.py

"""V2M_F3 — Single-mode comparison via confidence strips (refuge V2-M backbone).

V2-M counterpart to F3. Panels use the refuge_efficientnet_v2_m image-tower
backbone wherever the image stream is present. cd_solo_single is shared
(no image backbone), so the clinical floor is the same as in F3.

Panels (left → right):
    Clinical only          (cd_solo_single — backbone-independent)
    Image only             (refuge_v2m_baseline/img_solo_single_refuge_v2m)
    Concat fusion          (v2m_variants/single_bcd_concat_v2m)
    Pairwise fusion        (v2m_variants/single_bcd_pairwise_v2m)
    Gated fusion           (v2m_variants/single_bcd_gated_v2m)
    Hadamard fusion        (refuge_v2m_baseline/ensemble_single_refuge_v2m — baseline)

Re-run anytime:
    python -m v4.figures.V2M_F3_hyperfeature_ablation
"""
from __future__ import annotations

import warnings
warnings.filterwarnings("ignore")

from pathlib import Path

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

from v4.figures.util.loaders import REPO_ROOT, RESULTS_ROOT

OUT = Path(__file__).parent / "output" / "V2M_F3_single_mode_strips.png"

# ── Palette (matches v3 confidence_strips) ───────────────────────────────────
C_NORMAL    = "#78909C"
C_EARLY     = "#29B6F6"
C_MODERATE  = "#FFB300"
C_SEVERE    = "#E53935"
C_UNKNOWN   = "#BDBDBD"

SEV_LABELS = {
    "normal":   "Normal",
    "early":    "Glaucoma — early (VF_MD > −6)",
    "moderate": "Glaucoma — moderate (−12 to −6)",
    "severe":   "Glaucoma — severe (VF_MD < −12)",
    "unknown":  "Glaucoma — VF_MD not recorded",
}
SEV_COLORS = {
    "normal":   C_NORMAL,
    "early":    C_EARLY,
    "moderate": C_MODERATE,
    "severe":   C_SEVERE,
    "unknown":  C_UNKNOWN,
}

SEV_ORDER = ["normal", "unknown", "early", "moderate", "severe"]
SEV_ALPHA = {"normal": 0.40, "unknown": 0.35, "early": 0.55, "moderate": 0.70, "severe": 0.85}
SEV_SIZE  = {"normal": 6,    "unknown": 6,    "early": 8,    "moderate": 10, "severe": 12}

# ── Per-panel definitions: (label, results dir, eval_stage) ──────────────────
# Top row: single-modality reference runs (cd_solo is backbone-independent)
TOP_ROW = [
    ("Clinical only",   RESULTS_ROOT / "phase2_v4" / "cd_solo_single",                          "cd_fuse"),
    ("Image only",      RESULTS_ROOT / "refuge_v2m_baseline" / "img_solo_single_refuge_v2m",   "img_fuse"),
]

# Bottom row: L1 fusion bridge variants at refuge V2-M (eye-level img+cd ensembles)
BOTTOM_ROW = [
    ("Concat fusion",   RESULTS_ROOT / "v2m_variants" / "single_bcd_concat_v2m",                "nt"),
    ("Pairwise fusion", RESULTS_ROOT / "v2m_variants" / "single_bcd_pairwise_v2m",              "nt"),
    ("Gated fusion",    RESULTS_ROOT / "v2m_variants" / "single_bcd_gated_v2m",                 "nt"),
    ("Hadamard fusion", RESULTS_ROOT / "refuge_v2m_baseline" / "ensemble_single_refuge_v2m",   "nt"),
]

CLINICAL_DIR = REPO_ROOT / "Papila" / "ClinicalData"


# ── VFI loader (patient-level worst-eye severity, matches v3) ────────────────

def load_vfi() -> pd.DataFrame:
    od = pd.read_excel(CLINICAL_DIR / "patient_data_od.xlsx", header=1)
    os_= pd.read_excel(CLINICAL_DIR / "patient_data_os.xlsx", header=1)

    def _clean(df):
        df = df.copy()
        if "Patient ID" not in df.columns and "ID" in df.columns:
            df.rename(columns={"ID": "Patient ID"}, inplace=True)
        df["Patient ID"] = df["Patient ID"].astype(str).str.extract(r"(\d+)")[0].astype(int)
        df["Diagnosis"]  = pd.to_numeric(df["Diagnosis"], errors="coerce")
        df["VF_MD"]      = pd.to_numeric(df["VF_MD"], errors="coerce")
        # PAPILA: 0=Normal, 1=Glaucoma, 2=Suspect — keep only binary subjects
        df = df[df["Diagnosis"].isin([0, 1])].copy()
        return df[["Patient ID", "Diagnosis", "VF_MD"]]

    both = pd.concat([_clean(od), _clean(os_)], ignore_index=True)
    diag = both.groupby("Patient ID")["Diagnosis"].agg(lambda x: x.mode().iloc[0]).reset_index()
    vf   = both.groupby("Patient ID")["VF_MD"].min().reset_index()
    out  = diag.merge(vf, on="Patient ID").rename(
        columns={"Patient ID": "patient_id", "Diagnosis": "diagnosis", "VF_MD": "vf_md"}
    )

    def _sev(row):
        if int(row["diagnosis"]) == 0: return "normal"
        v = row["vf_md"]
        if pd.isna(v): return "unknown"
        if v > -6:    return "early"
        if v > -12:   return "moderate"
        return "severe"
    out["severity"] = out.apply(_sev, axis=1)
    return out


# ── Prediction pooler ────────────────────────────────────────────────────────

def collect_predictions(run_dir: Path, eval_stage: str) -> pd.DataFrame:
    """Pool test rows across reps × folds. Returns DataFrame with
    patient_id, y_true, prob_glaucoma, rep, fold.
    Applies softmax to the 2-class logits."""
    if not run_dir.exists():
        return pd.DataFrame()
    rows: list[dict] = []
    for rep in sorted(run_dir.glob("rep*")):
        fp = next(iter(rep.rglob("predictions.h5")), None)
        if fp is None: continue
        with h5py.File(fp, "r") as f:
            if eval_stage not in f: continue
            grp = f[eval_stage]
            logits = grp["logits"][:]
            y_true = grp["y_true"][:].astype(int)
            split  = grp["split"][:]
            eid0   = grp["entity_id_0"][:]
            n_folds, n_epochs, n_samples, n_heads, n_outputs = logits.shape
            if n_outputs != 2: continue
            ep, head = n_epochs - 1, n_heads - 1
            for fold in range(n_folds):
                labels = np.array([s.decode() if isinstance(s, bytes) else str(s) for s in split[fold]])
                test_mask = (labels == "test")
                if not test_mask.any(): continue
                lg = logits[fold, ep, test_mask, head, :]   # (n_test, 2)
                # softmax
                e = np.exp(lg - lg.max(axis=1, keepdims=True))
                p = e / e.sum(axis=1, keepdims=True)
                for k, idx in enumerate(np.where(test_mask)[0]):
                    rows.append({
                        "rep":             rep.name,
                        "fold":            fold,
                        "patient_id":      int(eid0[idx]),
                        "y_true":          int(y_true[idx]),
                        "prob_glaucoma":   float(p[k, 1]),
                    })
    return pd.DataFrame(rows)


# ── Panel render ─────────────────────────────────────────────────────────────

C_NORMAL_VIOLIN   = "#4c72b0"
C_GLAUCOMA_VIOLIN = "#c44e52"


def _draw_panel(ax, df: pd.DataFrame, vfi: pd.DataFrame, label: str):
    df = df.merge(vfi[["patient_id", "severity"]], on="patient_id", how="left")
    df["severity"] = df["severity"].fillna("unknown")

    rng = np.random.default_rng(42)
    x_pos = {0: 0.0, 1: 1.0}
    jitter_scale = 0.18

    # Violin density behind everything (per true class)
    data_by_class = [df.loc[df["y_true"] == cls, "prob_glaucoma"].values for cls in [0, 1]]
    if all(len(d) > 0 for d in data_by_class):
        vp = ax.violinplot(data_by_class, positions=[0, 1], widths=0.7,
                           showmedians=False, showextrema=False)
        for body, color in zip(vp["bodies"], [C_NORMAL_VIOLIN, C_GLAUCOMA_VIOLIN]):
            body.set_facecolor(color); body.set_alpha(0.30)
            body.set_edgecolor("none"); body.set_zorder(2)

    for sev in SEV_ORDER:
        mask = df["severity"] == sev
        if not mask.any(): continue
        sub = df[mask]
        jitter = rng.uniform(-jitter_scale, jitter_scale, len(sub))
        x = np.array([x_pos[int(v)] for v in sub["y_true"]]) + jitter
        ax.scatter(x, sub["prob_glaucoma"].values,
                   c=SEV_COLORS[sev], s=SEV_SIZE[sev],
                   alpha=SEV_ALPHA[sev], linewidths=0, zorder=3)

    # Median lines + TN/TP rate labels per class
    xtick_labels = []
    for cls, xc in x_pos.items():
        vals = df.loc[df["y_true"] == cls, "prob_glaucoma"]
        if not len(vals):
            xtick_labels.append("Normal" if cls == 0 else "Glaucoma"); continue
        med = float(np.median(vals))
        ax.plot([xc - jitter_scale - 0.04, xc + jitter_scale + 0.04],
                [med, med], color="#222", lw=2.0, zorder=5)
        if cls == 0:
            rate = (vals <= 0.5).mean() * 100
            xtick_labels.append(f"Normal\nTN {rate:.0f}%")
        else:
            rate = (vals > 0.5).mean() * 100
            xtick_labels.append(f"Glaucoma\nTP {rate:.0f}%")

    ax.axhline(0.5, color="#888", lw=1.2, ls="--", alpha=0.7, zorder=2)
    ax.set_xticks([0, 1]); ax.set_xticklabels(xtick_labels, fontsize=10)
    ax.set_ylim(-0.04, 1.04); ax.set_xlim(-0.55, 1.55)
    ax.set_title(label, fontsize=11, fontweight="bold")
    ax.grid(axis="y", alpha=0.3, zorder=1)

    # AUC across reps×folds (per-fold AUC averaged)
    fold_aucs = []
    for _, g in df.groupby(["rep", "fold"]):
        if g["y_true"].nunique() < 2: continue
        try: fold_aucs.append(roc_auc_score(g["y_true"], g["prob_glaucoma"]))
        except Exception: pass
    if fold_aucs:
        ax.text(0.66, 0.0,
                f"AUC = {np.mean(fold_aucs):.3f} ± {np.std(fold_aucs):.3f}",
                transform=ax.transAxes, ha="right", va="bottom",
                fontsize=9, color="#333",
                bbox=dict(facecolor="white", alpha=0.7, edgecolor="none", pad=2))


def render() -> None:
    vfi = load_vfi()
    top_dfs    = [(lbl, collect_predictions(p, s)) for lbl, p, s in TOP_ROW]
    bottom_dfs = [(lbl, collect_predictions(p, s)) for lbl, p, s in BOTTOM_ROW]

    for lbl, df in top_dfs + bottom_dfs:
        if len(df):
            print(f"  {lbl:<18s}  n_rows={len(df):>5d}  (pid={df['patient_id'].nunique()}, reps={df['rep'].nunique()})")
        else:
            print(f"  {lbl:<18s}  no data")

    n_cols = len(BOTTOM_ROW)
    fig = plt.figure(figsize=(4.4 * n_cols, 12))
    fig.patch.set_facecolor("#e8e8e8")
    fig.suptitle("L1 Fusion Comparison — refuge V2-M backbone", fontsize=14, fontweight="bold")

    gs = fig.add_gridspec(2, n_cols, hspace=0.30, wspace=0.15)

    # Top row: 2 reference panels at same width as bottom panels, centered.
    # In a 4-column bottom grid, that's columns 1 and 2.
    n_top = len(top_dfs)
    top_offset = (n_cols - n_top) // 2  # leading empty columns
    top_axes = []
    for i, (lbl, df) in enumerate(top_dfs):
        ax = fig.add_subplot(gs[0, top_offset + i])
        top_axes.append(ax)
        ax.set_facecolor("#e8e8e8")
        if len(df):
            _draw_panel(ax, df, vfi, lbl)
        else:
            ax.text(0.5, 0.5, "(pending)", ha="center", va="center",
                    fontsize=12, color="#888", transform=ax.transAxes)
            ax.set_xticks([]); ax.set_yticks([])
            ax.set_title(lbl, fontsize=11, fontweight="bold")

    # Bottom row: 4 fusion variants
    bottom_axes = []
    sharey = None
    for i, (lbl, df) in enumerate(bottom_dfs):
        ax = fig.add_subplot(gs[1, i], sharey=sharey)
        sharey = sharey or ax
        bottom_axes.append(ax)
        ax.set_facecolor("#e8e8e8")
        if len(df):
            _draw_panel(ax, df, vfi, lbl)
        else:
            ax.text(0.5, 0.5, "(pending)", ha="center", va="center",
                    fontsize=12, color="#888", transform=ax.transAxes)
            ax.set_xticks([]); ax.set_yticks([])
            ax.set_title(lbl, fontsize=11, fontweight="bold")

    top_axes[0].set_ylabel("Predicted P(Glaucoma)", fontsize=11)
    bottom_axes[0].set_ylabel("Predicted P(Glaucoma)", fontsize=11)
    for ax in bottom_axes[1:]:
        ax.set_yticklabels([])

    legend_patches = [mpatches.Patch(color=SEV_COLORS[s], label=SEV_LABELS[s])
                      for s in ["normal", "early", "moderate", "severe", "unknown"]]
    fig.legend(handles=legend_patches, fontsize=9,
               loc="lower center", ncol=len(legend_patches),
               framealpha=0.75, bbox_to_anchor=(0.5, -0.01))

    fig.tight_layout(rect=[0, 0.06, 1, 0.97])
    OUT.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(OUT, dpi=180, bbox_inches="tight")
    plt.close(fig)
    print(f"saved {OUT}")


if __name__ == "__main__":
    render()