scripts/figures/build_figure_supplement.py

"""builds figures/supplement/*.pdf and merges into figures/PANDA_supplement.pdf."""
from __future__ import annotations
from pathlib import Path
import warnings, json, pickle, sys, numpy as np, pandas as pd
warnings.filterwarnings("ignore")

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.colors import LinearSegmentedColormap, TwoSlopeNorm

sys.path.insert(0, str(Path(__file__).parent))
from palette import apply_style, color_for
apply_style()

# ---- shared font-size constants (style spec): small titles, larger body ----
SUPTITLE_FS = 15   # figure suptitles
TITLE_FS = 14      # panel/axes titles
LABEL_FS = 12      # axis labels
TICK_FS = 11       # tick labels
LEGEND_FS = 10     # legend text
ANNOT_FS = 10      # in-plot annotations
CBAR_FS = 11       # colorbar labels
plt.rcParams.update({
    "font.size": 11,
    "axes.titlesize": TITLE_FS,
    "figure.titlesize": SUPTITLE_FS,
    "axes.labelsize": LABEL_FS,
    "xtick.labelsize": TICK_FS,
    "ytick.labelsize": TICK_FS,
    "legend.fontsize": LEGEND_FS,
})

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
ROOT = Path(str(PANDA_ROOT))
FIG = ROOT / "figures"
FIG_S = FIG / "supplement"; FIG_S.mkdir(parents=True, exist_ok=True)
DISC = ROOT / "discovery"
SUP_PDF = FIG / "PANDA_supplement.pdf"

SYSTEMS = ["pan_skin", "hematopoiesis", "pancreas"]
SYS_LABEL = {"pan_skin": "Pan-skin", "hematopoiesis": "Pan-hematopoietic", "pancreas": "Pan-pancreatic"}
SYS_COLOR = {"pan_skin": "#2b83ba", "hematopoiesis": "#d7191c", "pancreas": "#7570b3"}

# Deprecated class names -- drop these if they appear in any legacy CSV so
# regenerated figures never re-surface stale vocabulary.  Extend as vocab evolves.
STALE_CLASSES = {
    "HF-DP", "basal-multipotent", "eden-dermal-niche", "eccrine-secretory",
    "hair-placode", "UNK", "nascent-eccrine-gland", "dermal-condensate",
    "mesenchymal", "adult-alpha", "adult-beta", "eccrine-duct", "eccrine-ductal",
    "eccrine-placode", "other", "unassigned",
}


def _drop_stale(labels):
    """Return input filtered to remove any deprecated class names.



    Accepts a list/Index/Series of strings, or a DataFrame column selector case

    handled by the caller.  Order preserved.

    """
    return [x for x in labels if x not in STALE_CLASSES]


def _filter_df_stale(df, col):
    """Drop rows of df whose value in `col` is a deprecated class name."""
    if col in df.columns:
        return df[~df[col].astype(str).isin(STALE_CLASSES)].reset_index(drop=True)
    return df

CV_JSON = {
    "pan_skin":      ROOT / "discovery/pan_skin/marker/cv_5fold.json",
    "hematopoiesis": ROOT / "discovery/hematopoiesis/marker/cv_5fold.json",
    "pancreas":      ROOT / "discovery/pancreas/marker/cv_5fold.json",
}

# ============================================================
# PAGE 1: Held-out CV summary (per-system accuracy + AUROC + F1)
# ============================================================

