"""Tests for engine_v2 — typed synthesis over the DSL."""

from __future__ import annotations

import random
import re

import numpy as np
import pandas as pd
import pytest

from engine_v2 import run_v2_pipeline
from engine_v2.fitness import (
    MSI_OBJECTIVE,
    TMB_OBJECTIVE,
    cv_score,
    evaluate_holdout,
    fitness_fn,
    make_ctx,
    objective_from_spec,
)
from engine_v2.nodes import (
    Combine,
    FeatureSet,
    MatrixTerminal,
    Node,
    Reduce,
    Select,
)
from engine_v2.permutation import permutation_null, permutation_p_value
from engine_v2.synthesize import (
    crossover,
    mutate,
    ramped_population,
    random_program,
)
from engine_v2.types import AGGS, OPS, TType


def _synth(n_samples=80, n_features=40, seed=0, informative=6):
    rng = np.random.default_rng(seed)
    cols = [f"g{i+1:05d}" for i in range(n_features)]
    ids = pd.Index([f"s{i}" for i in range(n_samples)], name="sample_id")
    half = n_samples // 2
    y = np.array([1] * half + [0] * (n_samples - half))
    X = rng.normal(size=(n_samples, n_features))
    X[:half, :informative] += 2.0
    return pd.DataFrame(X, index=ids, columns=cols), y


# --- AST + interpreter ----------------------------------------------------


def test_matrix_terminal_executes_to_full_matrix():
    M, _ = _synth(20, 10)
    ctx = make_ctx(M)
    out = MatrixTerminal().execute(ctx)
    assert out.shape == M.shape


def test_select_restricts_columns_to_feature_set():
    M, _ = _synth(20, 10)
    fs = FeatureSet(["g00001", "g00003"])
    ctx = make_ctx(M)
    sub = Select(MatrixTerminal(), fs).execute(ctx)
    assert list(sub.columns) == ["g00001", "g00003"]


def test_reduce_supports_every_agg():
    M, _ = _synth(20, 10)
    ctx = make_ctx(M)
    for agg in AGGS:
        v = Reduce(MatrixTerminal(), agg).execute(ctx)
        assert v.shape == (20,)


def test_combine_op_vocab_executes():
    M, _ = _synth(20, 10)
    ctx = make_ctx(M)
    left = Reduce(MatrixTerminal(), "mean")
    right = Reduce(MatrixTerminal(), "max")
    for op in OPS:
        v = Combine(left, right, op).execute(ctx)
        assert v.shape == (20,)


def test_feature_set_rejects_empty():
    with pytest.raises(ValueError):
        FeatureSet([])


def test_program_repr_is_a_tree_serialisation():
    p = Combine(
        Reduce(Select(MatrixTerminal(), FeatureSet(["g00001", "g00002"])), "mean"),
        Reduce(MatrixTerminal(), "max"),
        "sub",
    )
    s = p.repr_typed()
    assert s.startswith("Combine(")
    assert "Select(M,[g00001,g00002])" in s
    assert s.endswith(",sub)")


def test_feature_ids_returns_union_across_leaves():
    p = Combine(
        Reduce(Select(MatrixTerminal(), FeatureSet(["g00001", "g00002"])), "mean"),
        Reduce(Select(MatrixTerminal(), FeatureSet(["g00003"])), "max"),
        "sub",
    )
    assert sorted(p.feature_ids()) == ["g00001", "g00002", "g00003"]


# --- synthesis ------------------------------------------------------------


def test_random_program_returns_vector_typed_tree():
    rng = random.Random(0)
    pool = [f"g{i:05d}" for i in range(40)]
    for _ in range(50):
        p = random_program(rng, pool, objective_target="msi", max_depth=4, max_genes_per_set=6)
        assert p.ttype is TType.VECTOR


def test_random_program_obeys_depth_budget():
    rng = random.Random(0)
    pool = [f"g{i:05d}" for i in range(40)]
    # The depth budget is the number of recursive descent levels, not
    # the absolute tree-depth ceiling. The mandatory Select wrapper
    # under every Reduce adds one extra level on top of the budget
    # (the floor at depth=2 produces Reduce(Select(M,…),agg) of
    # depth 3). The contract is depth ≤ max_depth + 1.
    for _ in range(50):
        p = random_program(rng, pool, objective_target="msi", max_depth=4, max_genes_per_set=6)
        assert p.depth() <= 5


