api/app.py

"""FastAPI app for OncoDSL — backs both the Streamlit H2 tab (static artefacts)
and the React/Next.js "Lab" (live-streamed runs).

Endpoints split by surface:

**Legacy / Streamlit (single persisted MSI run from `scripts.run_h2`):**
- `GET  /health`        — liveness + which artefacts are on disk.
- `GET  /run`           — the persisted evolution log (anonymised).
- `GET  /result`        — the persisted result (anonymised).
- `POST /reveal`        — translate opaque IDs back to symbols.

**Lab (live runs in memory; objective spec drives the engine):**
- `POST /runs`          — start a GP run in a worker thread; return run_id.
- `GET  /runs/{id}`     — status + accumulated log so far (polling fallback).
- `GET  /runs/{id}/stream`  — SSE; per-generation events then a `done` sentinel.
- `GET  /runs/{id}/result`  — final result; 425 Too Early if still running.
- `POST /evaluate`      — reveal supplied IDs + score overlap with a named
                          reference gene set.

CORS is wide-open so the Next.js dev server on :3000 can hit any endpoint.
The reveal-side endpoints (`/reveal`, `/evaluate`) are the ONLY map-readers
— they never dump the whole map; only the supplied IDs are translated.

Start with:
    uvicorn api.app:app --reload
"""

from __future__ import annotations

import asyncio
import json
import logging
import threading
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal

import numpy as np
import pandas as pd
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from sse_starlette.sse import EventSourceResponse

from airgap import anonymise, reveal
from data_pipeline import schema
from dsl import Load
from engine import objective_from_spec, run_gp_pipeline_streaming
from engine_v2 import run_v2_pipeline_streaming
from engine_v2.fitness import objective_from_spec as v2_objective_from_spec
from validate.h1 import (
    IMMUNE_GENES,
    MMR_GENES,
    POSITIVE_LABEL,
    usable_msi_cohort,
)


log = logging.getLogger("oncodsl.api")
logging.basicConfig(level=logging.INFO)


import math


def _json_finite(value):
    """Recursively replace NaN / +inf / -inf with None — those tokens
    aren't valid JSON and break ``JSON.parse`` in the browser.

    The engines already floor invalid fitness to the objective's finite
    worst-case before reaching this point; this walker is a last-line
    safety net for anything else (e.g. an empty permutation array
    yielding a NaN mean)."""
    if isinstance(value, float):
        return value if math.isfinite(value) else None
    if isinstance(value, dict):
        return {k: _json_finite(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_json_finite(v) for v in value]
    if isinstance(value, tuple):
        return tuple(_json_finite(v) for v in value)
    return value


app = FastAPI(title="OncoDSL API", version="1.1")

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:3000",
        "http://127.0.0.1:3000",
        "*",
    ],
    allow_credentials=False,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)


H2_DIR = schema.PROCESSED_DIR / "h2"
EVOLUTION_PATH = H2_DIR / "evolution_log.json"
RESULT_PATH = H2_DIR / "result.json"


# Per-dataset reference gene sets for the /evaluate endpoint. Colorectal
# uses the existing MMR / immune sets from validate.h1; head & neck has
# HPV-detection markers (CDKN2A/p16 + RB1-pathway / E2F-target proxies
# that go UP when HPV-E7 disrupts RB1) and a small cell-cycle set.
# HNSC reference sets. ``p16`` is just the canonical surrogate marker
# (CDKN2A); ``cell_cycle`` is a standard E2F-target / proliferation core
# the literature uses, NOT cherry-picked to match any winner. HPV+
# tumours run their cell cycle high because HPV-E7 disrupts RB1 and
# releases E2F.
HPV_P16_GENES: list[str] = ["CDKN2A"]
HPV_CELL_CYCLE_GENES: list[str] = [
    "MCM2", "MCM3", "MCM4", "MCM5", "MCM6", "MCM7",
    "PCNA", "CDK1", "CCNE1", "CCNB1", "CDC6", "CDC20",
    "MKI67", "TOP2A", "RRM2", "TYMS", "FOXM1", "E2F1",
    "BUB1", "AURKB",
]

# Weight applied to the coherence prior (mean abs pairwise correlation
# over Select'd opaque columns) when RunRequest.coherence is true.
# Modest so separation still dominates; a single sharp gene can still
# win, but coordinated modules now have an edge.
COHERENCE_DEFAULT_WEIGHT = 0.10

# Immune-infiltration proxy used by the module-ranking "purity" flag.
# Per-patient mean expression of these standard cytotoxic-T-cell markers
# is a coarse proxy for immune infiltration; low proxy => fewer immune
# cells in the bulk biopsy => higher tumour purity. Curated once here
# so the validation layer has a single source of truth. Resolved to
# opaque IDs via the sealed map at module-ranking time (bounded reveal).
HPV_IMMUNE_PROXY_GENES: list[str] = ["CD8A", "GZMB", "PRF1", "CD3D", "CD2"]


REFERENCE_SETS_BY_DATASET: dict[str, dict[str, list[str]]] = {
    "coadread": {
        "MMR": list(MMR_GENES),
        "immune": list(IMMUNE_GENES),
    },
    "hnsc": {
        "p16": HPV_P16_GENES,
        "cell_cycle": HPV_CELL_CYCLE_GENES,
    },
}

# Legacy flat alias — old tests + the streamlit /evaluate path still
# read the colorectal sets directly through this name. New code should
# route by dataset via REFERENCE_SETS_BY_DATASET.
REFERENCE_SETS: dict[str, list[str]] = REFERENCE_SETS_BY_DATASET["coadread"]


# Valid (dataset, target) combinations. Anything outside this map is
# rejected at the /runs boundary before any work starts.
DATASET_TARGETS: dict[str, set[str]] = {
    "coadread": {"msi", "tmb", "none"},
    "hnsc": {"hpv", "none"},
}


# ---------------------------------------------------------------------------
# Legacy / Streamlit endpoints (unchanged)
# ---------------------------------------------------------------------------


def _load_json(path: Path):
    if not path.exists():
        raise HTTPException(
            status_code=503,
            detail=(
                f"{path.name} not found. Run: `python -m scripts.run_h2` "
                "to produce the H2 artefacts."
            ),
        )
    return json.loads(path.read_text())


class RevealRequest(BaseModel):
    gene_ids: list[str] = Field(..., description="Opaque feature IDs (^g\\d+$).")


class RevealResponse(BaseModel):
    symbols: list[str]


@app.get("/health")
def health() -> dict:
    return {
        "status": "ok",
        "artefacts": {
            "evolution_log": EVOLUTION_PATH.exists(),
            "result": RESULT_PATH.exists(),
        },
    }


@app.get("/run")
def get_run():
    return _load_json(EVOLUTION_PATH)


@app.get("/result")
def get_result():
    return _load_json(RESULT_PATH)


@app.post("/reveal", response_model=RevealResponse)
def post_reveal(req: RevealRequest) -> RevealResponse:
    try:
        symbols = reveal(req.gene_ids)
    except FileNotFoundError as exc:
        raise HTTPException(status_code=503, detail=str(exc))
    except KeyError as exc:
        raise HTTPException(status_code=400, detail=str(exc))
    return RevealResponse(symbols=symbols)


# ---------------------------------------------------------------------------
# Lab — cached data prep
# ---------------------------------------------------------------------------

_DATA_LOCK = threading.Lock()
_DATA_CACHE: dict[str, Any] = {}


