engine_v2/pipeline.py

"""engine_v2 top-level pipeline.

Inputs are name-blind (opaque feature IDs only). Two entry points:

- ``run_v2_pipeline`` — batch path returning ``(evolution_log, result)``.
- ``run_v2_pipeline_streaming`` — on_generation callback for SSE.

Steps: split TRAIN/TEST → optional prefilter → typed GP → winner
held-out → winner permutation null (winner held FIXED, target
shuffled). The baseline is dropped from v2 — its old "univariate
top-K" interpretation doesn't transfer cleanly to typed trees.
"""

from __future__ import annotations

import re
import time
from typing import Callable

import numpy as np
import pandas as pd

from engine.prefilter import top_n_features
from engine.split import make_split
# We reuse the v1 prefilter only via the BinaryAUROCObjective signal, so
# import lazily inside the prefilter branch — engine_v2 itself stays free
# of v1 fitness coupling.
from engine_v2.fitness import (
    V2Objective,
    evaluate_holdout,
    fitness_fn,
    make_ctx,
)
from engine_v2.gp import run_gp_v2
from engine_v2.nodes import ExecContext, Node
from engine_v2.permutation import (
    permutation_null,
    permutation_p_value,
    unsup_random_null,
)


_OPAQUE_ID_RE = re.compile(r"^g\d+$")


def _check_opaque_only(M: pd.DataFrame) -> None:
    bad = [c for c in M.columns if not _OPAQUE_ID_RE.match(str(c))]
    if bad:
        raise ValueError(
            "engine_v2: matrix columns must be opaque IDs (^g\\d+$); "
            f"got non-conforming columns e.g. {bad[:5]}"
        )


def _prefilter_pool(
    M_train: pd.DataFrame,
    y_train: np.ndarray,
    n: int,
    objective: V2Objective,
) -> list[str]:
    """Narrow the column pool. We dispatch on the objective so that:
       - MSI uses the binary-AUROC univariate ranking
         (engine.objectives.BinaryAUROCObjective).
       - TMB uses the absolute-Spearman ranking
         (engine.objectives.CorrelationObjective).
    """
    if objective.target == "msi":
        from engine.objectives import BinaryAUROCObjective
        obj = BinaryAUROCObjective()
    else:
        from engine.objectives import CorrelationObjective
        obj = CorrelationObjective(direction="neg")
    shortlist, _ = top_n_features(M_train, y_train, n=n, objective=obj)
    return shortlist