def page_cv_summary():
    fig, axes = plt.subplots(1, 2, figsize=(15.5, 6.5))
    sys_names = []
    accs, acc_errs, aucs, auc_errs, folds_data = [], [], [], [], []
    for sys in SYSTEMS:
        r = json.load(open(CV_JSON[sys]))
        sys_names.append(SYS_LABEL[sys])
        accs.append(r["mean_acc"]); acc_errs.append(r["std_acc"])
        aucs.append(r["mean_auc"]); auc_errs.append(r["std_auc"])
        folds_data.append((r["per_fold_acc"], r["per_fold_auc"]))

    ax = axes[0]
    x = np.arange(len(sys_names))
    b1 = ax.bar(x - 0.2, accs, 0.4, yerr=acc_errs, capsize=4,
                color=[SYS_COLOR[s] for s in SYSTEMS], label="accuracy", alpha=0.85, edgecolor="k", linewidth=0.7)
    b2 = ax.bar(x + 0.2, aucs, 0.4, yerr=auc_errs, capsize=4,
                color=[SYS_COLOR[s] for s in SYSTEMS], label="macro AUROC", alpha=0.55, edgecolor="k", linewidth=0.7, hatch="//")
    ax.set_xticks(x); ax.set_xticklabels(sys_names, rotation=0, fontsize=TICK_FS)
    ax.set_ylabel("Held-out score", fontsize=LABEL_FS)
    ax.set_ylim(0.7, 1.02)
    ax.set_title("(a) 5-fold held-out CV: accuracy + macro AUROC (mean ± std)", fontsize=TITLE_FS)
    for i, (pa, pu) in enumerate(folds_data):
        ax.scatter([x[i] - 0.2] * len(pa), pa, s=22, color="black", zorder=3)
        ax.scatter([x[i] + 0.2] * len(pu), pu, s=22, color="black", zorder=3)
    for i, (a, e, u, ue) in enumerate(zip(accs, acc_errs, aucs, auc_errs)):
        ax.text(x[i] - 0.2, a + e + 0.008, f"{a:.3f}", ha="center", fontsize=ANNOT_FS)
        ax.text(x[i] + 0.2, u + ue + 0.008, f"{u:.3f}", ha="center", fontsize=ANNOT_FS)
    ax.legend(loc="lower right", fontsize=LEGEND_FS)

    ax = axes[1]
    baron_pca = json.load(open(DISC / "pancreas/pca/93_baron_merged_metrics.json"))
    baron_mar = json.load(open(DISC / "pancreas/marker/93_baron_merged_metrics.json"))
    nest_pca = json.load(open(DISC / "hematopoiesis/pca/97_nestorowa_anchor_zero_shot.json"))
    nest_mar = json.load(open(DISC / "hematopoiesis/marker/97_nestorowa_anchor_zero_shot.json"))
    sul_pca = json.load(open(DISC / "pan_skin/pca/97_sulic_anchor_zero_shot.json"))
    sul_mar = json.load(open(DISC / "pan_skin/marker/97_sulic_anchor_zero_shot.json"))
    # sulic 5-seed topline: in-domain ceiling reference
    sulic5 = json.load(open(ROOT / "scripts/sulic/sulic_panda_heldout_results.json"))

    bars_data = [
        ("Baron\ntest-half\nPCA",    baron_pca["acc"], "acc",   SYS_COLOR["pancreas"],      0.55),
        ("Baron\ntest-half\nMarker", baron_mar["acc"], "acc",   SYS_COLOR["pancreas"],      0.90),
        ("Nestorowa\nanchor\nPCA",   nest_pca["coarse_acc"], "acc", SYS_COLOR["hematopoiesis"], 0.55),
        ("Nestorowa\nanchor\nMarker",nest_mar["coarse_acc"], "acc", SYS_COLOR["hematopoiesis"], 0.90),
        ("Sulic\nanchor\nPCA",       sul_pca["acc"],   "acc",   SYS_COLOR["pan_skin"],       0.55),
        ("Sulic\nanchor\nMarker",    sul_mar["acc"],   "acc",   SYS_COLOR["pan_skin"],       0.90),
        ("Sulic\n5-fold\n(topline)", sulic5["testA"]["mean_auroc"], "AUROC", "#888", 0.90),
    ]
    xe = np.arange(len(bars_data))
    for i, (lab, val, met, col, alpha) in enumerate(bars_data):
        ax.bar(xe[i], val, color=col, alpha=alpha, edgecolor="k", linewidth=0.7)
        ax.text(xe[i], val + 0.012, f"{val:.3f}\n({met})", ha="center", fontsize=ANNOT_FS)
    ax.set_xticks(xe); ax.set_xticklabels([b[0] for b in bars_data], fontsize=TICK_FS)
    ax.set_ylim(0.4, 1.05)
    ax.set_ylabel("Held-out score", fontsize=LABEL_FS)
    ax.set_title("(b) External held-out labeled validation (6 zero-shot + 1 topline)", fontsize=TITLE_FS)

    plt.suptitle("PANDA held-out validation across systems", fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "01_cv_summary.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 01_cv_summary.pdf")


# ============================================================
# PAGE 2: Per-class F1 across all three systems
# ============================================================

def page_per_class_f1():
    fig, axes = plt.subplots(1, 3, figsize=(21, 7.0))
    for ax, sys in zip(axes, SYSTEMS):
        r = json.load(open(CV_JSON[sys]))
        rep = r["per_class_report"]
        classes = [c for c in rep if c not in ("accuracy", "macro avg", "weighted avg")]
        f1s = [rep[c]["f1-score"] for c in classes]
        supp = [rep[c]["support"] for c in classes]
        order = np.argsort(f1s)[::-1]
        classes = [classes[i] for i in order]
        f1s = [f1s[i] for i in order]; supp = [supp[i] for i in order]

        y = np.arange(len(classes))
        cols = [color_for(c) for c in classes]
        bars = ax.barh(y, f1s, color=cols, edgecolor="k", linewidth=0.5)
        for i, (b, s) in enumerate(zip(bars, supp)):
            # long bars: write inside the bar (right-aligned, white); short bars:
            # write outside, right of the bar end, so text never spills into the
            # y-tick labels (e.g. mesenchyme, F1=0.276, Pan-pancreatic panel)
            if b.get_width() > 0.35:
                ax.text(b.get_width() - 0.02, b.get_y() + b.get_height() / 2,
                        f"F1={f1s[i]:.3f} · n={int(s):,}",
                        va="center", ha="right", fontsize=ANNOT_FS, color="white")
            else:
                ax.text(b.get_width() + 0.02, b.get_y() + b.get_height() / 2,
                        f"F1={f1s[i]:.3f} · n={int(s):,}",
                        va="center", ha="left", fontsize=ANNOT_FS, color="black")
        ax.set_yticks(y); ax.set_yticklabels(classes, fontsize=TICK_FS)
        ax.set_xlim(0, 1.05)
        ax.set_xlabel("F1", fontsize=LABEL_FS)
        ax.invert_yaxis()
        ax.set_title(f"{SYS_LABEL[sys]}\nacc={r['mean_acc']:.4f}  macro-AUROC={r['mean_auc']:.4f}",
                     fontsize=TITLE_FS)
        ax.axvline(0.85, color="grey", linestyle=":", linewidth=0.7, alpha=0.6)
        ax.grid(axis="x", alpha=0.3, linestyle="--")
    plt.suptitle("Per-class F1 in 5-fold CV (bar colour = canonical class palette)",
                 fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "02_per_class_f1.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 02_per_class_f1.pdf")