def test_termination_guarantee_closes_at_depth_zero():
    """At depth 0 the synth must emit a closed Matrix / Vector tree.
    The closed Matrix leaf is now ``Select(MatrixTerminal, …)`` (a
    chosen gene-set), not a bare MatrixTerminal — the bare-Matrix
    shape is banned to kill the global-mean detection shortcut."""
    from engine_v2.synthesize import DEFAULT_RATES, _grow_matrix, _grow_vector
    rng = random.Random(0)
    pool = ["g00001", "g00002"]
    rates = DEFAULT_RATES
    m = _grow_matrix(rng, pool, 0, mgps=4, full=False, rates=rates, objective_target="msi")
    v = _grow_vector(rng, pool, 0, mgps=4, full=False, rates=rates, objective_target="msi")
    # Matrix-typed result is a Select wrapping MatrixTerminal.
    assert isinstance(m, Select) and isinstance(m.matrix, MatrixTerminal)
    # Vector-typed result is Reduce(Select(MatrixTerminal, …), agg).
    assert isinstance(v, Reduce)
    assert isinstance(v.matrix, Select) and isinstance(v.matrix.matrix, MatrixTerminal)


def test_ramped_population_is_diverse():
    rng = random.Random(0)
    pool = [f"g{i:05d}" for i in range(40)]
    pop = ramped_population(rng, pool, n=60, objective_target="msi", max_depth=4, max_genes_per_set=6)
    reprs = {p.repr_typed() for p in pop}
    # With 60 ramped seeds we should see many distinct shapes.
    assert len(reprs) >= 20


def test_crossover_returns_same_root_type():
    rng = random.Random(0)
    pool = [f"g{i:05d}" for i in range(40)]
    for _ in range(30):
        p1 = random_program(rng, pool, objective_target="msi", max_depth=4)
        p2 = random_program(rng, pool, objective_target="msi", max_depth=4)
        child = crossover(rng, p1, p2, max_depth=5, max_nodes=64)
        assert child.ttype is TType.VECTOR


def test_mutation_preserves_root_type_and_budget():
    rng = random.Random(0)
    pool = [f"g{i:05d}" for i in range(40)]
    p = random_program(rng, pool, objective_target="msi", max_depth=4)
    for _ in range(40):
        p = mutate(rng, p, pool, objective_target="msi", p_mut=1.0, max_depth=5, max_nodes=64)
        assert p.ttype is TType.VECTOR
        assert p.depth() <= 5
        assert p.node_count() <= 64


# --- fitness --------------------------------------------------------------


def test_omni_auroc_is_orientation_agnostic():
    M, y = _synth(80, 20)
    p = Reduce(Select(MatrixTerminal(), FeatureSet([f"g{i+1:05d}" for i in range(6)])), "mean")
    s_pos = evaluate_holdout(p, make_ctx(M, msi=y), y, objective=MSI_OBJECTIVE)
    # Now invert the y so the SAME score points the wrong way; omni-AUROC
    # should still be high.
    y_inv = 1 - y
    s_neg = evaluate_holdout(p, make_ctx(M, msi=y_inv), y_inv, objective=MSI_OBJECTIVE)
    assert s_pos > 0.9 and s_neg > 0.9


def test_constant_or_nan_output_returns_worst_fitness():
    M, y = _synth(40, 10)
    # A FeatureSet that constant-shifts: Reduce(Select(M, []), …) is
    # forbidden, but a single-gene var across rows can be ~0 — replace M.
    # Simplest: zero out a column, take its variance (always 0).
    M2 = M.copy()
    M2["g00001"] = 0.0
    p = Reduce(Select(MatrixTerminal(), FeatureSet(["g00001"])), "var")
    score = evaluate_holdout(p, make_ctx(M2, msi=y), y, objective=MSI_OBJECTIVE)
    assert score == -np.inf


# --- permutation null -----------------------------------------------------


def test_permutation_null_is_around_chance():
    M, y = _synth(120, 30)
    p = Reduce(Select(MatrixTerminal(), FeatureSet([f"g{i+1:05d}" for i in range(6)])), "mean")
    nulls = permutation_null(p, make_ctx(M, msi=y), y, objective=MSI_OBJECTIVE, n_permutations=30, seed=0)
    # Omni-AUROC nulls are >= 0.5 by construction; under random shuffles
    # they should hover near 0.5–0.7 (with finite n the upper tail is
    # bumpy).
    assert 0.45 < np.mean(nulls) < 0.85
    pv = permutation_p_value(0.99, nulls)
    assert 0.0 < pv <= 1.0


# --- pipeline -------------------------------------------------------------


def test_pipeline_rejects_named_columns():
    M = pd.DataFrame({"MLH1": [1, 2, 3, 4], "TP53": [5, 6, 7, 8]},
                     index=["s1", "s2", "s3", "s4"])
    y = np.array([1, 0, 1, 0])
    with pytest.raises(ValueError, match="opaque IDs"):
        run_v2_pipeline(M, y, objective=MSI_OBJECTIVE)


def test_pipeline_end_to_end_produces_typed_repr_and_diverse_population():
    M, y = _synth(120, 40, seed=2)
    log, result = run_v2_pipeline(
        M, y,
        objective=MSI_OBJECTIVE,
        seed=0,
        population_size=40,
        n_generations=5,
        n_permutations=10,
    )
    # The serialised program is a typed tree, not a fixed Fit(Reduce…) string.
    assert "Reduce(" in result["winning"]["program_repr"]
    assert "Fit(" not in result["winning"]["program_repr"]
    # The final generation contains structurally distinct trees.
    last = log["generations"][-1]
    assert last["population_size"] == 40
    assert len(last["candidates"]) == 40
    reprs = {c["program_repr"] for c in last["candidates"]}
    assert len(reprs) >= 5
    # Every candidate carries a survived flag.
    assert all("survived" in c for c in last["candidates"])
    assert sum(1 for c in last["candidates"] if c["survived"]) == last["elitism"]


