scripts/figures/generate_paper_figures.py

"""main-text figures for PAPER.tex — cv per-class F1 bars, dahlin heatmap, veres stage stack.

Uses shared canonical palette (scripts/figures/palette.py) so the same class
gets the same color in every figure.
"""
from __future__ import annotations
from pathlib import Path
import warnings, json, sys
warnings.filterwarnings("ignore")

import numpy as np, pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from pathlib import Path as _P_root
ROOT = _P_root(__file__).resolve().parents[2]
ROOT_STR = str(ROOT)
FIG = Path(f"{ROOT_STR}/figures")
FIG.mkdir(exist_ok=True)

# shared canonical palette + style
sys.path.insert(0, str(_P_root(__file__).resolve().parent))
from palette import color_for, apply_style, CLASS_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
CBAR_LABEL_FS = 12    # colorbar labels
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,
})


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


def fig1_perclass_f1():
    """per-class F1 bars from 5-fold CV JSONs, one axis per system, class-colored."""
    panels = [
        ("Pan-skin", f"{ROOT_STR}/discovery/pan_skin/marker/cv_5fold.json"),
        ("Pan-hematopoiesis", f"{ROOT_STR}/discovery/hematopoiesis/marker/cv_5fold.json"),
        ("Pan-pancreas", f"{ROOT_STR}/discovery/pancreas/marker/cv_5fold.json"),
    ]
    fig, axes = plt.subplots(1, 3, figsize=(18, 6), constrained_layout=True)
    for i, (ax, (title, path)) in enumerate(zip(axes, panels)):
        r = json.load(open(path))
        rep = r["per_class_report"]
        # collect (class, f1, support)
        rows = [(c, rep[c]["f1-score"], int(rep[c]["support"]))
                for c in rep.keys()
                if c not in ("accuracy", "macro avg", "weighted avg")]
        # sort descending by F1 (best first at top)
        rows.sort(key=lambda t: t[1], reverse=True)
        classes = [t[0] for t in rows]
        f1s     = [t[1] for t in rows]
        supports = [t[2] for t in rows]

        colors = [color_for(c) for c in classes]
        y = np.arange(len(classes))
        bars = ax.barh(y, f1s, color=colors, edgecolor="white", linewidth=0.6)

        # per-bar n annotations (outside)
        for b, s in zip(bars, supports):
            ax.text(b.get_width() + 0.012, b.get_y() + b.get_height() / 2,
                    f"n={s:,}", va="center", fontsize=ANNOT_FS, color="#333")

        # scale
        ax.set_xlim(0, 1.22)
        ax.set_xticks([0.0, 0.25, 0.5, 0.75, 1.0])
        ax.set_xticklabels(["0.0", "0.25", "0.50", "0.75", "1.00"],
                           fontsize=TICK_FS)
        ax.set_xlabel("held-out F1", fontsize=LABEL_FS)

        # y-axis: shrink font slightly if many classes
        y_fs = TICK_FS + 1 if len(classes) <= 13 else TICK_FS
        ax.set_yticks(y)
        ax.set_yticklabels(classes, fontsize=y_fs)

        # panel title with n_cells; small acc/AUROC subtitle below
        n_cells = int(r["n_cells"])
        acc = r["mean_acc"]; auc = r["mean_auc"]
        ax.set_title(f"{title} (n_cells={n_cells:,})", fontsize=TITLE_FS, pad=32)
        ax.text(0.5, 1.01,
                f"acc = {acc:.3f}  ·  macro AUROC = {auc:.3f}",
                transform=ax.transAxes, ha="center", va="bottom",
                fontsize=ANNOT_FS, color="#555555")

        # F1=0.9 marker
        ax.axvline(0.9, color="#888888", linestyle="--", linewidth=1.1,
                   alpha=0.4, zorder=0)
        ax.grid(axis="x", alpha=0.25, linestyle=":", zorder=0)
        ax.invert_yaxis()
        _panel_letter(ax, "abc"[i], x=-0.32, y=1.02, fontsize=PANEL_FS)

    plt.savefig(FIG / "fig1_perclass_f1.pdf", bbox_inches="tight")
    plt.close()
    print(f"[fig1] wrote {FIG}/fig1_perclass_f1.pdf")


