engine_v2/synthesize.py

"""Typed program synthesis — init, crossover, mutation — full grammar.

Termination guarantee: every type closes at the depth floor.
- Matrix → MatrixTerminal
- Vector → Reduce(MatrixTerminal, agg)
- Scalar → Associate(Reduce(MatrixTerminal, "mean"), target, "spearman")
- Model  → (not a leaf — only via FitApply which itself is Vector)

The richer operators (Effect, Split, FitApply, Search) are introduced at
configurable rates so the population genuinely contains programs that
use them. ``Search`` is gated OFF by default (rate 0) — it's the
recursive operator and the prompt explicitly says "flag it and gate
behind a toggle" until we're satisfied with the cost.
"""

from __future__ import annotations

import random
from typing import Sequence

from engine_v2.nodes import (
    Associate,
    Combine,
    Effect,
    FeatureSet,
    FitApply,
    MatrixTerminal,
    Node,
    Reduce,
    Search,
    Select,
    Split,
)
from engine_v2.types import (
    AGGS,
    ASSOC_KINDS,
    ASSOC_TARGETS,
    DEFAULT_MAX_DEPTH,
    DEFAULT_MAX_GENES_PER_SET,
    MIN_GENES_PER_SET,
    OPS,
    PREDICATE_KINDS,
    SEARCH_MAX_K,
    TType,
)


# Init / mutation rates for the richer operators. Keep them modest:
# enough that the population genuinely contains them, low enough that
# the search space stays tractable.
DEFAULT_RATES = {
    "split": 0.10,         # Vector slot picks Split instead of Reduce/Combine.
    "fitapply": 0.10,      # Vector slot picks FitApply.
    "effect": 0.40,        # Scalar slot picks Effect (vs Associate).
    # Matrix slot picks Search instead of Select/M. ON at a modest
    # rate so Search appears in the population without dominating —
    # every Search runs an inner gene ranking, which is the
    # heaviest per-node operation. Caps stay (k ≤ 4 selected,
    # ≤ 200 candidate columns) so each Search's cost is bounded.
    "search": 0.05,
}


# ---------------------------------------------------------------------------
# Feature-set sampling
# ---------------------------------------------------------------------------


def _sample_feature_set(
    rng: random.Random,
    pool: Sequence[str],
    *,
    max_genes_per_set: int,
) -> FeatureSet:
    if not pool:
        raise ValueError("synth: empty gene pool")
    upper = min(max_genes_per_set, len(pool))
    size = rng.randint(MIN_GENES_PER_SET, max(MIN_GENES_PER_SET, upper))
    return FeatureSet(rng.sample(list(pool), size))


# ---------------------------------------------------------------------------
# Grow — typed recursive expansion
# ---------------------------------------------------------------------------


def _grow_matrix(rng: random.Random, pool, depth: int, *,
                 mgps: int, full: bool, rates: dict,
                 objective_target: str) -> Node:
    """Generate a Matrix-typed subtree with depth <= ``depth``.

    Every Matrix leaf is wrapped in a ``Select`` — a bare
    ``MatrixTerminal()`` would let a program score on the global
    expression vector with zero gene choice (a "whole-matrix mean"
    detector that conflates detection with gene discovery). The rule
    applies to ALL objectives: MSI/TMB/HPV/unsup. ``Search`` is gated
    off under unsupervised (the rate is already 0 by default but the
    gate makes it explicit).
    """
    is_unsup = objective_target == "none"
    if depth < 1:
        depth = 1
    if depth == 1:
        return Select(
            MatrixTerminal(),
            _sample_feature_set(rng, pool, max_genes_per_set=mgps),
        )
    # Matrix options: Select(Matrix, FeatureSet) | Search(Matrix, k).
    # The bare-MatrixTerminal fall-through is gone — it would have
    # produced a global-mean detector that scores well via bulk
    # expression rather than gene choice.
    r = rng.random()
    if not is_unsup and rates["search"] > 0 and r < rates["search"]:
        inner = _grow_matrix(
            rng, pool, depth - 1, mgps=mgps, full=full, rates=rates,
            objective_target=objective_target,
        )
        return Search(inner, k=rng.randint(2, SEARCH_MAX_K))
    inner = _grow_matrix(
        rng, pool, depth - 1, mgps=mgps, full=full, rates=rates,
        objective_target=objective_target,
    )
    return Select(inner, _sample_feature_set(rng, pool, max_genes_per_set=mgps))