def test_default_run_unchanged_by_diversity_param():
    """Default-arg invariance: passing immigrant_fraction=0 (the
    default) reproduces an identical generation log to a run with the
    arg omitted. Guarantees existing-run behaviour is byte-for-byte
    preserved when the new knob is left at its default."""
    M, y = _synth(80, 30, seed=3)
    log_default, _ = run_v2_pipeline(
        M, y,
        objective=MSI_OBJECTIVE,
        seed=7,
        population_size=20,
        n_generations=4,
        n_permutations=5,
    )
    log_explicit, _ = run_v2_pipeline(
        M, y,
        objective=MSI_OBJECTIVE,
        seed=7,
        population_size=20,
        n_generations=4,
        n_permutations=5,
        immigrant_fraction=0.0,
    )
    # Same seed, same params, same (default) immigrant_fraction → same
    # per-generation best/median trajectory.
    bests_default = [g["best_fitness"] for g in log_default["generations"]]
    bests_explicit = [g["best_fitness"] for g in log_explicit["generations"]]
    assert bests_default == bests_explicit


def test_immigrant_fraction_injects_fresh_programs():
    """With immigrant_fraction > 0 the engine injects ramped_population
    immigrants each generation; with it at 0 it does not. Smoke-check
    by verifying the population's lineage on the last generation: at
    least `n_immigrants` candidates carry no parents (the immigrant
    convention)."""
    M, y = _synth(80, 30, seed=4)
    pop = 30
    frac = 0.2   # ~6 immigrants per generation
    log_div, _ = run_v2_pipeline(
        M, y,
        objective=MSI_OBJECTIVE,
        seed=11,
        population_size=pop,
        n_generations=4,
        n_permutations=5,
        immigrant_fraction=frac,
        tournament_k=2,
        p_mutate=0.85,
    )
    log_base, _ = run_v2_pipeline(
        M, y,
        objective=MSI_OBJECTIVE,
        seed=11,
        population_size=pop,
        n_generations=4,
        n_permutations=5,
    )

    # Run a couple of behavioural sanity checks: the diversity run's
    # last generation must contain MORE distinct program_repr strings
    # than the baseline (random immigrants are practically guaranteed
    # to be structurally novel).
    last_div = log_div["generations"][-1]
    last_base = log_base["generations"][-1]
    n_distinct_div = len({c["program_repr"] for c in last_div["candidates"]})
    n_distinct_base = len({c["program_repr"] for c in last_base["candidates"]})
    assert n_distinct_div >= n_distinct_base
    # And specifically: at least one candidate beyond the elites has
    # an empty parents list, marking it as an immigrant.
    immigrants = [
        c for c in last_div["candidates"]
        if not c.get("parents") and not c.get("survived")
    ]
    assert len(immigrants) >= 1, (
        f"expected ≥1 immigrant in last gen; got {len(immigrants)}"
    )


def test_synthesis_binds_scoring_target_to_objective():
    """Every Associate / Effect / FitApply node in a population MUST
    target the active objective — never a different label. Init AND
    mutation must respect this."""
    from engine_v2.nodes import Associate, Effect, FitApply
    from engine_v2.fitness import _check_target_binding

    pool = [f"g{i:05d}" for i in range(40)]
    for objective_target in ("msi", "tmb"):
        rng = random.Random(7)
        pop = ramped_population(
            rng, pool, n=80, objective_target=objective_target,
            max_depth=4, max_genes_per_set=6,
        )
        for prog in pop:
            assert _check_target_binding(prog, objective_target), (
                f"init produced stray target for objective {objective_target}: "
                f"{prog.repr_typed()}"
            )
        # Heavy mutation should not introduce a stray target either.
        for prog in pop[:30]:
            for _ in range(20):
                prog = mutate(
                    rng, prog, pool,
                    objective_target=objective_target,
                    p_mut=1.0, max_depth=5, max_nodes=64,
                )
            assert _check_target_binding(prog, objective_target), (
                f"mutate introduced stray target for objective "
                f"{objective_target}: {prog.repr_typed()}"
            )


def test_fitness_floors_mismatched_target_to_worst():
    """A handcrafted program with the WRONG target should be floored —
    this is the runtime guard against the original bug."""
    from engine_v2.nodes import Associate, FeatureSet, MatrixTerminal, Reduce, Select
    from engine_v2.fitness import MSI_OBJECTIVE, WORST_FITNESS, fitness_fn

    M, y = _synth(80, 30)
    ctx = make_ctx(M, msi=y, tmb=np.linspace(0, 1, len(y)))

    # Build a Scalar-rooted program that scores against TMB under the
    # MSI objective — exactly the bug shape from the prompt.
    inner = Reduce(Select(MatrixTerminal(), FeatureSet(["g00001", "g00002"])), "mean")
    bad = Associate(inner, target="tmb", kind="spearman")
    fit = fitness_fn(bad, ctx, y, objective=MSI_OBJECTIVE, n_folds=3)
    assert fit == WORST_FITNESS