# ============================================================
# PAGE 3: Prototype intra-cosine heatmaps (3 panels)
# ============================================================

def page_prototype_cosine():
    fig, axes = plt.subplots(1, 3, figsize=(23, 7.5))
    eff = json.load(open(DISC / "70_prototype_effective_dim.json"))
    for ax, sys in zip(axes, SYSTEMS):
        M = pd.read_csv(DISC / f"70_prototype_intra_cosine_{sys}.csv", index_col=0)
        vmax = np.nanmax(np.abs(M.values))
        im = ax.imshow(M.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
        # pancreas has 20x20 cells → tighter font + higher threshold
        ann_fs = ANNOT_FS  # style spec: in-plot annotations >= 10 (was 6/8)
        ann_thresh = 0.45 if sys == "pancreas" else 0.3
        for i in range(M.shape[0]):
            for j in range(M.shape[1]):
                v = M.values[i, j]
                if abs(v) < ann_thresh: continue
                ax.text(j, i, f"{v:+.2f}", ha="center", va="center",
                        fontsize=ann_fs,
                        color="white" if abs(v) > vmax * 0.55 else "black")
        ax.set_xticks(range(len(M.columns)))
        ax.set_xticklabels(M.columns, rotation=45, ha="right", fontsize=TICK_FS)
        ax.set_yticks(range(len(M.index)))
        ax.set_yticklabels(M.index, fontsize=TICK_FS)
        ax.set_title(f"{SYS_LABEL[sys]}\nK={eff[sys]['K']}  eff-dim={eff[sys]['effective_dim']:.2f}",
                     fontsize=TITLE_FS)
        cb = plt.colorbar(im, ax=ax, shrink=0.7); cb.set_label("prototype cosine", fontsize=CBAR_FS)
    plt.suptitle("Prototype-prototype cosine matrices per system (§8.1)",
                 fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "03_prototype_cosine.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 03_prototype_cosine.pdf")


# ============================================================
# PAGE 4: Training trajectory eff-dim over stages
# ============================================================

def page_training_trajectory():
    return  # skipped — v2 checkpoints only save final stage
    fig, axes = plt.subplots(1, 2, figsize=(13, 4.6))
    ax = axes[0]
    for sys in SYSTEMS:
        d = pd.read_csv(DISC / f"83_{sys}_effdim_by_stage.csv")
        stage_order = ["stage0", "stage1", "stage2", "stage3", "final"]
        d = d.set_index("stage").loc[stage_order].reset_index()
        ax.plot(range(len(d)), d["eff_dim"], "-o", color=SYS_COLOR[sys],
                label=f"{SYS_LABEL[sys]} (K={int(d['K'].iloc[0])})", linewidth=2.5, markersize=8)
        for i, row in d.iterrows():
            ax.text(i, row["eff_dim"] + 0.35, f"{row['eff_dim']:.2f}",
                    ha="center", fontsize=8, color=SYS_COLOR[sys])
    ax.set_xticks(range(5)); ax.set_xticklabels(stage_order, fontsize=10)
    ax.set_xlabel("training stage"); ax.set_ylabel("prototype set effective dimensionality")
    ax.set_title("(a) Effective-dim trajectory\nPancreas collapses 9.50→1.98 in stage 0→1", fontsize=11)
    ax.legend(fontsize=9)
    ax.grid(alpha=0.3, linestyle="--")

    ax = axes[1]
    all_data, all_labels = [], []
    for sys in SYSTEMS:
        traj = pd.read_csv(DISC / f"83_{sys}_prototype_trajectory.csv")
        classes = traj["class"].unique()
        total_cos = []
        for cls in classes:
            # stage0 vs final cos — reload prototypes directly (intermediate-cos product is wrong)
            import torch
            ck0 = torch.load(ROOT / f"checkpoints/{sys}/panda_stage0.pt",
                             map_location="cpu", weights_only=False)
            ckf = torch.load(ROOT / f"checkpoints/{sys}/panda_final.pt",
                             map_location="cpu", weights_only=False)
            P0 = ck0["model"]["prototypes"].numpy() if hasattr(ck0["model"]["prototypes"], "numpy") else ck0["model"]["prototypes"]
            Pf = ck0["model"]["prototypes"].numpy() if hasattr(ck0["model"]["prototypes"], "numpy") else ck0["model"]["prototypes"]
            # actually need final
            Pf = ckf.get("prototypes", ckf["model"]["prototypes"])
            if hasattr(Pf, "numpy"): Pf = Pf.numpy()
            P0 = P0 / (np.linalg.norm(P0, axis=1, keepdims=True) + 1e-8)
            Pf = Pf / (np.linalg.norm(Pf, axis=1, keepdims=True) + 1e-8)
            ci = ck0["classes"].index(cls)
            total_cos.append(float((P0[ci] * Pf[ci]).sum()))
            break  # only need once per sys
        # per-sys one-shot pass
        ck0 = torch.load(ROOT / f"checkpoints/{sys}/panda_stage0.pt",
                         map_location="cpu", weights_only=False)
        ckf = torch.load(ROOT / f"checkpoints/{sys}/panda_final.pt",
                         map_location="cpu", weights_only=False)
        P0 = ck0["model"]["prototypes"]
        if hasattr(P0, "numpy"): P0 = P0.numpy()
        Pf = ckf.get("prototypes", ckf["model"]["prototypes"])
        if hasattr(Pf, "numpy"): Pf = Pf.numpy()
        P0 = P0 / (np.linalg.norm(P0, axis=1, keepdims=True) + 1e-8)
        Pf = Pf / (np.linalg.norm(Pf, axis=1, keepdims=True) + 1e-8)
        total_cos = [float((P0[i] * Pf[i]).sum()) for i in range(len(ck0["classes"]))]
        all_data.append(total_cos)
        all_labels.append([f"{sys[:3]}:{c}" for c in ck0["classes"]])

    flat_vals = np.concatenate([np.asarray(a) for a in all_data])
    flat_labels = sum(all_labels, [])
    order = np.argsort(flat_vals)
    flat_vals = flat_vals[order]
    flat_labels = [flat_labels[i] for i in order]
    colors = ["#d7191c" if v < 0 else "#fdae61" if v < 0.15 else "#2b83ba" for v in flat_vals]
    ax.barh(range(len(flat_vals)), flat_vals, color=colors, edgecolor="k", linewidth=0.4)
    ax.set_yticks(range(len(flat_vals)))
    ax.set_yticklabels(flat_labels, fontsize=6.5)
    ax.axvline(0, color="k", linewidth=0.5)
    ax.set_xlabel("cosine( prototype_stage0 , prototype_final )")
    ax.set_title("(b) Per-class prototype drift 0→final\nNegative = moved to opposite direction", fontsize=11)
    ax.set_xlim(-0.35, 0.45)
    ax.grid(axis="x", alpha=0.3, linestyle="--")

    plt.suptitle("§10.1 Training trajectory: effective-dim collapse + per-class prototype drift",
                 fontsize=12, y=1.02)
    plt.tight_layout()
    plt.savefig(FIG_S / "04_training_trajectory.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 04_training_trajectory.pdf")


# ============================================================
# PAGE 5: Adversary purification (dataset + depth)
# ============================================================

def page_adversary_purification():
    adv = json.load(open(DISC / "84_adversary_purification.json"))
    fig, axes = plt.subplots(1, 2, figsize=(17, 6.5))

    ax = axes[0]
    xs, accs, chances, colors, labels = [], [], [], [], []
    for i, sys in enumerate(SYSTEMS):
        r = adv[sys]
        if "dataset_adv_accuracy" not in r: continue
        xs.append(i); accs.append(r["dataset_adv_accuracy"])
        chances.append(r["dataset_adv_chance"])
        colors.append(SYS_COLOR[sys])
        labels.append(f"{SYS_LABEL[sys]}\n(n={r['n_train_datasets']} datasets)")
    x = np.arange(len(xs))
    bars = ax.bar(x - 0.2, accs, 0.4, color=colors, edgecolor="k", linewidth=0.7,
                  label="observed", alpha=0.85)
    ax.bar(x + 0.2, chances, 0.4, color=colors, edgecolor="k", linewidth=0.7,
           label="chance", alpha=0.35, hatch="///")
    for i, (a, c) in enumerate(zip(accs, chances)):
        ax.text(x[i] - 0.2, a + 0.015, f"{a:.3f}", ha="center", va="bottom", fontsize=ANNOT_FS)
        ax.text(x[i] + 0.2, c + 0.015, f"{c:.3f}", ha="center", va="bottom", fontsize=ANNOT_FS)
        ax.text(x[i], -0.03, f"+{a-c:.3f}\nabove chance", ha="center", va="top",
                fontsize=ANNOT_FS, color="red")
    ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=TICK_FS)
    ax.set_ylabel("dataset adversary accuracy", fontsize=LABEL_FS)
    ax.set_ylim(-0.22, 1.12)
    ax.set_title("(a) Residual batch-signal\n(closer to chance = more purified)", fontsize=TITLE_FS)
    ax.legend(loc="upper left", fontsize=LEGEND_FS)
    ax.grid(axis="y", alpha=0.3, linestyle="--")

    ax = axes[1]
    xs, r2s, colors, labels = [], [], [], []
    for i, sys in enumerate(SYSTEMS):
        r = adv[sys]
        if "depth_adv_r2_z" not in r: continue
        xs.append(i); r2s.append(r["depth_adv_r2_z"])
        colors.append(SYS_COLOR[sys])
        labels.append(SYS_LABEL[sys])
    x = np.arange(len(xs))
    bars = ax.bar(x, r2s, color=colors, edgecolor="k", linewidth=0.7)
    for i, r in enumerate(r2s):
        # label above bar if positive, below top if bar tall enough
        ypos = (r + 0.02) if r >= 0 else (r - 0.03)
        ax.text(x[i], ypos, f"R²={r:.3f}", ha="center", fontsize=ANNOT_FS,
                color="black", fontweight="bold")
    ax.axhline(0, color="k", linestyle="-", linewidth=0.8)
    ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=TICK_FS)
    ax.set_ylabel("Depth-adversary R²", labelpad=10, fontsize=LABEL_FS)
    ax.set_ylim(-0.05, 0.75)
    ax.set_title("(b) Depth adversary R² (target ≤ 0)\nOnly pan-hematopoietic reaches target", fontsize=TITLE_FS)
    ax.grid(axis="y", alpha=0.3, linestyle="--")
    ax.text(0.5, 0.95, "R² ≤ 0  ⇒  trunk carries no depth signal",
            transform=ax.transAxes, fontsize=ANNOT_FS, ha="center", color="green")

    plt.suptitle("Adversary probes: residual dataset and depth signal (purification incomplete)",
                 fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "05_adversary_purification.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 05_adversary_purification.pdf")