def fig3_dahlin_heatmap():
    """dahlin within-class module-score heatmap (Kit_W41 minus WT)."""
    p = Path(f"{ROOT_STR}/discovery/hematopoiesis/marker/57_pathway_analysis.csv")
    if not p.exists(): print(f"[fig3] {p} not found"); return
    df = pd.read_csv(p)
    pivot = df.pivot(index="module_name", columns="class", values="delta")
    pivot_p = df.pivot(index="module_name", columns="class", values="mannu_p_adj_bonferroni")
    row_order = ["Kit_signaling", "Kit_ligand", "MYC_targets", "Integrated_stress",
                 "Apoptosis_pro", "Apoptosis_anti", "Cell_cycle",
                 "Erythropoiesis_early", "Erythropoiesis_late",
                 "OXPHOS_ETC", "Glycolysis",
                 "Redox_glutathione", "LT_HSC_quiescence"]
    row_order = [r for r in row_order if r in pivot.index]
    col_order = ["LT-HSC", "MPP", "erythroid", "myeloid", "megakaryocyte",
                 "lymphoid", "basophil-mast", "monocyte", "macrophage"]
    col_order = [c for c in col_order if c in pivot.columns]
    P = pivot.loc[row_order, col_order]
    Pp = pivot_p.loc[row_order, col_order]

    fig, ax = plt.subplots(figsize=(11, 8), constrained_layout=True)
    vmax = np.nanmax(np.abs(P.values))
    im = ax.imshow(P.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")

    # cell annotations
    for i in range(P.shape[0]):
        for j in range(P.shape[1]):
            v = P.values[i, j]; pv = Pp.values[i, j]
            if np.isnan(v): continue
            star = "**" if pv < 1e-3 else "*" if pv < 0.05 else ""
            if abs(v) > 0.1:
                label = f"{v:+.2f}"
                if star: label = f"{label} {star}"
                color = "white" if abs(v) >= vmax * 0.6 else "black"
                ax.text(j, i, label, ha="center", va="center",
                        fontsize=ANNOT_FS, color=color)
            elif star:
                ax.text(j, i, star, ha="center", va="center",
                        fontsize=ANNOT_FS, color="black")

    # col labels (rotated 45)
    ax.set_xticks(range(len(col_order)))
    ax.set_xticklabels(col_order, rotation=45, ha="right", fontsize=TICK_FS)
    # row labels (module names)
    ax.set_yticks(range(len(row_order)))
    ax.set_yticklabels(row_order, fontsize=TICK_FS)

    # color the x tick labels (class names) using canonical palette
    for tl, cls in zip(ax.get_xticklabels(), col_order):
        tl.set_color(color_for(cls))
        tl.set_fontweight("bold")

    cbar = plt.colorbar(im, ax=ax, shrink=0.85, pad=0.02)
    cbar.set_label("Δ module score (Kit_W41 − WT)", fontsize=CBAR_LABEL_FS)
    cbar.ax.tick_params(labelsize=TICK_FS)

    ax.set_title("Dahlin: Kit-W41 vs WT within-class pathway module contrast",
                 fontsize=TITLE_FS, pad=14)
    fig.text(0.5, -0.01,
             "*  p<0.05     **  p<10$^{-3}$   (Mann–Whitney, Bonferroni)",
             ha="center", fontsize=ANNOT_FS, color="#555555")

    plt.savefig(FIG / "fig3_dahlin_heatmap.pdf", bbox_inches="tight")
    plt.close()
    print(f"[fig3] wrote {FIG}/fig3_dahlin_heatmap.pdf")


def fig4_sharon_stage_stack():
    """stacked-bar class fractions across veres stages 3-6, class-colored via palette."""
    import re
    p = Path(f"{ROOT_STR}/discovery/pancreas/marker/veres_predictions.csv")
    if not p.exists(): print(f"[fig4] {p} not found"); return
    pred = pd.read_csv(p)

    stg_re = re.compile(r"_S(\d)c_")
    stages = pred["cell_id"].astype(str).apply(
        lambda s: int(stg_re.search(s).group(1)) if stg_re.search(s) else np.nan)
    pred = pred.assign(stage=stages).dropna(subset=["stage"])
    pred["stage"] = pred["stage"].astype(int)
    n_staged = len(pred); n_pre = len(stages); n_drop = n_pre - n_staged
    print(f"[fig4] staged cells: {n_staged:,} (dropped {n_drop:,} primary-islet cells)")

    ct = (pred.groupby(["stage", "pred_label"]).size()
              .unstack("pred_label", fill_value=0))
    frac = ct.div(ct.sum(axis=1), axis=0)
    priority = ["pancreatic-progenitor", "proliferating",
                "endocrine-progenitor", "endocrine-progenitor-primed",
                "Fev-EP", "beta_progenitor", "beta",
                "alpha_progenitor", "alpha", "delta", "gamma", "epsilon",
                "acinar", "ductal", "exocrine",
                "mesenchyme", "endothelial", "immune"]
    present = list(frac.columns)
    ordered = [c for c in priority if c in present] + \
              [c for c in sorted(present, key=lambda x: -frac[x].sum())
               if c not in priority]
    frac = frac[ordered]

    fig, ax = plt.subplots(figsize=(12, 7), constrained_layout=True)
    bottom = np.zeros(frac.shape[0])
    x = np.arange(frac.shape[0])
    for cls in ordered:
        vals = frac[cls].values
        ax.bar(x, vals, bottom=bottom, label=cls,
               color=color_for(cls), edgecolor="white", linewidth=0.6)
        bottom += vals

    for xi, s in zip(x, frac.index):
        n_stage = int(ct.loc[s].sum())
        ax.text(xi, 1.03, f"n = {n_stage:,}", ha="center", va="bottom",
                fontsize=ANNOT_FS, color="#222222")

    ax.set_xticks(x)
    ax.set_xticklabels([f"Stage {int(s)}" for s in frac.index], fontsize=TICK_FS)
    ax.set_xlabel("Veres protocol stage", fontsize=LABEL_FS)
    ax.set_ylabel("Predicted class fraction", fontsize=LABEL_FS)
    ax.set_title("Veres 2019 pancreatic differentiation", fontsize=TITLE_FS,
                 pad=12)

    # legend reversed so it reads top-of-stack first (matches the visual)
    handles, labels = ax.get_legend_handles_labels()
    ax.legend(handles[::-1], labels[::-1],
              bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=LEGEND_FS,
              frameon=False, title="Predicted class",
              title_fontsize=LEGEND_TITLE_FS,
              handlelength=1.6, borderaxespad=0.4)
    ax.set_ylim(0, 1.12)
    ax.set_yticks([0.0, 0.25, 0.5, 0.75, 1.0])
    ax.tick_params(axis="y", labelsize=TICK_FS)
    plt.savefig(FIG / "fig4_veres_stage_stack.pdf", bbox_inches="tight")
    plt.close()
    print(f"[fig4] wrote {FIG}/fig4_veres_stage_stack.pdf "
          f"({len(ordered)} classes over stages {list(frac.index)})")


if __name__ == "__main__":
    fig1_perclass_f1()
    fig3_dahlin_heatmap()
    fig4_sharon_stage_stack()
    print(f"\nAll figures in {FIG}/")