"""regenerate fig5 (dingwall UMAP, 2 panels) + fig6 (3-panel multi UMAP).

Uses shared canonical palette (scripts/figures/palette.py) so a class gets
the same color in every figure. Reuses project()/do_umap()/_load_dahlin_raw()/
_load_veres() from build_pca_vs_marker_umaps.py.

Caches expensive UMAP embeddings to figures/_cache_*.npz for reuse.

Fig 6c restricts Veres to the 12,297 held-out slice
(data/corpus/pancreas/held_out_labeled/veres_GSE114412_test.h5ad).

outputs:
  figures/fig5_dingwall_umap.pdf
  figures/fig6_multi_umap.pdf
"""
from __future__ import annotations
from pathlib import Path
import warnings, sys, numpy as np, pandas as pd
warnings.filterwarnings("ignore")
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
import anndata as ad

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.mkdir(exist_ok=True)

sys.path.insert(0, str(ROOT / "scripts/figures"))
from build_pca_vs_marker_umaps import (
    project, do_umap, _load_dahlin_raw, _load_veres,
)
from palette import color_for, apply_style, GENOTYPE_COLORS, STAGE_COLORS
apply_style()

# ---- shared style constants (user style spec) ----
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
LEGEND_TITLE_FS = 11  # legend titles
ANNOT_FS = 11         # in-plot annotations
PANEL_FS = 13         # panel letters (a)/(b)/(c)

# override rc defaults from apply_style() so implicit sizes also conform
plt.rcParams.update({
    "axes.titlesize": TITLE_FS,
    "axes.labelsize": LABEL_FS,
    "xtick.labelsize": TICK_FS,
    "ytick.labelsize": TICK_FS,
    "legend.fontsize": LEGEND_FS,
    "legend.title_fontsize": LEGEND_TITLE_FS,
    "figure.titlesize": SUPTITLE_FS,
})

RS = 42

# stage color for the primary-islet bucket (not in ordinal STAGE_COLORS)
# medium grey so the legend swatch is visible on white backgrounds
_ISLET_GRAY = "#707070"

# darker grey for the "other" bucket in fig5 panel (a) — the light #e5e5e5
# used previously was nearly invisible in legend swatches on white paper
_OTHER_GREY = "#909090"

# Dingwall genotype mapping (verified against GSE220977 metadata)
DINGWALL_CKO = {"GSM6833482", "GSM6833483"}
DINGWALL_WT  = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"}


def _panel_letter(ax, letter, x=-0.08, y=1.03, fontsize=PANEL_FS):
    ax.text(x, y, f"({letter})", transform=ax.transAxes,
            fontsize=fontsize, fontweight="bold", va="bottom", ha="left")


def _dominant_classes(P, min_frac=0.005):
    """return the classes making up >= min_frac of the cells (in count order)."""
    from collections import Counter
    n = len(P)
    counts = Counter(P.tolist())
    return [c for c, k in counts.most_common() if k / n >= min_frac]