def _prepare_lab_data(
    target: str, dataset: str = "coadread",
) -> tuple[pd.DataFrame, np.ndarray, pd.DataFrame, dict[str, np.ndarray]]:
    """Load + anonymise once per (dataset, target). Cached for the
    process lifetime.

    Returns ``(M, y, clinical, extra_labels)``. Clinical only contains
    named fields (stage, age) — never gene symbols. ``extra_labels``
    carries any OTHER named targets so engine_v2 programs can compute
    post-hoc / Associate against them.
    """
    cache_key = f"{dataset}:{target}"
    with _DATA_LOCK:
        if cache_key in _DATA_CACHE:
            entry = _DATA_CACHE[cache_key]
            return entry["M"], entry["y"], entry["clinical"], entry["extra_labels"]

        if dataset == "coadread" and target == "msi":
            cohort = usable_msi_cohort(Load("processed"))
            M = anonymise(cohort.expression)
            y = (cohort.labels["msi_status"] == POSITIVE_LABEL).astype(int).values
            tmb = pd.to_numeric(cohort.labels["tmb"], errors="coerce")
            extra = {"tmb": tmb.reindex(M.index).astype(float).values}
            clinical = cohort.clinical.reindex(M.index)[["stage", "age"]].copy()
        elif dataset == "coadread" and target == "tmb":
            cohort = Load("processed")
            tmb = pd.to_numeric(cohort.labels["tmb"], errors="coerce")
            usable = tmb.notna() & (~cohort.expression.isna().any(axis=1))
            ids = cohort.sample_ids[usable]
            M = anonymise(cohort.expression.loc[ids])
            y = tmb.loc[ids].astype(float).values
            msi = (cohort.labels["msi_status"]
                   .reindex(ids)
                   .eq(POSITIVE_LABEL)
                   .astype(int)
                   .values)
            extra = {"msi": msi}
            clinical = cohort.clinical.reindex(ids)[["stage", "age"]].copy()
        elif dataset == "coadread" and target == "none":
            # Unsupervised: keep every patient with complete expression,
            # regardless of MSI / TMB. y is None at engine time; the
            # labels travel alongside (NaN where missing) so the worker
            # can compute the post-hoc alignment AFTER the GP finishes.
            cohort = Load("processed")
            usable = ~cohort.expression.isna().any(axis=1)
            ids = cohort.sample_ids[usable]
            M = anonymise(cohort.expression.loc[ids])
            y = None  # type: ignore[assignment]
            msi_status = cohort.labels["msi_status"].reindex(ids)
            msi_arr = np.where(
                msi_status.eq(POSITIVE_LABEL),
                1.0,
                np.where(msi_status.eq("MSS"), 0.0, np.nan),
            )
            tmb_arr = (
                pd.to_numeric(cohort.labels["tmb"], errors="coerce")
                .reindex(ids)
                .astype(float)
                .values
            )
            extra = {"msi": msi_arr, "tmb": tmb_arr}
            clinical = cohort.clinical.reindex(ids)[["stage", "age"]].copy()
        elif dataset == "hnsc" and target in ("hpv", "none"):
            # Head & neck. Load the HNSC processed dir directly — Load's
            # else-branch handles any directory containing the two
            # parquets, so no Load() rewrite needed.
            cohort = Load(schema.HNSC_PROCESSED_DIR)
            hpv_status = cohort.labels.get("hpv_status")
            if hpv_status is None:
                raise ValueError(
                    "HNSC processed cohort is missing the hpv_status label. "
                    "Re-run: python -m data_pipeline.build_hnsc"
                )
            if target == "hpv":
                usable = hpv_status.isin(["HPV+", "HPV-"]) & (
                    ~cohort.expression.isna().any(axis=1)
                )
                ids = cohort.sample_ids[usable]
                M = anonymise(cohort.expression.loc[ids])
                y = (hpv_status.reindex(ids) == "HPV+").astype(int).values
                extra = {}
                clin_cols = [
                    c
                    for c in [
                        "stage", "age", "sex", "race",
                        "tissue_site", "icd_o_3_site", "is_oropharynx",
                    ]
                    if c in cohort.clinical.columns
                ]
                clinical = cohort.clinical.reindex(ids)[clin_cols].copy()
            else:
                # Unsupervised: every patient with complete expression;
                # carry the HPV label aside for the post-hoc alignment.
                usable = ~cohort.expression.isna().any(axis=1)
                ids = cohort.sample_ids[usable]
                M = anonymise(cohort.expression.loc[ids])
                y = None  # type: ignore[assignment]
                hpv_arr = np.where(
                    hpv_status.reindex(ids).eq("HPV+"),
                    1.0,
                    np.where(hpv_status.reindex(ids).eq("HPV-"), 0.0, np.nan),
                )
                extra = {"hpv": hpv_arr}
                clin_cols = [
                    c
                    for c in [
                        "stage", "age", "sex", "race",
                        "tissue_site", "icd_o_3_site", "is_oropharynx",
                    ]
                    if c in cohort.clinical.columns
                ]
                clinical = cohort.clinical.reindex(ids)[clin_cols].copy()
        else:
            raise ValueError(f"Unknown (dataset, target): ({dataset!r}, {target!r})")

        _DATA_CACHE[cache_key] = {
            "M": M,
            "y": np.asarray(y) if y is not None else None,
            "clinical": clinical,
            "extra_labels": extra,
        }
        entry = _DATA_CACHE[cache_key]
        return entry["M"], entry["y"], entry["clinical"], entry["extra_labels"]


# ---------------------------------------------------------------------------
# Lab — run store + worker thread
# ---------------------------------------------------------------------------


@dataclass
class Run:
    id: str
    objective_spec: dict
    params: dict
    engine: str = "v1"   # "v1" or "v2"
    dataset: str = "coadread"   # "coadread" | "hnsc"
    status: str = "running"   # "running" | "done" | "error"
    coherence: bool = False    # "Prefer coordinated gene modules" prior was on
    diversity: bool = False    # "Maintain diversity" knob was on
    # Per-run DSL injection-rate overrides (e.g. {"search": 0.0,
    # "scalar_share": 0.30}). Missing keys keep their DEFAULT_RATES
    # value. None ⇒ behaviour byte-for-byte unchanged.
    rates_override: dict[str, float] | None = None
    log: list[dict] = field(default_factory=list)
    result: dict | None = None
    error: str | None = None
    queue: asyncio.Queue | None = None
    loop: asyncio.AbstractEventLoop | None = None
    lock: threading.Lock = field(default_factory=threading.Lock)


RUN_STORE: dict[str, Run] = {}


class ObjectiveSpecModel(BaseModel):
    target: str
    metric: str
    direction: str | None = None


class RunParamsModel(BaseModel):
    generations: int = Field(30, ge=2, le=1000)
    population: int = Field(150, ge=5, le=3000)
    genes_per_set: int = Field(8, ge=2, le=8)
    max_sets: int = Field(2, ge=1, le=2)
    lambda_size: float = Field(0.005, ge=0.0, le=0.1, alias="lambda")
    seed: int = 42
    # ``None`` means "no prefilter — sample from the full opaque-ID set".
    prefilter_n: int | None = Field(default=None, ge=10, le=20000)
    permutations: int = Field(200, ge=2, le=2000)

    model_config = {"populate_by_name": True}


class RunRequest(BaseModel):
    objective_spec: ObjectiveSpecModel
    params: RunParamsModel
    engine: Literal["v1", "v2"] = "v1"
    # Which dataset's cohort to load. "coadread" (default) for the
    # original TCGA colorectal study; "hnsc" for the head & neck study
    # that ships the HPV+/HPV- objective. The engine / DSL / params /
    # live view are identical across datasets; only the loader and the
    # valid (dataset, target) combinations swap.
    dataset: str = "coadread"
    # Peel-off chain: a list of prior run_ids whose winners' full-
    # cohort scores should be linearly residualised out of the
    # expression matrix before this run starts. Works for ANY
    # objective; all priors must share the same (dataset, target) as
    # the new run. The chain is in-memory only — valid within a
    # single server session.
    residualize_against: list[str] | None = None
    # "Prefer coordinated gene modules" prior. When true, fitness =
    # separation + COHERENCE_WEIGHT * mean-abs-pairwise-corr over the
    # Select'd opaque columns. Names no gene or pathway. Default off,
    # so existing runs are unchanged.
    coherence: bool = False
    # "Maintain diversity" knob. When true, the worker lowers selection
    # pressure (tournament_k 3 → 2), raises mutation (p_mutate 0.7 →
    # 0.85), and injects ~10% random immigrants each generation. The
    # search-internal change has no effect on the airgap or the wire
    # shape beyond this additive field; default false preserves the
    # current run behaviour byte-for-byte.
    diversity: bool = False
    # Per-run DSL injection-rate overrides. Keys we accept (all
    # optional): ``split``, ``effect``, ``fitapply``, ``search``, and
    # ``scalar_share``. Each is a float in [0, 1]. Missing keys fall
    # back to engine_v2.synthesize.DEFAULT_RATES / its scalar_share
    # default (0.20), so an empty/missing override reproduces current
    # behaviour byte-for-byte. ``search: 0.0`` is the replacement for
    # the old ``enable_search: False`` toggle. Airgap-safe — search-
    # internal hyperparameters, no gene identities.
    rates_override: dict[str, float] | None = None


class RunResponse(BaseModel):
    run_id: str


class EvaluateRequest(BaseModel):
    gene_ids: list[str]
    reference_set: str
    # Which dataset's reference-set vocabulary to look up `reference_set`
    # against. Backward-compat: legacy callers omit this and default to
    # colorectal (MMR / immune).
    dataset: str = "coadread"
    # Active objective. When set to one with a single-gene diagnostic
    # (hpv → hpv_rank, tmb → tmb_rank), /evaluate attaches each revealed
    # gene's individual rank + metric to the response. Optional / used
    # only on the reveal side; the engine never sees it.
    target: str | None = None


class EvaluateRevealedRow(BaseModel):
    id: str
    symbol: str
    matched: bool
    # Single-gene rank diagnostic (populated only for (dataset, target)
    # pairs where one exists — HNSC/HPV, coadread/TMB).
    rank: int | None = None
    total: int | None = None
    single_gene_metric: float | None = None   # AUROC or signed Spearman
    metric_kind: str | None = None            # "auroc" | "spearman"


class EvaluateResponse(BaseModel):
    revealed: list[EvaluateRevealedRow]
    overlap_count: int
    reference_set: str


