scripts/capability_ratio_test.py

"""Capability test — can the blind engine discover a planted gene ratio?

Self-contained known-answer capability check. Plants a synthetic binary
label defined by the balance between TWO real, positively-correlated
genes (a log-ratio direction) — rigged so NEITHER gene helps on its
own — and asks the engine to rediscover it blind. If the engine's
winner uses exactly those two genes, combined multiplicatively /
divisively, we've shown the DSL can compose real two-gene interactions.

Isolation
---------
This is ONE file under scripts/. It:
- READS the existing processed colorectal matrix via ``dsl.Load``.
- ANONYMISES a small panel via ``airgap.anonymise`` (reuses the
  existing sealed map; no re-seal).
- Runs ``engine_v2.run_v2_pipeline`` (never modified).
- Reveals the winner's genes ONCE at the end via ``airgap.reveal``
  — bounded to the winner's opaque IDs (same discipline /evaluate uses).
It writes NOTHING to disk, adds no dataset, no API route, no UI, and
does not touch ``engine_v2/``, ``dsl/``, ``api/``, or ``web/``. Deleting
this file leaves zero trace.

Why a small panel
-----------------
Two individually-uninformative genes give the search NO univariate
gradient, so finding them among ~20,000 columns is an impossible
needle-hunt — that tests search-at-scale, not composition. To test
COMPOSITION fairly, we restrict the search to a tiny controlled panel:
the two planted genes + ~50 random decoys. Restricting the panel and
forcing Vector programs is fair framing (the engine still has to find
WHICH two genes among the panel AND HOW to combine them), not a hint
that hands over the answer.

Run:
    python -m scripts.capability_ratio_test
    python -m scripts.capability_ratio_test --pair GENE_A GENE_B
    python -m scripts.capability_ratio_test --seeds 1 3 7 --pop 200 --gens 40
"""

from __future__ import annotations

import argparse
import sys
from dataclasses import dataclass
from typing import Iterable, Sequence

import numpy as np
import pandas as pd

from airgap import anonymise, reveal
from dsl import Load
from engine_v2 import run_v2_pipeline
from engine_v2.fitness import V2Objective


# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------


def _zscore(x: np.ndarray) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    mu = float(np.mean(x))
    sd = float(np.std(x, ddof=0))
    if sd == 0.0 or not np.isfinite(sd):
        return np.zeros_like(x)
    return (x - mu) / sd


def _omni_auroc(score: np.ndarray, y: np.ndarray) -> float:
    """Orientation-agnostic AUROC of a 1-D score vs a binary label.
    Rank-sum formula; returns 0.5 for degenerate inputs."""
    s = np.asarray(score, dtype=float)
    yy = np.asarray(y, dtype=int)
    finite = np.isfinite(s)
    if int(finite.sum()) < 2:
        return 0.5
    s = s[finite]
    yy = yy[finite]
    n_pos = int((yy == 1).sum())
    n_neg = int((yy == 0).sum())
    if n_pos == 0 or n_neg == 0:
        return 0.5
    ranks = pd.Series(s).rank(method="average").to_numpy()
    S_pos = ranks[yy == 1].sum()
    auroc = (S_pos - n_pos * (n_pos + 1) / 2.0) / (n_pos * n_neg)
    return float(max(auroc, 1.0 - auroc))