def _grow_vector(rng: random.Random, pool, depth: int, *,
                 mgps: int, full: bool, rates: dict,
                 objective_target: str,
                 in_split: bool = False) -> Node:
    """Generate a Vector-typed subtree with depth <= ``depth``.

    Operators that score against a label (FitApply) are constructed
    with ``target = objective_target`` — the engine never picks its
    own target. See ``_check_target_binding`` in ``fitness.py``.
    """
    is_unsup = objective_target == "none"
    if depth < 2:
        depth = 2
    if depth == 2:
        # Depth-floor: every objective wraps in Select so no Reduce
        # ever sits on a bare MatrixTerminal (would be a global-mean
        # shortcut that detects without choosing genes).
        leaf: Node = Select(
            MatrixTerminal(),
            _sample_feature_set(rng, pool, max_genes_per_set=mgps),
        )
        return Reduce(leaf, rng.choice(AGGS))
    r = rng.random()
    # Split has ONE-level rule — never inside another Split.
    if not in_split and r < rates["split"]:
        inner = _grow_vector(
            rng, pool, depth - 1,
            mgps=mgps, full=full, rates=rates,
            objective_target=objective_target, in_split=True,
        )
        # Under unsup the structure must come from the gene-based score,
        # not the known clinical axis (stage_late) — force "score".
        predicate = (
            "score" if is_unsup else rng.choice(PREDICATE_KINDS)
        )
        return Split(inner, predicate=predicate)
    if r < rates["split"] + rates["fitapply"]:
        inner = _grow_vector(
            rng, pool, depth - 1,
            mgps=mgps, full=full, rates=rates,
            objective_target=objective_target, in_split=in_split,
        )
        # FitApply's target is BOUND to the active objective, never picked.
        return FitApply(inner, target=objective_target)
    if full or rng.random() < 0.5:
        return Combine(
            _grow_vector(rng, pool, depth - 1, mgps=mgps, full=full, rates=rates,
                         objective_target=objective_target, in_split=in_split),
            _grow_vector(rng, pool, depth - 1, mgps=mgps, full=full, rates=rates,
                         objective_target=objective_target, in_split=in_split),
            rng.choice(OPS),
        )
    return Reduce(
        _grow_matrix(
            rng, pool, depth - 1, mgps=mgps, full=full, rates=rates,
            objective_target=objective_target,
        ),
        rng.choice(AGGS),
    )


def _grow_scalar(rng: random.Random, pool, depth: int, *,
                 mgps: int, full: bool, rates: dict,
                 objective_target: str) -> Node:
    """Generate a Scalar-typed subtree (Associate or Effect).

    The target is BOUND to the active objective — the engine picks the
    kind (pearson / spearman) and whether to use Effect (adjusted) or
    Associate (raw), never the target.
    """
    if depth < 3:
        depth = 3
    inner_depth = max(2, depth - 1)
    inner = _grow_vector(rng, pool, inner_depth,
                         mgps=mgps, full=full, rates=rates,
                         objective_target=objective_target)
    kind = rng.choice(ASSOC_KINDS)
    if rng.random() < rates["effect"]:
        return Effect(inner, target=objective_target, kind=kind)
    return Associate(inner, target=objective_target, kind=kind)


def random_program(
    rng: random.Random,
    pool: Sequence[str],
    *,
    objective_target: str,
    max_depth: int = DEFAULT_MAX_DEPTH,
    max_genes_per_set: int = DEFAULT_MAX_GENES_PER_SET,
    full: bool | None = None,
    rates: dict | None = None,
    return_type: TType = TType.VECTOR,
) -> Node:
    """A random tree of the requested return type. ``return_type=Vector``
    is the engine's normal root — but the GP also explores Scalar-rooted
    programs (Associate / Effect) since A3 allows Scalar outputs. The
    objective's target is bound; the engine never chooses it.
    """
    rates = rates or DEFAULT_RATES
    if full is None:
        full = rng.random() < 0.5
    if return_type is TType.SCALAR:
        return _grow_scalar(rng, pool, max_depth,
                            mgps=max_genes_per_set, full=full, rates=rates,
                            objective_target=objective_target)
    return _grow_vector(rng, pool, max_depth,
                        mgps=max_genes_per_set, full=full, rates=rates,
                        objective_target=objective_target)