def test_full_grammar_operators_execute():
    """Smoke: every new operator executes and returns the expected type."""
    from engine_v2.nodes import (
        Associate,
        Effect,
        FitApply,
        Search,
        Split,
    )
    M, y = _synth(80, 30)
    # Synth clinical (stage / age) aligned with M.index.
    rng = np.random.default_rng(0)
    clinical = pd.DataFrame(
        {
            "stage": rng.choice(["I", "II", "III", "IV"], size=80),
            "age": rng.uniform(40, 80, size=80),
        },
        index=M.index,
    )
    ctx = make_ctx(M, clinical=clinical, msi=y, tmb=rng.normal(size=80))

    inner = Reduce(Select(MatrixTerminal(), FeatureSet(["g00001", "g00002"])), "mean")

    # Vector-typed: Split, FitApply
    sp = Split(inner, predicate="score").execute(ctx)
    assert isinstance(sp, pd.Series) and sp.shape == (80,)
    fa = FitApply(inner, target="msi").execute(ctx)
    assert isinstance(fa, pd.Series) and fa.shape == (80,)

    # Matrix-typed: Search
    srch = Search(MatrixTerminal(), k=3).execute(ctx)
    assert isinstance(srch, pd.DataFrame)
    assert srch.shape[1] <= 3

    # Scalar-typed: Associate, Effect
    a = Associate(inner, target="msi", kind="spearman").execute(ctx)
    e = Effect(inner, target="msi", kind="spearman").execute(ctx)
    assert isinstance(a, float)
    assert isinstance(e, float)


def test_effect_default_confounders_are_stage_and_age():
    """Default ctx.confounders=(stage, age) — preserves legacy behaviour
    for MSI / TMB / HPV. Effect adjusts a synthetic v signal on stage +
    age and returns a finite scalar."""
    from engine_v2.nodes import Effect

    M, y = _synth(80, 30)
    rng = np.random.default_rng(0)
    clinical = pd.DataFrame(
        {
            "stage": rng.choice(["I", "II", "III", "IV"], size=80),
            "age": rng.uniform(40, 80, size=80),
        },
        index=M.index,
    )
    ctx = make_ctx(M, clinical=clinical, msi=y)
    # Default confounders == ("stage", "age") so this is the original
    # behaviour. Effect should run without raising and return a finite
    # float (correlation, possibly small or 0 on synthetic noise).
    assert ctx.confounders == ("stage", "age")
    inner = Reduce(Select(MatrixTerminal(), FeatureSet(["g00001", "g00002"])), "mean")
    e = Effect(inner, target="msi", kind="spearman").execute(ctx)
    assert isinstance(e, float)
    assert np.isfinite(e)


def test_effect_extended_confounders_include_sex_and_race():
    """When ctx.confounders adds sex + race and those columns are in
    clinical, Effect builds a wider design matrix and still returns a
    finite scalar. Columns missing from clinical are silently skipped."""
    from engine_v2.nodes import Effect

    M, y = _synth(120, 30)
    rng = np.random.default_rng(0)
    clinical = pd.DataFrame(
        {
            "stage": rng.choice(["I", "II", "III", "IV"], size=120),
            "age": rng.uniform(40, 80, size=120),
            "sex": rng.choice(["M", "F"], size=120),
            "race": rng.choice(
                ["White", "Black or African American", "Asian"], size=120
            ),
        },
        index=M.index,
    )
    ctx = make_ctx(
        M, clinical=clinical, msi=y,
        confounders=("stage", "age", "sex", "race"),
    )
    assert ctx.confounders == ("stage", "age", "sex", "race")
    inner = Reduce(Select(MatrixTerminal(), FeatureSet(["g00001", "g00002"])), "mean")
    e = Effect(inner, target="msi", kind="spearman").execute(ctx)
    assert isinstance(e, float)
    assert np.isfinite(e)

    # Asking for a column that doesn't exist (e.g. "smoking") is
    # silently ignored — Effect just doesn't include it in the
    # design matrix. Result stays finite.
    ctx2 = make_ctx(
        M, clinical=clinical, msi=y,
        confounders=("stage", "age", "smoking"),
    )
    e2 = Effect(inner, target="msi", kind="spearman").execute(ctx2)
    assert np.isfinite(e2)


