engine_v2/gp.py

"""Typed GP loop for engine_v2.

Init via ramped half-and-half. Tournament selection. Elitism. Subtree
crossover with strict type matching. Subtree + point mutation. A single
seed governs all RNG (init, selection, crossover, mutation, FeatureSet
sampling). The fitness cache is keyed by ``program.repr_typed()`` so
elites that carry across generations don't recompute.
"""

from __future__ import annotations

import random
from typing import Callable, Sequence

import numpy as np
import pandas as pd

from engine_v2.fitness import V2Objective, fitness_fn
from engine_v2.nodes import ExecContext, Node
from engine_v2.synthesize import (
    crossover,
    mutate,
    ramped_population,
)


def _assign_id(program: Node, generation: int, idx: int) -> str:
    return f"g{generation}c{idx}"


def _tournament_select(
    rng: random.Random,
    population: list[Node],
    fitnesses: list[float],
    k: int,
) -> int:
    contenders = rng.sample(range(len(population)), k)
    return max(contenders, key=lambda i: fitnesses[i])


def run_gp_v2(
    ctx_train: ExecContext,
    y_train: np.ndarray | None,
    pool: Sequence[str],
    *,
    objective: V2Objective,
    population_size: int = 150,
    n_generations: int = 30,
    tournament_k: int = 3,
    elitism: int = 5,
    p_mutate: float = 0.7,
    lambda_size: float = 0.005,
    cv_folds: int = 5,
    seed: int = 42,
    max_depth: int = 4,
    max_genes_per_set: int = 8,
    max_nodes: int = 64,
    coherence_weight: float = 0.0,
    immigrant_fraction: float = 0.0,
    rates_override: dict | None = None,
    scalar_share_override: float | None = None,
    on_generation: Callable[[dict], None] | None = None,
) -> tuple[list[dict], Node, float]:
    """Returns ``(generation_log, winner, winner_fitness)``.

    Each ``generation_log`` entry is a complete snapshot of the
    population: every candidate with its id, program_repr, fitness,
    survived flag, n_nodes, and depth. Top-12 are also exposed inline
    via ``top_candidates`` for the SSE stream's bandwidth budget.
    """
    py_rng = random.Random(seed)

    # Per-objective synthesis overrides (e.g. UNSUP forbids Scalar roots
    # and Associate / Effect / FitApply — Vector-only programs). The
    # caller may also pass `rates_override` (e.g. to gate Search off
    # per-run); those keys win over the objective's overrides which
    # win over DEFAULT_RATES inside synthesize. ``scalar_share`` lives
    # on its own param (it's a single float, not a per-operator rate);
    # an explicit `scalar_share_override` wins over the objective's
    # value.
    overrides = objective.synthesis_overrides() if hasattr(objective, "synthesis_overrides") else {}
    objective_rates = overrides.get("rates") or {}
    objective_scalar_share = overrides.get("scalar_share")
    effective_scalar_share = (
        scalar_share_override
        if scalar_share_override is not None
        else objective_scalar_share
    )
    merged_rates_override: dict | None
    if rates_override is None and not objective_rates:
        merged_rates_override = None
    else:
        merged_rates_override = {
            **(objective_rates or {}),
            **(rates_override or {}),
        }

    population: list[Node] = ramped_population(
        py_rng, pool,
        n=population_size,
        objective_target=objective.target,
        max_depth=max_depth,
        max_genes_per_set=max_genes_per_set,
        **(
            {"rates": merged_rates_override} if merged_rates_override is not None else {}
        ),
        **(
            {"scalar_share": effective_scalar_share}
            if effective_scalar_share is not None else {}
        ),
    )
    ids: list[str] = [_assign_id(p, 0, i) for i, p in enumerate(population)]
    parents: list[list[str]] = [[] for _ in population]

    fitness_cache: dict[str, float] = {}

    def evaluate(pop: list[Node]) -> list[float]:
        out = []
        for prog in pop:
            sig = prog.repr_typed()
            if sig not in fitness_cache:
                fitness_cache[sig] = fitness_fn(
                    prog, ctx_train, y_train,
                    objective=objective,
                    lambda_size=lambda_size,
                    n_folds=cv_folds,
                    random_state=seed,
                    coherence_weight=coherence_weight,
                )
            out.append(fitness_cache[sig])
        return out

    worst = objective.worst_score()

    def _finite(x: float) -> float:
        return float(x) if np.isfinite(x) else worst

    log: list[dict] = []
    best_overall: Node | None = None
    best_overall_fit = -np.inf

    for gen in range(n_generations):
        fitnesses = evaluate(population)
        ranked = sorted(
            range(len(population)),
            key=lambda i: fitnesses[i],
            reverse=True,
        )
        elite_ix = set(ranked[:elitism])

        candidates_payload = [
            {
                "id": ids[i],
                "program_repr": population[i].repr_typed(),
                "fitness": _finite(fitnesses[i]),
                "n_nodes": int(population[i].node_count()),
                "depth": int(population[i].depth()),
                "gene_ids": list(population[i].feature_ids()),
                "parents": list(parents[i]),
                "survived": i in elite_ix,
            }
            for i in ranked
        ]
        top_payload = candidates_payload[:12]
        # Median computed over the finite values only, so a generation
        # where many programs are degenerate doesn't collapse to -inf.
        finite_vals = [f for f in fitnesses if np.isfinite(f)]
        median_fit = float(np.median(finite_vals)) if finite_vals else worst

        entry = {
            "generation": gen,
            "best_fitness": _finite(fitnesses[ranked[0]]),
            "median_fitness": median_fit,
            "elitism": elitism,
            "population_size": len(population),
            "top_candidates": top_payload,
            "candidates": candidates_payload,
        }
        log.append(entry)
        if on_generation is not None:
            on_generation(entry)

        if fitnesses[ranked[0]] > best_overall_fit:
            best_overall_fit = fitnesses[ranked[0]]
            best_overall = population[ranked[0]]

        if gen == n_generations - 1:
            break

        # ----- Build next generation -----
        new_pop: list[Node] = []
        new_ids: list[str] = []
        new_parents: list[list[str]] = []

        # Elites carry over unchanged.
        for ei, idx in enumerate(ranked[:elitism]):
            new_pop.append(population[idx])  # share the node — it's evaluated.
            new_ids.append(_assign_id(population[idx], gen + 1, ei))
            new_parents.append([ids[idx]])

        i_offset = elitism

        # Random immigrants — fresh programs drawn from ramped_population
        # using the same grammar / objective / depth constraints as init.
        # They displace offspring slots (never elites), counter premature
        # convergence, and are gated to immigrant_fraction > 0 so existing
        # runs are byte-for-byte unchanged. Rounding: round(frac * pop),
        # capped so it never pushes past the budget after elites.
        n_immigrants = (
            int(round(immigrant_fraction * population_size))
            if immigrant_fraction > 0.0 else 0
        )
        n_immigrants = max(0, min(n_immigrants, population_size - len(new_pop)))
        if n_immigrants > 0:
            fresh = ramped_population(
                py_rng, pool,
                n=n_immigrants,
                objective_target=objective.target,
                max_depth=max_depth,
                max_genes_per_set=max_genes_per_set,
                **(
                    {"rates": merged_rates_override} if merged_rates_override is not None else {}
                ),
                **(
                    {"scalar_share": effective_scalar_share}
                    if effective_scalar_share is not None else {}
                ),
            )
            for child in fresh:
                new_pop.append(child)
                new_ids.append(_assign_id(child, gen + 1, i_offset))
                new_parents.append([])   # immigrant — no parents
                i_offset += 1

        while len(new_pop) < population_size:
            i1 = _tournament_select(py_rng, population, fitnesses, tournament_k)
            i2 = _tournament_select(py_rng, population, fitnesses, tournament_k)
            child = crossover(
                py_rng, population[i1], population[i2],
                max_depth=max_depth, max_nodes=max_nodes,
            )
            child = mutate(
                py_rng, child, pool,
                objective_target=objective.target,
                p_mut=p_mutate,
                max_depth=max_depth,
                max_genes_per_set=max_genes_per_set,
                max_nodes=max_nodes,
                **(
                    {"rates": merged_rates_override}
                    if merged_rates_override is not None else {}
                ),
            )
            new_pop.append(child)
            new_ids.append(_assign_id(child, gen + 1, i_offset))
            new_parents.append([ids[i1], ids[i2]])
            i_offset += 1

        population, ids, parents = new_pop, new_ids, new_parents

    assert best_overall is not None
    return log, best_overall, float(best_overall_fit)