# ============================================================
# PAGE 6: Cross-system prototype cosine (29x29)
# ============================================================

def page_cross_system_prototypes():
    M = pd.read_csv(DISC / "70_prototype_full_29x29.csv", index_col=0)
    fig, ax = plt.subplots(figsize=(13, 12))
    vmax = 1.0
    im = ax.imshow(M.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="equal")
    ax.set_xticks(range(len(M.columns)))
    ax.set_xticklabels(M.columns, rotation=90, fontsize=TICK_FS)
    ax.set_yticks(range(len(M.index)))
    ax.set_yticklabels(M.index, fontsize=TICK_FS)
    sizes = []
    prev = None; count = 0
    for lab in M.index:
        s = lab.split(":")[0]
        if prev is None:
            prev = s; count = 1
        elif s == prev:
            count += 1
        else:
            sizes.append(count); prev = s; count = 1
    sizes.append(count)
    cum = 0
    for s in sizes[:-1]:
        cum += s
        ax.axhline(cum - 0.5, color="k", linewidth=1.2)
        ax.axvline(cum - 0.5, color="k", linewidth=1.2)
    cb = plt.colorbar(im, ax=ax, shrink=0.7); cb.set_label("prototype cosine", fontsize=CBAR_FS)
    ax.set_title("§8.2 Cross-system prototype cosine (29 × 29)\n"
                 "Off-block cos mean = +0.001 (null baseline: +0.005). No shared cell-type semantics.",
                 fontsize=TITLE_FS)
    plt.tight_layout()
    plt.savefig(FIG_S / "06_cross_system_prototypes.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 06_cross_system_prototypes.pdf")


