tests/test_engine.py

"""Smoke tests for the GP engine. Synthetic data only — no network, no parquet."""

from __future__ import annotations

import random

import numpy as np
import pandas as pd
import pytest

from engine.baseline import make_baseline
from engine.fitness import cv_auroc, fitness_fn, program_state
from engine.gp import run_gp
from engine.permutation import permutation_null, permutation_p_value
from engine.pipeline import run_gp_pipeline
from engine.prefilter import top_n_features
from engine.program import (
    MAX_SET_SIZE,
    Program,
    crossover,
    mutate,
    random_program,
)
from engine.split import make_split


# --- synthetic data --------------------------------------------------------


def _synthetic_matrix(
    n_samples: int = 80,
    n_features: int = 200,
    n_informative: int = 10,
    seed: int = 0,
) -> tuple[pd.DataFrame, np.ndarray]:
    rng = np.random.default_rng(seed)
    half = n_samples // 2
    y = np.array([1] * half + [0] * (n_samples - half))
    cols = [f"g{i+1:05d}" for i in range(n_features)]
    sample_ids = pd.Index([f"s{i}" for i in range(n_samples)], name="sample_id")
    X = rng.normal(loc=0.0, scale=1.0, size=(n_samples, n_features))
    # Inject signal in the first `n_informative` features.
    X[:half, :n_informative] += 2.5
    M = pd.DataFrame(X, index=sample_ids, columns=cols)
    return M, y


# --- prefilter --------------------------------------------------------------


def test_prefilter_surfaces_the_informative_features():
    M, y = _synthetic_matrix(n_samples=100, n_features=200, n_informative=10)
    top, scores = top_n_features(M, y, n=20)
    informative = {f"g{i+1:05d}" for i in range(10)}
    overlap = informative & set(top)
    assert len(overlap) >= 8   # the prefilter should grab most/all of them


# --- program operators -----------------------------------------------------


def test_random_program_well_formed():
    rng = random.Random(0)
    pool = [f"g{i:05d}" for i in range(50)]
    p = random_program(rng, pool)
    assert 1 <= len(p.feature_sets) <= 2
    for fs in p.feature_sets:
        assert 2 <= len(fs) <= MAX_SET_SIZE
        assert all(g in pool for g in fs)
    # Gene IDs should be unique across sets.
    assert len(p.gene_ids) == len(set(p.gene_ids))


def test_mutate_preserves_invariants():
    rng = random.Random(1)
    pool = [f"g{i:05d}" for i in range(50)]
    p = random_program(rng, pool)
    for _ in range(20):
        p = mutate(rng, p, pool, p_mut=1.0)
    assert 1 <= len(p.feature_sets) <= 2
    for fs in p.feature_sets:
        assert 2 <= len(fs) <= MAX_SET_SIZE
    assert len(p.gene_ids) == len(set(p.gene_ids))


def test_crossover_produces_valid_child():
    rng = random.Random(2)
    pool = [f"g{i:05d}" for i in range(50)]
    p1 = random_program(rng, pool)
    p2 = random_program(rng, pool)
    child = crossover(rng, p1, p2)
    assert child.feature_sets
    assert len(child.gene_ids) == len(set(child.gene_ids))
    assert child.parents == [p1.program_id, p2.program_id]


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


def test_program_state_shape():
    M, _ = _synthetic_matrix(n_samples=20, n_features=10)
    p = Program(feature_sets=[["g00001", "g00002"], ["g00005", "g00006", "g00007"]])
    state = program_state(M, p)
    assert state.shape == (20, 2)


def test_cv_auroc_strong_for_informative_program():
    M, y = _synthetic_matrix(n_samples=120, n_features=50, n_informative=8)
    p = Program(feature_sets=[[f"g{i+1:05d}" for i in range(8)]])
    score = cv_auroc(M, y, p, n_folds=5, random_state=0)
    assert score > 0.95


