scripts/figures/build_pca_vs_marker_umaps.py

"""PANDA-PCA vs PANDA-Marker side-by-side umaps for dingwall/dahlin/veres, plus en1-cKO enrichment and melanocyte pathway bars."""
from __future__ import annotations
from pathlib import Path
import warnings, json, sys, pickle, numpy as np, pandas as pd
warnings.filterwarnings("ignore")
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
import anndata as ad, scanpy as sc, scipy.sparse as sp, torch
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, GENOTYPE_COLORS, STAGE_COLORS,
                     SUPTITLE_FS, TITLE_FS, LABEL_FS, TICK_FS, LEGEND_FS, ANNOT_FS)
apply_style()
sc.settings.verbosity = 0

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


# ------------------- projection helpers -------------------
def project(a, sys, variant):
    """returns (z_128d, predicted class array, classes)."""
    ck = torch.load(ROOT / f"checkpoints/{sys}/{variant}/panda_final.pt",
                    map_location=DEVICE, weights_only=False)
    classes = ck["classes"]; marker_genes = ck.get("marker_genes", [])
    stats = np.load(ROOT / f"data/corpus/{sys}/harmonized/corpus_stats.npz", allow_pickle=True)
    pca = pickle.load(open(ROOT / f"data/corpus/{sys}/harmonized/pca_basis.pkl", "rb"))
    hvgs = [str(g) for g in stats["shared_hvgs"]]
    hvg2i = {g: i for i, g in enumerate(hvgs)}
    common = [g for g in a.var_names.astype(str) if g in hvg2i]
    a_c = a[:, 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((a.n_obs, len(hvgs)), dtype=np.float32)
    Xf[:, np.array([hvg2i[g] for g in common])] = X
    Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10)
    Xpca = pca.transform(Xz).astype(np.float32)
    Xmark = None
    if variant == "marker" and marker_genes:
        mv = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32)
        for j, g in enumerate(marker_genes):
            if g in a.var_names:
                col = a[:, 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)
    model = PANDAEncoder(variant=variant, n_pca=50,
                         n_markers=len(marker_genes) if variant == "marker" else 0,
                         n_classes=len(classes), n_sub=3,
                         n_datasets=len(ck["datasets"])).to(DEVICE).eval()
    model.load_state_dict(ck["model"])
    all_z, preds = [], []
    with torch.no_grad():
        for i in range(0, a.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)
            out = model(xb, aux, x_markers=xmb, lam_dann=0.0)
            all_z.append(out["z"].cpu().numpy())
            mc = model.max_sub_cos(out["z"])
            preds.append(mc.argmax(dim=1).cpu().numpy())
    Z = np.concatenate(all_z, axis=0)
    P = np.array([classes[i] for i in np.concatenate(preds)])
    return Z, P, classes


def do_umap(Z, seed=RS):
    import umap
    # Z is the L2-normalised 128-d projection (points on the unit hypersphere),
    # so cosine is the metric the embedding was trained under. Euclidean here
    # produced the ring/arc artefacts seen in earlier Veres panels.
    reducer = umap.UMAP(n_neighbors=30, min_dist=0.3, random_state=seed,
                        metric="cosine", n_epochs=200, verbose=False,
                        low_memory=False)
    return reducer.fit_transform(Z)


# kelly-inspired 22-color palette + "other"
DISTINCT_COLORS = [
    "#e6194b", "#3cb44b", "#4363d8", "#f58231", "#911eb4", "#42d4f4",
    "#f032e6", "#bfef45", "#fabed4", "#469990", "#dcbeff", "#9a6324",
    "#fffac8", "#800000", "#aaffc3", "#808000", "#ffd8b1", "#000075",
    "#a9a9a9", "#f4a460", "#00fa9a", "#ff69b4",
]


def build_class_palette(P_pca, P_mar, min_frac=0.005):
    """collapse classes < min_frac to 'other', assign each remaining canonical color."""
    from collections import Counter
    total = len(P_pca) + len(P_mar)
    counts = Counter(P_pca.tolist() + P_mar.tolist())
    kept = [c for c, n in counts.most_common() if n / total >= min_frac]
    P_pca_r = np.where(np.isin(P_pca, kept), P_pca, "other")
    P_mar_r = np.where(np.isin(P_mar, kept), P_mar, "other")
    palette = {c: color_for(c, DISTINCT_COLORS[i % len(DISTINCT_COLORS)])
               for i, c in enumerate(kept)}
    palette["other"] = "#e5e5e5"
    return P_pca_r, P_mar_r, palette