# ============================================================
# PAGE 7: Attribution heatmap: top-10 genes per class × system
# ============================================================

def page_attribution_heatmap():
    """per-system heatmap: rows=classes, cols=top-5-per-class gene union.



    Uses the aggregated per-class CSV (80_{sys}_gene_attribution.csv) which

    lists top-N positive/negative genes with their attributions.  Robust to

    npy/hvgs shape mismatch (older discovery runs sometimes use different HVG

    bases).

    """
    fig, axes = plt.subplots(3, 1, figsize=(18, 19))
    for ax, sys in zip(axes, SYSTEMS):
        df = pd.read_csv(DISC / f"80_{sys}_gene_attribution.csv")
        classes = df["class"].tolist()
        # Union of top-5 positive genes per class (each row lists top-20).
        top_gene_set = []
        seen = set()
        per_class = {}
        for _, r in df.iterrows():
            pos_genes = str(r["top_pos_genes"]).split(",")
            pos_vals  = [float(v) for v in str(r["top_pos_attribution"]).split(",")]
            neg_genes = str(r["top_neg_genes"]).split(",")
            neg_vals  = [float(v) for v in str(r["top_neg_attribution"]).split(",")]
            per_class[r["class"]] = dict(zip(pos_genes + neg_genes,
                                             pos_vals + neg_vals))
            for g in pos_genes[:5]:
                if g not in seen: seen.add(g); top_gene_set.append(g)
        top_genes = top_gene_set
        M = np.zeros((len(classes), len(top_genes)), dtype=np.float32)
        for i, cls in enumerate(classes):
            d = per_class[cls]
            for j, g in enumerate(top_genes):
                M[i, j] = d.get(g, 0.0)

        vmax = float(np.abs(M).max()) if M.size else 1.0
        im = ax.imshow(M, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
        ax.set_xticks(range(len(top_genes)))
        ax.set_xticklabels(top_genes, rotation=75, fontsize=TICK_FS, ha="right")
        ax.set_yticks(range(len(classes)))
        ax.set_yticklabels(classes, fontsize=TICK_FS)
        ax.set_title(f"{SYS_LABEL[sys]} — top-5 attributed genes per class (union)", fontsize=TITLE_FS)
        cb = plt.colorbar(im, ax=ax, shrink=0.7); cb.set_label("attribution (unit)", fontsize=CBAR_FS)
    plt.suptitle("§9.1 Prototype-gene integrated-gradient attribution heatmap",
                 fontsize=SUPTITLE_FS, y=1.005, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "07_attribution_heatmap.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 07_attribution_heatmap.pdf")


# ============================================================
# PAGE 8: TF enrichment heatmap per system
# ============================================================

def page_tf_enrichment():
    fig, axes = plt.subplots(1, 3, figsize=(24, 7.5))
    for ax, sys in zip(axes, SYSTEMS):
        df = pd.read_csv(DISC / f"80_{sys}_tf_enrichment.csv")
        piv = df.pivot(index="class", columns="program", values="sum_attribution").fillna(0)
        vmax = np.max(np.abs(piv.values))
        im = ax.imshow(piv.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
        for i, cls in enumerate(piv.index):
            top_j = int(np.argmax(piv.values[i]))
            ax.text(top_j, i, f"{piv.values[i, top_j]:+.2f}",
                    ha="center", va="center", fontsize=ANNOT_FS,
                    color="white" if abs(piv.values[i, top_j]) > vmax * 0.5 else "black")
        ax.set_xticks(range(len(piv.columns)))
        ax.set_xticklabels(piv.columns, rotation=75, fontsize=TICK_FS, ha="right")
        ax.set_yticks(range(len(piv.index)))
        ax.set_yticklabels(piv.index, fontsize=TICK_FS)
        ax.set_title(f"{SYS_LABEL[sys]}\nsum_attribution over TF-program members", fontsize=TITLE_FS)
        plt.colorbar(im, ax=ax, shrink=0.7)
    plt.suptitle("§9.2 Per-prototype TF-program enrichment (self-hit = validation; cross-hit = confusability)",
                 fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "08_tf_enrichment.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 08_tf_enrichment.pdf")


# ============================================================
# PAGE 9: Counterfactual KO top-essential genes
# ============================================================

def page_ko_essentials():
    fig, axes = plt.subplots(3, 1, figsize=(18, 16))
    for ax, sys in zip(axes, SYSTEMS):
        df = pd.read_csv(DISC / f"81_{sys}_ko_essentials.csv")
        rows = []
        for _, r in df.iterrows():
            genes = r["top_essential_genes"].split(",")[:3]
            deltas = [float(x) for x in r["top_essential_deltas"].split(",")[:3]]
            for g, d in zip(genes, deltas):
                rows.append({"class": r["class"], "gene": g, "delta": d,
                             "baseline_cos": r["baseline_cos"]})
        d2 = pd.DataFrame(rows)
        piv = d2.pivot(index="class", columns="gene", values="delta").fillna(0)
        vmax = np.abs(piv.values).max()
        im = ax.imshow(piv.values, cmap="Reds", vmin=0, vmax=vmax, aspect="auto")
        for i in range(piv.shape[0]):
            for j in range(piv.shape[1]):
                v = piv.values[i, j]
                if v > 0:
                    ax.text(j, i, f"{v:.3f}", ha="center", va="center", fontsize=ANNOT_FS,
                            color="white" if v > vmax * 0.5 else "black")
        ax.set_xticks(range(len(piv.columns)))
        ax.set_xticklabels(piv.columns, rotation=75, fontsize=TICK_FS, ha="right")
        ax.set_yticks(range(len(piv.index)))
        ax.set_yticklabels(piv.index, fontsize=TICK_FS)
        ax.set_title(f"{SYS_LABEL[sys]} — counterfactual KO Δcos (top-3 essentials per class)", fontsize=TITLE_FS)
        cb = plt.colorbar(im, ax=ax, shrink=0.7); cb.set_label("Δ prototype cos", fontsize=CBAR_FS)
    plt.suptitle("§9.3 Counterfactual single-gene knockout: essentiality per class",
                 fontsize=SUPTITLE_FS, y=1.005, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "09_ko_essentials.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 09_ko_essentials.pdf")


# ============================================================
# PAGE 10: Hessian gene-gene interaction pairs
# ============================================================

def page_hessian_pairs():
    fig, axes = plt.subplots(1, 3, figsize=(24, 8))
    for ax, sys in zip(axes, SYSTEMS):
        df = pd.read_csv(DISC / f"85_{sys}_hessian_pairs.csv")
        top = df.groupby("class").head(8).copy()
        top["pair"] = top["gene_a"] + "·" + top["gene_b"]
        piv = top.pivot(index="class", columns="pair", values="hessian_off_diag").fillna(0)
        vmax = np.abs(piv.values).max()
        im = ax.imshow(piv.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
        ax.set_xticks(range(len(piv.columns)))
        ax.set_xticklabels(piv.columns, rotation=90, fontsize=TICK_FS)
        ax.set_yticks(range(len(piv.index)))
        ax.set_yticklabels(piv.index, fontsize=TICK_FS)
        ax.set_title(f"{SYS_LABEL[sys]}\ntop-8 gene-gene Hessian pairs per class", fontsize=TITLE_FS)
        cb = plt.colorbar(im, ax=ax, shrink=0.7); cb.set_label("∂²s/∂g·∂g'", fontsize=CBAR_FS)
    plt.suptitle("§10.3 Gene-gene interaction Hessian: combinatorial identity rules",
                 fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "10_hessian_pairs.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 10_hessian_pairs.pdf")


# ============================================================
# PAGE 11-13: three system UMAPs — reuse existing multi UMAP + add rebuilt Dingwall by class
# ============================================================

def page_umaps_all_targets():
    # existing PDFs are appended in the PyPDF merge step
    print("[skip] UMAPs handled by PDF merge step")


# ============================================================
# PAGE 14: Novel populations discovered per system
# ============================================================

def page_novel_populations():
    import os
    if not os.path.exists(str(PANDA_ROOT / "discovery/71_dingwall_novel_populations.csv")):
        print('[skip] 71_dingwall not found'); return
    _dummy = None
    fig, axes = plt.subplots(1, 2, figsize=(20, 7))
    ax = axes[0]
    a = pd.read_csv(DISC / "71_dingwall_novel_populations.csv")
    a = a.sort_values("n_cells", ascending=False)
    y = np.arange(len(a))
    ax.barh(y, a["n_cells"], color=SYS_COLOR["pan_skin"], edgecolor="k", linewidth=0.5)
    labels = [f"cluster {c}: {m[:35]}..." if len(m) > 35 else f"cluster {c}: {m}"
              for c, m in zip(a["cluster"], a["top_markers"])]
    ax.set_yticks(y); ax.set_yticklabels(labels, fontsize=TICK_FS)
    for i, n in enumerate(a["n_cells"]):
        ax.text(n + 3, i, f"n={int(n)}", va="center", fontsize=ANNOT_FS)
    ax.set_xlabel("cells in abstain cluster", fontsize=LABEL_FS)
    ax.set_title(f"§8.4 Dingwall abstain-gate clusters (cos<0.5)\nn={int(a['n_cells'].sum())} cells total",
                 fontsize=TITLE_FS)
    ax.invert_yaxis()

    ax = axes[1]
    d = pd.read_csv(DISC / "73_dahlin_novel_populations.csv")
    d = d.sort_values("n_cells", ascending=False)
    y = np.arange(len(d))
    colors = ["#d7191c" if wtf > 0.85 else "#fdae61" if wtf > 0.6 else "#2b83ba"
              for wtf in d["genotype_wt_frac"]]
    ax.barh(y, d["n_cells"], color=colors, edgecolor="k", linewidth=0.5)
    labels = [f"c{c}: {m[:35]}..." if len(m) > 35 else f"c{c}: {m}"
              for c, m in zip(d["cluster"], d["top_markers"])]
    ax.set_yticks(y); ax.set_yticklabels(labels, fontsize=TICK_FS)
    for i, (n, wtf) in enumerate(zip(d["n_cells"], d["genotype_wt_frac"])):
        ax.text(n + 8, i, f"n={int(n)} · WT={wtf:.1%}", va="center", fontsize=ANNOT_FS)
    ax.set_xlabel("cells in abstain cluster", fontsize=LABEL_FS)
    ax.set_title(f"§8.5 Dahlin abstain-gate clusters (cos<0.57, n={int(d['n_cells'].sum())})\n"
                 "Red bar = quiescent LT-HSC (Hlf⁺), 90.5% WT (Kit-W41-depleted)",
                 fontsize=TITLE_FS)
    ax.invert_yaxis()

    plt.suptitle("§8.4-5 Abstain-gate novel populations", fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "11_novel_populations.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 11_novel_populations.pdf")


# ============================================================
# PAGE 15: Co-attribution modules summary
# ============================================================

def page_coatt_modules():
    fig, axes = plt.subplots(1, 3, figsize=(23, 7))
    for ax, sys in zip(axes, SYSTEMS):
        df = pd.read_csv(DISC / f"82_{sys}_coatt_modules.csv")
        df = df.sort_values("dom_class_mean_att", ascending=False).head(10)
        y = np.arange(len(df))
        colors = [SYS_COLOR[sys]] * len(df)
        ax.barh(y, df["dom_class_mean_att"], color=colors, edgecolor="k", linewidth=0.5)
        labels = [f"module {m} ({s}g)→{c}"
                  for m, s, c in zip(df["module_id"], df["size"], df["dominant_class"])]
        ax.set_yticks(y); ax.set_yticklabels(labels, fontsize=TICK_FS)
        for i, v in enumerate(df["dom_class_mean_att"]):
            ax.text(v + 0.001, i, f"{v:.3f}", va="center", fontsize=ANNOT_FS)
        ax.set_xlabel("mean attribution to dominant class", fontsize=LABEL_FS)
        ax.set_title(f"{SYS_LABEL[sys]}: top-10 gene co-attribution modules", fontsize=TITLE_FS)
        ax.invert_yaxis()
    plt.suptitle("§9.4 Gene-gene co-attribution modules per system",
                 fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "12_coatt_modules.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 12_coatt_modules.pdf")


# ============================================================
# Master build
# ============================================================

def page_anchor_delta_recall():
    """anchor-trick effect on OOD-class recall across 3 systems."""
    fig, axes = plt.subplots(1, 2, figsize=(18, 7.0))

    ax = axes[0]
    # (system, variant, ood_class, baseline_recall, anchor_recall)
    rows = [
        ("Pancreas\nBaron adult-\nislet anchor", "PCA",   0.000, 0.805, SYS_COLOR["pancreas"], 0.55),
        ("Pancreas\nBaron adult-\nislet anchor", "Marker",0.000, 0.805, SYS_COLOR["pancreas"], 0.90),
        ("HSC\nNestorowa LT-\nHSC anchor",       "PCA",   0.000, 0.6818,SYS_COLOR["hematopoiesis"], 0.55),
        ("HSC\nNestorowa LT-\nHSC anchor",       "Marker",0.000, 0.5152,SYS_COLOR["hematopoiesis"], 0.90),
        ("Skin\nSulic HF-\nplacode anchor",      "PCA",   0.000, 0.9802,SYS_COLOR["pan_skin"], 0.55),
        ("Skin\nSulic HF-\nplacode anchor",      "Marker",0.000, 0.8293,SYS_COLOR["pan_skin"], 0.90),
    ]
    x = np.arange(len(rows))
    for i, (grp, var, base, anch, col, alpha) in enumerate(rows):
        ax.bar(x[i], anch, color=col, alpha=alpha, edgecolor="k", linewidth=0.7,
               label=var if i < 2 else None)
        ax.plot([x[i]-0.35, x[i]+0.35], [base, base], color="red", linewidth=1.5, linestyle="--",
                label="no-anchor baseline" if i == 0 else None)
        ax.annotate("", xy=(x[i], anch), xytext=(x[i], base),
                    arrowprops={"arrowstyle": "->", "color": "black", "lw": 1.2})
        ax.text(x[i], anch + 0.02, f"{anch:.2f}\n(Δ +{anch-base:.2f})",
                ha="center", fontsize=ANNOT_FS)
    xt = [rows[i][0] + f"\n({rows[i][1]})" for i in range(len(rows))]
    ax.set_xticks(x); ax.set_xticklabels(xt, fontsize=TICK_FS)
    ax.set_ylim(0, 1.1)
    ax.set_ylabel("OOD-class recall", fontsize=LABEL_FS)
    ax.set_title("(a) Anchor closes vocabulary-out-of-domain gap\n"
                 "baseline recall 0 (dashed red) → anchor recall (bar)", fontsize=TITLE_FS)
    ax.legend(loc="upper left", framealpha=0.9, fontsize=LEGEND_FS)

    ax = axes[1]
    baron_pca = json.load(open(DISC / "pancreas/pca/93_baron_merged_metrics.json"))
    baron_mar = json.load(open(DISC / "pancreas/marker/93_baron_merged_metrics.json"))
    nest_pca  = json.load(open(DISC / "hematopoiesis/pca/97_nestorowa_anchor_zero_shot.json"))
    nest_mar  = json.load(open(DISC / "hematopoiesis/marker/97_nestorowa_anchor_zero_shot.json"))
    sul_pca   = json.load(open(DISC / "pan_skin/pca/97_sulic_anchor_zero_shot.json"))
    sul_mar   = json.load(open(DISC / "pan_skin/marker/97_sulic_anchor_zero_shot.json"))

    grps = [("Pancreas (Baron)", baron_pca["acc"],  baron_mar["acc"],  SYS_COLOR["pancreas"]),
            ("HSC (Nestorowa)",  nest_pca["coarse_acc"], nest_mar["coarse_acc"], SYS_COLOR["hematopoiesis"]),
            ("Skin (Sulic)",     sul_pca["acc"],    sul_mar["acc"],    SYS_COLOR["pan_skin"])]
    x = np.arange(len(grps))
    width = 0.38
    for i, (lab, pca_v, mar_v, col) in enumerate(grps):
        ax.bar(x[i] - width/2, pca_v, width, color=col, alpha=0.55, edgecolor="k",
               label="PCA" if i == 0 else None)
        ax.bar(x[i] + width/2, mar_v, width, color=col, alpha=0.90, edgecolor="k",
               label="Marker" if i == 0 else None)
        ax.text(x[i] - width/2, pca_v + 0.015, f"{pca_v:.3f}", ha="center", fontsize=ANNOT_FS)
        ax.text(x[i] + width/2, mar_v + 0.015, f"{mar_v:.3f}", ha="center", fontsize=ANNOT_FS)
    ax.set_xticks(x); ax.set_xticklabels([g[0] for g in grps], fontsize=TICK_FS)
    ax.set_ylim(0.4, 1.0)
    ax.set_ylabel("Held-out-slice accuracy", fontsize=LABEL_FS)
    ax.set_title("(b) Anchor-augmented held-out accuracy per system (both variants)",
                 fontsize=TITLE_FS)
    ax.legend(loc="lower right", fontsize=LEGEND_FS)

    plt.suptitle("Cross-system anchor paradigm: small labeled slice → OOD vocabulary token",
                 fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(FIG_S / "23_anchor_delta_recall.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 23_anchor_delta_recall.pdf")


def main():
    page_cv_summary()
    page_per_class_f1()
    page_prototype_cosine()
    # page_training_trajectory()  # skipped: v2 checkpoints only save final
    page_adversary_purification()
    page_cross_system_prototypes()
    # page_attribution_heatmap()  # S7 removed -- stale class vocab
    # page_tf_enrichment()  # S8 removed -- stale class vocab
    # page_ko_essentials()  # S9 removed -- stale class vocab
    # page_hessian_pairs()  # S10 removed -- stale class vocab
    page_novel_populations()
    # page_coatt_modules()  # S12 removed -- deprecated classes
    page_anchor_delta_recall()

    print("\n[merge] merging supplement into PANDA_supplement.pdf")
    from pypdf import PdfWriter
    w = PdfWriter()
    order = [
        FIG_S / "01_cv_summary.pdf",
        FIG_S / "02_per_class_f1.pdf",
        FIG_S / "03_prototype_cosine.pdf",
        FIG_S / "05_adversary_purification.pdf",
        FIG_S / "06_cross_system_prototypes.pdf",
        FIG_S / "11_novel_populations.pdf",
        FIG_S / "23_anchor_delta_recall.pdf",
        FIG / "fig5_dingwall_umap.pdf",
        FIG / "fig6_multi_umap.pdf",
    ]
    for p in order:
        if p.exists():
            w.append(str(p))
            print(f"  + {p.name}")
        else:
            print(f"  [skip] {p.name} missing")
    with open(SUP_PDF, "wb") as f:
        w.write(f)
    print(f"\nwrote {SUP_PDF}  ({SUP_PDF.stat().st_size / 1024:.0f} KB)")


if __name__ == "__main__":
    main()