def _compute_unsup_posthoc(
    result: dict,
    M: pd.DataFrame,
    clinical: pd.DataFrame | None,
    extra_labels: dict[str, np.ndarray],
    *,
    seed: int,
    test_size: float,
) -> dict:
    """Post-hoc alignment for an unsupervised run.

    The engine was blind to labels during search. Now, on the SAME
    held-out subset the engine evaluated against, look up the labels
    (kept aside in ``extra_labels``) and report:

    - ``msi_auroc``: orientation-agnostic AUROC of the winner's scores
      vs the MSI binary label. ``None`` if either class has fewer than
      10 samples in the held-out subset (too small to be meaningful).
    - ``tmb_abs_spearman``: ``|spearman(scores, tmb)|`` on held-out
      patients with non-NaN TMB. ``None`` if fewer than 10 have TMB.

    This step lives in the reveal/evaluate side of the airgap (api) —
    NEVER in engine_v2.
    """
    from sklearn.metrics import roc_auc_score
    from scipy.stats import spearmanr

    winning = result.get("winning", {})
    scores_raw: list[float | None] = winning.get("holdout_scores") or []
    ids: list[str] = winning.get("holdout_sample_ids") or []
    if not scores_raw or not ids or len(scores_raw) != len(ids):
        return {"msi_auroc": None, "tmb_abs_spearman": None, "n_holdout": 0}

    scores_arr = np.array(
        [np.nan if v is None else float(v) for v in scores_raw], dtype=float
    )
    if not np.isfinite(scores_arr).all() or float(np.nanstd(scores_arr)) == 0.0:
        return {"msi_auroc": None, "tmb_abs_spearman": None, "n_holdout": len(ids)}

    # Map held-out sample IDs to positions in the original M to slice the
    # saved label arrays.
    id_to_pos = {sid: i for i, sid in enumerate(M.index)}
    pos = np.array([id_to_pos[sid] for sid in ids if sid in id_to_pos])
    aligned = pos.shape[0] == len(ids)
    if not aligned:
        return {"msi_auroc": None, "tmb_abs_spearman": None, "n_holdout": len(ids)}

    out: dict = {"n_holdout": int(len(ids))}

    msi_full = extra_labels.get("msi")
    if msi_full is not None and len(msi_full) == len(M):
        msi_held = np.asarray(msi_full, dtype=float)[pos]
        mask = np.isfinite(msi_held)
        n_pos = int(((msi_held == 1.0) & mask).sum())
        n_neg = int(((msi_held == 0.0) & mask).sum())
        if n_pos >= 10 and n_neg >= 10:
            try:
                auroc = float(roc_auc_score(msi_held[mask].astype(int), scores_arr[mask]))
                out["msi_auroc"] = float(max(auroc, 1.0 - auroc))
            except Exception:
                out["msi_auroc"] = None
        else:
            out["msi_auroc"] = None
        out["n_msi_held"] = n_pos + n_neg
    else:
        out["msi_auroc"] = None

    tmb_full = extra_labels.get("tmb")
    if tmb_full is not None and len(tmb_full) == len(M):
        tmb_held = np.asarray(tmb_full, dtype=float)[pos]
        mask = np.isfinite(tmb_held)
        if int(mask.sum()) >= 10:
            corr, _ = spearmanr(scores_arr[mask], tmb_held[mask])
            out["tmb_abs_spearman"] = float(abs(corr)) if np.isfinite(corr) else None
        else:
            out["tmb_abs_spearman"] = None
        out["n_tmb_held"] = int(mask.sum())
    else:
        out["tmb_abs_spearman"] = None

    # HNSC: hpv label travels alongside via extra_labels for unsup runs.
    # Same orientation-agnostic AUROC machinery as msi.
    hpv_full = extra_labels.get("hpv")
    if hpv_full is not None and len(hpv_full) == len(M):
        hpv_held = np.asarray(hpv_full, dtype=float)[pos]
        mask = np.isfinite(hpv_held)
        n_pos = int(((hpv_held == 1.0) & mask).sum())
        n_neg = int(((hpv_held == 0.0) & mask).sum())
        if n_pos >= 10 and n_neg >= 10:
            try:
                auroc = float(roc_auc_score(hpv_held[mask].astype(int), scores_arr[mask]))
                out["hpv_auroc"] = float(max(auroc, 1.0 - auroc))
            except Exception:
                out["hpv_auroc"] = None
        else:
            out["hpv_auroc"] = None
        out["n_hpv_held"] = n_pos + n_neg
    else:
        out["hpv_auroc"] = None

    return out


def _push(run: Run, event_name: str, data: dict | str) -> None:
    if run.loop is None or run.queue is None:
        return
    payload = (event_name, data)
    try:
        run.loop.call_soon_threadsafe(run.queue.put_nowait, payload)
    except RuntimeError:
        # Event loop already closed (e.g. between TestClient requests).
        # The log already captured the event under the lock; subscribers
        # that connect via /stream will see it via the replay path.
        pass


def _assemble_residualize_df(
    prior_ids: list[str], required_target: str, required_dataset: str = "coadread",
) -> pd.DataFrame:
    """Build a sample-id-indexed DataFrame of prior axis scores for the
    peel-off chain. One column per prior axis (named ``axis_<run_id>``).
    Rows where ANY prior scored ``None`` (non-finite at run time) are
    kept in the per-axis Series but dropped from the design matrix
    downstream via the pipeline's NaN filter. Validation here raises
    ValueError on missing / wrong-(dataset,target) / no-scores priors —
    caller surfaces as HTTP 400. The chain works for every objective;
    all priors must share the same (dataset, target) as the new run.
    """
    columns: dict[str, pd.Series] = {}
    for prior_id in prior_ids:
        prior = RUN_STORE.get(prior_id)
        if prior is None:
            raise ValueError(f"unknown prior run_id: {prior_id!r}")
        with prior.lock:
            if prior.status != "done":
                raise ValueError(
                    f"prior run {prior_id!r} is not done (status={prior.status!r})"
                )
            if prior.objective_spec.get("target") != required_target:
                raise ValueError(
                    f"prior run {prior_id!r} has target "
                    f"{prior.objective_spec.get('target')!r}; "
                    f"peel-off chain requires {required_target!r}"
                )
            if prior.dataset != required_dataset:
                raise ValueError(
                    f"prior run {prior_id!r} is on dataset "
                    f"{prior.dataset!r}; peel-off chain requires "
                    f"{required_dataset!r}"
                )
            result = prior.result or {}
        winning = result.get("winning", {}) or {}
        scores = winning.get("full_scores") or []
        ids = winning.get("full_sample_ids") or []
        if not scores or not ids or len(scores) != len(ids):
            raise ValueError(
                f"prior run {prior_id!r} has no full-cohort scores stored"
            )
        # _json_finite() turns non-finite values into None on the wire;
        # build a Series with NaN where missing so the pipeline's NaN
        # filter drops those patients during residualisation.
        col_values = [
            float(v) if v is not None else float("nan") for v in scores
        ]
        columns[f"axis_{prior_id}"] = pd.Series(col_values, index=ids)
    df = pd.DataFrame(columns)
    return df