def _align_priors(
    M: pd.DataFrame, residualize_scores: pd.DataFrame,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Drop cohort rows whose prior axis scores are missing or
    non-finite — they can't be projected. Returns (M, priors) on the
    same valid index. Defining the cohort this way is NOT leakage; it
    just says "we can only score patients we have priors for"."""
    aligned = residualize_scores.apply(pd.to_numeric, errors="coerce")
    aligned = aligned.reindex(M.index)
    valid = aligned.notna().all(axis=1) & np.isfinite(aligned).all(axis=1)
    if not bool(valid.all()):
        M = M.loc[valid]
        aligned = aligned.loc[valid]
    return M, aligned


def _fit_residualise_beta(
    M_train: pd.DataFrame, priors_train: pd.DataFrame,
) -> np.ndarray | None:
    """Vectorised OLS of every column of M on ``[intercept, *priors]``,
    fit using TRAIN rows ONLY. Returns the (1+k, g) coefficient matrix,
    or None when the train slice is too small to fit. The key
    invariant: these coefficients depend on TRAIN rows only — they're
    later applied to the full cohort (train + test) so the test rows
    are never used to choose the projection. This is what makes
    held-out AUROC honest under peel-off."""
    if M_train.shape[0] < 5:
        return None
    n = M_train.shape[0]
    P = np.column_stack(
        [np.ones(n), priors_train.to_numpy(dtype=float)],
    )
    Y = M_train.to_numpy(dtype=float)
    beta, *_ = np.linalg.lstsq(P, Y, rcond=None)
    return beta


def _apply_residualise(
    M: pd.DataFrame, priors: pd.DataFrame, beta: np.ndarray,
) -> pd.DataFrame:
    """Apply a previously-fit residualisation projection to ``M``.
    Test rows are residualised using train-fit ``beta`` — never their
    own. ``priors`` must be aligned to ``M.index``; columns must match
    the priors used at fit time."""
    n = len(M)
    P = np.column_stack([np.ones(n), priors.to_numpy(dtype=float)])
    Y = M.to_numpy(dtype=float)
    resid = Y - P @ beta
    return pd.DataFrame(resid, index=M.index, columns=M.columns)


def _make_full_ctx(
    M: pd.DataFrame,
    clinical: pd.DataFrame | None,
) -> ExecContext:
    """Full-cohort ExecContext for executing the winner across every
    patient (so the iterative-discovery chain can residualise against
    the per-patient scores on the next run).

    Labels are deliberately empty for both supervised and unsupervised
    peel-off. For supervised winners that contain ``FitApply(...,
    target)``, executing with empty labels causes FitApply to
    short-circuit to its raw inner vector (the pre-fit score) — which
    is what we want as the residualisation target: it's monotonic
    with the LR-fitted prediction in the single-1D-input case so the
    chain is consistent, and it sidesteps any test-row leakage from
    re-fitting on the full cohort here.
    """
    return ExecContext(
        M=M,
        clinical=(
            clinical if clinical is not None
            else pd.DataFrame(index=M.index)
        ),
        labels={},
    )


def _full_cohort_winner_scores(
    winner: Node,
    M: pd.DataFrame,
    clinical: pd.DataFrame | None,
) -> tuple[list[float | None], list[str]]:
    """Re-execute the winner on the full cohort and return per-patient
    scores (finite-guarded) + their sample-id labels. Used to feed the
    iterative-discovery chain's residualisation step on the NEXT run.
    Currently only emitted for the unsupervised objective."""
    full_ctx = _make_full_ctx(M, clinical)
    sample_ids = [str(s) for s in M.index]
    try:
        out = winner.execute(full_ctx)
    except Exception:
        return [], sample_ids
    if not isinstance(out, pd.Series):
        return [], sample_ids
    scores: list[float | None] = []
    for v in out.values:
        try:
            f = float(v)
        except (TypeError, ValueError):
            scores.append(None)
            continue
        scores.append(f if np.isfinite(f) else None)
    return scores, sample_ids


def _build_ctxs(
    M: pd.DataFrame,
    split,
    *,
    primary_target_name: str,
    clinical: pd.DataFrame | None,
    extra_labels: dict[str, np.ndarray] | None,
    confounders: tuple[str, ...] = ("stage", "age"),
) -> tuple[ExecContext, ExecContext]:
    """Build TRAIN and TEST ExecContexts from a split + optional clinical
    + optional other-target labels. The primary y is keyed by the
    objective's target name.

    For the unsupervised objective (target='none'), every label is
    stripped — the engine literally cannot see msi / tmb during search.
    The held-out labels are recovered in api/_worker for the post-hoc
    alignment check.
    """
    M_train = M.loc[split.train_ids]
    M_test = M.loc[split.test_ids]
    clin_train = clinical.loc[split.train_ids] if clinical is not None else None
    clin_test = clinical.loc[split.test_ids] if clinical is not None else None

    is_unsup = primary_target_name == "none"

    def slice_labels(side_ids, side_y) -> dict[str, np.ndarray]:
        if is_unsup:
            return {}   # airgap: no labels to the engine during unsup search
        labels: dict[str, np.ndarray] = {primary_target_name: side_y}
        if extra_labels:
            id_to_pos = {sid: i for i, sid in enumerate(M.index)}
            pos = np.array([id_to_pos[sid] for sid in side_ids])
            for k, v in extra_labels.items():
                if k == primary_target_name:
                    continue
                arr = np.asarray(v)
                if len(arr) == len(M):
                    labels[k] = arr[pos]
        return labels

    ctx_train = ExecContext(
        M=M_train,
        clinical=clin_train if clin_train is not None else pd.DataFrame(index=M_train.index),
        labels=slice_labels(split.train_ids, split.y_train),
        confounders=confounders,
    )
    ctx_test = ExecContext(
        M=M_test,
        clinical=clin_test if clin_test is not None else pd.DataFrame(index=M_test.index),
        labels=slice_labels(split.test_ids, split.y_test),
        confounders=confounders,
    )
    if is_unsup:
        assert ctx_train.labels == {} and ctx_test.labels == {}, (
            "unsup ExecContext must carry no labels — engine airgap"
        )
    return ctx_train, ctx_test


def run_v2_pipeline(
    M: pd.DataFrame,
    y: np.ndarray | None,
    *,
    objective: V2Objective,
    seed: int = 42,
    test_size: float = 0.3,
    prefilter_n: int | None = None,
    population_size: int = 200,
    n_generations: int = 40,
    n_permutations: int = 200,
    cv_folds: int = 5,
    tournament_k: int = 3,
    elitism: int = 5,
    p_mutate: float = 0.7,
    lambda_size: float = 0.005,
    max_depth: int = 4,
    max_genes_per_set: int = 8,
    max_nodes: int = 64,
    clinical: pd.DataFrame | None = None,
    extra_labels: dict[str, np.ndarray] | None = None,
    residualize_scores: pd.DataFrame | None = None,
    coherence_weight: float = 0.0,
    confounders: tuple[str, ...] = ("stage", "age"),
    immigrant_fraction: float = 0.0,
    rates_override: dict | None = None,
    scalar_share_override: float | None = None,
) -> tuple[dict, dict]:
    """Full engine_v2 pipeline.

    ``clinical`` (stage / age, indexed like M) and ``extra_labels``
    (e.g. include "tmb" when the objective targets "msi", and vice
    versa) feed the full-DSL operators (Split / Effect / Associate /
    FitApply) without compromising the airgap — genes stay opaque.
    """
    _check_opaque_only(M)

    n_genes_input = int(M.shape[1])
    nan_cols = M.columns[M.isna().any(axis=0)]
    if len(nan_cols):
        M = M.drop(columns=nan_cols)
    # Peel-off chain: align M to patients with prior-axis scores
    # BEFORE the split (defining the cohort isn't leakage). The actual
    # residualisation projection is fit AFTER the split, on TRAIN rows
    # only, then applied to the full M — so test features are
    # residualised with train-fit coefficients, never their own. This
    # is what makes held-out AUROC honest under peel-off.
    priors_aligned: pd.DataFrame | None = None
    if residualize_scores is not None and len(residualize_scores.columns) > 0:
        M, priors_aligned = _align_priors(M, residualize_scores)
    if clinical is not None:
        clinical = clinical.reindex(M.index)

    is_unsup = objective.target == "none"
    split_y = y if y is not None else np.zeros(len(M), dtype=float)
    split = make_split(
        M.index, split_y, test_size=test_size, random_state=seed,
        stratify=objective.binary,
    )

    # Fit residualisation on TRAIN rows only; apply to the full M so
    # train + test features sit in the same residualised space without
    # the test rows ever being seen by the projection fit.
    if priors_aligned is not None:
        beta = _fit_residualise_beta(
            M.loc[split.train_ids], priors_aligned.loc[split.train_ids],
        )
        if beta is not None:
            M = _apply_residualise(M, priors_aligned, beta)
    ctx_train, ctx_test = _build_ctxs(
        M, split,
        primary_target_name=objective.target,
        clinical=clinical,
        extra_labels=extra_labels,
        confounders=confounders,
    )

    # Prefilter is target-driven; for unsup there's no target, fall back
    # to the full opaque pool.
    pool = (
        _prefilter_pool(ctx_train.M, split.y_train, prefilter_n, objective)
        if (prefilter_n is not None and not is_unsup)
        else list(ctx_train.M.columns)
    )

    gp_y_train = None if is_unsup else split.y_train
    gp_y_test = None if is_unsup else split.y_test

    t0 = time.time()
    log, winner, winner_cv_fitness = run_gp_v2(
        ctx_train, gp_y_train, pool,
        objective=objective,
        population_size=population_size,
        n_generations=n_generations,
        tournament_k=tournament_k,
        elitism=elitism,
        p_mutate=p_mutate,
        lambda_size=lambda_size,
        coherence_weight=coherence_weight,
        immigrant_fraction=immigrant_fraction,
        rates_override=rates_override,
        scalar_share_override=scalar_share_override,
        cv_folds=cv_folds,
        seed=seed,
        max_depth=max_depth,
        max_genes_per_set=max_genes_per_set,
        max_nodes=max_nodes,
    )
    gp_seconds = time.time() - t0

    winner_holdout = evaluate_holdout(
        winner, ctx_test, gp_y_test,
        objective=objective, ctx_train=ctx_train,
    )
    # Re-execute the winner to extract per-patient held-out scores so the
    # API worker can compute post-hoc alignment (unsup) without parsing
    # the program_repr back into a Node. Falls back to empty if execution
    # was degenerate.
    try:
        _w_scores_series = winner.execute(ctx_test)
        if isinstance(_w_scores_series, pd.Series):
            winner_holdout_scores = [
                float(v) if np.isfinite(v) else None for v in _w_scores_series.values
            ]
        else:
            winner_holdout_scores = []
    except Exception:
        winner_holdout_scores = []
    holdout_sample_ids = [str(s) for s in ctx_test.M.index]

    # Full-cohort scores: re-execute the winner on every patient
    # (train+test combined, labels stripped) so the iterative-
    # discovery chain can residualise against this axis on the next
    # run. Emitted for every objective now — MSI/HPV/TMB can enumerate
    # axes too via Find next axis.
    full_scores, full_sample_ids = _full_cohort_winner_scores(
        winner, M, clinical,
    )

    if is_unsup:
        rates = objective.synthesis_overrides().get("rates")
        nulls = unsup_random_null(
            ctx_test, pool,
            objective=objective,
            n_permutations=n_permutations,
            rates=rates,
            max_depth=max_depth,
            max_genes_per_set=max_genes_per_set,
            seed=seed,
            ctx_train=ctx_train,
        )
    else:
        nulls = permutation_null(
            winner, ctx_test, split.y_test,
            objective=objective,
            n_permutations=n_permutations,
            seed=seed,
        )
    p_value = permutation_p_value(winner_holdout, nulls)

    winner_repr = winner.repr_typed()
    winner_gene_ids = list(dict.fromkeys(winner.feature_ids()))   # dedup keep-order

    evolution_log = {
        "engine": "v2",
        "run": {
            "seed": int(seed),
            "objective_spec": objective.to_dict(),
            "fitness_label": objective.fitness_label(),
            "params": {
                "population_size": population_size,
                "n_generations": n_generations,
                "tournament_k": tournament_k,
                "elitism": elitism,
                "p_mutate": p_mutate,
                "lambda_size": lambda_size,
                "test_size": test_size,
                "prefilter_n": prefilter_n,
                "cv_folds": cv_folds,
                "n_permutations": n_permutations,
                "max_depth": max_depth,
                "max_genes_per_set": max_genes_per_set,
                "max_nodes": max_nodes,
            },
            "n_train": int(len(split.train_ids)),
            "n_test": int(len(split.test_ids)),
            "n_genes": int(M.shape[1]),
            "n_genes_input": n_genes_input,
            "n_genes_dropped_nan": int(len(nan_cols)),
            "prefilter_N": None if prefilter_n is None else int(prefilter_n),
            "prefilter_note": (
                "Prefilter off: typed GP samples FeatureSets from the full "
                "opaque-ID column set."
                if prefilter_n is None
                else f"Prefilter on: typed GP samples FeatureSets from the "
                     f"top-{prefilter_n} univariate features (TRAIN only)."
            ),
            "gp_seconds": round(gp_seconds, 2),
        },
        "generations": log,
    }

    worst = objective.worst_score()
    finite_nulls = [n for n in nulls if np.isfinite(n)]

    def _f(x: float) -> float:
        return float(x) if np.isfinite(x) else worst

    result = {
        "engine": "v2",
        "objective_spec": objective.to_dict(),
        "fitness_label": objective.fitness_label(),
        "winning": {
            "id": "winner",
            "program_repr": winner_repr,
            "gene_ids": winner_gene_ids,
            "n_nodes": int(winner.node_count()),
            "depth": int(winner.depth()),
            "cv_fitness": _f(winner_cv_fitness),
            "holdout_score": _f(winner_holdout),
            "holdout_auroc": _f(winner_holdout),
            "permutation_p": (
                float(p_value) if np.isfinite(p_value) else 1.0
            ),
            "holdout_scores": winner_holdout_scores,
            "holdout_sample_ids": holdout_sample_ids,
            "full_scores": full_scores,
            "full_sample_ids": full_sample_ids,
        },
        "permutation_summary": {
            "n_permutations": n_permutations,
            "null_kind": (
                "random_vector_programs" if is_unsup
                else "winner_fixed_target_shuffle"
            ),
            "null_score_mean": (
                float(np.mean(finite_nulls)) if finite_nulls else worst
            ),
            "null_score_p95": (
                float(np.quantile(finite_nulls, 0.95))
                if finite_nulls else worst
            ),
        },
    }
    return evolution_log, result


def run_v2_pipeline_streaming(
    M: pd.DataFrame,
    y: np.ndarray | None,
    *,
    objective: V2Objective,
    on_generation: Callable[[dict], None],
    seed: int = 42,
    test_size: float = 0.3,
    prefilter_n: int | None = None,
    population_size: int = 200,
    n_generations: int = 40,
    n_permutations: int = 200,
    cv_folds: int = 5,
    tournament_k: int = 3,
    elitism: int = 5,
    p_mutate: float = 0.7,
    lambda_size: float = 0.005,
    max_depth: int = 4,
    max_genes_per_set: int = 8,
    max_nodes: int = 64,
    clinical: pd.DataFrame | None = None,
    extra_labels: dict[str, np.ndarray] | None = None,
    residualize_scores: pd.DataFrame | None = None,
    coherence_weight: float = 0.0,
    confounders: tuple[str, ...] = ("stage", "age"),
    immigrant_fraction: float = 0.0,
    rates_override: dict | None = None,
    scalar_share_override: float | None = None,
) -> dict:
    """Streaming variant — calls ``on_generation(entry)`` each generation
    and returns the final result dict. The caller stores the full
    population in its own log via the callback."""
    _check_opaque_only(M)
    n_genes_input = int(M.shape[1])
    nan_cols = M.columns[M.isna().any(axis=0)]
    if len(nan_cols):
        M = M.drop(columns=nan_cols)
    # Peel-off chain: align M to patients with prior-axis scores
    # BEFORE the split (defining the cohort isn't leakage). The actual
    # residualisation projection is fit AFTER the split, on TRAIN rows
    # only, then applied to the full M — so test features are
    # residualised with train-fit coefficients, never their own. This
    # is what makes held-out AUROC honest under peel-off.
    priors_aligned: pd.DataFrame | None = None
    if residualize_scores is not None and len(residualize_scores.columns) > 0:
        M, priors_aligned = _align_priors(M, residualize_scores)
    if clinical is not None:
        clinical = clinical.reindex(M.index)

    is_unsup = objective.target == "none"
    split_y = y if y is not None else np.zeros(len(M), dtype=float)
    split = make_split(
        M.index, split_y, test_size=test_size, random_state=seed,
        stratify=objective.binary,
    )

    # Fit residualisation on TRAIN rows only; apply to the full M so
    # train + test features sit in the same residualised space without
    # the test rows ever being seen by the projection fit.
    if priors_aligned is not None:
        beta = _fit_residualise_beta(
            M.loc[split.train_ids], priors_aligned.loc[split.train_ids],
        )
        if beta is not None:
            M = _apply_residualise(M, priors_aligned, beta)
    ctx_train, ctx_test = _build_ctxs(
        M, split,
        primary_target_name=objective.target,
        clinical=clinical,
        extra_labels=extra_labels,
        confounders=confounders,
    )
    pool = (
        _prefilter_pool(ctx_train.M, split.y_train, prefilter_n, objective)
        if (prefilter_n is not None and not is_unsup)
        else list(ctx_train.M.columns)
    )

    gp_y_train = None if is_unsup else split.y_train
    gp_y_test = None if is_unsup else split.y_test

    t0 = time.time()
    _log, winner, winner_cv_fitness = run_gp_v2(
        ctx_train, gp_y_train, pool,
        objective=objective,
        population_size=population_size,
        n_generations=n_generations,
        tournament_k=tournament_k,
        elitism=elitism,
        p_mutate=p_mutate,
        lambda_size=lambda_size,
        coherence_weight=coherence_weight,
        immigrant_fraction=immigrant_fraction,
        rates_override=rates_override,
        scalar_share_override=scalar_share_override,
        cv_folds=cv_folds,
        seed=seed,
        max_depth=max_depth,
        max_genes_per_set=max_genes_per_set,
        max_nodes=max_nodes,
        on_generation=on_generation,
    )
    gp_seconds = time.time() - t0

    winner_holdout = evaluate_holdout(
        winner, ctx_test, gp_y_test,
        objective=objective, ctx_train=ctx_train,
    )
    # Re-execute the winner to extract per-patient held-out scores so the
    # API worker can compute post-hoc alignment (unsup) without parsing
    # the program_repr back into a Node. Falls back to empty if execution
    # was degenerate.
    try:
        _w_scores_series = winner.execute(ctx_test)
        if isinstance(_w_scores_series, pd.Series):
            winner_holdout_scores = [
                float(v) if np.isfinite(v) else None for v in _w_scores_series.values
            ]
        else:
            winner_holdout_scores = []
    except Exception:
        winner_holdout_scores = []
    holdout_sample_ids = [str(s) for s in ctx_test.M.index]

    # Full-cohort scores: feeds the peel-off chain's residualisation on
    # the next run. Emitted for every objective now — MSI/HPV/TMB can
    # enumerate axes too.
    full_scores, full_sample_ids = _full_cohort_winner_scores(
        winner, M, clinical,
    )

    if is_unsup:
        rates = objective.synthesis_overrides().get("rates")
        nulls = unsup_random_null(
            ctx_test, pool,
            objective=objective,
            n_permutations=n_permutations,
            rates=rates,
            max_depth=max_depth,
            max_genes_per_set=max_genes_per_set,
            seed=seed,
            ctx_train=ctx_train,
        )
    else:
        nulls = permutation_null(
            winner, ctx_test, split.y_test,
            objective=objective,
            n_permutations=n_permutations,
            seed=seed,
        )
    p_value = permutation_p_value(winner_holdout, nulls)
    worst = objective.worst_score()
    finite_nulls = [n for n in nulls if np.isfinite(n)]

    def _f(x: float) -> float:
        return float(x) if np.isfinite(x) else worst

    return {
        "engine": "v2",
        "objective_spec": objective.to_dict(),
        "fitness_label": objective.fitness_label(),
        "winning": {
            "id": "winner",
            "program_repr": winner.repr_typed(),
            "gene_ids": list(dict.fromkeys(winner.feature_ids())),
            "n_nodes": int(winner.node_count()),
            "depth": int(winner.depth()),
            "cv_fitness": _f(winner_cv_fitness),
            "holdout_score": _f(winner_holdout),
            "holdout_auroc": _f(winner_holdout),
            "permutation_p": (
                float(p_value) if np.isfinite(p_value) else 1.0
            ),
            "holdout_scores": winner_holdout_scores,
            "holdout_sample_ids": holdout_sample_ids,
            "full_scores": full_scores,
            "full_sample_ids": full_sample_ids,
        },
        "permutation_summary": {
            "n_permutations": n_permutations,
            "null_kind": (
                "random_vector_programs" if is_unsup
                else "winner_fixed_target_shuffle"
            ),
            "null_score_mean": (
                float(np.mean(finite_nulls)) if finite_nulls else worst
            ),
            "null_score_p95": (
                float(np.quantile(finite_nulls, 0.95))
                if finite_nulls else worst
            ),
        },
        "run_meta": {
            "seed": int(seed),
            "n_genes": int(M.shape[1]),
            "n_genes_input": n_genes_input,
            "n_genes_dropped_nan": int(len(nan_cols)),
            "prefilter_N": None if prefilter_n is None else int(prefilter_n),
            "gp_seconds": round(gp_seconds, 2),
        },
    }