def fig5_dingwall():
    cache_p = FIG / "_cache_dingwall_full_marker.npz"
    if cache_p.exists():
        print(f"[fig5] reusing cache {cache_p.name}", flush=True)
        c = np.load(cache_p, allow_pickle=True)
        emb, P, gt = c["emb"], c["P"], c["genotype"]
        n = len(emb)
    else:
        print("[fig5] loading dingwall raw ...", flush=True)
        raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad")
        gt = np.where(raw.obs["sample"].astype(str).isin(list(DINGWALL_CKO)), "En1-cKO",
             np.where(raw.obs["sample"].astype(str).isin(list(DINGWALL_WT)), "WT", "other"))
        n = raw.n_obs
        print(f"[fig5] projecting {n:,} cells (PANDA-Marker) ...", flush=True)
        Z, P, _ = project(raw, "pan_skin", "marker")
        print(f"[fig5] UMAP on {n:,} projections ...", flush=True)
        emb = do_umap(Z)
        np.savez(cache_p, emb=emb, P=P, genotype=gt)

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

    # (a) PANDA-predicted class
    ax = axes[0]
    kept = _dominant_classes(P, min_frac=0.005)
    P_r = np.where(np.isin(P, kept), P, "other")
    # plot "other" first so kept classes render on top
    m_other = P_r == "other"
    if m_other.sum() > 0:
        ax.scatter(emb[m_other, 0], emb[m_other, 1], s=6, alpha=0.6,
                   c=_OTHER_GREY, label=f"other (n={int(m_other.sum()):,})",
                   linewidths=0, rasterized=True)
    for cls in kept:
        m = P_r == cls
        if m.sum() == 0: continue
        ax.scatter(emb[m, 0], emb[m, 1], s=6, alpha=0.6,
                   c=color_for(cls),
                   label=f"{cls} (n={int(m.sum()):,})", linewidths=0,
                   rasterized=True)
    ax.set_title("Dingwall skin — PANDA-Marker predicted class",
                 fontsize=TITLE_FS, pad=8)
    ax.set_xlabel("UMAP-1", fontsize=LABEL_FS)
    ax.set_ylabel("UMAP-2", fontsize=LABEL_FS)
    ax.set_xticks([]); ax.set_yticks([])
    ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=LEGEND_FS,
              markerscale=1.6, frameon=False,
              title=f"Predicted class (n={n:,})",
              title_fontsize=LEGEND_TITLE_FS,
              handletextpad=0.5, borderaxespad=0.4)
    _panel_letter(ax, "a")

    # (b) En1 genotype
    ax = axes[1]
    for g in ["other", "WT", "En1-cKO"]:  # cKO last so it plots on top
        m = gt == g
        if m.sum() == 0: continue
        ax.scatter(emb[m, 0], emb[m, 1], s=6, alpha=0.6,
                   c=GENOTYPE_COLORS.get(g, "#bbbbbb"),
                   label=f"{g} (n={int(m.sum()):,})",
                   linewidths=0, rasterized=True)
    ax.set_title("Dingwall skin — En1 genotype", fontsize=TITLE_FS, pad=8)
    ax.set_xlabel("UMAP-1", fontsize=LABEL_FS)
    ax.set_ylabel("UMAP-2", fontsize=LABEL_FS)
    ax.set_xticks([]); ax.set_yticks([])
    ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=LEGEND_FS,
              markerscale=1.6, frameon=False,
              title="En1 genotype", title_fontsize=LEGEND_TITLE_FS,
              handletextpad=0.5, borderaxespad=0.4)
    _panel_letter(ax, "b")

    plt.savefig(FIG / "fig5_dingwall_umap.pdf", bbox_inches="tight", dpi=200)
    plt.close()
    print(f"[fig5] wrote {FIG}/fig5_dingwall_umap.pdf", flush=True)


def _get_dingwall_cached():
    cache_p = FIG / "_cache_dingwall_full_marker.npz"
    if cache_p.exists():
        c = np.load(cache_p, allow_pickle=True)
        return c["emb"], c["genotype"]
    # otherwise recompute
    print("[fig6] recomputing dingwall ...", flush=True)
    raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad")
    gt_s = np.where(raw.obs["sample"].astype(str).isin(list(DINGWALL_CKO)), "En1-cKO",
           np.where(raw.obs["sample"].astype(str).isin(list(DINGWALL_WT)), "WT", "other"))
    Z_s, P_s, _ = project(raw, "pan_skin", "marker")
    emb_s = do_umap(Z_s)
    np.savez(cache_p, emb=emb_s, P=P_s, genotype=gt_s)
    return emb_s, gt_s


def _get_dahlin_cached():
    cache_p = FIG / "_cache_dahlin_marker.npz"
    if cache_p.exists():
        print(f"[fig6] reusing {cache_p.name}", flush=True)
        c = np.load(cache_p, allow_pickle=True)
        return c["emb"], c["genotype"]
    print("[fig6] loading dahlin raw ...", flush=True)
    a_d = _load_dahlin_raw()
    gt_d = a_d.obs["genotype"].astype(str).values
    print(f"[fig6] projecting dahlin ({a_d.n_obs:,}) ...", flush=True)
    Z_d, _, _ = project(a_d, "hematopoiesis", "marker")
    print(f"[fig6] UMAP dahlin ...", flush=True)
    emb_d = do_umap(Z_d)
    np.savez(cache_p, emb=emb_d, genotype=gt_d)
    return emb_d, gt_d


def _get_veres_heldout_cached():
    """Project + UMAP on the 12,297 held-out Veres slice only."""
    cache_p = FIG / "_cache_veres_heldout_marker.npz"
    if cache_p.exists():
        print(f"[fig6] reusing {cache_p.name}", flush=True)
        c = np.load(cache_p, allow_pickle=True)
        return c["emb"], c["stage"]

    print("[fig6] loading veres raw + held-out obs list ...", flush=True)
    heldout_p = ROOT / "data/corpus/pancreas/held_out_labeled/veres_GSE114412_test.h5ad"
    if heldout_p.exists():
        heldout_raw = set(ad.read_h5ad(heldout_p).obs_names.astype(str).tolist())
        heldout_names = {n[6:] if n.startswith("veres_") else n for n in heldout_raw}
        heldout_names |= heldout_raw
    else:
        print(f"[fig6] WARN: held-out file missing; using ALL veres cells")
        heldout_names = None

    a_v = _load_veres()
    if heldout_names is not None:
        keep = np.array([str(n) in heldout_names for n in a_v.obs_names])
        print(f"[fig6] restricting veres to held-out: {keep.sum():,}/{a_v.n_obs:,}", flush=True)
        a_v = a_v[keep].copy()

    stage_col = "Stage" if "Stage" in a_v.obs.columns else "stage"
    stage = pd.to_numeric(a_v.obs[stage_col], errors="coerce").fillna(-1).astype(int).values
    st_str = np.array([str(s) if s > 0 else "islet" for s in stage])
    print(f"[fig6] projecting veres held-out ({a_v.n_obs:,}) ...", flush=True)
    Z_v, _, _ = project(a_v, "pancreas", "marker")
    print(f"[fig6] UMAP veres held-out ...", flush=True)
    emb_v = do_umap(Z_v)
    np.savez(cache_p, emb=emb_v, stage=st_str)
    return emb_v, st_str