def _worker(
    run: Run,
    M: pd.DataFrame,
    y: np.ndarray | None,
    clinical: pd.DataFrame | None = None,
    extra_labels: dict[str, np.ndarray] | None = None,
    residualize_against: list[str] | None = None,
    coherence: bool = False,
    diversity: bool = False,
    rates_override: dict[str, float] | None = None,
) -> None:
    """Run the GP pipeline; stream events; capture errors."""
    try:
        p = run.params

        # Peel-off chain: assemble prior-axis scores so the pipeline
        # residualises M before its train/test split. Works for every
        # objective; all priors must share the run's (dataset, target).
        residualize_df: pd.DataFrame | None = None
        if residualize_against and run.engine == "v2":
            residualize_df = _assemble_residualize_df(
                residualize_against,
                run.objective_spec.get("target", ""),
                required_dataset=run.dataset,
            )

        # "Prefer coordinated gene modules" prior. Picked here so the
        # constant is in one place; UI toggles the boolean.
        coherence_weight = COHERENCE_DEFAULT_WEIGHT if coherence else 0.0

        # "Maintain diversity" knob. Off (default): preserve current
        # behaviour byte-for-byte (k=3, p_mutate=0.7, no immigrants).
        # On: lower selection pressure + raise mutation + inject ~10%
        # random immigrants per generation. Tuned constants live here
        # so the UI just flips a bool.
        tournament_k = 2 if diversity else 3
        p_mutate = 0.85 if diversity else 0.7
        immigrant_fraction = 0.10 if diversity else 0.0

        # Per-run DSL injection-rate overrides. The UI exposes one
        # rate per OPTIONAL operator (Split / Effect / FitApply /
        # Search) plus a Scalar-share field; mandatory operators
        # (Select / Reduce / Combine; Associate is the leftover of
        # the Scalar branch) get no input. The wire shape is a flat
        # dict where ``scalar_share`` is a sibling key alongside the
        # rates; we split it back out here because synthesize takes
        # ``rates`` and ``scalar_share`` as separate parameters. We
        # overlay the rates on top of DEFAULT_RATES so synthesize
        # sees a complete dict (its inner code indexes rates["split"]
        # etc. by key without a fallback). When the caller sends
        # nothing, both args go through as None and synthesize falls
        # back to its DEFAULT_RATES / 0.20 scalar_share — behaviour
        # byte-for-byte unchanged.
        merged_rates_override: dict | None = None
        scalar_share_override: float | None = None
        if rates_override:
            from engine_v2.synthesize import DEFAULT_RATES as _DR
            merged_rates_override = {**_DR}
            for k, v in rates_override.items():
                try:
                    f = float(v)
                except (TypeError, ValueError):
                    continue
                # Clamp every value to [0, 1] so a typo can't break
                # synthesize's rng comparisons.
                f = max(0.0, min(1.0, f))
                if k == "scalar_share":
                    scalar_share_override = f
                else:
                    merged_rates_override[k] = f
            # If the user touched only scalar_share, leave the rates
            # at their unmodified defaults — no override needed.
            if merged_rates_override == _DR:
                merged_rates_override = None

        def on_gen(event: dict) -> None:
            with run.lock:
                run.log.append(event)
            # Light SSE payload: trim per-gen to top-12 for bandwidth.
            stream_event = dict(event)
            if "candidates" in stream_event:
                stream_event = {
                    **stream_event,
                    "candidates": stream_event["top_candidates"]
                    if "top_candidates" in stream_event
                    else stream_event["candidates"][:12],
                }
            _push(run, "generation", _json_finite(stream_event))

        if run.engine == "v2":
            objective_v2 = v2_objective_from_spec(run.objective_spec)
            # Dataset-aware Effect confounders. Default (coadread / MSI /
            # TMB) stays (stage, age) so existing runs are byte-for-byte
            # unchanged. HNSC adds sex + race when those columns are
            # present in `clinical` — Effect skips any column missing
            # from the frame. Never includes a gene identity.
            confounders = ("stage", "age")
            if run.dataset == "hnsc" and clinical is not None:
                extras = [
                    c for c in ("sex", "race") if c in clinical.columns
                ]
                confounders = ("stage", "age", *extras)
            result = run_v2_pipeline_streaming(
                M, y,
                objective=objective_v2,
                on_generation=on_gen,
                seed=p["seed"],
                prefilter_n=p["prefilter_n"],
                population_size=p["population"],
                n_generations=p["generations"],
                n_permutations=p["permutations"],
                lambda_size=p["lambda_size"],
                max_genes_per_set=p["genes_per_set"],
                tournament_k=tournament_k,
                p_mutate=p_mutate,
                clinical=clinical,
                extra_labels=extra_labels,
                residualize_scores=residualize_df,
                coherence_weight=coherence_weight,
                confounders=confounders,
                immigrant_fraction=immigrant_fraction,
                rates_override=merged_rates_override,
                scalar_share_override=scalar_share_override,
            )
            if objective_v2.target == "none":
                result["posthoc"] = _compute_unsup_posthoc(
                    result, M, clinical, extra_labels or {},
                    seed=p["seed"], test_size=0.3,
                )
        else:
            objective = objective_from_spec(run.objective_spec)
            result = run_gp_pipeline_streaming(
                M, y,
                objective=objective,
                on_generation=on_gen,
                seed=p["seed"],
                prefilter_n=p["prefilter_n"],
                population_size=p["population"],
                n_generations=p["generations"],
                n_permutations=p["permutations"],
            )
        safe_result = _json_finite(result)
        with run.lock:
            run.result = safe_result
            run.status = "done"
        _push(run, "done", safe_result)
    except Exception as exc:  # noqa: BLE001 — surface to the client
        log.exception("Run %s failed", run.id)
        with run.lock:
            run.status = "error"
            run.error = f"{type(exc).__name__}: {exc}"
        _push(run, "error", run.error)


# ---------------------------------------------------------------------------
# Lab endpoints
# ---------------------------------------------------------------------------


@app.post("/runs", response_model=RunResponse)
async def post_runs(req: RunRequest) -> RunResponse:
    spec = req.objective_spec.model_dump(exclude_none=True)
    engine_choice = req.engine
    try:
        if engine_choice == "v2":
            v2_objective_from_spec(spec)
        else:
            objective_from_spec(spec)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc))

    target = spec["target"]
    dataset = req.dataset

    # Validate dataset + objective combination — fail fast at the API
    # boundary, before any data prep.
    if dataset not in DATASET_TARGETS:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Unknown dataset {dataset!r}. "
                f"Supported: {sorted(DATASET_TARGETS)}."
            ),
        )
    if target not in DATASET_TARGETS[dataset]:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Objective {target!r} is not valid for dataset {dataset!r}. "
                f"Valid targets for {dataset!r}: "
                f"{sorted(DATASET_TARGETS[dataset])}."
            ),
        )

    # Validate the peel-off chain BEFORE we spin a worker. Priors must
    # exist, be done, share the same (dataset, target) as the new run,
    # and carry full-cohort scores. Works for every objective; gated
    # to v2 because v1 doesn't persist full_scores.
    residualize_against = req.residualize_against or None
    if residualize_against:
        if engine_choice != "v2":
            raise HTTPException(
                status_code=400,
                detail="residualize_against requires engine='v2'.",
            )
        try:
            _assemble_residualize_df(
                residualize_against, target, required_dataset=dataset,
            )
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc))

    try:
        M, y, clinical, extra_labels = _prepare_lab_data(target, dataset)
    except Exception as exc:
        raise HTTPException(status_code=503, detail=f"Data prep failed: {exc}")

    run_id = uuid.uuid4().hex[:12]
    run = Run(
        id=run_id,
        objective_spec=spec,
        params=req.params.model_dump(by_alias=False),
        engine=engine_choice,
        dataset=dataset,
        coherence=bool(req.coherence),
        diversity=bool(req.diversity),
        rates_override=dict(req.rates_override) if req.rates_override else None,
    )
    run.loop = asyncio.get_running_loop()
    run.queue = asyncio.Queue()
    RUN_STORE[run_id] = run

    threading.Thread(
        target=_worker,
        args=(run, M, y, clinical, extra_labels),
        kwargs={
            "residualize_against": residualize_against,
            "coherence": bool(req.coherence),
            "diversity": bool(req.diversity),
            "rates_override": (
                dict(req.rates_override) if req.rates_override else None
            ),
        },
        daemon=True,
    ).start()
    return RunResponse(run_id=run_id)


@app.get("/runs/{run_id}")
def get_run_by_id(run_id: str) -> dict:
    run = RUN_STORE.get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail=f"Unknown run_id {run_id}")
    with run.lock:
        # Lightweight summary log: drop the full per-gen population to keep
        # the payload small. Use /runs/{id}/population/{gen} for the full
        # population of any generation.
        summary_log = [
            {
                k: v for k, v in entry.items()
                if k != "candidates"   # full population lives in the per-gen endpoint
            }
            for entry in run.log
        ]
        return _json_finite({
            "id": run.id,
            "engine": run.engine,
            "dataset": run.dataset,
            "coherence": bool(run.coherence),
            "diversity": bool(run.diversity),
            "rates_override": run.rates_override,
            "objective_spec": run.objective_spec,
            "params": run.params,
            "status": run.status,
            "error": run.error,
            "n_generations_seen": len(run.log),
            "generations_persisted": len(run.log),
            "log": summary_log,
        })


@app.get("/runs/{run_id}/population/{generation}")
def get_run_population(run_id: str, generation: int) -> dict:
    """Full population of one generation (opaque IDs only).

    SSE streams keep the top-12 to stay light; this endpoint serves the
    rest of the candidates on demand for the tile grid.
    """
    run = RUN_STORE.get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail=f"Unknown run_id {run_id}")
    with run.lock:
        n = len(run.log)
        if generation < 0 or generation >= n:
            raise HTTPException(
                status_code=404,
                detail=(
                    f"Generation {generation} not persisted; "
                    f"{n} available (0..{max(n - 1, 0)})."
                ),
            )
        entry = run.log[generation]
        return _json_finite({
            "run_id": run.id,
            "engine": run.engine,
            "generation": entry["generation"],
            "best_fitness": entry["best_fitness"],
            "median_fitness": entry["median_fitness"],
            "elitism": entry["elitism"],
            "population_size": entry.get("population_size", len(entry["candidates"])),
            "candidates": list(entry["candidates"]),
        })


@app.get("/runs/{run_id}/result")
def get_run_result(run_id: str) -> dict:
    run = RUN_STORE.get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail=f"Unknown run_id {run_id}")
    with run.lock:
        if run.status == "error":
            raise HTTPException(
                status_code=500,
                detail=run.error or "run failed",
            )
        if run.status != "done" or run.result is None:
            raise HTTPException(status_code=425, detail="run still running")
        return _json_finite(run.result)


@app.get("/runs/{run_id}/stream")
async def stream_run(run_id: str):
    run = RUN_STORE.get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail=f"Unknown run_id {run_id}")

    async def event_source():
        # 1) Replay everything we already have on the log (for late subscribers).
        with run.lock:
            snapshot = list(run.log)
            already_done = run.status in ("done", "error")
            final_result = run.result
            final_error = run.error
        for entry in snapshot:
            yield {"event": "generation", "data": json.dumps(entry)}
        if already_done:
            if run.status == "done" and final_result is not None:
                yield {"event": "done", "data": json.dumps(final_result)}
            elif run.status == "error":
                yield {"event": "error", "data": json.dumps({"detail": final_error})}
            return

        # 2) Subscribe to new events. We don't deduplicate — late subscribers
        # accept that the very-first events past the snapshot may overlap by
        # one generation at most (acceptable for MVP).
        assert run.queue is not None
        while True:
            try:
                event_name, payload = await asyncio.wait_for(
                    run.queue.get(), timeout=30.0,
                )
            except asyncio.TimeoutError:
                yield {"event": "ping", "data": "{}"}
                continue
            yield {"event": event_name, "data": json.dumps(payload)
                   if not isinstance(payload, str) else json.dumps({"detail": payload})}
            if event_name in ("done", "error"):
                return

    return EventSourceResponse(event_source())