def test_fitness_penalises_size():
    M, y = _synthetic_matrix(n_samples=120, n_features=50, n_informative=8)
    small = Program(feature_sets=[[f"g{i+1:05d}" for i in range(2)]])
    big = Program(feature_sets=[[f"g{i+1:05d}" for i in range(8)]])
    small_fit = fitness_fn(M, y, small, lambda_size=0.01)
    big_fit = fitness_fn(M, y, big, lambda_size=0.01)
    # both should be high; size penalty makes big at least slightly closer.
    assert (big_fit + 0.01 * (big.n_genes - small.n_genes)) >= small_fit - 0.05


# --- GP loop ---------------------------------------------------------------


def test_run_gp_improves_or_holds_fitness_over_generations():
    M, y = _synthetic_matrix(n_samples=120, n_features=80, n_informative=10)
    pool, _ = top_n_features(M, y, n=40)
    log, winner, fit = run_gp(
        M, y, pool,
        population_size=20, n_generations=8, seed=0,
    )
    assert len(log) == 8
    assert log[-1]["best_fitness"] >= log[0]["best_fitness"] - 1e-9


# --- baseline / permutation ------------------------------------------------


def test_baseline_uses_first_k():
    shortlist = [f"g{i:05d}" for i in range(20)]
    b = make_baseline(shortlist, k=5)
    assert b.feature_sets == [shortlist[:5]]
    assert b.program_id == "baseline"


def test_permutation_null_is_around_chance():
    M, y = _synthetic_matrix(n_samples=120, n_features=100, n_informative=10)
    split = make_split(M.index, y, test_size=0.3, random_state=0)
    nulls = permutation_null(
        M.loc[split.train_ids], split.y_train,
        M.loc[split.test_ids], split.y_test,
        n_permutations=20, seed=0,
    )
    assert 0.3 < np.mean(nulls) < 0.7
    p = permutation_p_value(0.99, nulls)
    assert 0.0 < p <= 1.0


# --- pipeline / airgap enforcement -----------------------------------------


def test_pipeline_rejects_named_columns():
    M = pd.DataFrame(
        {"MLH1": [1.0, 2.0, 3.0, 4.0], "TP53": [5.0, 6.0, 7.0, 8.0]},
        index=["s1", "s2", "s3", "s4"],
    )
    y = np.array([1, 0, 1, 0])
    with pytest.raises(ValueError, match="opaque IDs"):
        run_gp_pipeline(M, y)


def test_pipeline_end_to_end_on_synthetic():
    M, y = _synthetic_matrix(n_samples=120, n_features=80, n_informative=10)
    evolution_log, result = run_gp_pipeline(
        M, y,
        seed=0,
        prefilter_n=30,
        population_size=15,
        n_generations=4,
        n_permutations=10,
    )
    assert "winning" in result
    assert result["winning"]["holdout_auroc"] > 0.8
    assert 0.0 < result["winning"]["permutation_p"] <= 1.0
    assert len(evolution_log["generations"]) == 4
    # Top-of-rank candidates carry the survived flag.
    cand0 = evolution_log["generations"][0]["candidates"][0]
    assert cand0["survived"] is True


# --- objective dispatch + correlation path --------------------------------


def _synthetic_continuous(
    n_samples: int = 120,
    n_features: int = 100,
    n_informative: int = 8,
    seed: int = 0,
) -> tuple[pd.DataFrame, np.ndarray]:
    """Build a name-blind matrix where the first `n_informative` features
    NEGATIVELY track a continuous y (good test for CorrelationObjective
    with direction='neg')."""
    rng = np.random.default_rng(seed)
    cols = [f"g{i+1:05d}" for i in range(n_features)]
    sample_ids = pd.Index([f"s{i}" for i in range(n_samples)], name="sample_id")
    y = rng.normal(loc=10.0, scale=3.0, size=n_samples)
    X = rng.normal(loc=5.0, scale=1.0, size=(n_samples, n_features))
    # Negatively correlated informative cols: feature = -y + noise
    for j in range(n_informative):
        X[:, j] = -0.8 * y + rng.normal(scale=0.5, size=n_samples)
    M = pd.DataFrame(X, index=sample_ids, columns=cols)
    return M, y