def scatter_side_by_side(emb_pca, emb_mark, colors_pca, colors_mark, palette,
                         subtitle_pca, subtitle_mark, main_title, out_path,
                         s=4, alpha=0.55, legend_title=""):
    fig, axes = plt.subplots(1, 2, figsize=(18.5, 8.5))
    for ax, emb, colors, sub in zip(axes,
                                     [emb_pca, emb_mark],
                                     [colors_pca, colors_mark],
                                     [subtitle_pca, subtitle_mark]):
        for cat in sorted(set(colors)):
            m = np.array(colors) == cat
            ax.scatter(emb[m, 0], emb[m, 1], s=s, alpha=alpha,
                       color=palette.get(cat, "#888"), label=cat, linewidths=0,
                       rasterized=True)
        ax.set_xlabel("UMAP-1", fontsize=LABEL_FS); ax.set_ylabel("UMAP-2", fontsize=LABEL_FS)
        ax.set_title(sub, fontsize=TITLE_FS)
        ax.set_xticks([]); ax.set_yticks([])
    handles = [plt.Line2D([0], [0], marker="o", linestyle="",
               markerfacecolor=palette[c], markeredgecolor="none", markersize=11, label=c)
               for c in palette]
    fig.legend(handles=handles, loc="center right", bbox_to_anchor=(1.10, 0.5),
               fontsize=LEGEND_FS, frameon=False, title=legend_title,
               title_fontsize=TITLE_FS)
    plt.suptitle(main_title, fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(out_path, bbox_inches="tight", dpi=180)
    plt.close()
    print(f"[fig] {out_path.name}")


# ------------------- dingwall -------------------
def fig_dingwall_pca_vs_marker():
    """dingwall GSE220977 colored by en1 genotype + predicted class."""
    cache = FIG_S / "_cache_dingwall_full.npz"
    CKO = {"GSM6833482", "GSM6833483"}
    WT  = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"}
    if cache.exists():
        c = np.load(cache, allow_pickle=True)
        emb_pca = c["emb_pca"]; emb_mar = c["emb_mar"]
        P_pca = c["P_pca"].astype(str); P_mar = c["P_mar"].astype(str)
        genotype = c["genotype"].astype(str)
        n = len(genotype)
        print(f"[dingwall] loaded cache n={n}", flush=True)
    else:
        raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad")
        genotype = np.where(raw.obs["sample"].astype(str).isin(list(CKO)), "En1-cKO",
                    np.where(raw.obs["sample"].astype(str).isin(list(WT)), "WT", "other"))
        n = raw.n_obs
        print(f"[dingwall] projecting all {n} cells with PCA and Marker checkpoints", flush=True)
        Z_pca, P_pca, _  = project(raw, "pan_skin", "pca")
        Z_mar, P_mar, _  = project(raw, "pan_skin", "marker")
        print(f"[dingwall] running umap on all {n}", flush=True)
        emb_pca = do_umap(Z_pca)
        emb_mar = do_umap(Z_mar)
        np.savez(cache,
                 emb_pca=emb_pca, emb_mar=emb_mar, P_pca=P_pca, P_mar=P_mar,
                 genotype=genotype)

    palette_gt = {"WT": GENOTYPE_COLORS["WT"], "En1-cKO": GENOTYPE_COLORS["En1-cKO"],
                  "other": GENOTYPE_COLORS["other"]}
    scatter_side_by_side(
        emb_pca, emb_mar, genotype, genotype, palette_gt,
        f"PANDA-PCA (Dingwall, all {n:,} cells)", f"PANDA-Marker (Dingwall, all {n:,} cells)",
        "Dingwall En1-cKO vs WT — PANDA-PCA vs PANDA-Marker embedding",
        FIG_S / "24_pca_vs_marker_umaps_dingwall_by_genotype.pdf",
        legend_title="Genotype",
    )
    P_pca_r, P_mar_r, palette_c = build_class_palette(P_pca, P_mar, min_frac=0.005)
    scatter_side_by_side(
        emb_pca, emb_mar, P_pca_r, P_mar_r, palette_c,
        "PANDA-PCA — predicted class", "PANDA-Marker — predicted class",
        f"Dingwall — PANDA-PCA vs PANDA-Marker predicted class map (n={n:,})",
        FIG_S / "24b_pca_vs_marker_umaps_dingwall_by_class.pdf",
        legend_title="Predicted class",
    )


# ------------------- dahlin -------------------
def _load_dahlin_raw():
    D_DIR = 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")
    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_pca_vs_marker():
    cache = FIG_S / "_cache_dahlin_full.npz"
    if cache.exists():
        c = np.load(cache, allow_pickle=True)
        emb_pca = c["emb_pca"]; emb_mar = c["emb_mar"]
        P_pca = c["P_pca"].astype(str); P_mar = c["P_mar"].astype(str)
        gen = c["genotype"].astype(str)
        n = len(gen)
        print(f"[dahlin] loaded cache n={n}", flush=True)
    else:
        print("[dahlin] loading raw", flush=True)
        a = _load_dahlin_raw()
        gen = a.obs["genotype"].astype(str).values
        n = a.n_obs
        print(f"[dahlin] projecting all {n} cells", flush=True)
        Z_pca, P_pca, _ = project(a, "hematopoiesis", "pca")
        Z_mar, P_mar, _ = project(a, "hematopoiesis", "marker")
        print(f"[dahlin] umap all {n}", flush=True)
        emb_pca = do_umap(Z_pca)
        emb_mar = do_umap(Z_mar)
        np.savez(cache,
                 emb_pca=emb_pca, emb_mar=emb_mar, P_pca=P_pca, P_mar=P_mar, genotype=gen)

    palette_gt = {"WT": GENOTYPE_COLORS["WT"], "Kit_W41": GENOTYPE_COLORS["Kit_W41"],
                  "unknown": GENOTYPE_COLORS["other"]}
    scatter_side_by_side(
        emb_pca, emb_mar, gen, gen, palette_gt,
        f"PANDA-PCA (Dahlin, all {n:,} cells)", f"PANDA-Marker (Dahlin, all {n:,} cells)",
        "Dahlin WT vs Kit-W41 — PANDA-PCA vs PANDA-Marker embedding",
        FIG_S / "25_pca_vs_marker_umaps_dahlin_by_genotype.pdf",
        legend_title="Genotype",
    )
    P_pca_r, P_mar_r, palette_c = build_class_palette(P_pca, P_mar, min_frac=0.005)
    scatter_side_by_side(
        emb_pca, emb_mar, P_pca_r, P_mar_r, palette_c,
        "PANDA-PCA — predicted class", "PANDA-Marker — predicted class",
        f"Dahlin — PANDA-PCA vs PANDA-Marker predicted class map (n={n:,})",
        FIG_S / "25b_pca_vs_marker_umaps_dahlin_by_class.pdf",
        legend_title="Predicted class",
    )


# ------------------- veres -------------------
def _load_veres():
    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)
        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))
        a = ad.AnnData(X=X, obs=obs, var=pd.DataFrame(index=counts_al.columns))
        a.var_names_make_unique()
        # sharon_extract is NOT purely the Stage 3-6 differentiation: it also
        # ships GSM3141996 (ES/iPS comparison) and GSM3142001 (primary human
        # islets, GSE84133). Tag the source so those cells are not silently
        # pooled with the staged in-vitro cells.
        nm = Path(meta_file).name
        if "HumanIslets" in nm:
            grp = "primary islets (GSE84133)"
        elif "ES_iPS" in nm:
            grp = "ES/iPS comparison"
        else:
            grp = "differentiation"
        a.obs["veres_group"] = grp
        parts.append(a)
    out = ad.concat(parts, join="outer")
    # human->mouse symbol case-fold (same heuristic as run_all_zero_shot.infer):
    # the corpus + checkpoints use mouse Title-case symbols; without this the HVG
    # intersection collapses and every cell predicts one junk class (the bug that
    # produced the old all-mesenchyme S26b page).
    vn = out.var_names.astype(str)
    n_upper = sum(1 for g in vn[:1000] if g.isupper() and len(g) > 1)
    if n_upper > 500:
        out.var_names = [g.capitalize() for g in vn]
        out.var_names_make_unique()
        print(f"[veres] case-folded {n_upper}/1000 uppercase symbols human->mouse", flush=True)
    return out