def _find_pair(
    expr: pd.DataFrame,
    *,
    min_r: float = 0.85,
    single_gene_band: float = 0.07,
    ratio_min: float = 0.9,
    rng: np.random.Generator,
    n_seed_genes: int = 200,
    top_expressed_frac: float = 0.6,
    var_percentile: int = 50,
    max_r: float = 0.99,
) -> tuple[str, str, dict] | None:
    """Find a positively-correlated pair (A, B) that satisfies all the
    setup constraints:
      (a) Pearson r in [min_r, max_r] — enough shared variance that
          zA − zB cancels most of the single-gene information (near-
          perfect r is degenerate and yields no signal at all);
      (b) single-gene AUROC(A) and AUROC(B) inside 0.5 ± single_gene_band
          against the planted label (individually ~useless);
      (c) AUROC(zA − zB) ≥ ratio_min (the ratio is strong).

    Strategy: pick a random "seed" gene from a variance-filtered pool,
    rank the rest of the pool by |Pearson r| against it, and check the
    top few candidates. Random-pair search is too slow because pairs
    with r ≥ 0.85 are rare in a 20k-gene matrix.
    """
    means = expr.mean(axis=0)
    stds = expr.std(axis=0, ddof=0)
    keep = (
        (means >= means.quantile(1.0 - top_expressed_frac))
        & (stds >= stds.quantile(var_percentile / 100.0))
    )
    pool = list(expr.columns[keep])
    if len(pool) < 20:
        return None
    pool_arr = expr[pool].to_numpy(dtype=float)
    # Z-score each column once; Pearson r between two z-scored cols is
    # just their normalised dot product / n.
    pool_z = (pool_arr - pool_arr.mean(axis=0)) / np.where(
        pool_arr.std(axis=0, ddof=0) > 0, pool_arr.std(axis=0, ddof=0), 1.0,
    )
    n_samples = pool_arr.shape[0]

    seed_ixs = list(range(len(pool)))
    rng.shuffle(seed_ixs)
    n_seed = min(int(n_seed_genes), len(seed_ixs))
    for si in seed_ixs[:n_seed]:
        zs = pool_z[:, si]
        corrs = (pool_z.T @ zs) / n_samples
        # Rank by descending correlation (positive first).
        ranked = np.argsort(-corrs)
        for pi in ranked:
            if int(pi) == si:
                continue
            r = float(corrs[int(pi)])
            if not np.isfinite(r) or r < min_r or r > max_r:
                # ranked is sorted desc; once we drop below min_r no
                # further candidate for this seed can qualify.
                if r < min_r:
                    break
                continue
            A = pool[si]
            B = pool[int(pi)]
            xa = expr[A].to_numpy(dtype=float)
            xb = expr[B].to_numpy(dtype=float)
            zA = _zscore(xa)
            zB = _zscore(xb)
            signal = zA - zB
            med = float(np.median(signal))
            y = (signal > med).astype(int)
            auroc_A = _omni_auroc(xa, y)
            auroc_B = _omni_auroc(xb, y)
            if abs(auroc_A - 0.5) > single_gene_band:
                continue
            if abs(auroc_B - 0.5) > single_gene_band:
                continue
            ratio_auroc = _omni_auroc(signal, y)
            if ratio_auroc < ratio_min:
                continue
            return (
                A, B,
                {
                    "r": r,
                    "auroc_A": auroc_A,
                    "auroc_B": auroc_B,
                    "ratio_auroc": ratio_auroc,
                    "median_signal": med,
                    "n_positive": int(y.sum()),
                    "n_negative": int((1 - y).sum()),
                },
            )
    return None


def _build_panel_columns(
    expr: pd.DataFrame,
    *,
    keep_pair: tuple[str, str],
    n_decoys: int,
    rng: np.random.Generator,
) -> list[str]:
    """Return {A, B} ∪ a random sample of n_decoys other genes."""
    other = [c for c in expr.columns if c not in keep_pair]
    if len(other) <= n_decoys:
        decoys = other
    else:
        idx = rng.choice(len(other), size=n_decoys, replace=False)
        decoys = [other[int(i)] for i in idx]
    # Deterministic column order so anonymise assigns stable opaque IDs
    # across seeds within a single run.
    return sorted(list(keep_pair) + list(decoys))


def _panel_controls(panel_expr: pd.DataFrame, y: np.ndarray) -> dict:
    """Print-ready control diagnostics on the panel: best single-gene
    AUROC and the AUROC of the plain mean across all panel genes.
    Both should sit near 0.5 for the composition claim to be meaningful."""
    aurocs = {
        col: _omni_auroc(panel_expr[col].to_numpy(dtype=float), y)
        for col in panel_expr.columns
    }
    best_col = max(aurocs, key=lambda k: aurocs[k])
    best = aurocs[best_col]
    mean_score = panel_expr.mean(axis=1).to_numpy(dtype=float)
    mean_auroc = _omni_auroc(mean_score, y)
    return {
        "aurocs": aurocs,
        "best_single_col": best_col,
        "best_single_auroc": float(best),
        "panel_mean_auroc": float(mean_auroc),
    }


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------