# Cache the heavy per-gene rank lookup keyed by (dataset, target). The
# diagnostic computes ranks for ALL ~20k genes once; subsequent /evaluate
# calls do a dict lookup.
_GENE_RANK_LOOKUPS: dict[tuple[str, str], dict[str, dict] | None] = {}


def _gene_rank_lookup(dataset: str, target: str) -> dict[str, dict] | None:
    """Return a ``{symbol: {rank, total, metric, metric_kind}}`` map for
    the active (dataset, target) pair if a single-gene diagnostic exists;
    None otherwise. Cached for the process lifetime."""
    key = (dataset, target)
    if key in _GENE_RANK_LOOKUPS:
        return _GENE_RANK_LOOKUPS[key]

    out: dict[str, dict] | None
    if dataset == "hnsc" and target == "hpv":
        from validate.hpv_rank import (
            DEFAULT_SEED,
            DEFAULT_TEST_SIZE,
            _auroc_per_column,
            _hpv_cohort_named,
            _train_slice,
        )

        X, y = _hpv_cohort_named(None)
        X_train, y_train = _train_slice(
            X, y, seed=DEFAULT_SEED, test_size=DEFAULT_TEST_SIZE,
        )
        aurocs = _auroc_per_column(X_train, y_train)
        valid = aurocs.dropna()
        ranks = valid.rank(method="min", ascending=False).astype(int)
        n_genes = int(len(valid))
        out = {
            str(sym): {
                "rank": int(ranks.loc[sym]),
                "total": n_genes,
                "metric": float(valid.loc[sym]),
                "metric_kind": "auroc",
            }
            for sym in valid.index
        }
    elif dataset == "coadread" and target == "tmb":
        from validate.tmb_rank import _spearman_per_column, _tmb_cohort_named

        X, y = _tmb_cohort_named(None)
        corr = _spearman_per_column(X, y)
        valid = corr.dropna()
        # Ascending so rank 1 = most-negative (the TMB objective rewards
        # the most-negative association). Matches /diagnostic/tmb-rank.
        ranks = valid.rank(method="min", ascending=True).astype(int)
        n_genes = int(len(valid))
        out = {
            str(sym): {
                "rank": int(ranks.loc[sym]),
                "total": n_genes,
                "metric": float(valid.loc[sym]),
                "metric_kind": "spearman",
            }
            for sym in valid.index
        }
    else:
        out = None

    _GENE_RANK_LOOKUPS[key] = out
    return out


@app.post("/evaluate", response_model=EvaluateResponse)
def post_evaluate(req: EvaluateRequest) -> EvaluateResponse:
    if req.dataset not in REFERENCE_SETS_BY_DATASET:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Unknown dataset {req.dataset!r}. "
                f"Supported: {sorted(REFERENCE_SETS_BY_DATASET)}."
            ),
        )
    dataset_sets = REFERENCE_SETS_BY_DATASET[req.dataset]
    if req.reference_set not in dataset_sets:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Unknown reference_set {req.reference_set!r} for dataset "
                f"{req.dataset!r}. Supported: {sorted(dataset_sets)}"
            ),
        )
    try:
        symbols = reveal(req.gene_ids)
    except FileNotFoundError as exc:
        raise HTTPException(status_code=503, detail=str(exc))
    except KeyError as exc:
        raise HTTPException(status_code=400, detail=str(exc))
    ref_set = set(dataset_sets[req.reference_set])
    # Single-gene rank lookup is optional — None for (MSI, unsup) and
    # any (dataset, target) without a diagnostic.
    rank_lookup = (
        _gene_rank_lookup(req.dataset, req.target) if req.target else None
    )
    rows: list[EvaluateRevealedRow] = []
    for gid, sym in zip(req.gene_ids, symbols):
        row_kwargs = {
            "id": gid,
            "symbol": sym,
            "matched": sym in ref_set,
        }
        if rank_lookup is not None and sym in rank_lookup:
            r = rank_lookup[sym]
            row_kwargs["rank"] = r["rank"]
            row_kwargs["total"] = r["total"]
            row_kwargs["single_gene_metric"] = r["metric"]
            row_kwargs["metric_kind"] = r["metric_kind"]
        rows.append(EvaluateRevealedRow(**row_kwargs))
    overlap = sum(1 for r in rows if r.matched)
    return EvaluateResponse(
        revealed=rows,
        overlap_count=overlap,
        reference_set=req.reference_set,
    )


# ---------------------------------------------------------------------------
# Diagnostic: where MMR / IMMUNE genes land under the TMB target. Belongs to
# the reveal side of the airgap (named genes; lives in validate/).
# ---------------------------------------------------------------------------

_TMB_RANK_CACHE: dict[str, object] = {}


@app.get("/diagnostic/tmb-rank")
def get_tmb_rank_diagnostic() -> dict:
    from dataclasses import asdict

    from validate.tmb_rank import tmb_rank_diagnostic

    cached = _TMB_RANK_CACHE.get("payload")
    if cached is not None:
        return cached  # type: ignore[return-value]

    d = tmb_rank_diagnostic()
    payload = _json_finite({
        "cohort": d.cohort,
        "n_samples": d.n_samples,
        "n_genes": d.n_genes,
        "mmr": [asdict(r) for r in d.mmr_rows],
        "immune": [asdict(r) for r in d.immune_rows],
        "top_negative": [asdict(r) for r in d.top_negative],
    })
    _TMB_RANK_CACHE["payload"] = payload
    return payload


_HPV_RANK_CACHE: dict[str, object] = {}

# /diagnostic/full-rank cache keyed by (dataset, target).
_FULL_RANK_CACHE: dict[tuple[str, str], dict] = {}


def _compute_full_rank(dataset: str, target: str) -> dict | None:
    """Single-gene rank of every OPAQUE column on the engine's TRAIN
    split. Binary targets (msi, hpv) → orientation-agnostic AUROC;
    TMB → signed Spearman. Returns the opaque-only payload or None
    when no single-gene ranking applies (unsupervised).

    Lives entirely on the API side; the engine never sees gene names
    because the matrix is already anonymised before _prepare_lab_data
    returns it. The wire payload carries opaque IDs only — the
    frontend reveals symbols for the winner + reference genes via the
    single-ID /evaluate path."""
    if target == "none":
        return None
    M, y, _clinical, _extra = _prepare_lab_data(target, dataset)
    if y is None:
        return None
    from engine.split import make_split

    split = make_split(
        M.index,
        np.asarray(y),
        test_size=0.3,
        random_state=42,
        stratify=(target in ("msi", "hpv")),
    )
    X_train = M.loc[split.train_ids]
    y_train = np.asarray(split.y_train)

    metric_kind: str
    if target in ("msi", "hpv"):
        # Same Mann-Whitney / rank-sum AUROC formula as validate/hpv_rank.
        n = X_train.shape[0]
        n_pos = int((y_train == 1).sum())
        n_neg = int((y_train == 0).sum())
        if n_pos == 0 or n_neg == 0:
            return None
        R = X_train.rank(axis=0).to_numpy(dtype=float)
        pos_mask = (y_train == 1)
        S_pos = R[pos_mask].sum(axis=0)
        auroc = (S_pos - n_pos * (n_pos + 1) / 2.0) / (n_pos * n_neg)
        std = X_train.std(axis=0, ddof=0).to_numpy(dtype=float)
        omni = np.maximum(auroc, 1.0 - auroc)
        omni = np.where(std == 0, np.nan, omni)
        scores = pd.Series(omni, index=X_train.columns)
        ascending = False  # highest AUROC = rank 1
        metric_kind = "auroc"
        n_pos_out, n_neg_out = n_pos, n_neg
    elif target == "tmb":
        # Signed Spearman, ascending — most negative = rank 1.
        from validate.tmb_rank import _spearman_per_column

        scores = _spearman_per_column(X_train, y_train)
        ascending = True
        metric_kind = "spearman"
        n_pos_out = n_neg_out = 0
    else:
        return None

    valid = scores.dropna()
    ranks = valid.rank(method="min", ascending=ascending).astype(int)
    n_genes = int(len(valid))
    # Order rows by rank ascending so rank 1 sits first.
    order = valid.sort_values(ascending=ascending)
    rows = [
        {
            "opaque_id": str(opq),
            "score": float(valid.loc[opq]),
            "rank": int(ranks.loc[opq]),
        }
        for opq in order.index
    ]
    # Reference marks: the dataset's reference-set genes, with their
    # opaque IDs revealed up-front so the rank track can pin them.
    # Tiny lookup against the sealed map (a few dozen symbols max);
    # the full map never crosses the wire.
    reference_marks: list[dict] = []
    try:
        from airgap.seal import _read_sealed

        id_to_sym = _read_sealed().get("id_to_symbol", {})
        sym_to_id = {v: k for k, v in id_to_sym.items()}
        for set_name, syms in REFERENCE_SETS_BY_DATASET[dataset].items():
            for sym in syms:
                opq = sym_to_id.get(sym)
                if opq is None or opq not in ranks.index:
                    continue
                reference_marks.append({
                    "opaque_id": str(opq),
                    "symbol": str(sym),
                    "set_name": set_name,
                    "rank": int(ranks.loc[opq]),
                    "score": float(valid.loc[opq]),
                })
    except Exception:
        pass

    return {
        "dataset": dataset,
        "target": target,
        "metric_kind": metric_kind,
        "n_samples": int(X_train.shape[0]),
        "n_pos": int(n_pos_out),
        "n_neg": int(n_neg_out),
        "n_genes": n_genes,
        "seed": 42,
        "test_size": 0.3,
        "ranks": rows,
        "reference_marks": reference_marks,
    }