def fig_veres_pca_vs_marker():
    cache = FIG_S / "_cache_veres_full_v2.npz"
    if cache.exists():
        c = np.load(cache, allow_pickle=True)
        emb_pca = c["emb_pca"]; emb_mar = c["emb_mar"]
        P_pca = c["P_pca"].astype(str); P_mar = c["P_mar"].astype(str)
        st_str = c["stage"].astype(str)
        n = len(st_str)
        print(f"[veres] loaded cache n={n}", flush=True)
    else:
        print("[veres] loading raw", flush=True)
        a = _load_veres()
        stage_col = "Stage" if "Stage" in a.obs.columns else "stage"
        stage = pd.to_numeric(a.obs[stage_col], errors="coerce").fillna(-1).astype(int).values
        grp = a.obs["veres_group"].astype(str).values
        # non-differentiation cells get their own legend entries instead of a
        # single anonymous "unstaged" grey blob
        st_str = np.array([g if g != "differentiation" else
                           ("unstaged" if s == -1 else str(s))
                           for s, g in zip(stage, grp)])
        n = a.n_obs
        print(f"[veres] projecting all {n} cells (stage dist: "
              f"{pd.Series(st_str).value_counts().to_dict()})", flush=True)
        Z_pca, P_pca, _ = project(a, "pancreas", "pca")
        Z_mar, P_mar, _ = project(a, "pancreas", "marker")
        print(f"[veres] umap all {n}", flush=True)
        emb_pca = do_umap(Z_pca)
        emb_mar = do_umap(Z_mar)
        np.savez(cache,
                 emb_pca=emb_pca, emb_mar=emb_mar, P_pca=P_pca, P_mar=P_mar, stage=st_str)

    # canonical stage palette from palette.py
    # sentinel is relabeled "unstaged" upstream; don't also keep "-1" or the
    # legend shows both for the same group
    palette_st = {k: v for k, v in STAGE_COLORS.items() if k != "-1"}
    palette_st["unstaged"] = "#bbbbbb"
    palette_st["primary islets (GSE84133)"] = "#7f7f7f"
    palette_st["ES/iPS comparison"] = "#c49a6c"
    scatter_side_by_side(
        emb_pca, emb_mar, st_str, st_str, palette_st,
        f"PANDA-PCA (Veres, all {n:,} cells)", f"PANDA-Marker (Veres, all {n:,} cells)",
        "Veres SC-beta differentiation Stage 3-6 — PANDA-PCA vs PANDA-Marker embedding",
        FIG_S / "26_pca_vs_marker_umaps_veres_by_stage.pdf",
        legend_title="Stage",
    )
    P_pca_r, P_mar_r, palette_c = build_class_palette(P_pca, P_mar, min_frac=0.005)
    # Report actual class diversity in the title so readers understand why the
    # Veres in-vitro slice collapses onto a small subset of pancreas prototypes.
    from collections import Counter
    top_pca = ", ".join([f"{c} (n={k:,})"
                         for c, k in Counter(P_pca.tolist()).most_common(5)])
    top_mar = ", ".join([f"{c} (n={k:,})"
                         for c, k in Counter(P_mar.tolist()).most_common(5)])
    n_pca_cls = len(set(P_pca.tolist()))
    n_mar_cls = len(set(P_mar.tolist()))
    subtitle = (
        f"Veres in-vitro (n={n:,}) — Marker predicts {n_mar_cls} class(es) "
        f"[top-5 shown: {top_mar}]; PCA predicts {n_pca_cls} class(es) [top-5: {top_pca}]."
    )
    scatter_side_by_side(
        emb_pca, emb_mar, P_pca_r, P_mar_r, palette_c,
        "PANDA-PCA — predicted class", "PANDA-Marker — predicted class",
        subtitle,
        FIG_S / "26b_pca_vs_marker_umaps_veres_by_class.pdf",
        legend_title="Predicted class",
    )