@dataclass
class SeedResult:
    seed: int
    holdout: float
    winner_gene_ids: list[str]
    winner_symbols: list[str]
    program_repr: str
    passed_genes: bool
    passed_score: bool
    passed_interaction: bool
    interaction_ops: list[str]


_INTERACTION_TOKENS = ("protected_div", "mul", "sub")


def _has_interaction(repr_str: str) -> tuple[bool, list[str]]:
    ops = [t for t in _INTERACTION_TOKENS if t in repr_str]
    combined = "Combine(" in repr_str
    return (combined and len(ops) > 0, ops)


def main(argv: Sequence[str] | None = None) -> int:
    ap = argparse.ArgumentParser(
        description=(
            "Blind capability test: plant a two-gene ratio, run engine_v2 on "
            "a small panel, and check whether the winner recovers the ratio."
        ),
    )
    ap.add_argument(
        "--pair", nargs=2, metavar=("GENE_A", "GENE_B"), default=None,
        help="Pin the planted pair (reproducibility); otherwise auto-picked.",
    )
    ap.add_argument("--seeds", type=int, nargs="+", default=[1, 3, 7])
    ap.add_argument("--pop", type=int, default=200)
    ap.add_argument("--gens", type=int, default=40)
    ap.add_argument("--perms", type=int, default=100)
    ap.add_argument("--n-decoys", type=int, default=50)
    ap.add_argument("--pair-seed", type=int, default=0,
                    help="RNG seed for the pair search + decoy sampling.")
    ap.add_argument("--min-r", type=float, default=0.85,
                    help="Minimum Pearson r for the planted pair.")
    ap.add_argument("--single-band", type=float, default=0.07,
                    help="Allowed |single-gene AUROC − 0.5| for the pair.")
    ap.add_argument("--ratio-min", type=float, default=0.9,
                    help="Minimum ratio AUROC before the pair counts.")
    ap.add_argument("--pass-score", type=float, default=0.80,
                    help="Held-out AUROC required to pass the score check.")
    args = ap.parse_args(argv or sys.argv[1:])

    print("=" * 78)
    print("Capability test — can the engine discover a planted two-gene RATIO?")
    print("=" * 78)
    print(
        "This script is a known-answer engine unit test. It plants a "
        "synthetic\ntarget defined by the balance between two REAL, "
        "positively-correlated\ngenes and asks the engine to rediscover the "
        "interaction BLIND (on a\nsmall controlled panel — the two planted "
        "genes + a handful of decoys)."
    )
    print("-" * 78)

    # --- 1. Load NAMED expression matrix ----------------------------------
    cohort = Load("processed")
    expr = cohort.expression.dropna(axis=0, how="any")
    # Keep only samples with fully complete expression across all genes.
    n_samples = int(expr.shape[0])
    n_genes = int(expr.shape[1])
    print(f"Loaded processed matrix : {n_samples} samples × {n_genes} genes")

    rng = np.random.default_rng(int(args.pair_seed))

    # --- 2/3. Pick correlated pair + build planted target -----------------
    if args.pair:
        A, B = args.pair
        missing = [g for g in (A, B) if g not in expr.columns]
        if missing:
            print(f"\nABORT: --pair genes not in the matrix: {missing}")
            return 1
        xa = expr[A].to_numpy(dtype=float)
        xb = expr[B].to_numpy(dtype=float)
        r = float(np.corrcoef(xa, xb)[0, 1])
        zA = _zscore(xa)
        zB = _zscore(xb)
        signal = zA - zB
        med = float(np.median(signal))
        y = (signal > med).astype(int)
        stats = {
            "r": r,
            "auroc_A": _omni_auroc(xa, y),
            "auroc_B": _omni_auroc(xb, y),
            "ratio_auroc": _omni_auroc(signal, y),
            "median_signal": med,
            "n_positive": int(y.sum()),
            "n_negative": int((1 - y).sum()),
        }
        # Loud warning if the user's pair doesn't satisfy the meaningful-
        # setup constraints, but still run — the user asked for it.
        if abs(stats["auroc_A"] - 0.5) > args.single_band \
           or abs(stats["auroc_B"] - 0.5) > args.single_band:
            print(
                f"\nWARN: single-gene AUROCs (A={stats['auroc_A']:.3f}, "
                f"B={stats['auroc_B']:.3f}) are outside 0.5 ± "
                f"{args.single_band:.2f}. The pair isn't individually-"
                f"uninformative; the composition claim will be weakened."
            )
        if stats["ratio_auroc"] < args.ratio_min:
            print(
                f"\nWARN: ratio AUROC {stats['ratio_auroc']:.3f} < "
                f"{args.ratio_min:.2f}. Weak planted signal — the engine may "
                f"legitimately fail to find it."
            )
    else:
        found = _find_pair(
            expr,
            min_r=args.min_r,
            single_gene_band=args.single_band,
            ratio_min=args.ratio_min,
            rng=rng,
        )
        if found is None:
            print(
                f"\nABORT: could not find a pair with r >= {args.min_r}, "
                f"single-gene AUROCs within 0.5 ± {args.single_band}, and "
                f"ratio AUROC >= {args.ratio_min} in the search budget. "
                f"Try --min-r 0.5, --single-band 0.09, or --pair-seed."
            )
            return 1
        A, B, stats = found
        xa = expr[A].to_numpy(dtype=float)
        xb = expr[B].to_numpy(dtype=float)
        zA = _zscore(xa)
        zB = _zscore(xb)
        signal = zA - zB
        med = stats["median_signal"]
        y = (signal > med).astype(int)

    print("\n" + "-" * 78)
    print("PLANTED PAIR")
    print("-" * 78)
    print(f"  A                 : {A}")
    print(f"  B                 : {B}")
    print(f"  Pearson r(A, B)   : {stats['r']:+.3f}   (correlated ↑)")
    print(f"  AUROC(A) alone    : {stats['auroc_A']:.3f}   (individually ~0.5)")
    print(f"  AUROC(B) alone    : {stats['auroc_B']:.3f}   (individually ~0.5)")
    print(f"  AUROC(zA − zB)    : {stats['ratio_auroc']:.3f}   (planted ratio ≫)")
    print(f"  y distribution    : {stats['n_positive']} positive / "
          f"{stats['n_negative']} negative")

    # --- 4/5. Build panel + control diagnostics ---------------------------
    panel_cols = _build_panel_columns(
        expr,
        keep_pair=(A, B),
        n_decoys=int(args.n_decoys),
        rng=rng,
    )
    panel_expr = expr[panel_cols].copy()
    controls = _panel_controls(panel_expr, y)
    print("\n" + "-" * 78)
    print("PANEL CONTROLS")
    print("-" * 78)
    print(
        f"  Panel size          : {len(panel_cols)} genes "
        f"(A + B + {len(panel_cols) - 2} random decoys)"
    )
    print(
        f"  Best single AUROC   : {controls['best_single_auroc']:.3f}   "
        f"(nothing wins alone; column: {controls['best_single_col']})"
    )
    print(
        f"  Panel-mean AUROC    : {controls['panel_mean_auroc']:.3f}   "
        f"(averaging doesn't work)"
    )
    print(
        "  → Any high held-out score MUST come from composition, not "
        "single-gene or plain-mean signal."
    )

    # --- 6. Anonymise + run blind for each seed ---------------------------
    # Anonymise reuses the existing sealed map (no re-seal). The pipeline
    # sees only opaque IDs.
    M = anonymise(panel_expr)
    y_arr = np.asarray(y, dtype=int)

    # Pinned known-answer opaque IDs so the grader can compare without
    # revealing anything upstream. reveal() is called ONCE at the end
    # per seed and only for the winner's IDs — bounded.
    sym_to_id = {sym: cid for cid, sym in zip(M.columns, panel_cols)
                 if sym in (A, B)}
    # M's columns are already sorted; map A/B via panel_cols → M.columns.
    # Rebuild the mapping carefully: anonymise sorts by opaque ID after
    # renaming, so the sample-order-preserving mapping is symbol → id.
    _sealed_symbols = list(panel_cols)  # sorted
    # We don't need the mapping to run the engine — only to verify the
    # answer. Reveal handles that at the end.

    objective = V2Objective(target="msi", binary=True)

    print("\n" + "-" * 78)
    print("BLIND RUNS")
    print("-" * 78)
    print(
        f"  Seeds               : {list(args.seeds)}\n"
        f"  Population × gens   : {args.pop} × {args.gens}\n"
        f"  scalar_share        : 0.00 (force Vector programs)\n"
        f"  prefilter           : off (small panel; must not univariate-filter)\n"
        f"  coherence           : off (composition, not co-expression)\n"
        f"  diversity           : ON (k=2, p_mutate=0.85, immigrants 10%)\n"
        f"  Objective           : binary AUROC "
        f"(target='msi' as the carrier for the planted label)"
    )
    print(
        "  Airgap              : the engine sees only opaque IDs; reveal is "
        "bounded to the WINNER's IDs, once per seed, at the end."
    )

    results: list[SeedResult] = []
    for seed in args.seeds:
        print(f"\n[seed={seed}] running …", flush=True)
        _log, result = run_v2_pipeline(
            M, y_arr,
            objective=objective,
            seed=int(seed),
            prefilter_n=None,
            population_size=int(args.pop),
            n_generations=int(args.gens),
            n_permutations=int(args.perms),
            tournament_k=2,
            p_mutate=0.85,
            immigrant_fraction=0.10,
            coherence_weight=0.0,
            scalar_share_override=0.0,
        )
        winning = result.get("winning", {}) or {}
        winner_ids = list(winning.get("gene_ids") or [])
        winner_syms = reveal(winner_ids) if winner_ids else []
        program_repr = str(winning.get("program_repr") or "")
        holdout = float(winning.get("holdout_score") or 0.0)
        genes_ok = set(winner_syms) >= {A, B}
        score_ok = holdout >= float(args.pass_score)
        inter_ok, ops = _has_interaction(program_repr)
        passed = genes_ok and score_ok and inter_ok
        results.append(SeedResult(
            seed=int(seed),
            holdout=holdout,
            winner_gene_ids=winner_ids,
            winner_symbols=winner_syms,
            program_repr=program_repr,
            passed_genes=genes_ok,
            passed_score=score_ok,
            passed_interaction=inter_ok,
            interaction_ops=ops,
        ))
        print(f"  program            : {program_repr}")
        print(f"  revealed genes     : {winner_syms}")
        print(
            f"  held-out AUROC     : {holdout:.3f}   "
            f"(pass ≥ {args.pass_score:.2f}: {'YES' if score_ok else 'no'})"
        )
        print(
            f"  gene check         : winner ⊇ {{A, B}}? "
            f"{'YES' if genes_ok else 'no'} — got {sorted(set(winner_syms))}"
        )
        print(
            f"  interaction check  : Combine + {'/'.join(_INTERACTION_TOKENS)}? "
            f"{'YES' if inter_ok else 'no'}   ops seen: {ops}"
        )
        print(f"  seed verdict       : {'PASS' if passed else 'FAIL'}")

    # --- 7. Overall verdict ----------------------------------------------
    n_pass = sum(1 for r in results if r.passed_genes and r.passed_score and r.passed_interaction)
    n_total = len(results)
    print("\n" + "=" * 78)
    print(f"OVERALL: recovered the interaction in {n_pass} / {n_total} seeds")
    print("=" * 78)
    if n_total > 0 and n_pass >= max(1, (n_total + 1) // 2):
        print(
            "PASS. The engine can discover a genuine two-gene interaction "
            "blind — so\nwhen a real target (HPV) yields only averages, that's "
            "because the biology\ndoesn't need a ratio, not because the engine "
            "can't build one."
        )
        return 0
    else:
        print(
            "FAIL. The engine did NOT recover the planted interaction even on "
            "a small\npanel — a real limitation to fix (raise Combine rate / "
            "diversity / budget)\nbefore claiming the DSL composes."
        )
        return 2


if __name__ == "__main__":
    sys.exit(main())