@app.get("/diagnostic/full-rank")
def get_full_rank_diagnostic(dataset: str, target: str) -> dict:
    """Opaque-only single-gene ranking of every column on the engine's
    TRAIN split. Powers the Lab's Result rank track + browsable
    ranking list. Airgap-clean: no gene names in the payload."""
    key = (dataset, target)
    cached = _FULL_RANK_CACHE.get(key)
    if cached is not None:
        return cached
    try:
        payload = _compute_full_rank(dataset, target)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc))
    if payload is None:
        raise HTTPException(
            status_code=404,
            detail=(
                f"No single-gene ranking applies for "
                f"(dataset={dataset!r}, target={target!r})."
            ),
        )
    payload_safe = _json_finite(payload)
    _FULL_RANK_CACHE[key] = payload_safe
    return payload_safe


@app.get("/diagnostic/hpv-rank")
def get_hpv_rank_diagnostic() -> dict:
    """HNSC-side mirror of /diagnostic/tmb-rank.

    Ranks every gene by single-gene orientation-agnostic AUROC vs the
    HPV+/HPV− label on the engine's TRAIN split. Reports where
    CDKN2A (p16) and each cell-cycle gene sit, plus the top single-
    gene separators. Lives in validate/ (named-genes allowed); the
    engine never imports this path."""
    from dataclasses import asdict

    from validate.hpv_rank import hpv_rank_diagnostic

    cached = _HPV_RANK_CACHE.get("payload")
    if cached is not None:
        return cached  # type: ignore[return-value]

    d = hpv_rank_diagnostic()
    payload = _json_finite({
        "cohort": d.cohort,
        "seed": d.seed,
        "n_samples": d.n_samples,
        "n_pos": d.n_pos,
        "n_neg": d.n_neg,
        "n_genes": d.n_genes,
        "p16": [asdict(r) for r in d.p16_rows],
        "cell_cycle": [asdict(r) for r in d.cell_cycle_rows],
        "top_separators": [asdict(r) for r in d.top_separators],
    })
    _HPV_RANK_CACHE["payload"] = payload
    return payload


# ---------------------------------------------------------------------------
# Module ranking — harvest distinct gene_ids sets from a run's persisted
# population, score each on the run's held-out split as a group, sort by
# combined held-out AUROC. Opaque-only payload; symbols are revealed
# per-module on demand via /evaluate.
# ---------------------------------------------------------------------------