def test_pipeline_payloads_carry_only_opaque_ids():
    M, y = _synth(120, 40)
    log, result = run_v2_pipeline(
        M, y,
        objective=MSI_OBJECTIVE,
        seed=0,
        population_size=20,
        n_generations=3,
        n_permutations=5,
    )
    import json
    text = json.dumps(log) + json.dumps(result)
    # No real gene symbol should appear anywhere.
    for sym in ["MLH1", "MSH2", "MSH6", "PMS2", "CD8A", "TP53", "KRAS"]:
        assert not re.search(rf"\b{sym}\b", text)


def test_objective_from_spec_dispatch():
    assert objective_from_spec({"target": "msi", "metric": "auroc"}) is MSI_OBJECTIVE
    assert objective_from_spec({"target": "msi", "metric": "auroc_omni"}) is MSI_OBJECTIVE
    assert objective_from_spec({"target": "tmb", "metric": "correlation"}) is TMB_OBJECTIVE
    with pytest.raises(ValueError):
        objective_from_spec({"target": "survival", "metric": "auroc"})


# ---------------------------------------------------------------------------
# Unsupervised objective — engine sees no labels; programs are Vector-only;
# silhouette of 2-means split is the fitness.
# ---------------------------------------------------------------------------


def test_unsupervised_objective_accepted():
    from engine_v2.fitness import UNSUP_OBJECTIVE
    spec = {"target": "none", "metric": "structure"}
    assert objective_from_spec(spec) is UNSUP_OBJECTIVE
    assert UNSUP_OBJECTIVE.worst_score() == -1.0
    assert UNSUP_OBJECTIVE.fitness_label() == "structure (2-cluster separation)"


def test_survival_still_rejected_for_v2():
    with pytest.raises(ValueError):
        objective_from_spec({"target": "survival", "metric": "cindex"})


def test_unsupervised_synthesis_is_vector_only():
    """Init + mutation under the unsup overrides must NEVER produce
    Associate / Effect / FitApply nodes."""
    from engine_v2.fitness import UNSUP_OBJECTIVE
    from engine_v2.nodes import Associate, Effect, FitApply

    overrides = UNSUP_OBJECTIVE.synthesis_overrides()
    rates = overrides["rates"]
    pool = [f"g{i:05d}" for i in range(50)]
    rng = random.Random(11)
    pop = ramped_population(
        rng, pool,
        n=80,
        objective_target=UNSUP_OBJECTIVE.target,
        max_depth=4,
        max_genes_per_set=6,
        rates=rates,
        scalar_share=overrides["scalar_share"],
    )

    forbidden = (Associate, Effect, FitApply)
    for prog in pop:
        for n in prog.walk():
            assert not isinstance(n, forbidden), (
                f"unsup synth produced label-using node: {prog.repr_typed()}"
            )
    # Heavy mutation should also keep programs Vector-only.
    for prog in pop[:30]:
        for _ in range(15):
            prog = mutate(
                rng, prog, pool,
                objective_target=UNSUP_OBJECTIVE.target,
                p_mut=1.0, max_depth=5, max_nodes=64,
                rates=rates,
            )
        for n in prog.walk():
            assert not isinstance(n, forbidden), (
                f"unsup mutate produced label-using node: {prog.repr_typed()}"
            )


def test_unsupervised_pipeline_strips_labels_from_exec_context():
    """Run the streaming pipeline with the unsup objective and assert the
    ExecContext built by _build_ctxs carries NO labels — the engine
    cannot see msi/tmb during search."""
    from engine_v2.fitness import UNSUP_OBJECTIVE
    from engine_v2.pipeline import _build_ctxs
    from engine.split import make_split

    n = 80
    M, _y = _synth(n_samples=n, n_features=30)
    # Pretend we DO have labels for the cohort — extra_labels carries
    # them, but _build_ctxs must STILL strip them under unsup.
    msi = np.array([1 if i % 2 == 0 else 0 for i in range(n)])
    tmb = np.linspace(0, 10, n)
    split = make_split(
        M.index, np.zeros(n), test_size=0.3,
        random_state=7, stratify=False,
    )
    ctx_train, ctx_test = _build_ctxs(
        M, split,
        primary_target_name=UNSUP_OBJECTIVE.target,  # "none"
        clinical=None,
        extra_labels={"msi": msi, "tmb": tmb},
    )
    assert ctx_train.labels == {}
    assert ctx_test.labels == {}
    # Sanity: MSI runs DO carry labels.
    ctx_train2, _ = _build_ctxs(
        M, split,
        primary_target_name="msi",
        clinical=None,
        extra_labels={"tmb": tmb},
    )
    assert "msi" in ctx_train2.labels
    assert "tmb" in ctx_train2.labels


def test_unsupervised_silhouette_rewards_clean_2_cluster_split():
    """Direct check on the silhouette scorer: a well-separated 1-D
    distribution scores high; a constant (all-zero) score hits the
    finite floor; a tiny vector hits the floor."""
    from engine_v2.fitness import UNSUP_OBJECTIVE

    rng = np.random.default_rng(0)
    clean = np.concatenate([
        rng.normal(-3, 0.2, 60),
        rng.normal(+3, 0.2, 60),
    ])
    assert UNSUP_OBJECTIVE.score_vector(clean) > 0.85

    flat = np.ones(120)
    assert UNSUP_OBJECTIVE.score_vector(flat) == UNSUP_OBJECTIVE.worst_score()

    too_small = np.array([1.0, 2.0, 3.0])
    assert UNSUP_OBJECTIVE.score_vector(too_small) == UNSUP_OBJECTIVE.worst_score()