def fig6_multi():
    """3-panel multi-system UMAP colored by biology-of-interest label."""
    fig, axes = plt.subplots(1, 3, figsize=(22, 8), constrained_layout=True)

    # ---- (a) Dingwall (skin, En1 genotype) ----
    emb_s, gt_s = _get_dingwall_cached()
    ax = axes[0]
    for g in ["other", "WT", "En1-cKO"]:
        m = gt_s == g
        if m.sum() == 0: continue
        ax.scatter(emb_s[m, 0], emb_s[m, 1], s=5, alpha=0.55,
                   c=GENOTYPE_COLORS.get(g, "#bbbbbb"),
                   label=f"{g} (n={int(m.sum()):,})",
                   linewidths=0, rasterized=True)
    ax.set_title("(a) Dingwall skin — En1 genotype", fontsize=TITLE_FS, pad=8)
    ax.set_xlabel("UMAP-1", fontsize=LABEL_FS)
    ax.set_ylabel("UMAP-2", fontsize=LABEL_FS)
    ax.set_xticks([]); ax.set_yticks([])
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.02),
              ncol=3, frameon=False, fontsize=LEGEND_FS, markerscale=2.0,
              handletextpad=0.4, columnspacing=1.2)

    # ---- (b) Dahlin (HSC, Kit genotype) ----
    emb_d, gt_d = _get_dahlin_cached()
    ax = axes[1]
    # dahlin uses "unknown" for cells outside labeled samples
    for g in ["unknown", "WT", "Kit_W41"]:
        m = gt_d == g
        if m.sum() == 0: continue
        c = GENOTYPE_COLORS.get(g, "#bbbbbb") if g != "unknown" else "#bbbbbb"
        ax.scatter(emb_d[m, 0], emb_d[m, 1], s=5, alpha=0.55,
                   c=c, label=f"{g} (n={int(m.sum()):,})",
                   linewidths=0, rasterized=True)
    ax.set_title("(b) Dahlin hematopoiesis — Kit genotype",
                 fontsize=TITLE_FS, pad=8)
    ax.set_xlabel("UMAP-1", fontsize=LABEL_FS)
    ax.set_ylabel("UMAP-2", fontsize=LABEL_FS)
    ax.set_xticks([]); ax.set_yticks([])
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.02),
              ncol=3, frameon=False, fontsize=LEGEND_FS, markerscale=2.0,
              handletextpad=0.4, columnspacing=1.2)

    # ---- (c) Veres held-out (pancreas, Stage) ----
    emb_v, st_str = _get_veres_heldout_cached()
    ax = axes[2]
    stage_order = ["islet", "3", "4", "5", "6"]
    for s in stage_order:
        m = st_str == s
        if m.sum() == 0: continue
        color = STAGE_COLORS[s] if s in STAGE_COLORS else _ISLET_GRAY
        label = (f"Stage {s} (n={int(m.sum()):,})" if s != "islet"
                 else f"islet (n={int(m.sum()):,})")
        ax.scatter(emb_v[m, 0], emb_v[m, 1], s=5, alpha=0.55,
                   c=color, label=label, linewidths=0, rasterized=True)
    ax.set_title(f"(c) Veres pancreas held-out (n={len(st_str):,}) — protocol stage",
                 fontsize=TITLE_FS, pad=8)
    ax.set_xlabel("UMAP-1", fontsize=LABEL_FS)
    ax.set_ylabel("UMAP-2", fontsize=LABEL_FS)
    ax.set_xticks([]); ax.set_yticks([])
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.02),
              ncol=5, frameon=False, fontsize=LEGEND_FS, markerscale=2.0,
              handletextpad=0.4, columnspacing=1.2)

    plt.savefig(FIG / "fig6_multi_umap.pdf", bbox_inches="tight", dpi=180)
    plt.close()
    print(f"[fig6] wrote {FIG}/fig6_multi_umap.pdf", flush=True)


if __name__ == "__main__":
    fig5_dingwall()
    fig6_multi()