def _compute_module_ranking(run: "Run") -> dict:
    """Build the module-ranking payload for a completed run. Reproduces
    the run's exact train/test split via the persisted ``full_sample_ids``
    / ``holdout_sample_ids`` so "held-out" actually is held-out — and
    aggregates each module's genes with the parameter-free MEAN (matches
    Reduce(mean)), so there's no model to fit and nothing to leak.
    """
    from sklearn.metrics import roc_auc_score
    from scipy.stats import spearmanr

    with run.lock:
        target = run.objective_spec.get("target", "")
        dataset = run.dataset
        result = run.result
        log = list(run.log)
        coherence_flag = bool(run.coherence)

    if target == "none":
        raise HTTPException(
            status_code=400,
            detail=(
                "Module ranking requires a supervised target "
                "(msi, hpv, or tmb)."
            ),
        )

    winning = (result or {}).get("winning", {}) or {}
    full_ids: list[str] = list(winning.get("full_sample_ids") or [])
    test_ids: list[str] = list(winning.get("holdout_sample_ids") or [])
    if not full_ids or not test_ids:
        raise HTTPException(
            status_code=400,
            detail=(
                "Run is missing held-out sample IDs — module ranking "
                "needs a v2 run with persisted full / holdout sample IDs."
            ),
        )

    M_all, y_all, clinical, _extra = _prepare_lab_data(target, dataset)
    if y_all is None:
        raise HTTPException(
            status_code=400,
            detail="No labels available for this target — cannot rank modules.",
        )

    # Slice to the cohort the run actually used (peel-off may have
    # trimmed patients).
    keep_full = [sid for sid in full_ids if sid in M_all.index]
    if len(keep_full) != len(full_ids):
        raise HTTPException(
            status_code=500,
            detail="Persisted sample IDs no longer align with the cohort matrix.",
        )
    M = M_all.loc[keep_full]
    y_full = pd.Series(np.asarray(y_all), index=M_all.index).loc[keep_full]
    clinical_sub = (
        clinical.reindex(keep_full).copy() if clinical is not None else None
    )

    test_set = set(test_ids)
    train_ids_seq = [sid for sid in keep_full if sid not in test_set]
    if not train_ids_seq:
        raise HTTPException(
            status_code=500,
            detail="Could not reconstruct train split from persisted IDs.",
        )
    M_train = M.loc[train_ids_seq]
    y_train = y_full.loc[train_ids_seq].to_numpy()
    test_ids_present = [sid for sid in test_ids if sid in M.index]
    M_test = M.loc[test_ids_present]
    y_test = y_full.loc[test_ids_present].to_numpy()

    is_binary = target in ("msi", "hpv")
    if is_binary:
        n_pos_train = int((y_train == 1).sum())
        n_neg_train = int((y_train == 0).sum())
        if n_pos_train == 0 or n_neg_train == 0:
            raise HTTPException(
                status_code=400,
                detail=(
                    "Train split has only one class — module ranking impossible."
                ),
            )

    # Per-gene single-gene rank lookup on TRAIN (same metric as the
    # /diagnostic/full-rank endpoint, recomputed locally so we don't
    # need the cache to be warm).
    if is_binary:
        R = M_train.rank(axis=0).to_numpy(dtype=float)
        pos_mask = (y_train == 1)
        S_pos = R[pos_mask].sum(axis=0)
        gene_auroc = (S_pos - n_pos_train * (n_pos_train + 1) / 2.0) / (
            n_pos_train * n_neg_train
        )
        omni = np.maximum(gene_auroc, 1.0 - gene_auroc)
        std = M_train.std(axis=0, ddof=0).to_numpy(dtype=float)
        gene_score = np.where(std == 0, np.nan, omni)
        scores_series = pd.Series(gene_score, index=M_train.columns)
        ascending = False
        metric_kind = "auroc"
    else:   # tmb
        from validate.tmb_rank import _spearman_per_column

        scores_series = _spearman_per_column(M_train, y_train)
        ascending = True
        metric_kind = "spearman"

    valid_scores = scores_series.dropna()
    ranks_series = valid_scores.rank(method="min", ascending=ascending).astype(int)
    n_genes_valid = int(len(valid_scores))

    # Harvest modules: distinct unordered gene_ids sets across every
    # persisted generation. engine_v2 stores a flat gene_ids list per
    # candidate (engine_v2/gp.py: list(population[i].feature_ids())).
    # For each distinct gene-set we also track the MAX GP fitness
    # observed for any candidate carrying that set — this is the
    # "engine's own preference" signal exposed to the merged Groups
    # table as a sort lens (alongside the Combined-AUROC / Coherence
    # / Synergy re-score lenses).
    seen: set[frozenset[str]] = set()
    modules: list[list[str]] = []
    gp_fitness_by_key: dict[frozenset[str], float] = {}
    # Argmax program_repr per gene-set — the actual tree of the
    # candidate that earned ``gp_fitness_by_key[key]``. Carries
    # through to the merged table's expanded row so the user can
    # see the engine's preferred shape for that set.
    best_program_repr_by_key: dict[frozenset[str], str] = {}
    for entry in log:
        for cand in entry.get("candidates", []) or []:
            ids = list(dict.fromkeys(cand.get("gene_ids") or []))
            if len(ids) < 2:
                continue
            key = frozenset(ids)
            # Track max GP fitness for every candidate that surfaces
            # this gene-set, regardless of dedupe; record the
            # corresponding program_repr too.
            fit = cand.get("fitness")
            if fit is not None:
                try:
                    f = float(fit)
                    if math.isfinite(f):
                        prev = gp_fitness_by_key.get(key)
                        if prev is None or f > prev:
                            gp_fitness_by_key[key] = f
                            repr_str = cand.get("program_repr")
                            if isinstance(repr_str, str) and repr_str:
                                best_program_repr_by_key[key] = repr_str
                except (TypeError, ValueError):
                    pass
            if key in seen:
                continue
            seen.add(key)
            modules.append(ids)

    # Reference-set membership: resolve the dataset's reference symbols
    # to opaque IDs ONCE, then intersect each module against them. This
    # is a bounded reveal of a known small set (a few dozen symbols
    # max) — the same pattern /diagnostic/full-rank uses for
    # ``reference_marks``. Only a boolean membership crosses the wire
    # in the collapsed list — no module gene NAMES.
    ref_opaque_by_set: dict[str, set[str]] = {}
    sym_to_id: dict[str, str] = {}
    try:
        from airgap.seal import _read_sealed

        id_to_sym = _read_sealed().get("id_to_symbol", {})
        sym_to_id = {v: k for k, v in id_to_sym.items()}
        for set_name, syms in REFERENCE_SETS_BY_DATASET.get(dataset, {}).items():
            ref_opaque_by_set[set_name] = {
                sym_to_id[s] for s in syms if s in sym_to_id
            }
    except Exception:
        ref_opaque_by_set = {}
        sym_to_id = {}

    # ------------------------------------------------------------------
    # Confounder-survival subgroups (HNSC/HPV only, for now)
    # ------------------------------------------------------------------
    # Site-stratified subgroup: held-out patients with is_oropharynx=True.
    # HPV+ HNSC concentrates in the oropharynx (~68% in this cohort vs
    # ~5% of HPV-), so a module that just marks oropharynx tissue will
    # collapse here even though it looks great on the full cohort.
    orop_test_mask: pd.Series | None = None
    if (
        dataset == "hnsc"
        and clinical_sub is not None
        and "is_oropharynx" in clinical_sub.columns
        and target == "hpv"
    ):
        orop_series = clinical_sub.reindex(test_ids_present)["is_oropharynx"]
        orop_test_mask = orop_series.fillna(False).astype(bool)

    # Immune-infiltration proxy: bottom tertile of mean(CD8A, GZMB, PRF1,
    # CD3D, CD2) on TEST = high-purity subset. A module whose "HPV
    # signal" is really an immune-composition artifact will collapse
    # here. Lives in the validation/API layer — engine never sees
    # these symbols.
    high_purity_test_mask: pd.Series | None = None
    proxy_opaque_used: list[str] = []
    if dataset == "hnsc" and len(sym_to_id) > 0 and target == "hpv":
        proxy_opaque = [
            sym_to_id[s] for s in HPV_IMMUNE_PROXY_GENES if s in sym_to_id
        ]
        proxy_opaque = [g for g in proxy_opaque if g in M.columns]
        if len(proxy_opaque) >= 2 and len(M_test) >= 9:
            proxy_opaque_used = proxy_opaque
            proxy = M_test[proxy_opaque].mean(axis=1)
            threshold = proxy.quantile(1.0 / 3.0)
            high_purity_test_mask = (proxy <= threshold)

    def _combined_in_subset(
        valid_ids: list[str], mask: pd.Series | None,
    ) -> tuple[float | None, int, int, int]:
        """Return (combined_holdout, n, n_pos, n_neg) over the masked
        subset. None when the subgroup is too small or degenerate."""
        if mask is None:
            return None, 0, 0, 0
        sub_ids = mask.index[mask.values.astype(bool)]
        if len(sub_ids) == 0:
            return None, 0, 0, 0
        sub_scores = (
            M_test.loc[sub_ids][valid_ids].mean(axis=1).to_numpy(dtype=float)
        )
        sub_y = y_full.loc[sub_ids].to_numpy()
        if is_binary:
            n_pos = int((sub_y == 1).sum())
            n_neg = int((sub_y == 0).sum())
            if n_pos < 5 or n_neg < 5:
                return None, int(len(sub_ids)), n_pos, n_neg
            std_sub = float(np.nanstd(sub_scores))
            if not np.isfinite(std_sub) or std_sub == 0.0:
                return None, int(len(sub_ids)), n_pos, n_neg
            try:
                a = float(roc_auc_score(sub_y.astype(int), sub_scores))
                return float(max(a, 1.0 - a)), int(len(sub_ids)), n_pos, n_neg
            except Exception:
                return None, int(len(sub_ids)), n_pos, n_neg
        else:
            if len(sub_ids) < 10:
                return None, int(len(sub_ids)), 0, 0
            try:
                c, _ = spearmanr(sub_scores, sub_y)
                return (
                    float(abs(c)) if np.isfinite(c) else None,
                    int(len(sub_ids)), 0, 0,
                )
            except Exception:
                return None, int(len(sub_ids)), 0, 0

    # Tolerance for the survival flag: the stratified AUROC must stay
    # within this margin of the full-cohort combined AUROC to count as
    # "survives". 0.05 is a forgiving margin given the much smaller
    # subgroup sample sizes.
    SURVIVE_TOLERANCE = 0.05

    out_modules: list[dict] = []
    for ids in modules:
        valid_ids = [g for g in ids if g in M.columns]
        if len(valid_ids) < 2:
            continue

        # Combined per-patient score = mean over the module's genes.
        # Parameter-free; matches Reduce(mean) and carries no leakable
        # model.
        per_pat_test = M_test[valid_ids].mean(axis=1).to_numpy(dtype=float)

        if is_binary:
            std_test = float(np.nanstd(per_pat_test))
            if not np.isfinite(std_test) or std_test == 0.0:
                combined: float = float("nan")
            else:
                try:
                    auroc = float(
                        roc_auc_score(y_test.astype(int), per_pat_test)
                    )
                    combined = float(max(auroc, 1.0 - auroc))
                except Exception:
                    combined = float("nan")
        else:
            try:
                corr, _ = spearmanr(per_pat_test, y_test)
                combined = float(abs(corr)) if np.isfinite(corr) else float("nan")
            except Exception:
                combined = float("nan")

        # Coherence = mean absolute pairwise correlation on TRAIN (same
        # quantity the coherence prior rewards).
        sub_train = M_train[valid_ids].to_numpy(dtype=float)
        try:
            corr_mat = np.corrcoef(sub_train.T)
            m = corr_mat.shape[0] if corr_mat.ndim == 2 else 0
            if m < 2:
                coherence_val: float = float("nan")
            else:
                iu = np.triu_indices(m, k=1)
                vals = corr_mat[iu]
                coherence_val = float(np.nanmean(np.abs(vals)))
        except Exception:
            coherence_val = float("nan")

        per_gene = []
        for g in valid_ids:
            if g in ranks_series.index:
                per_gene.append({
                    "id": g,
                    "single_gene_metric": float(valid_scores.loc[g]),
                    "rank": int(ranks_series.loc[g]),
                    "total": n_genes_valid,
                })
            else:
                per_gene.append({
                    "id": g,
                    "single_gene_metric": None,
                    "rank": None,
                    "total": n_genes_valid,
                })

        ref_sets = [
            name for name, opqs in ref_opaque_by_set.items()
            if any(g in opqs for g in valid_ids)
        ]

        # Confounder-survival flags (HNSC/HPV only — None elsewhere).
        # A module that separates HPV beyond marking oropharynx tissue
        # or immune-composition artifact stays close to its full-cohort
        # AUROC within the stratified subgroup.
        orop_auroc, orop_n, orop_pos, orop_neg = _combined_in_subset(
            valid_ids, orop_test_mask,
        )
        pur_auroc, pur_n, pur_pos, pur_neg = _combined_in_subset(
            valid_ids, high_purity_test_mask,
        )

        def _survives(sub: float | None) -> bool | None:
            if sub is None:
                return None
            if combined is None or not math.isfinite(combined):
                return None
            return bool(sub + SURVIVE_TOLERANCE >= combined)

        out_modules.append({
            "gene_ids": valid_ids,
            "size": len(valid_ids),
            "combined_holdout": combined,
            "coherence": coherence_val,
            "ref_sets": ref_sets,
            "per_gene": per_gene,
            # Engine's own preference: the max GP fitness observed for
            # any candidate carrying this gene-set in the persisted
            # population. Lets the merged "Groups the engine explored"
            # table re-sort by what the SEARCH preferred without
            # discarding the alternate re-score lenses.
            "gp_fitness": gp_fitness_by_key.get(frozenset(valid_ids)),
            # The actual tree of the candidate that earned the
            # gp_fitness above (argmax over the persisted population
            # for this gene-set). Opaque-safe: program_repr is built
            # from opaque IDs only — no gene names.
            "best_program_repr": best_program_repr_by_key.get(
                frozenset(valid_ids),
            ),
            # Survival in held-out OROPHARYNX subgroup. None outside
            # HNSC/HPV (no oropharynx flag) or when subgroup too small.
            "combined_holdout_oropharynx": orop_auroc,
            "n_holdout_oropharynx": orop_n,
            "n_pos_oropharynx": orop_pos,
            "n_neg_oropharynx": orop_neg,
            "survives_site": _survives(orop_auroc),
            # Survival in held-out HIGH-PURITY (bottom-tertile immune-
            # infiltration proxy) subgroup. None when proxy genes can't
            # be resolved or subgroup too small.
            "combined_holdout_highpurity": pur_auroc,
            "n_holdout_highpurity": pur_n,
            "n_pos_highpurity": pur_pos,
            "n_neg_highpurity": pur_neg,
            "survives_purity": _survives(pur_auroc),
        })

    def _sort_key(m: dict) -> float:
        v = m["combined_holdout"]
        if v is None:
            return float("-inf")
        try:
            f = float(v)
        except (TypeError, ValueError):
            return float("-inf")
        return f if math.isfinite(f) else float("-inf")

    out_modules.sort(key=_sort_key, reverse=True)

    return _json_finite({
        "run_id": run.id,
        "target": target,
        "dataset": dataset,
        "metric_kind": metric_kind,
        "coherence": coherence_flag,
        "n_modules": len(out_modules),
        "n_train": int(len(train_ids_seq)),
        "n_test": int(len(test_ids_present)),
        "subgroups": {
            # Site-stratified subgroup metadata. None when not
            # applicable (non-HNSC, or HNSC without is_oropharynx).
            "site": (
                {
                    "kind": "oropharynx",
                    "n": int(orop_test_mask.sum()) if orop_test_mask is not None else 0,
                    "tolerance": SURVIVE_TOLERANCE,
                }
                if orop_test_mask is not None else None
            ),
            "purity": (
                {
                    "kind": "high_purity_bottom_tertile",
                    "n": (
                        int(high_purity_test_mask.sum())
                        if high_purity_test_mask is not None else 0
                    ),
                    "n_proxy_genes": len(proxy_opaque_used),
                    "tolerance": SURVIVE_TOLERANCE,
                }
                if high_purity_test_mask is not None else None
            ),
        },
        "modules": out_modules,
    })