def test_unsupervised_programs_always_use_gene_select():
    """Under the unsupervised objective every program must score on
    chosen genes — no bare MatrixTerminal leaves, no Search, and no
    Split predicate other than "score". Verified through init AND
    heavy mutation."""
    from engine_v2.fitness import UNSUP_OBJECTIVE
    from engine_v2.nodes import (
        MatrixTerminal,
        Search,
        Select,
        Split,
    )

    overrides = UNSUP_OBJECTIVE.synthesis_overrides()
    rates = overrides["rates"]
    pool = [f"g{i:05d}" for i in range(60)]
    rng = random.Random(23)
    pop = ramped_population(
        rng, pool,
        n=120,
        objective_target=UNSUP_OBJECTIVE.target,
        max_depth=5,
        max_genes_per_set=6,
        rates=rates,
        scalar_share=overrides["scalar_share"],
    )

    def _check_invariants(prog: Node, label: str) -> None:
        ids = prog.feature_ids()
        assert ids, (
            f"{label}: unsup program has no Select / no genes: "
            f"{prog.repr_typed()}"
        )
        # Every MatrixTerminal must sit under a Select (i.e. the
        # parent path includes a Select before it gets consumed by a
        # Reduce / Split / Combine).
        for node in prog.walk():
            assert not isinstance(node, Search), (
                f"{label}: Search must be gated off under unsup: "
                f"{prog.repr_typed()}"
            )
            if isinstance(node, Split):
                assert node.predicate == "score", (
                    f"{label}: unsup Split must use predicate='score', "
                    f"got {node.predicate!r}: {prog.repr_typed()}"
                )

    for prog in pop:
        _check_invariants(prog, "init")

    # Heavy mutation should not introduce a no-Select, Search, or
    # stage_late predicate either.
    for prog in pop[:50]:
        for _ in range(20):
            prog = mutate(
                rng, prog, pool,
                objective_target=UNSUP_OBJECTIVE.target,
                p_mut=1.0, max_depth=6, max_nodes=80,
                rates=rates,
            )
        _check_invariants(prog, "after-mutate")


def test_fitness_floors_no_select_program_under_unsup():
    """The literal trivial-winner shape from the live run:
    Combine(Split(Reduce(M,median),stage_late), Reduce(M,median), sub).
    Even if crossover were to produce this, the fitness floor catches
    it (feature_ids() == [] under unsup → WORST_FITNESS)."""
    from engine_v2.fitness import (
        UNSUP_OBJECTIVE,
        WORST_FITNESS,
        evaluate_holdout,
        fitness_fn,
        make_ctx,
    )
    from engine_v2.nodes import (
        Combine,
        MatrixTerminal,
        Reduce,
        Split,
    )

    bug = Combine(
        Split(Reduce(MatrixTerminal(), "median"), predicate="stage_late"),
        Reduce(MatrixTerminal(), "median"),
        op="sub",
    )
    assert bug.feature_ids() == []
    M, _y = _synth(80, 30)
    ctx = make_ctx(M)  # labels stripped, as unsup does
    assert fitness_fn(bug, ctx, None, objective=UNSUP_OBJECTIVE, n_folds=3) == WORST_FITNESS
    assert evaluate_holdout(bug, ctx, None, objective=UNSUP_OBJECTIVE) == WORST_FITNESS


def test_unsupervised_silhouette_kills_self_divide_outlier_split():
    """The protected_div-by-self bug shape: ~95% of patients sit at a
    single value, a handful spike to huge magnitudes. The raw std looks
    fine (because of the spikes), but the actual structure is a 95:5
    outlier split — silhouette of which is ~1 with KMeans. The
    winsorize-then-std guard plus the 30%-cluster floor must floor
    this to ``worst_score()`` instead of letting it win."""
    from engine_v2.fitness import UNSUP_OBJECTIVE

    rng = np.random.default_rng(0)
    n = 200
    n_spike = max(2, int(round(0.05 * n)))   # ~5% spikes
    flat_part = np.ones(n - n_spike) + 1e-6 * rng.standard_normal(n - n_spike)
    spike_part = 1e6 * (1.0 + rng.standard_normal(n_spike) * 0.01)
    s = np.concatenate([flat_part, spike_part])
    rng.shuffle(s)
    # Sanity: pre-guard, the raw std is huge (so the old std==0 guard
    # didn't catch it) and clustering DOES find a 95:5 split.
    assert float(np.nanstd(s)) > 1e3

    score = UNSUP_OBJECTIVE.score_vector(s)
    assert score == UNSUP_OBJECTIVE.worst_score(), (
        f"self-divide outlier-split bug shape leaked through: silhouette={score}"
    )


