engine_v2/fitness.py

"""Fitness for engine_v2 programs — handles every output type.

Three objective shapes:

- **MSI separation** — fitness is **orientation-agnostic AUROC**:
  ``max(AUROC, 1 - AUROC)``. A perfectly inverted score is still
  rewarded; the engine doesn't need to learn a sign convention to
  separate the two classes.

- **TMB correlation** — fitness is the **signed strength of the
  negative correlation** between the output Vector and TMB:
  ``-spearman(score, TMB)``. Higher = score goes DOWN as TMB goes UP.

A valid program outputs Vector, Scalar, or Model:
- Vector → apply the objective metric.
- Scalar (from Associate / Effect) → THE statistic IS the fitness; we
  re-evaluate it held-out via the program's recipe rather than relying
  on the train-fold value.
- Model → not exposed as a root; FitApply auto-applies and yields a Vector.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import numpy as np
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.metrics import roc_auc_score, silhouette_score
from sklearn.model_selection import KFold, StratifiedKFold
from scipy.stats import spearmanr

from engine_v2.nodes import (
    Associate,
    Effect,
    ExecContext,
    FitApply,
    Node,
)
from engine_v2.types import TType


WORST_FITNESS = -np.inf

# Max gap a single-gene-monotonic program is allowed to have between
# its reported held-out AUROC and its constituent gene's single-gene
# held-out AUROC. Anything beyond this is treated as contaminated
# (e.g. residualise-before-split leak) and floored to WORST_FITNESS.
# Tolerance covers tie-breaking + numerical drift in the AUROC.
LEAKAGE_TOLERANCE = 0.02


@dataclass
class V2Objective:
    target: Literal["msi", "tmb", "none", "hpv"]
    binary: bool

    def fitness_label(self) -> str:
        if self.target in ("msi", "hpv"):
            return "max(AUROC,1−AUROC)"
        if self.target == "tmb":
            return "−spearman(score,TMB)"
        return "structure (2-cluster separation)"

    def worst_score(self) -> float:
        if self.target in ("msi", "hpv"):
            return 0.5
        if self.target == "tmb":
            return 0.0
        return -1.0  # unsupervised: silhouette floor

    def to_dict(self) -> dict:
        if self.target == "msi":
            return {"target": "msi", "metric": "auroc_omni"}
        if self.target == "tmb":
            return {"target": "tmb", "metric": "correlation", "direction": "neg"}
        if self.target == "hpv":
            return {"target": "hpv", "metric": "auroc_omni"}
        return {"target": "none", "metric": "structure"}

    def synthesis_overrides(self) -> dict:
        """Per-objective synthesis defaults. For unsupervised runs we
        forbid Scalar roots and disable Associate / Effect / FitApply
        entirely (none of them have a target to bind against), so the
        synthesised population is Vector-only.
        """
        if self.target == "none":
            return {
                "rates": {
                    "split": 0.10,
                    "fitapply": 0.0,
                    "effect": 0.0,
                    "search": 0.0,
                },
                "scalar_share": 0.0,
            }
        return {}

    def score_vector(
        self, scores: np.ndarray, y: np.ndarray | None = None
    ) -> float:
        if self.target == "none":
            return self._score_silhouette(scores)
        if self.target in ("msi", "hpv"):
            auroc = float(roc_auc_score(y, scores))
            return max(auroc, 1.0 - auroc)
        corr, _ = spearmanr(scores, y)
        if np.isnan(corr):
            return 0.0
        return -float(corr)

    def _score_oos_silhouette(
        self, train_scores: np.ndarray, test_scores: np.ndarray
    ) -> float:
        """Out-of-sample silhouette.

        Winsorize bounds, mean/std, and KMeans centres are fit on TRAIN
        scores ONLY. TEST scores are clipped + standardised with the
        train statistics, assigned to the train centres via
        ``KMeans.predict``, and the silhouette is computed on the TEST
        points. No test leakage. Guards live on the TEST side:
        ``n_test ≥ 20``, near-constant after clipping → ``worst_score()``,
        ≥30% per cluster.

        Used by ``cv_score`` (per fold) and ``evaluate_holdout`` (final
        held-out) so a program that overfits silhouette in-sample no
        longer wins selection and the reported held-out is true
        generalisation.
        """
        tr = np.asarray(train_scores, dtype=float)
        te = np.asarray(test_scores, dtype=float)
        n_train = tr.shape[0]
        n_test = te.shape[0]
        if n_train < 20 or n_test < 20:
            return self.worst_score()
        if not np.isfinite(tr).all() or not np.isfinite(te).all():
            return self.worst_score()
        # Fit on train only.
        lo = float(np.nanpercentile(tr, 2.5))
        hi = float(np.nanpercentile(tr, 97.5))
        tr_c = np.clip(tr, lo, hi)
        mean_w = float(np.nanmean(tr_c))
        std_w = float(np.nanstd(tr_c))
        if std_w < 1e-9:
            return self.worst_score()
        z_tr = (tr_c - mean_w) / std_w
        try:
            km = KMeans(n_clusters=2, n_init=10, random_state=0)
            km.fit(z_tr.reshape(-1, 1))
        except Exception:
            return self.worst_score()
        # Apply train statistics to test, assign to train centres.
        z_te = (np.clip(te, lo, hi) - mean_w) / std_w
        if float(np.nanstd(z_te)) < 1e-9:
            return self.worst_score()
        try:
            labels_te = km.predict(z_te.reshape(-1, 1))
        except Exception:
            return self.worst_score()
        min_cluster = max(10, int(np.ceil(0.30 * n_test)))
        counts = np.bincount(labels_te, minlength=2)
        if counts[0] < min_cluster or counts[1] < min_cluster:
            return self.worst_score()
        try:
            return float(silhouette_score(z_te.reshape(-1, 1), labels_te))
        except Exception:
            return self.worst_score()

    def _score_silhouette(self, scores: np.ndarray) -> float:
        """Silhouette of a 2-means split on standardised 1-D scores.

        Range [-1, 1]: higher = cleaner two-group separation. Degenerate
        outputs are floored to ``worst_score()``:

        - The raw ``std == 0`` guard misses ``protected_div(x, x)``-style
          programs that produce ~1 for most patients and a few huge
          values — a "perfect" 98:2 outlier split with silhouette ~1
          aligning with nothing. To catch it we WINSORIZE first (clip
          to the 2.5–97.5 percentile range), which collapses the
          outlier blow-up to near-constant; the std-on-winsorized
          guard then fires and floors the score. We don't rank-
          transform because that would erase a genuine bimodal gap.
        - We also raise the per-cluster floor from 10% to 30% of n,
          which would catch the same 98:2 split downstream if the
          winsorize check ever misses.
        """
        s = np.asarray(scores, dtype=float)
        n = s.shape[0]
        if n < 20 or not np.isfinite(s).all():
            return self.worst_score()
        # Winsorize then recompute mean / std on the clipped values.
        lo = float(np.nanpercentile(s, 2.5))
        hi = float(np.nanpercentile(s, 97.5))
        sw = np.clip(s, lo, hi)
        mean_w = float(np.nanmean(sw))
        std_w = float(np.nanstd(sw))
        if std_w < 1e-9:
            return self.worst_score()
        z = (sw - mean_w) / std_w
        try:
            km = KMeans(n_clusters=2, n_init=10, random_state=0)
            labels = km.fit_predict(z.reshape(-1, 1))
        except Exception:
            return self.worst_score()
        min_cluster = max(10, int(np.ceil(0.30 * n)))
        counts = np.bincount(labels, minlength=2)
        if counts[0] < min_cluster or counts[1] < min_cluster:
            return self.worst_score()
        try:
            return float(silhouette_score(z.reshape(-1, 1), labels))
        except Exception:
            return self.worst_score()

    def score_scalar(self, value: float) -> float:
        """A Scalar program output (Associate / Effect) IS the fitness,
        once oriented to the objective's direction.

        For MSI we use |scalar| because either sign separates a binary
        label equally well; for TMB we want negative correlation, so
        fitness = -scalar. For unsupervised the Scalar path is never
        produced (rates=0, scalar_share=0); the safety floor catches
        any stray case.
        """
        if not np.isfinite(value):
            return self.worst_score()
        if self.target in ("msi", "hpv"):
            return float(abs(value))
        if self.target == "tmb":
            return float(-value)
        return self.worst_score()


MSI_OBJECTIVE = V2Objective(target="msi", binary=True)
TMB_OBJECTIVE = V2Objective(target="tmb", binary=False)
UNSUP_OBJECTIVE = V2Objective(target="none", binary=False)
HPV_OBJECTIVE = V2Objective(target="hpv", binary=True)


def _single_gene_monotonic(program: Node) -> str | None:
    """If the program is single-gene monotonic — exactly one gene
    selected, no Combine/Split branches (i.e. its output is a
    monotonic function of one column) — return that gene's opaque
    ID. Otherwise None. Used by the held-out leakage guard: such a
    program MUST score AUROC equal to its constituent gene's
    single-gene AUROC; if it doesn't, the held-out is contaminated.
    """
    from engine_v2.nodes import Combine, Split

    ids = list(dict.fromkeys(program.feature_ids()))
    if len(ids) != 1:
        return None
    for n in program.walk():
        if isinstance(n, (Combine, Split)):
            return None
    return ids[0]


def _single_gene_auroc(
    gene_id: str, ctx: ExecContext, y: np.ndarray,
) -> float | None:
    """Orientation-agnostic single-gene AUROC on ``ctx``'s rows. Used
    by the leakage guard for binary objectives only."""
    if gene_id not in ctx.M.columns:
        return None
    x = ctx.M[gene_id].to_numpy(dtype=float)
    finite = np.isfinite(x) & np.isfinite(np.asarray(y, dtype=float))
    if int(finite.sum()) < 4:
        return None
    yf = np.asarray(y)[finite].astype(int)
    if len(np.unique(yf)) < 2:
        return None
    try:
        auroc = float(roc_auc_score(yf, x[finite]))
    except Exception:
        return None
    return max(auroc, 1.0 - auroc)


def _has_bare_matrix_reduce(program: Node) -> bool:
    """Return True if any Reduce in the program operates directly on a
    bare MatrixTerminal — i.e. a global reduction over the WHOLE
    expression matrix with no Select / Search restricting the columns.
    The synthesiser bans this shape for every objective; this is the
    runtime airbag in case a stray tree slips through crossover or a
    replay from disk."""
    from engine_v2.nodes import MatrixTerminal, Reduce as _Reduce

    for n in program.walk():
        if isinstance(n, _Reduce) and isinstance(n.matrix, MatrixTerminal):
            return True
    return False


def _check_target_binding(program: Node, expected_target: str) -> bool:
    """Every node that scores against a label (Associate / Effect /
    FitApply) MUST target the active objective. Returns True on a clean
    program; False if any node carries a stray target — the program is
    then floored to WORST_FITNESS by the caller, never silently let
    onto the score board.
    """
    for n in program.walk():
        if isinstance(n, (Associate, Effect, FitApply)):
            if n.target != expected_target:
                return False
    return True


def objective_from_spec(spec: dict) -> V2Objective:
    target = spec.get("target")
    metric = spec.get("metric")
    if target == "msi" and metric in ("auroc", "auroc_omni"):
        return MSI_OBJECTIVE
    if target == "tmb" and metric == "correlation":
        return TMB_OBJECTIVE
    if target == "none" and metric == "structure":
        return UNSUP_OBJECTIVE
    if target == "hpv" and metric in ("auroc", "auroc_omni"):
        return HPV_OBJECTIVE
    raise ValueError(
        f"engine_v2: unsupported objective spec {spec!r}. "
        "Supported: msi+auroc, tmb+correlation, none+structure, hpv+auroc."
    )


# ---------------------------------------------------------------------------
# Context construction
# ---------------------------------------------------------------------------


def make_ctx(
    M: pd.DataFrame,
    *,
    clinical: pd.DataFrame | None = None,
    msi: np.ndarray | None = None,
    tmb: np.ndarray | None = None,
    confounders: tuple[str, ...] = ("stage", "age"),
) -> ExecContext:
    labels: dict[str, np.ndarray] = {}
    if msi is not None:
        labels["msi"] = np.asarray(msi)
    if tmb is not None:
        labels["tmb"] = np.asarray(tmb)
    if clinical is None:
        clinical = pd.DataFrame(index=M.index)
    return ExecContext(
        M=M, clinical=clinical, labels=labels, confounders=confounders,
    )


# ---------------------------------------------------------------------------
# Output extraction — runs the program on a context, normalises to a
# per-patient score (or a Scalar)
# ---------------------------------------------------------------------------


def _vector_from_program(program: Node, ctx: ExecContext) -> np.ndarray | None:
    """Execute a program and return its per-patient score as a 1-D numpy
    array — or None if the result is degenerate."""
    try:
        out = program.execute(ctx)
    except Exception:
        return None
    if isinstance(out, pd.Series):
        arr = np.asarray(out.values, dtype=float)
    elif isinstance(out, (float, int, np.floating, np.integer)):
        return None   # callers handle Scalar via _scalar_from_program
    else:
        return None
    if arr.shape[0] != ctx.M.shape[0]:
        return None
    if not np.isfinite(arr).all():
        return None
    if np.nanstd(arr) == 0.0:
        return None
    return arr


def _scalar_from_program(program: Node, ctx: ExecContext) -> float | None:
    try:
        out = program.execute(ctx)
    except Exception:
        return None
    if isinstance(out, (float, int, np.floating, np.integer)):
        v = float(out)
        return v if np.isfinite(v) else None
    return None


# ---------------------------------------------------------------------------
# Held-out + CV scoring
# ---------------------------------------------------------------------------


def evaluate_holdout(
    program: Node,
    ctx: ExecContext,
    y: np.ndarray | None,
    *,
    objective: V2Objective,
    ctx_train: ExecContext | None = None,
) -> float:
    """Single-shot evaluation: program is re-executed against the held-out
    context (its own ctx.M / ctx.labels), and the objective metric is
    applied. Returns the objective's score.

    For the unsupervised objective ``y`` is ignored. When ``ctx_train``
    is passed (the pipeline always does), the silhouette is computed
    OUT-OF-SAMPLE: cluster centres fit on the train context, test scores
    assigned to those centres and scored. Falls back to in-sample
    ``_score_silhouette`` only when ``ctx_train`` is None (kept for
    backward-compat with old test callers).

    When ``ctx_train`` is provided, the test ctx is wrapped so
    ``FitApply`` fits its model on train labels and applies the frozen
    model to the test inputs — never fits on test labels.
    """
    # Safety floor: every objective requires gene-based programs.
    # A no-Select tree has no genes; a Reduce over a bare
    # MatrixTerminal is a global-mean shortcut. Both score the worst.
    if not program.feature_ids():
        return WORST_FITNESS
    if _has_bare_matrix_reduce(program):
        return WORST_FITNESS
    # Honest FitApply: wrap the test ctx with fit_ctx=ctx_train so
    # FitApply fits on train labels and applies on test. Identity for
    # programs without FitApply nodes.
    if ctx_train is not None and ctx.fit_ctx is None:
        ctx = ExecContext(
            M=ctx.M, clinical=ctx.clinical, labels=ctx.labels,
            fit_ctx=ctx_train,
        )
    if program.ttype is TType.SCALAR:
        v = _scalar_from_program(program, ctx)
        if v is None:
            return WORST_FITNESS
        return objective.score_scalar(v)
    scores = _vector_from_program(program, ctx)
    if scores is None:
        return WORST_FITNESS
    if objective.target == "none":
        if ctx_train is None:
            # Backward-compat fallback for callers that don't have a
            # train ctx (older tests). Pipeline always supplies it.
            return objective.score_vector(scores)
        train_scores = _vector_from_program(program, ctx_train)
        if train_scores is None:
            return WORST_FITNESS
        return objective._score_oos_silhouette(train_scores, scores)
    out = objective.score_vector(scores, y)
    # Leakage guard for binary AUROC objectives. A single-gene-
    # monotonic program's held-out AUROC MUST equal that gene's
    # single-gene held-out AUROC on the same test rows (modest
    # tolerance for ties/numerical drift). If it doesn't, the held-out
    # features are contaminated (e.g. residualise-before-split leak)
    # — floor to worst. Skip for multi-gene / Combine / Split programs
    # where genuine synergy beyond single genes is legitimate.
    if objective.binary and objective.target in ("msi", "hpv"):
        sg = _single_gene_monotonic(program)
        if sg is not None and y is not None:
            honest = _single_gene_auroc(sg, ctx, y)
            if honest is not None and out > honest + LEAKAGE_TOLERANCE:
                return WORST_FITNESS
    return out


def cv_score(
    program: Node,
    ctx: ExecContext,
    y: np.ndarray | None,
    *,
    objective: V2Objective,
    n_folds: int = 5,
    random_state: int = 0,
) -> float:
    """k-fold CV on the rows of ``ctx.M``. For Scalar-rooted programs we
    re-execute against the fold's slice; for Vector-rooted we slice the
    output. For the unsupervised objective ``y`` is ignored — the
    structure score is intrinsic to the per-patient Vector."""
    n = ctx.M.shape[0]
    if n < n_folds + 2:
        return evaluate_holdout(program, ctx, y, objective=objective)
    splitter: object
    if objective.binary:
        splitter = StratifiedKFold(
            n_splits=n_folds, shuffle=True, random_state=random_state,
        )
        fold_iter = splitter.split(np.zeros((n, 1)), y)
    else:
        splitter = KFold(n_splits=n_folds, shuffle=True, random_state=random_state)
        fold_iter = splitter.split(np.zeros((n, 1)))

    is_scalar = program.ttype is TType.SCALAR
    if not is_scalar:
        # Compute once on the full training context — slicing is fine
        # because Vector outputs are per-row.
        scores = _vector_from_program(program, ctx)
        if scores is None:
            return WORST_FITNESS

    is_unsup = objective.target == "none"

    out: list[float] = []
    for tr, te in fold_iter:
        if is_scalar:
            # Re-execute on the fold's slice so e.g. Effect re-residualises.
            sub_ctx = _slice_ctx(ctx, te)
            v = _scalar_from_program(program, sub_ctx)
            if v is None:
                return WORST_FITNESS
            out.append(objective.score_scalar(v))
        else:
            if is_unsup:
                # Train-fit / test-eval per fold: programs that overfit
                # in-sample silhouette no longer win selection.
                out.append(
                    objective._score_oos_silhouette(
                        scores[tr], scores[te],  # type: ignore[index]
                    )
                )
            else:
                out.append(objective.score_vector(scores[te], y[te]))  # type: ignore[index]
    return float(np.mean(out))


def _slice_ctx(ctx: ExecContext, idx: np.ndarray) -> ExecContext:
    """Slice every row-aligned field of the context by integer index."""
    M = ctx.M.iloc[idx]
    clin = ctx.clinical.iloc[idx] if not ctx.clinical.empty else ctx.clinical
    labels = {k: np.asarray(v)[idx] for k, v in ctx.labels.items()}
    return ExecContext(M=M, clinical=clin, labels=labels)


def _coherence_score(program: Node, ctx: ExecContext) -> float:
    """Generic 'prefer coordinated gene modules' bias.

    Returns the mean absolute pairwise correlation among the program's
    Select'd opaque columns on the training matrix. Names no gene or
    pathway — purely structural. Programs with fewer than 2 distinct
    genes get a coherence of 0 (neutral). Range [0, 1]; higher = more
    coordinated module.
    """
    ids = list(dict.fromkeys(program.feature_ids()))
    if len(ids) < 2:
        return 0.0
    cols = [c for c in ids if c in ctx.M.columns]
    if len(cols) < 2:
        return 0.0
    try:
        corr = ctx.M[cols].corr().abs().to_numpy()
    except Exception:
        return 0.0
    n = corr.shape[0]
    if n < 2:
        return 0.0
    iu = np.triu_indices(n, k=1)
    vals = corr[iu]
    if vals.size == 0:
        return 0.0
    mean = float(np.nanmean(vals))
    return mean if np.isfinite(mean) else 0.0


def fitness_fn(
    program: Node,
    ctx: ExecContext,
    y: np.ndarray | None,
    *,
    objective: V2Objective,
    lambda_size: float = 0.005,
    n_folds: int = 5,
    random_state: int = 0,
    coherence_weight: float = 0.0,
) -> float:
    # The active objective binds the scoring target; a program that
    # carries any other label as its target is treated as degenerate.
    # For unsupervised runs the program is Vector-only (no target-bearing
    # nodes), so the guard is a no-op.
    if not _check_target_binding(program, objective.target):
        return WORST_FITNESS
    # Every objective requires gene-based programs — a Reduce over the
    # whole matrix is a detection shortcut that scores via bulk
    # expression rather than gene choice. Two floors:
    #   (a) no Select anywhere → no genes chosen at all.
    #   (b) any Reduce on a bare MatrixTerminal → global-mean shortcut.
    if not program.feature_ids():
        return WORST_FITNESS
    if _has_bare_matrix_reduce(program):
        return WORST_FITNESS
    base = cv_score(
        program, ctx, y,
        objective=objective, n_folds=n_folds, random_state=random_state,
    )
    if base == WORST_FITNESS:
        return WORST_FITNESS
    out = base - lambda_size * program.node_count()
    if coherence_weight > 0.0:
        out += coherence_weight * _coherence_score(program, ctx)
    return out