def ramped_population(
    rng: random.Random,
    pool: Sequence[str],
    *,
    n: int,
    objective_target: str,
    max_depth: int = DEFAULT_MAX_DEPTH,
    max_genes_per_set: int = DEFAULT_MAX_GENES_PER_SET,
    rates: dict | None = None,
    scalar_share: float = 0.20,
) -> list[Node]:
    """Ramped half-and-half across depths, with ``scalar_share`` of the
    population rooted at a Scalar (Associate / Effect)."""
    out: list[Node] = []
    depths = list(range(2, max_depth + 1)) or [2]
    for i in range(n):
        d = depths[i % len(depths)]
        full = (i // len(depths)) % 2 == 0
        rt = TType.SCALAR if rng.random() < scalar_share else TType.VECTOR
        out.append(random_program(
            rng, pool, objective_target=objective_target,
            max_depth=d, max_genes_per_set=max_genes_per_set,
            full=full, rates=rates, return_type=rt,
        ))
    return out


# ---------------------------------------------------------------------------
# Crossover — swap subtrees of matching return type
# ---------------------------------------------------------------------------


def _enumerate(parent: Node, slot: TType) -> list[tuple[Node, "_Cursor"]]:
    targets: list[tuple[Node, _Cursor]] = []
    if parent.ttype is slot:
        targets.append((parent, _Cursor.root(parent, slot)))
    _walk_for_replacement(parent, slot, targets)
    return targets


def _walk_for_replacement(parent: Node, slot: TType,
                          out: list[tuple[Node, "_Cursor"]]) -> None:
    # Single-child carriers
    one_child_attrs = {
        Select: "matrix",
        Reduce: "matrix",
        Split: "inner",
        Associate: "inner",
        Effect: "inner",
        FitApply: "inner",
        Search: "matrix",
    }
    for cls, attr in one_child_attrs.items():
        if isinstance(parent, cls):
            child: Node = getattr(parent, attr)
            if child.ttype is slot:
                out.append((child, _Cursor.field(parent, attr, slot)))
            _walk_for_replacement(child, slot, out)
            return

    if isinstance(parent, Combine):
        for attr in ("left", "right"):
            ch = getattr(parent, attr)
            if ch.ttype is slot:
                out.append((ch, _Cursor.field(parent, attr, slot)))
            _walk_for_replacement(ch, slot, out)
        return
    # MatrixTerminal has no children.


class _Cursor:
    def __init__(self, applier, slot):
        self._apply = applier
        self.slot = slot

    def apply(self, new_node: Node) -> Node:
        return self._apply(new_node)

    @staticmethod
    def root(root: Node, slot: TType) -> "_Cursor":
        def apply(new_node: Node) -> Node:
            return new_node
        return _Cursor(apply, slot)

    @staticmethod
    def field(parent: Node, name: str, slot: TType) -> "_Cursor":
        def apply(new_node: Node) -> Node:
            setattr(parent, name, new_node)
            return new_node
        return _Cursor(apply, slot)


def crossover(
    rng: random.Random,
    p1: Node,
    p2: Node,
    *,
    max_depth: int = DEFAULT_MAX_DEPTH,
    max_nodes: int = 64,
) -> Node:
    from copy import deepcopy

    child = deepcopy(p1)
    candidate_slots = (
        _enumerate(child, TType.VECTOR)
        + _enumerate(child, TType.MATRIX)
        + _enumerate(child, TType.SCALAR)
    )
    options: list[tuple[TType, _Cursor, Node]] = []
    seen: set[int] = set()
    for sub, cur in candidate_slots:
        if id(sub) in seen:
            continue
        seen.add(id(sub))
        options.append((sub.ttype, cur, sub))
    rng.shuffle(options)
    for slot_t, cur, _sub in options:
        donor_slots = _enumerate(p2, slot_t)
        if not donor_slots:
            continue
        donor_sub, _ = rng.choice(donor_slots)
        donor_copy = deepcopy(donor_sub)
        replaced = cur.apply(donor_copy)
        if cur.slot is child.ttype and id(replaced) is id(donor_copy):
            new_root = donor_copy
        else:
            new_root = child
        if new_root.depth() <= max_depth and new_root.node_count() <= max_nodes:
            return new_root
        child = deepcopy(p1)
    return deepcopy(p1)


# ---------------------------------------------------------------------------
# Mutation
# ---------------------------------------------------------------------------


def mutate(
    rng: random.Random,
    program: Node,
    pool: Sequence[str],
    *,
    objective_target: str,
    p_mut: float = 0.7,
    max_depth: int = DEFAULT_MAX_DEPTH,
    max_genes_per_set: int = DEFAULT_MAX_GENES_PER_SET,
    max_nodes: int = 64,
    rates: dict | None = None,
) -> Node:
    """Subtree + point mutation. The objective's target is BOUND — it is
    never a point-mutation spot, and freshly-grown subtrees inherit the
    same binding."""
    from copy import deepcopy

    if rng.random() > p_mut:
        return program
    rates = rates or DEFAULT_RATES

    program = deepcopy(program)

    # ----- Subtree mutation -----
    if rng.random() < 0.5:
        candidates = (
            _enumerate(program, TType.VECTOR)
            + _enumerate(program, TType.MATRIX)
            + _enumerate(program, TType.SCALAR)
        )
        seen: set[int] = set()
        unique = []
        for sub, cur in candidates:
            if id(sub) in seen:
                continue
            seen.add(id(sub))
            unique.append((sub, cur))
        if not unique:
            return program
        original = deepcopy(program)
        sub, cur = rng.choice(unique)
        sub_budget = 3 if cur.slot is TType.VECTOR else 2
        if cur.slot is TType.SCALAR:
            new_sub = _grow_scalar(
                rng, pool, 3,
                mgps=max_genes_per_set, full=False, rates=rates,
                objective_target=objective_target,
            )
        elif cur.slot is TType.MATRIX:
            new_sub = _grow_matrix(
                rng, pool, sub_budget,
                mgps=max_genes_per_set, full=False, rates=rates,
                objective_target=objective_target,
            )
        else:
            new_sub = _grow_vector(
                rng, pool, sub_budget,
                mgps=max_genes_per_set, full=False, rates=rates,
                objective_target=objective_target,
            )
        new_root = cur.apply(new_sub)
        if cur.slot is program.ttype and id(new_root) is id(new_sub):
            program = new_sub
        if program.depth() > max_depth or program.node_count() > max_nodes:
            return original
        return program

    # ----- Point mutation -----
    # NOTE: the target field on Associate / Effect / FitApply is NOT a
    # point-mutation spot — the engine must not flip a program's target
    # mid-run.
    spots: list[tuple[str, object]] = []
    for n in program.walk():
        if isinstance(n, Reduce):
            spots.append(("agg", n))
        elif isinstance(n, Combine):
            spots.append(("op", n))
        elif isinstance(n, Select):
            spots.append(("gene", n.features))
        elif isinstance(n, Associate) or isinstance(n, Effect):
            spots.append(("scalar_kind", n))
        elif isinstance(n, Split):
            # Under unsupervised the only legal predicate is "score" —
            # don't emit a mutation spot that could flip it to a clinical
            # variable (stage_late).
            if objective_target != "none":
                spots.append(("predicate", n))
        elif isinstance(n, Search):
            spots.append(("search_k", n))
    if not spots:
        return program
    kind, target = rng.choice(spots)
    if kind == "agg":
        target.agg = rng.choice(  # type: ignore[attr-defined]
            [a for a in AGGS if a != target.agg] or list(AGGS)
        )
    elif kind == "op":
        target.op = rng.choice(  # type: ignore[attr-defined]
            [o for o in OPS if o != target.op] or list(OPS)
        )
    elif kind == "gene":
        fs = target  # FeatureSet
        if not pool or not fs.ids:  # type: ignore[attr-defined]
            return program
        idx = rng.randrange(len(fs.ids))   # type: ignore[attr-defined]
        replacements = [g for g in pool if g not in fs.ids]  # type: ignore[attr-defined]
        if not replacements:
            return program
        fs.ids[idx] = rng.choice(replacements)  # type: ignore[attr-defined]
    elif kind == "scalar_kind":
        # Flip only the correlation kind (pearson/spearman) — never the
        # target. The target is the objective's, by construction.
        target.kind = rng.choice(  # type: ignore[attr-defined]
            [k for k in ASSOC_KINDS if k != target.kind] or list(ASSOC_KINDS)
        )
    elif kind == "predicate":
        target.predicate = rng.choice(  # type: ignore[attr-defined]
            [p for p in PREDICATE_KINDS if p != target.predicate]
            or list(PREDICATE_KINDS)
        )
    elif kind == "search_k":
        target.k = max(2, min(SEARCH_MAX_K, target.k + rng.choice([-1, 1])))  # type: ignore[attr-defined]
    return program