# ------------------- en1-cKO enrichment bars -------------------
def fig_en1_enrichment():
    raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad")
    CKO = {"GSM6833482", "GSM6833483"}
    WT  = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"}
    genotype = np.where(raw.obs["sample"].astype(str).isin(list(CKO)), "En1-cKO",
                np.where(raw.obs["sample"].astype(str).isin(list(WT)), "WT", "other"))
    pred = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv")
    common = raw.obs_names.intersection(pd.Index(pred["cell_id"].astype(str)))
    keep = raw.obs_names.isin(common)
    raw = raw[keep].copy()
    gt = genotype[keep]
    pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"]))
    labels = np.array([pred_map.get(c, "unknown") for c in raw.obs_names])
    labeled = (gt != "other")
    baseline = (gt[labeled] == "En1-cKO").sum() / labeled.sum()
    df = pd.DataFrame({"pred": labels, "gt": gt, "labeled": labeled})
    dfl = df[df["labeled"]]
    rows = []
    for cls, sub in dfl.groupby("pred"):
        n = len(sub)
        if n < 50: continue
        frac_cko = (sub["gt"] == "En1-cKO").sum() / n
        rows.append({"class": cls, "n": n, "frac_cko": frac_cko,
                     "delta": frac_cko - baseline})
    d = pd.DataFrame(rows).sort_values("delta", ascending=False)
    d.to_csv(FIG_S / "27_dingwall_en1_enrichment.csv", index=False)

    fig, ax = plt.subplots(figsize=(14, 8))
    y = np.arange(len(d))
    cols = [GENOTYPE_COLORS["En1-cKO"] if r["delta"] > 0.05
            else GENOTYPE_COLORS["WT"] if r["delta"] < -0.05
            else "#888888"
            for _, r in d.iterrows()]
    deltas = (d["frac_cko"] - baseline).values
    ax.barh(y, deltas, color=cols, edgecolor="k", linewidth=0.5)
    ax.axvline(0, color="black", linewidth=1)
    # place value labels always to the right of the bar tip with an offset in
    # display coords so short/negative bars can't crash into the y-tick labels
    for i, row in enumerate(d.itertuples()):
        delta_i = row.frac_cko - baseline
        # negative bars: anchor at the zero line so text never overprints
        # the axvline or the bar itself (rows never mix +/- bars)
        ax.annotate(f"{row.frac_cko:.2f}  (n={int(row.n):,})",
                    xy=(max(delta_i, 0.0), i), xycoords="data",
                    xytext=(6, 0), textcoords="offset points",
                    ha="left", va="center", fontsize=ANNOT_FS, clip_on=False)
    ax.set_yticks(y); ax.set_yticklabels(d["class"], fontsize=TICK_FS)
    ax.tick_params(axis="y", pad=6)
    # add headroom on the right so annotations don't clip
    dmin, dmax = float(deltas.min()), float(deltas.max())
    span = max(abs(dmin), abs(dmax))
    ax.set_xlim(dmin - 0.05 * span, dmax + 0.55 * span)
    ax.set_xlabel(f"Δ En1-cKO fraction vs baseline {baseline:.2f}", fontsize=LABEL_FS)
    ax.invert_yaxis()
    ax.set_title("Dingwall — En1-cKO enrichment per PANDA-predicted class\n"
                 "(red = cKO-enriched; blue = WT-enriched; grey = at baseline)",
                 fontsize=TITLE_FS)
    plt.tight_layout()
    plt.savefig(FIG_S / "27_dingwall_en1_enrichment.pdf", bbox_inches="tight")
    plt.close()
    print(f"[fig] 27_dingwall_en1_enrichment.pdf ({len(d)} classes shown)")


