scripts/figures/build_umaps_and_discovery.py

"""UMAPs + discovery figures, writes figures/supplement/13+ and merges into 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
from matplotlib.patches import Patch
import anndata as ad
import scipy.sparse as sp
import scanpy as sc
import torch
import umap as _umap

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
sys.path.insert(0, str(PANDA_ROOT))
sys.path.insert(0, str(Path(__file__).parent))
from panda import PANDAEncoder
from palette import apply_style, color_for
apply_style()

ROOT = Path(str(PANDA_ROOT))
FIG = ROOT / "figures"
FIG_S = FIG / "supplement"; FIG_S.mkdir(parents=True, exist_ok=True)
DISC = ROOT / "discovery"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

RS = 42
SAMPLE_N = 8000  # per panel


from palette import CLASS_COLORS as _CLASS_COLORS
# canonical palette, defer to color_for(); keep "other" as light grey
CLASS_PALETTE = dict(_CLASS_COLORS)
CLASS_PALETTE["other"] = "#c8c8c8"

# classes retired from the canonical class list — filter out of plots
DEPRECATED_CLASSES = {"HF-DP", "eccrine-duct"}


def get_projection_from_ckpt(adata_target, sys, shared_hvgs, mu, sig, pca):
    ck = torch.load(ROOT / f"checkpoints/{sys}/marker/panda_final.pt", map_location=DEVICE, weights_only=False)
    classes = ck["classes"]
    marker_genes = ck.get("marker_genes", [])
    n_markers = len(marker_genes)
    model = PANDAEncoder(variant="marker" if n_markers else "pca",
                         n_pca=50, n_markers=n_markers, n_sub=3,
                         n_classes=len(classes),
                         n_datasets=len(ck["datasets"])).to(DEVICE).eval()
    model.load_state_dict(ck["model"])
    protos = ck["prototypes"]
    protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8)

    hvg2i = {g: i for i, g in enumerate(shared_hvgs)}
    common = [g for g in adata_target.var_names.astype(str) if g in hvg2i]
    a_c = adata_target[:, common].copy()
    sc.pp.normalize_total(a_c, target_sum=1e4); sc.pp.log1p(a_c)
    X = a_c.X.toarray().astype(np.float32) if sp.issparse(a_c.X) else a_c.X.astype(np.float32)
    Xf = np.zeros((adata_target.n_obs, len(shared_hvgs)), dtype=np.float32)
    cols = np.array([hvg2i[g] for g in common])
    Xf[:, cols] = X
    Xz = np.clip((Xf - mu.astype(np.float32)) / sig.astype(np.float32), -10, 10)
    Xpca = pca.transform(Xz).astype(np.float32)

    Xmark = None
    if n_markers:
        mv = np.zeros((adata_target.n_obs, n_markers), dtype=np.float32)
        for j, g in enumerate(marker_genes):
            if g in adata_target.var_names:
                col = adata_target[:, g].X
                if sp.issparse(col): col = col.toarray()
                mv[:, j] = col.flatten().astype(np.float32)
        mmu = mv.mean(axis=0, keepdims=True); msig = mv.std(axis=0, keepdims=True) + 1e-6
        Xmark = np.clip((mv - mmu) / msig, -5, 5).astype(np.float32)

    all_z = []
    with torch.no_grad():
        for i in range(0, adata_target.n_obs, 4096):
            xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE)
            xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None
            aux = torch.zeros(len(xb), 2, device=DEVICE)
            all_z.append(model(xb, aux, x_markers=xmb, lam_dann=0.0)["z"].cpu().numpy())
    Z = np.concatenate(all_z)
    Zn = Z / (np.linalg.norm(Z, axis=1, keepdims=True) + 1e-8)
    cos = Zn @ protos.T
    pred = np.array([classes[i] for i in cos.argmax(axis=1)])
    max_cos = cos.max(axis=1)
    return Z, pred, max_cos, classes


def do_umap(Z, seed=RS):
    return _umap.UMAP(n_neighbors=30, min_dist=0.3, random_state=seed,
                     metric="cosine", n_components=2).fit_transform(Z)


def scatter_by_cat(ax, emb, categories, palette=None, alpha=0.55, s=3, legend_title=""):
    cats = sorted(pd.unique(categories))
    for cat in cats:
        m = np.asarray(categories) == cat
        color = palette.get(cat, "#999999") if palette else None
        ax.scatter(emb[m, 0], emb[m, 1], s=s, alpha=alpha, c=color,
                   label=f"{cat} (n={int(m.sum())})", edgecolors="none")
    ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2")
    leg = ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left",
                    fontsize=7.5, markerscale=6, frameon=False, title=legend_title)
    leg.get_title().set_fontsize(9)
    ax.tick_params(axis="both", labelsize=8)


# =========================================================================
# Fig 13: Dingwall 2-panel UMAP (predicted class + En1 genotype)
# =========================================================================

def fig_dingwall_umap():
    """Dingwall UMAP with PANDA-Marker predicted class + En1 genotype.



    Prefers the existing PCA-vs-Marker cache (_cache_dingwall_full.npz which

    already holds emb_mar, P_mar, genotype for all 25,800 cells). Falls back

    to the 50_aldrich_projections.h5ad if the cache is missing.

    """
    cache = FIG_S / "_cache_dingwall_full.npz"
    if cache.exists():
        c = np.load(cache, allow_pickle=True)
        emb = np.asarray(c["emb_mar"])
        pred = c["P_mar"].astype(str)
        genotype = c["genotype"].astype(str)
        rng = np.random.default_rng(RS)
        idx = rng.choice(len(emb), size=min(SAMPLE_N, len(emb)), replace=False)
        emb, pred, genotype = emb[idx], pred[idx], genotype[idx]
        total_n = int(c["genotype"].shape[0])
    else:
        p = ad.read_h5ad(ROOT / "discovery/pan_skin/marker/50_aldrich_projections.h5ad")
        Z = np.asarray(p.obsm["Z_projection"])
        pred = (p.obs["pred_bbse_label"] if "pred_bbse_label" in p.obs
                else p.obs["pred_label"]).astype(str).values
        genotype = p.obs["genotype"].astype(str).values
        rng = np.random.default_rng(RS)
        idx = rng.choice(len(Z), size=min(SAMPLE_N, len(Z)), replace=False)
        Z_s = Z[idx]
        emb = do_umap(Z_s)
        pred, genotype = pred[idx], genotype[idx]
        total_n = int(len(Z))

    # drop deprecated classes from the class panel
    keep_c = ~np.isin(pred, list(DEPRECATED_CLASSES))
    emb_c = emb[keep_c]; pred_c = pred[keep_c]

    fig, axes = plt.subplots(1, 2, figsize=(15, 6.2))
    scatter_by_cat(axes[0], emb_c, pred_c, palette=CLASS_PALETTE,
                   legend_title="predicted class")
    axes[0].set_title(f"Dingwall En1-cKO skin (n={total_n:,} total; {len(emb):,} shown)\n"
                      "PANDA-Marker predicted class", fontsize=11)
    gcolors = {"WT": "#2b83ba", "En1-cKO": "#d7191c",
               "unknown": "#999999", "other": "#bbbbbb"}
    scatter_by_cat(axes[1], emb, genotype, palette=gcolors, alpha=0.4, s=3,
                   legend_title="En1 genotype")
    axes[1].set_title("Dingwall En1-cKO skin\ncoloured by En1 genotype", fontsize=11)

    plt.suptitle("UMAP of PANDA's 128-d projection: Dingwall held-out target",
                 fontsize=13, y=1.00)
    plt.tight_layout()
    plt.savefig(FIG_S / "13_dingwall_umap.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 13_dingwall_umap.pdf")


# =========================================================================
# Fig 14: Dahlin 3-panel UMAP (predicted class + Kit genotype + LT-HSC cluster)
# =========================================================================

def fig_dahlin_umap():
    # placeholder; real impl in fig_dahlin_umap_full below (cache stored only emb+gt)
    from scripts.analysis.__init__ import dummy  # noqa
    return


def _load_dahlin_raw():
    from pathlib import Path as _P
    D_DIR = _P(str(PANDA_ROOT / "data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract"))
    GT = {"SIGAB1":"WT","SIGAC1":"WT","SIGAD1":"WT","SIGAF1":"WT","SIGAG1":"WT",
          "SIGAH1":"WT","SIGAG8":"Kit_W41","SIGAH8":"Kit_W41"}
    parts = []
    for f in sorted(D_DIR.glob("*.txt.gz")):
        sample = f.name.split("_")[1].split(".")[0]
        df = pd.read_csv(f, sep="\t", compression="gzip", index_col=0)
        X = sp.csr_matrix(df.values.T.astype(np.float32))
        obs = pd.DataFrame(index=[f"{sample}_{bc}" for bc in df.columns.astype(str)])
        obs["sample"] = sample; obs["genotype"] = GT.get(sample, "unknown")
        var = pd.DataFrame(index=df.index.astype(str))
        parts.append(ad.AnnData(X=X, obs=obs, var=var))
    a = ad.concat(parts, join="outer", label="_batch")
    import mygene
    mg = mygene.MyGeneInfo()
    res = mg.querymany(a.var_names.astype(str).tolist(), scopes="ensembl.gene",
                       fields="symbol", species="mouse", verbose=False)
    id2sym = {r["query"]: r["symbol"] for r in res if "symbol" in r}
    syms = pd.Series(a.var_names.astype(str)).map(id2sym).values
    keep = pd.notna(syms)
    a = a[:, keep].copy(); a.var_names = syms[keep]; a.var_names_make_unique()
    return a


def fig_dahlin_umap_full():
    """Dahlin UMAP — 2-panel (PANDA-Marker predicted class + Kit genotype).



    Prefers pca-vs-marker cache which holds emb_mar/P_mar/genotype for all

    61,122 cells. If a legacy cache with emb/pred/gt/max_cos is present, use

    it and render the 3-panel view (with a confidence colorbar).

    """
    cache = FIG_S / "_cache_dahlin_full.npz"
    if not cache.exists():
        stats = np.load(ROOT / "data/corpus/hematopoiesis/harmonized/corpus_stats.npz",
                        allow_pickle=True)
        shared_hvgs = [str(g) for g in stats["shared_hvgs"]]
        pca = pickle.load(open(ROOT / "data/corpus/hematopoiesis/harmonized/pca_basis.pkl", "rb"))
        a = _load_dahlin_raw()
        print(f"[dahlin] {a.shape}", flush=True)
        rng = np.random.default_rng(RS)
        idx = rng.choice(a.n_obs, size=min(SAMPLE_N, a.n_obs), replace=False)
        a_sub = a[idx].copy()
        Z, pred, mc, classes = get_projection_from_ckpt(a_sub, "hematopoiesis",
                                                        shared_hvgs, stats["mean"], stats["std"], pca)
        emb = do_umap(Z)
        gt = a_sub.obs["genotype"].values
        np.savez(cache, emb=emb, gt=np.asarray(gt, dtype=object),
                 pred=np.asarray(pred, dtype=object), max_cos=mc)
        _keys = {"emb", "pred", "gt", "max_cos"}
    c = np.load(cache, allow_pickle=True)
    keys = set(c.files)
    if {"emb_mar", "P_mar", "genotype"}.issubset(keys):
        emb = np.asarray(c["emb_mar"])
        pred = c["P_mar"].astype(str)
        gt = c["genotype"].astype(str)
        total_n = len(emb)
        mc = None
    else:
        emb = c["emb"]; gt = c["gt"].astype(str); pred = c["pred"].astype(str)
        mc = c["max_cos"] if "max_cos" in keys else None
        total_n = len(emb)

    # drop deprecated
    keep_c = ~np.isin(pred, list(DEPRECATED_CLASSES))
    emb_c, pred_c = emb[keep_c], pred[keep_c]

    n_panels = 3 if mc is not None else 2
    fig, axes = plt.subplots(1, n_panels, figsize=(7.0 * n_panels, 6.5))

    scatter_by_cat(axes[0], emb_c, pred_c, palette=CLASS_PALETTE,
                   legend_title="predicted class")
    axes[0].set_title(f"Dahlin Kit-mutant HSPCs (n={total_n:,})\n"
                      "PANDA-Marker predicted lineage class", fontsize=11)
    gcolors = {"WT": "#2b83ba", "Kit_W41": "#d7191c",
               "unknown": "#999999", "other": "#bbbbbb"}
    scatter_by_cat(axes[1], emb, gt, palette=gcolors, legend_title="Kit genotype")
    axes[1].set_title("Dahlin coloured by Kit genotype", fontsize=11)

    if mc is not None:
        ax = axes[2]
        sc_plot = ax.scatter(emb[:, 0], emb[:, 1], c=mc, cmap="viridis",
                             vmin=0.4, vmax=1.0, s=3, alpha=0.65, edgecolors="none")
        ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2")
        ax.tick_params(axis="both", labelsize=8)
        plt.colorbar(sc_plot, ax=ax, shrink=0.75, label="max prototype cosine")
        ax.set_title("Prototype-cosine confidence\n"
                     "(low cos → abstain-gate flagged)", fontsize=11)

    plt.suptitle("Dahlin UMAP — PANDA-Marker zero-shot on Kit-W41 (§8.5)",
                 fontsize=13, y=1.00)
    plt.tight_layout()
    plt.savefig(FIG_S / "14_dahlin_umap.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 14_dahlin_umap.pdf")


# =========================================================================
# Fig 15: Veres 3-panel UMAP (predicted class + stage + confidence)
# =========================================================================

def _load_veres_stages():
    SHARON_DIR = ROOT / "data/corpus/pancreas/held_out_unlabeled/sharon_extract"
    parts = []
    for meta_file in sorted(SHARON_DIR.glob("*.cell_metadata.tsv.gz")):
        counts_file = str(meta_file).replace("cell_metadata", "processed_counts")
        if not Path(counts_file).exists(): continue
        meta = pd.read_csv(meta_file, sep="\t", compression="gzip")
        counts = pd.read_csv(counts_file, sep="\t", compression="gzip", index_col=0)
        counts.columns = [c[0].upper() + c[1:].lower() if len(c) > 1 else c
                          for c in counts.columns.astype(str)]
        counts = counts.T.groupby(level=0).sum().T
        obs = meta.set_index("library.barcode")
        obs = obs.loc[obs.index.intersection(counts.index)]
        counts_al = counts.loc[obs.index]
        X = sp.csr_matrix(counts_al.values.astype(np.float32))
        obs["dataset"] = "veres"
        var = pd.DataFrame({"gene_symbol": counts_al.columns}, index=counts_al.columns)
        a = ad.AnnData(X=X, obs=obs, var=var); a.var_names_make_unique()
        parts.append(a)
    return ad.concat(parts, join="outer", label="_batch")


def fig_veres_umap_full():
    cache = FIG_S / "_cache_veres_full.npz"
    if cache.exists():
        c = np.load(cache, allow_pickle=True)
        emb = c["emb"]; pred = c["pred"]; stage = c["stage"]; mc = c["max_cos"]
    else:
        stats = np.load(ROOT / "data/corpus/pancreas/harmonized/corpus_stats.npz",
                        allow_pickle=True)
        shared_hvgs = [str(g) for g in stats["shared_hvgs"]]
        pca = pickle.load(open(ROOT / "data/corpus/pancreas/harmonized/pca_basis.pkl", "rb"))
        a = _load_veres_stages()
        stage_num = pd.to_numeric(a.obs["Stage"], errors="coerce")
        keep = stage_num.notna().values
        a = a[keep].copy(); a.obs["Stage_int"] = stage_num[keep].astype(int).values
        rng = np.random.default_rng(RS)
        idx = rng.choice(a.n_obs, size=min(SAMPLE_N, a.n_obs), replace=False)
        a_sub = a[idx].copy()
        Z, pred, mc, classes = get_projection_from_ckpt(a_sub, "pancreas",
                                                        shared_hvgs, stats["mean"], stats["std"], pca)
        emb = do_umap(Z)
        stage = a_sub.obs["Stage_int"].values
        np.savez(cache, emb=emb, pred=np.asarray(pred, dtype=object),
                 stage=stage, max_cos=mc)

    fig, axes = plt.subplots(1, 3, figsize=(21, 6.5))
    scatter_by_cat(axes[0], emb, pred, palette=CLASS_PALETTE, legend_title="predicted class")
    axes[0].set_title("Veres hPSC-directed pancreatic differentiation (57,297 total; 8,000 shown)\n"
                      "PANDA-predicted endocrine class", fontsize=11)
    stage_colors = {3: "#fdae61", 4: "#f8b0d1", 5: "#7570b3", 6: "#d7191c"}
    scatter_by_cat(axes[1], emb, stage, palette=stage_colors, legend_title="protocol stage")
    axes[1].set_title("Coloured by directed-differentiation stage\n"
                      "(3 → 4 → 5 → 6 = hPSC → SC-β target)", fontsize=11)

    ax = axes[2]
    sc_plot = ax.scatter(emb[:, 0], emb[:, 1], c=mc, cmap="viridis", vmin=0.4, vmax=1.0,
                          s=3, alpha=0.65, edgecolors="none")
    ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2")
    ax.tick_params(axis="both", labelsize=8)
    plt.colorbar(sc_plot, ax=ax, shrink=0.75, label="max prototype cosine")
    ax.set_title("Prototype-cosine confidence\n(low cos = zero-shot ambiguity)", fontsize=11)

    plt.suptitle("Veres UMAP — cross-species + cross-platform + in vitro triple shift (§8)",
                 fontsize=13, y=1.00)
    plt.tight_layout()
    plt.savefig(FIG_S / "15_veres_umap.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 15_veres_umap.pdf")


# =========================================================================
# Fig 16: Dingwall En1-cKO discovery evidence (class enrichment + melanocyte volcano)
# =========================================================================

def fig_dingwall_discovery():
    """Two-panel Dingwall En1-cKO discovery evidence.



    Falls back to the compact per-class enrichment CSV (27_*.csv, columns:

    class,n,frac_cko,delta) and the pathway module CSV (28_*.csv) when the

    older 53_*/56_* CSVs are not present.

    """
    fig, axes = plt.subplots(1, 2, figsize=(15, 5.5))

    ax = axes[0]
    p53 = ROOT / "discovery/pan_skin/marker/53_en1_cko_class_enrichment.csv"
    if p53.exists():
        enr = pd.read_csv(p53)
        cls_col, lfc_col, p_col = "class", "log2_fold_enrich_cKO_vs_WT", "fisher_pvalue"
    else:
        enr = pd.read_csv(ROOT / "discovery/pan_skin/marker/27_dingwall_en1_enrichment.csv")
        cls_col, lfc_col, p_col = "class", "delta", None  # delta already signed
    # drop deprecated HF-DP class (removed from canonical vocabulary)
    enr = enr[~enr[cls_col].isin(["HF-DP"])].copy()
    enr = enr.sort_values(lfc_col)
    y = np.arange(len(enr))
    colors = ["#d7191c" if v > 0 else "#2b83ba" for v in enr[lfc_col]]
    ax.barh(y, enr[lfc_col], color=colors, edgecolor="k", linewidth=0.5)
    xmax = max(abs(enr[lfc_col].min()), abs(enr[lfc_col].max())) * 1.3
    for i, (_, r) in enumerate(enr.iterrows()):
        if p_col is not None and p_col in r:
            pv = r[p_col]
            star = " ***" if pv < 1e-3 else " *" if pv < 0.05 else ""
            ax.text(xmax, i, f"p={pv:.1e}{star}",
                    ha="left", va="center", fontsize=8, color="black")
        else:
            ax.text(xmax, i, f"n={int(r['n']):,}",
                    ha="left", va="center", fontsize=8, color="black")
    ax.set_xlim(-xmax * 1.05, xmax * 1.9)
    ax.axvline(0, color="k", linewidth=0.6)
    ax.set_yticks(y); ax.set_yticklabels(enr[cls_col], fontsize=9)
    xlabel = ("log2 fold-change (cKO/WT)" if p_col is not None
              else "Δ En1-cKO fraction vs corpus baseline")
    ax.set_xlabel(xlabel, fontsize=10)
    ax.set_title("(a) Dingwall class enrichment (cKO vs WT)\n"
                 "Blue = depleted in cKO, red = enriched", fontsize=10)
    ax.grid(axis="x", alpha=0.3, linestyle="--")

    ax = axes[1]
    p56 = ROOT / "discovery/pan_skin/marker/56_melanocyte_pathways.csv"
    if p56.exists():
        pw = pd.read_csv(p56).sort_values("delta_cKO_minus_WT")
        vcol, pcol, ncol = "delta_cKO_minus_WT", "MannU_p", "pathway"
    else:
        pw = pd.read_csv(ROOT / "discovery/pan_skin/marker/28_melanocyte_pathway_modules.csv")
        pw = pw.sort_values("delta")
        vcol, pcol, ncol = "delta", "p", "module"
    y = np.arange(len(pw))
    colors = ["#d7191c" if d > 0 else "#2b83ba" for d in pw[vcol]]
    ax.barh(y, pw[vcol], color=colors, edgecolor="k", linewidth=0.5)
    xmax = max(abs(pw[vcol].min()), abs(pw[vcol].max())) * 1.3
    for i, (_, r) in enumerate(pw.iterrows()):
        pv = r[pcol]
        star = (" ***" if pv < 1e-10 else " **" if pv < 1e-3 else
                " *" if pv < 0.05 else "")
        ax.text(xmax, i, f"p={pv:.1e}{star}",
                ha="left", va="center", fontsize=8)
    ax.set_xlim(-xmax * 1.05, xmax * 1.9)
    ax.axvline(0, color="k", linewidth=0.6)
    labels = [str(s).split(" (")[0] for s in pw[ncol]]
    ax.set_yticks(y); ax.set_yticklabels(labels, fontsize=9)
    ax.set_xlabel("Δ module score (cKO − WT)", fontsize=10)
    ax.set_title("(b) Dingwall melanocyte pathway modules\n"
                 "Mann-Whitney U within melanocyte class", fontsize=10)
    ax.grid(axis="x", alpha=0.3, linestyle="--")

    plt.suptitle("§4 Dingwall En1-cKO mechanistic evidence", fontsize=13, y=1.02)
    plt.tight_layout()
    plt.savefig(FIG_S / "16_dingwall_discovery.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 16_dingwall_discovery.pdf")


# =========================================================================
# Fig 17: Dahlin discovery evidence (LT-HSC depletion + Kit_signaling module)
# =========================================================================

def fig_dahlin_discovery():
    fig, axes = plt.subplots(1, 2, figsize=(15, 5.5))

    ax = axes[0]
    d = pd.read_csv(DISC / "73_dahlin_novel_populations.csv")
    d = d.sort_values("genotype_wt_frac", ascending=False)
    def label(row):
        top = row["top_markers"].split(",")[0]
        return f"c{row['cluster']}:{top}⁺ (n={row['n_cells']})"
    labels = [label(r) for _, r in d.iterrows()]
    y = np.arange(len(d))
    colors = ["#d7191c" if wtf > 0.85 else "#fdae61" if wtf > 0.7 else "#2b83ba"
              for wtf in d["genotype_wt_frac"]]
    ax.barh(y, d["genotype_wt_frac"], color=colors, edgecolor="k", linewidth=0.5)
    ax.axvline(0.60, color="green", linestyle="--", linewidth=1.2, label="whole-corpus baseline (~60% WT)")
    for i, wtf in enumerate(d["genotype_wt_frac"]):
        ax.text(min(wtf + 0.015, 1.05), i, f"{wtf:.1%}", va="center", fontsize=8)
    ax.set_yticks(y); ax.set_yticklabels(labels, fontsize=8)
    ax.set_xlim(0, 1.15)
    ax.set_xlabel("fraction WT")
    ax.set_title("(a) Kit-W41 depletes quiescent LT-HSC (Hlf⁺, 90.5% WT)\n"
                 "Abstain-gate substates ranked by WT fraction", fontsize=10)
    ax.legend(fontsize=8, loc="lower right")
    ax.grid(axis="x", alpha=0.3, linestyle="--")

    ax = axes[1]
    ms = pd.read_csv(ROOT / "discovery/hematopoiesis/marker/67_dahlin_module_scores.csv")
    piv = ms.pivot(index="module", columns="class", values="delta_Kit_minus_WT")
    piv_p = ms.pivot(index="module", columns="class", values="MannU_p")
    row_order = ["Kit_signaling", "MYC_targets", "Integrated_stress",
                 "Apoptosis_pro", "Apoptosis_anti", "Cell_cycle", "Erythroid_dev"]
    row_order = [r for r in row_order if r in piv.index]
    col_order = ["MPP", "erythroid", "myeloid", "megakaryocyte", "lymphoid"]
    col_order = [c for c in col_order if c in piv.columns]
    P = piv.loc[row_order, col_order]; Pp = piv_p.loc[row_order, col_order]
    vmax = np.nanmax(np.abs(P.values))
    im = ax.imshow(P.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
    for i in range(P.shape[0]):
        for j in range(P.shape[1]):
            v = P.values[i, j]; p = Pp.values[i, j]
            if np.isnan(v): continue
            star = "***" if p < 1e-10 else "**" if p < 1e-3 else "*" if p < 0.05 else ""
            ax.text(j, i, f"{v:+.3f}\n{star}", ha="center", va="center",
                    fontsize=8, color="white" if abs(v) > vmax * 0.55 else "black")
    ax.set_xticks(range(len(col_order))); ax.set_xticklabels(col_order, rotation=30, ha="right")
    ax.set_yticks(range(len(row_order))); ax.set_yticklabels(row_order)
    plt.colorbar(im, ax=ax, label="Δ module score (Kit-W41 − WT)")
    ax.set_title("(b) Dahlin within-class module Δ\n"
                 "Kit_signaling ↓ + ISR ↑ + Apoptosis_pro erythroid ↓ (p=6.6e-123)", fontsize=10)

    plt.suptitle("§6 Dahlin Kit-W41 mechanistic evidence", fontsize=13, y=1.02)
    plt.tight_layout()
    plt.savefig(FIG_S / "17_dahlin_discovery.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 17_dahlin_discovery.pdf")


# =========================================================================
# Fig 18: Veres discovery evidence (stage stack + alpha-vs-beta TF axis)
# =========================================================================

def fig_veres_discovery():
    fig, axes = plt.subplots(1, 2, figsize=(15, 5.5))

    ax = axes[0]
    st = pd.read_csv(ROOT / "discovery/pancreas/marker/64_sharon_class_per_stage.csv", index_col=0)
    st.columns = st.columns.astype(float).astype(int)
    order = ["alpha", "delta", "gamma", "beta", "acinar", "ductal",
             "endocrine-progenitor", "endothelial", "other", "immune"]
    order = [c for c in order if c in st.index]
    st2 = st.loc[order]
    bottom = np.zeros(st2.shape[1])
    for cls in order:
        vals = st2.loc[cls].values
        ax.bar(st2.columns, vals, bottom=bottom, label=cls,
               color=CLASS_PALETTE.get(cls, "#999999"), edgecolor="k", linewidth=0.4)
        bottom += vals
    ax.set_xticks(st2.columns); ax.set_xticklabels([f"Stage {int(s)}" for s in st2.columns])
    ax.set_ylabel("PANDA-predicted class fraction")
    a6 = float(st.loc["alpha", 6]); b6 = float(st.loc["beta", 6])
    ax.text(0.98, 0.98, f"Stage 6:\nα = {a6:.1%}\nβ = {b6:.1%}",
            transform=ax.transAxes, fontsize=10, ha="right", va="top",
            bbox=dict(boxstyle="round", facecolor="white", alpha=0.9))
    ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=8)
    ax.set_ylim(0, 1.05)
    ax.set_title("(a) Veres SC-β protocol produces SC-α, not SC-β\nInefficient differentiation at Stage 6", fontsize=10)

    ax = axes[1]
    de = pd.read_csv(ROOT / "discovery/pancreas/marker/65_sharon_stage6_alpha_vs_beta.csv")
    # rows: (up_in, gene, logfc, padj)
    for updir, color in [("alpha", "#d7191c"), ("beta", "#2b83ba")]:
        sub = de[de["up_in"] == updir]
        neg_log10p = -np.log10(np.clip(sub["padj"].values, 1e-320, 1))
        sign = 1 if updir == "alpha" else -1
        ax.scatter(sign * sub["logfc"], neg_log10p, s=15, alpha=0.55, c=color,
                   label=f"up in SC-{updir}")
        top = sub.nsmallest(8, "padj")
        # skip labels landing within epsilon of an already-placed one so
        # dense clusters (e.g. Cpe/Rpl13a) don't overprint; deterministic
        placed = []
        for _, r in top.iterrows():
            lx = sign * r["logfc"]; ly = -np.log10(max(r["padj"], 1e-320)) + 3
            if any(abs(lx - px) < 0.25 and abs(ly - py) < 12 for px, py in placed):
                continue
            placed.append((lx, ly))
            ax.text(lx, ly, r["gene"], fontsize=8, ha="center", color=color)
    ax.axvline(0, color="k", linewidth=0.5)
    ax.set_xlabel("log2 FC (SC-α ← 0 → SC-β)")
    ax.set_ylabel("−log10 padj")
    ax.set_title("(b) Veres Stage-6 SC-α vs SC-β DE\n"
                 "Arx/Irx2 vs Nkx6-1/Mnx1/Neurod1 TF axis", fontsize=10)
    ax.legend(fontsize=9)
    ax.grid(alpha=0.3, linestyle="--")

    plt.suptitle("§7 Veres SC-β / SC-α mechanistic evidence", fontsize=13, y=1.02)
    plt.tight_layout()
    plt.savefig(FIG_S / "18_veres_discovery.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 18_veres_discovery.pdf")


# =========================================================================
# Fig 19: HSC myeloid combinatorial identity network
# =========================================================================

def fig_myeloid_network():
    df = pd.read_csv(DISC / "85_hematopoiesis_hessian_pairs.csv")
    my = df[df["class"] == "myeloid"].head(20).copy()

    fig, ax = plt.subplots(figsize=(12, 10.5))
    genes = list(pd.unique(pd.concat([my["gene_a"], my["gene_b"]])))
    n = len(genes)
    # circular node layout
    theta = np.linspace(0, 2 * np.pi, n, endpoint=False)
    pos = {g: (np.cos(t), np.sin(t)) for g, t in zip(genes, theta)}

    # edge width ∝ |H|
    max_h = my["abs_h"].max()
    for _, r in my.iterrows():
        x1, y1 = pos[r["gene_a"]]; x2, y2 = pos[r["gene_b"]]
        lw = 4 * r["abs_h"] / max_h
        alpha = min(0.85, 0.3 + 0.6 * r["abs_h"] / max_h)
        ax.plot([x1, x2], [y1, y2], color="#d7191c", lw=lw, alpha=alpha, zorder=1)

    for g in genes:
        x, y = pos[g]
        ax.scatter(x, y, s=420, c=color_for("myeloid"), edgecolor="k", linewidth=1, zorder=2)
        ax.text(x, y + 0.09, g, ha="center", fontsize=12, zorder=3,
                fontweight="bold")
    ax.set_xlim(-1.35, 1.35); ax.set_ylim(-1.25, 1.25)
    ax.set_aspect("equal"); ax.axis("off")
    ax.set_title("§10.3 Pan-hematopoietic myeloid prototype:\n"
                 "combinatorial identity via macrophage antimicrobial network\n"
                 "(top-20 Hessian pairs, edge width ∝ |∂²s/∂g·∂g'|)", fontsize=18)
    plt.tight_layout()
    plt.savefig(FIG_S / "19_myeloid_network.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 19_myeloid_network.pdf")


# =========================================================================
# Fig 20: HF placode Wnt/EDAR module co-attribution network
# =========================================================================

def fig_placode_wnt_module():
    modules = pd.read_csv(DISC / "82_pan_skin_coatt_modules.csv")
    # HF-placode Wnt/EDAR module was module 5 in earlier output
    hf_mods = modules[modules["dominant_class"] == "HF-placode"]
    if len(hf_mods) == 0:
        print("[skip] no HF-placode modules")
        return
    mod = hf_mods.iloc[0]
    genes = mod["member_genes"].split(",")[:20]

    # tight figsize — module contains only 3 genes, no need for 12x10 inch box
    fig, ax = plt.subplots(figsize=(6.5, 6.5))
    n = len(genes)
    theta = np.linspace(0, 2 * np.pi, n, endpoint=False)
    # place small modules further from center
    radius = 0.55 if n <= 3 else 1.0
    pos = {g: (radius * np.cos(t), radius * np.sin(t)) for g, t in zip(genes, theta)}

    canonical = {"Ptch2", "Lef1", "Edar", "Wnt6", "Wnt7b", "Bmp7", "Tfap2b", "Tfap2a"}
    hf_col = color_for("HF-placode")
    for g in genes:
        x, y = pos[g]
        col = hf_col if g in canonical else "#2b83ba"
        ax.scatter(x, y, s=520, c=col, edgecolor="k", linewidth=1, zorder=2)
        ax.text(x, y + 0.12, g, ha="center", fontsize=13, zorder=3,
                fontweight="bold" if g in canonical else "normal",
                color=hf_col if g in canonical else "black")

    # all-pair edges — every member is internally co-attributed
    for i, g1 in enumerate(genes):
        for g2 in genes[i+1:]:
            x1, y1 = pos[g1]; x2, y2 = pos[g2]
            ax.plot([x1, x2], [y1, y2], color="#888", lw=0.6, alpha=0.4, zorder=1)

    ax.set_xlim(-1.1, 1.1); ax.set_ylim(-1.1, 1.1)
    ax.set_aspect("equal"); ax.axis("off")
    genes_str = ", ".join(genes)
    ax.set_title(f"Pan-skin HF-placode co-attribution triangle\n"
                 f"module id={int(mod['module_id'])}, size={int(mod['size'])} genes: {genes_str}",
                 fontsize=13)
    plt.tight_layout()
    plt.savefig(FIG_S / "20_placode_wnt_module.pdf", bbox_inches="tight")
    plt.close()
    print("[fig] 20_placode_wnt_module.pdf")


# =========================================================================
# Master merge
# =========================================================================

def merge_pdf():
    from pypdf import PdfWriter
    SUP_PDF = FIG / "PANDA_supplement.pdf"
    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 / "13_dingwall_umap.pdf",
        FIG_S / "14_dahlin_umap.pdf",
        FIG_S / "15_veres_umap.pdf",
        FIG_S / "16_dingwall_discovery.pdf",
        FIG_S / "17_dahlin_discovery.pdf",
        FIG_S / "18_veres_discovery.pdf",
        FIG_S / "19_myeloid_network.pdf",
        FIG_S / "20_placode_wnt_module.pdf",
    ]
    w = PdfWriter()
    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"wrote {SUP_PDF} ({SUP_PDF.stat().st_size/1024:.0f} KB)")


def main():
    # UMAP figures first — they cache and are slow
    fig_dingwall_umap()
    fig_dahlin_umap_full()
    fig_veres_umap_full()
    fig_dingwall_discovery()
    fig_dahlin_discovery()
    fig_veres_discovery()
    fig_myeloid_network()
    fig_placode_wnt_module()
    merge_pdf()


if __name__ == "__main__":
    main()