@app.get("/runs/{run_id}/modules")
def get_run_modules(run_id: str) -> dict:
    """Ranked list of coordinated gene modules harvested from a run's
    persisted population. Each module = a distinct candidate's
    ``gene_ids`` set (≥2 genes), scored by combined held-out AUROC on
    the run's exact train/test split — opaque-only on the wire.

    Returns 425 if the run is still running, 404 if unknown, 400 for
    unsupervised runs (no target to evaluate against)."""
    run = RUN_STORE.get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail=f"Unknown run_id {run_id}")
    with run.lock:
        if run.status == "error":
            raise HTTPException(
                status_code=500, detail=run.error or "run failed",
            )
        if run.status != "done" or run.result is None:
            raise HTTPException(status_code=425, detail="run still running")
    return _compute_module_ranking(run)


# DSL operators recognised by the operator-usage endpoint. The token
# we match in ``program_repr`` is ``Name(`` so e.g. ``Select(`` won't
# collide with the bare ``M`` MatrixTerminal. ``FitApply`` was named
# this way in engine_v2; the user-visible label spells it "Fit/Apply".
_OPERATOR_TOKENS: list[tuple[str, str]] = [
    # (token_in_program_repr,  human_label_for_the_tile)
    ("Select(",    "Select"),
    ("Reduce(",    "Reduce"),
    ("Combine(",   "Combine"),
    ("Split(",     "Split"),
    ("Associate(", "Associate"),
    ("Effect(",    "Effect"),
    ("FitApply(",  "Fit/Apply"),
    ("Search(",    "Search"),
]


_OPERATOR_USAGE_CACHE: dict[str, dict] = {}


def _compute_operator_usage(run: "Run") -> dict:
    """Walk every candidate in every persisted generation; count how
    often each DSL operator token appears in the ``program_repr`` and
    in how many candidate-instances it appears at least once.

    Persistent elites are counted once per generation they appear in
    (matching the "Generations × Population" grid the user sees).
    Opaque-safe: operator keywords + integer counts only — no gene
    IDs or symbols ever leave this function.
    """
    with run.lock:
        log = list(run.log)
    totals: dict[str, int] = {label: 0 for _, label in _OPERATOR_TOKENS}
    programs_using: dict[str, int] = {label: 0 for _, label in _OPERATOR_TOKENS}
    n_candidates = 0
    n_generations = len(log)
    for entry in log:
        cands = entry.get("candidates", []) or []
        for cand in cands:
            n_candidates += 1
            repr_str = cand.get("program_repr") or ""
            if not repr_str:
                continue
            for token, label in _OPERATOR_TOKENS:
                # Counts every occurrence — a Reduce-inside-Reduce
                # program contributes twice to "Reduce" total uses.
                n_in_prog = repr_str.count(token)
                if n_in_prog == 0:
                    continue
                totals[label] += n_in_prog
                programs_using[label] += 1

    operators = [
        {
            "name": label,
            "total_uses": totals[label],
            "programs_using": programs_using[label],
        }
        for _, label in _OPERATOR_TOKENS
    ]
    return _json_finite({
        "run_id": run.id,
        "n_generations": int(n_generations),
        "n_candidates": int(n_candidates),
        "operators": operators,
    })


@app.get("/runs/{run_id}/operator-usage")
def get_run_operator_usage(run_id: str) -> dict:
    """How often each DSL operator (Select, Reduce, Combine, Split,
    Associate, Effect, Fit/Apply, Search) was used across every
    candidate program in every persisted generation — the
    "Generations × Population" grid. Opaque-safe by construction:
    operator keywords + integer counts only.

    Returns 425 if the run is still running, 404 if unknown."""
    run = RUN_STORE.get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail=f"Unknown run_id {run_id}")
    with run.lock:
        if run.status == "error":
            raise HTTPException(
                status_code=500, detail=run.error or "run failed",
            )
        if run.status != "done" or run.result is None:
            raise HTTPException(status_code=425, detail="run still running")
    cached = _OPERATOR_USAGE_CACHE.get(run_id)
    if cached is not None:
        return cached
    payload = _compute_operator_usage(run)
    _OPERATOR_USAGE_CACHE[run_id] = payload
    return payload


# ---------------------------------------------------------------------------
# Independent-cohort transfer test (GSE65858).
# ---------------------------------------------------------------------------
# Airgap discipline: this endpoint is the ONLY place the winning
# program's genes cross from the blind side to the named GSE65858
# cohort. We reveal ONLY the winner's opaque IDs (bounded — same
# discipline as /evaluate), pass those symbols to
# ``validate.transfer_gse65858.transfer_score``, and return the
# resulting AUROC + p + the same list of symbols we sent in. GSE65858's
# full gene list NEVER crosses back. The sealed map is only opened for
# the winner's genes, not dumped. Cached per run for the process
# lifetime.
_TRANSFER_CACHE: dict[str, dict] = {}


@app.get("/runs/{run_id}/transfer")
def get_run_transfer(run_id: str) -> dict:
    """External-cohort transfer test: score the HNSC/HPV winner on
    GSE65858 (GEO independent cohort, ~270 head & neck tumours,
    Illumina HumanHT-12 v4 microarray) and return AUROC + permutation
    p. Only defined for the HNSC/HPV objective.

    Returns 425 while the run is still going, 404 if unknown, 400 for
    non-HPV/HNSC runs."""
    run = RUN_STORE.get(run_id)
    if run is None:
        raise HTTPException(status_code=404, detail=f"Unknown run_id {run_id}")
    with run.lock:
        if run.status == "error":
            raise HTTPException(
                status_code=500, detail=run.error or "run failed",
            )
        if run.status != "done" or run.result is None:
            raise HTTPException(status_code=425, detail="run still running")
        target = run.objective_spec.get("target", "")
        dataset = run.dataset
        result = run.result

    if not (dataset == "hnsc" and target == "hpv"):
        raise HTTPException(
            status_code=400,
            detail=(
                "Transfer validation is only defined for the HNSC/HPV "
                "objective."
            ),
        )

    cached = _TRANSFER_CACHE.get(run_id)
    if cached is not None:
        return cached

    winning = (result or {}).get("winning", {}) or {}
    winner_ids = list(winning.get("gene_ids") or [])
    if not winner_ids:
        raise HTTPException(
            status_code=400,
            detail="Winner has no gene_ids to transfer.",
        )

    # Bounded reveal: only the winner's genes cross to the named side.
    try:
        symbols = reveal(winner_ids)
    except FileNotFoundError as exc:
        raise HTTPException(status_code=503, detail=str(exc))
    except KeyError as exc:
        raise HTTPException(status_code=400, detail=str(exc))
    symbols = [s for s in symbols if isinstance(s, str) and s]

    # Lazy-import the reveal-side validation module so the API import
    # graph doesn't pull it in eagerly (same discipline as the other
    # diagnostics).
    try:
        from validate.transfer_gse65858 import transfer_score
    except Exception as exc:  # noqa: BLE001 — module import failure
        raise HTTPException(
            status_code=503,
            detail=f"Transfer module could not be imported: {exc}",
        )

    try:
        score_out = transfer_score(symbols)
    except FileNotFoundError as exc:
        # Cohort parquets missing — the user hasn't built GSE65858
        # yet. Surface a clear 503 so the UI can show the graceful
        # "couldn't validate externally" note.
        raise HTTPException(status_code=503, detail=str(exc))

    payload = _json_finite({
        "run_id": run.id,
        "cohort": "GSE65858",
        "platform": "Illumina HumanHT-12 v4 microarray",
        "source": "GEO",
        "n_cohort": int(score_out.get("n", 0)),
        # These match the transfer_score dict verbatim; the only gene
        # NAMES here are the winner's own revealed symbols (found +
        # missing), same as /evaluate.
        "auroc": score_out.get("auroc"),
        "p": score_out.get("p"),
        "n": score_out.get("n"),
        "n_pos": score_out.get("n_pos"),
        "n_neg": score_out.get("n_neg"),
        "n_found": score_out.get("n_found"),
        "n_missing": score_out.get("n_missing"),
        "found_symbols": list(score_out.get("found_symbols") or []),
        "missing_symbols": list(score_out.get("missing_symbols") or []),
    })
    _TRANSFER_CACHE[run_id] = payload
    return payload