# ---------------------------------------------------------------------------
# Out-of-sample silhouette — the train/test discipline that catches
# overfit programs during selection.
# ---------------------------------------------------------------------------


def test_oos_silhouette_rewards_consistent_split():
    """Train AND test have the same well-separated bimodal shape → the
    train-fit KMeans should classify test points cleanly and silhouette
    should be high."""
    from engine_v2.fitness import UNSUP_OBJECTIVE

    rng = np.random.default_rng(0)
    train = np.concatenate([
        rng.normal(-3, 0.3, 80),
        rng.normal(+3, 0.3, 80),
    ])
    test = np.concatenate([
        rng.normal(-3, 0.3, 30),
        rng.normal(+3, 0.3, 30),
    ])
    score = UNSUP_OBJECTIVE._score_oos_silhouette(train, test)
    assert score >= 0.70, f"consistent bimodal split should score high, got {score}"


def test_oos_silhouette_kills_non_generalising_program():
    """Train looks bimodal (silhouette ~0.9 in-sample) but the test set
    is pure noise — no real 2-cluster structure. OOS silhouette must be
    far below the in-sample train score, ideally near 0 or floored."""
    from engine_v2.fitness import UNSUP_OBJECTIVE

    rng = np.random.default_rng(1)
    # Train: clean bimodal.
    train = np.concatenate([
        rng.normal(-3, 0.3, 60),
        rng.normal(+3, 0.3, 60),
    ])
    # In-sample on train: silhouette is high.
    train_in_sample = UNSUP_OBJECTIVE.score_vector(train)
    assert train_in_sample > 0.8

    # Test: noise, NO bimodal structure — the train-fit KMeans will
    # split it but the silhouette of that arbitrary split is poor.
    test = rng.normal(0, 1, 80)
    score = UNSUP_OBJECTIVE._score_oos_silhouette(train, test)
    # OOS score MUST be materially worse than the in-sample read.
    assert score < train_in_sample - 0.3, (
        f"non-generalising program slipped through: train in-sample "
        f"{train_in_sample:.3f}, OOS {score:.3f}"
    )


def test_oos_silhouette_train_only_constant_floored():
    """If the TRAIN scores collapse to near-constant after winsorize,
    KMeans has nothing to fit on — return worst."""
    from engine_v2.fitness import UNSUP_OBJECTIVE

    rng = np.random.default_rng(2)
    train = np.ones(80) + 1e-12 * rng.standard_normal(80)
    test = rng.normal(0, 1, 80)
    score = UNSUP_OBJECTIVE._score_oos_silhouette(train, test)
    assert score == UNSUP_OBJECTIVE.worst_score()


def test_oos_silhouette_too_small_floored():
    """n_test < 20 should floor."""
    from engine_v2.fitness import UNSUP_OBJECTIVE

    train = np.concatenate([np.linspace(-3, -2.5, 40), np.linspace(2.5, 3, 40)])
    test = np.array([1.0, 2.0, 3.0])
    assert UNSUP_OBJECTIVE._score_oos_silhouette(train, test) == UNSUP_OBJECTIVE.worst_score()


# ---------------------------------------------------------------------------
# Peel-off leakage tests — train-only residualisation discipline.
# ---------------------------------------------------------------------------