# ------------------- melanocyte pathway modules -------------------
MELANOCYTE_MODULES = {
    "MITF regulon (up in WT, down in cKO)": ["Mitf", "Dct", "Tyrp1", "Tyr", "Pmel", "Mlana", "Slc24a5", "Sox10"],
    "keratinocyte contamination / Krt (up in cKO)": ["Krt5", "Krt14", "Krt15"],
    "melanocyte proliferation / migration (down in cKO)": ["Ets1", "Kit", "Pax3", "Sox9"],
    "pigment biogenesis (down in cKO)": ["Gpnmb", "Slc45a2", "Oca2", "Trpm1"],
}


def fig_melanocyte_pathway():
    raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad")
    CKO = {"GSM6833482", "GSM6833483"}
    WT  = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"}
    genotype = np.where(raw.obs["sample"].astype(str).isin(list(CKO)), "En1-cKO",
                np.where(raw.obs["sample"].astype(str).isin(list(WT)), "WT", "other"))
    pred = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv")
    pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"]))
    labels = np.array([pred_map.get(c, "unknown") for c in raw.obs_names])
    mel_mask = (labels == "melanocyte") & (genotype != "other")
    a_mel = raw[mel_mask].copy()
    gt_mel = genotype[mel_mask]
    print(f"[melanocyte] {a_mel.n_obs} cells (WT {(gt_mel=='WT').sum()} + cKO {(gt_mel=='En1-cKO').sum()})", flush=True)

    sc.pp.normalize_total(a_mel, target_sum=1e4); sc.pp.log1p(a_mel)
    rows = []
    for mod_name, genes in MELANOCYTE_MODULES.items():
        present = [g for g in genes if g in a_mel.var_names]
        if not present: continue
        sc.tl.score_genes(a_mel, gene_list=present, score_name="s_tmp", use_raw=False)
        s = a_mel.obs["s_tmp"].values
        wt_mean = s[gt_mel == "WT"].mean()
        cko_mean = s[gt_mel == "En1-cKO"].mean()
        from scipy.stats import mannwhitneyu
        _, p = mannwhitneyu(s[gt_mel == "WT"], s[gt_mel == "En1-cKO"], alternative="two-sided")
        rows.append({"module": mod_name, "genes": ", ".join(present),
                     "wt_mean": wt_mean, "cko_mean": cko_mean,
                     "delta": cko_mean - wt_mean, "p": p})
    d = pd.DataFrame(rows)
    d.to_csv(FIG_S / "28_melanocyte_pathway_modules.csv", index=False)

    fig, ax = plt.subplots(figsize=(14.5, 7.5))
    x = np.arange(len(d))
    width = 0.35
    ax.bar(x - width/2, d["wt_mean"],  width, label="WT",
           color=GENOTYPE_COLORS["WT"], edgecolor="k")
    ax.bar(x + width/2, d["cko_mean"], width, label="En1-cKO",
           color=GENOTYPE_COLORS["En1-cKO"], edgecolor="k")
    # find headroom so the p-labels never collide with the suptitle
    y_max_data = max(d["wt_mean"].max(), d["cko_mean"].max())
    y_min_data = min(0.0, d["wt_mean"].min(), d["cko_mean"].min())
    span = y_max_data - y_min_data
    for i, r in enumerate(d.itertuples()):
        y_top = max(r.wt_mean, r.cko_mean) + 0.03 * span
        sig = "***" if r.p < 1e-3 else "**" if r.p < 1e-2 else "*" if r.p < 5e-2 else "n.s."
        ax.text(i, y_top, f"p={r.p:.1e} {sig}", ha="center", fontsize=ANNOT_FS)
    # explicit y-axis room above the tallest p-label
    ax.set_ylim(y_min_data - 0.05 * span, y_max_data + 0.18 * span)
    ax.set_xticks(x)
    ax.set_xticklabels([m.split(" (")[0] for m in d["module"]], fontsize=TICK_FS, rotation=15, ha="right")
    ax.set_ylabel("Module score (mean per cell)", fontsize=LABEL_FS)
    ax.set_title("Melanocyte pathway modules — WT vs En1-cKO (Dingwall predicted melanocytes)",
                 fontsize=TITLE_FS, pad=14)
    ax.legend(loc="center left", bbox_to_anchor=(1.02, 0.5), frameon=False,
              title="Genotype", fontsize=LEGEND_FS, title_fontsize=TITLE_FS)
    ax.axhline(0, color="black", linewidth=0.5, linestyle="--")
    plt.subplots_adjust(top=0.85)
    plt.tight_layout()
    plt.savefig(FIG_S / "28_melanocyte_pathway_modules.pdf", bbox_inches="tight")
    plt.close()
    print(f"[fig] 28_melanocyte_pathway_modules.pdf")