def test_correlation_objective_prefilter_finds_informative():
    from engine.objectives import CorrelationObjective
    from engine.prefilter import top_n_features

    M, y = _synthetic_continuous(n_samples=150, n_features=200, n_informative=10)
    obj = CorrelationObjective(direction="neg")
    top, scores = top_n_features(M, y, n=20, objective=obj)
    informative = {f"g{i+1:05d}" for i in range(10)}
    assert len(informative & set(top)) >= 8


def test_correlation_objective_pipeline_end_to_end():
    from engine.objectives import CorrelationObjective

    M, y = _synthetic_continuous(n_samples=180, n_features=120, n_informative=10)
    evolution_log, result = run_gp_pipeline(
        M, y,
        objective=CorrelationObjective(direction="neg"),
        seed=0,
        prefilter_n=40,
        population_size=20,
        n_generations=4,
        n_permutations=10,
    )
    # Direction-adjusted score should be solidly positive on this signal.
    assert result["winning"]["holdout_score"] > 0.5
    assert result["objective_spec"]["target"] == "tmb"
    assert result["objective_spec"]["direction"] == "neg"
    assert evolution_log["run"]["fitness_label"] == "|spearman|"


def test_pipeline_unknown_objective_rejected_by_spec_builder():
    from engine.objectives import objective_from_spec
    with pytest.raises(ValueError):
        objective_from_spec({"target": "survival", "metric": "auroc"})


def test_pipeline_no_prefilter_samples_full_pool():
    """With prefilter_n=None the GP must be able to reach genes a tight
    prefilter would have excluded."""
    from engine.gp import run_gp
    from engine.objectives import BinaryAUROCObjective

    # 60 features; the informative ones are at indices 50..59 — they would
    # never make a top-10 univariate prefilter because of how we build the
    # synthetic noise, but they MUST be reachable when prefilter is off.
    rng = np.random.default_rng(2)
    n_samples = 80
    n_features = 60
    cols = [f"g{i+1:05d}" for i in range(n_features)]
    sample_ids = pd.Index([f"s{i}" for i in range(n_samples)], name="sample_id")
    y = np.array([1] * 40 + [0] * 40)
    X = rng.normal(loc=0.0, scale=1.0, size=(n_samples, n_features))
    # Mild signal in late features only — visible enough to be picked when
    # the GP can reach them, but no single one would dominate a prefilter.
    for j in range(50, 60):
        X[:40, j] += 0.6
    M = pd.DataFrame(X, index=sample_ids, columns=cols)

    log_off, result_off = run_gp_pipeline(
        M, y,
        seed=0,
        prefilter_n=None,
        population_size=30,
        n_generations=6,
        n_permutations=5,
    )
    assert log_off["run"]["prefilter_N"] is None
    # The GP saw the full pool (n_features) — confirmed by the n_genes
    # field in the run header.
    assert log_off["run"]["n_genes"] == n_features
    # And it can pick from late features; check at least one late feature
    # appears anywhere in the top-K candidates over the run.
    late = {f"g{i+1:05d}" for i in range(50, 60)}
    seen = {g for gen in log_off["generations"]
            for c in gen["candidates"] for g in c["gene_ids"]}
    assert late & seen, (
        "GP with prefilter off never sampled any late-index feature — "
        "init/mutation are not drawing from the full pool"
    )


def test_streaming_pipeline_emits_one_event_per_generation():
    from engine.pipeline import run_gp_pipeline_streaming

    M, y = _synthetic_matrix(n_samples=120, n_features=60, n_informative=8)
    received: list[dict] = []
    result = run_gp_pipeline_streaming(
        M, y,
        on_generation=lambda e: received.append(e),
        seed=0,
        prefilter_n=20,
        population_size=12,
        n_generations=3,
        n_permutations=5,
    )
    assert [e["generation"] for e in received] == [0, 1, 2]
    # Every event has the lab-shape candidates.
    for e in received:
        assert "best_fitness" in e and "median_fitness" in e
        assert all("survived" in c for c in e["candidates"])
    # The final result is the same shape as the batch path.
    assert "winning" in result and "permutation_p" in result["winning"]