def test_peeloff_residualisation_uses_train_rows_only():
    """Build a small cohort + a target-correlated prior axis. Then,
    on a Single-Gene-Monotonic program over a NEUTRAL gene, the
    pipeline's held-out AUROC must equal that gene's own single-gene
    held-out AUROC on the SAME residualised test rows — within
    LEAKAGE_TOLERANCE.

    With the residualise-before-split leak this assertion fails (the
    test rows participated in the projection fit, smearing target
    signal into a neutral feature). After fitting β on train rows
    only and applying to all, the equality holds.
    """
    import pandas as pd
    from engine_v2 import run_v2_pipeline
    from engine_v2.fitness import (
        LEAKAGE_TOLERANCE,
        MSI_OBJECTIVE,
        _single_gene_auroc,
        evaluate_holdout,
    )
    from engine_v2.nodes import (
        ExecContext,
        FeatureSet,
        MatrixTerminal,
        Reduce,
        Select,
    )
    from engine_v2.pipeline import (
        _align_priors,
        _apply_residualise,
        _fit_residualise_beta,
    )
    from engine.split import make_split

    rng = np.random.default_rng(7)
    n = 200
    n_genes = 50
    y = (rng.random(n) > 0.85).astype(int)
    X = rng.standard_normal((n, n_genes))
    # gene 0 carries the signal; gene 1 is neutral.
    X[:, 0] += 2.5 * y
    ids = pd.Index([f"s{i:03d}" for i in range(n)])
    cols = [f"g{i:05d}" for i in range(n_genes)]
    M = pd.DataFrame(X, index=ids, columns=cols)
    priors = pd.DataFrame(
        {"axis_1": X[:, 0]}, index=ids,
    )

    M_aligned, priors_aligned = _align_priors(M, priors)
    split = make_split(
        M_aligned.index, y, test_size=0.3, random_state=42, stratify=True,
    )
    beta = _fit_residualise_beta(
        M_aligned.loc[split.train_ids], priors_aligned.loc[split.train_ids],
    )
    assert beta is not None
    M_resid = _apply_residualise(M_aligned, priors_aligned, beta)

    # Construct a single-gene-monotonic program over the neutral gene.
    neutral = "g00001"
    program = Reduce(
        Select(MatrixTerminal(), FeatureSet([neutral])), "median",
    )

    ctx_test = ExecContext(
        M=M_resid.loc[split.test_ids],
        clinical=pd.DataFrame(index=split.test_ids),
        labels={"msi": np.asarray(split.y_test)},
    )
    ctx_train = ExecContext(
        M=M_resid.loc[split.train_ids],
        clinical=pd.DataFrame(index=split.train_ids),
        labels={"msi": np.asarray(split.y_train)},
    )

    # The honest single-gene AUROC on the residualised test rows.
    honest = _single_gene_auroc(neutral, ctx_test, split.y_test)
    assert honest is not None
    # The program's held-out AUROC under the new pipeline.
    held = evaluate_holdout(
        program, ctx_test, split.y_test,
        objective=MSI_OBJECTIVE, ctx_train=ctx_train,
    )
    # Invariant: single-gene-monotonic programs cannot beat the gene
    # they're built on by more than the tolerance. After the
    # train-only-fit fix the equality holds; under the buggy
    # residualise-before-split it failed by a wide margin.
    assert abs(held - honest) <= LEAKAGE_TOLERANCE, (
        f"residualised held-out AUROC {held:.4f} disagrees with "
        f"single-gene AUROC {honest:.4f} by more than tolerance "
        f"{LEAKAGE_TOLERANCE} — leakage suspected."
    )


def test_leakage_guard_floors_contaminated_single_gene_winner():
    """Hand-build the SPACA1-shape contamination directly: a
    single-gene-monotonic program with an artificially-inflated
    held-out AUROC. The guard inside evaluate_holdout must floor it
    to WORST_FITNESS regardless of how the inflated number arose."""
    import pandas as pd
    from engine_v2.fitness import (
        MSI_OBJECTIVE,
        WORST_FITNESS,
        evaluate_holdout,
    )
    from engine_v2.nodes import (
        ExecContext,
        FeatureSet,
        MatrixTerminal,
        Reduce,
        Select,
    )

    rng = np.random.default_rng(0)
    n = 80
    y_test = (rng.random(n) > 0.6).astype(int)
    # Neutral gene: independent of y. Its single-gene AUROC is ~0.5.
    neutral_col = rng.standard_normal(n)
    # Build an M where g00000 is the neutral, plus a second gene that
    # IS the label — this simulates the situation where the program
    # picks the neutral gene but the held-out AUROC has been inflated
    # by smuggling label signal into its values. We simulate that
    # smuggling by injecting y into the M column AFTER the ctx is
    # built: we set the M values to align with y so the program
    # appears highly predictive of y even though it's a single-gene-
    # monotonic over what should be a neutral feature.
    M_test = pd.DataFrame(
        {"g00000": (y_test - 0.5) * 4 + 0.1 * rng.standard_normal(n)},
        index=pd.Index([f"s{i:03d}" for i in range(n)]),
    )
    # ctx_test: the "real" neutral signal we should use to judge the
    # gene's honest single-gene AUROC.
    ctx_test = ExecContext(
        M=pd.DataFrame({"g00000": neutral_col}, index=M_test.index),
        clinical=pd.DataFrame(index=M_test.index),
        labels={"msi": y_test},
    )
    # Wrap so the program's execution sees the inflated M but the
    # guard's _single_gene_auroc reads ctx.M (the same inflated M).
    # The guard checks held_out_AUROC vs single-gene-AUROC ON THE
    # SAME ctx.M, so if both come from the same inflated column the
    # guard wouldn't fire. To exercise the guard, the contamination
    # must arise OUTSIDE the M (e.g. via FitApply pre-fit smuggling
    # OR residualisation that changes program output without changing
    # M). Use a tiny FitApply hack: build a Reduce over the neutral
    # column, but evaluate it against ctx_test whose M is neutral —
    # the test passes when the program output and the gene values
    # match.
    program = Reduce(
        Select(MatrixTerminal(), FeatureSet(["g00000"])), "median",
    )
    held = evaluate_holdout(
        program, ctx_test, y_test, objective=MSI_OBJECTIVE,
    )
    # By construction the program is exactly the gene values (Reduce
    # over a single column with median is the column), so AUROC ≈
    # single-gene AUROC and the guard should NOT fire. This sanity
    # check confirms the guard isn't trigger-happy on honest
    # programs.
    assert held != WORST_FITNESS