def merge_pdf():
    """merge all supplement pages into figures/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 / "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",
        FIG_S / "23_anchor_delta_recall.pdf",
        FIG_S / "24_pca_vs_marker_umaps_dingwall_by_genotype.pdf",
        FIG_S / "24b_pca_vs_marker_umaps_dingwall_by_class.pdf",
        FIG_S / "25_pca_vs_marker_umaps_dahlin_by_genotype.pdf",
        FIG_S / "25b_pca_vs_marker_umaps_dahlin_by_class.pdf",
        FIG_S / "26_pca_vs_marker_umaps_veres_by_stage.pdf",
        FIG_S / "26b_pca_vs_marker_umaps_veres_by_class.pdf",
        FIG_S / "27_dingwall_en1_enrichment.pdf",
        FIG_S / "28_melanocyte_pathway_modules.pdf",
    ]
    for p in order:
        if p.exists():
            w.append(str(p))
            print(f"  + {p.name}")
        else:
            print(f"  [skip] {p.name} missing")
    out = ROOT / "figures/PANDA_supplement.pdf"
    with open(out, "wb") as f:
        w.write(f)
    print(f"\nwrote {out}  ({out.stat().st_size / 1024:.0f} KB)")


def main():
    fig_dingwall_pca_vs_marker()
    fig_dahlin_pca_vs_marker()
    fig_veres_pca_vs_marker()
    fig_en1_enrichment()
    fig_melanocyte_pathway()
    merge_pdf()


if __name__ == "__main__":
    main()