"""GP main loop with evolution logging.

Seeded RNG, tournament selection, elitism, crossover, mutation. We log the
top-K candidates per generation (not the full population) so the persisted
artefact stays small and the Streamlit replay can scrub instantly. Fitness
results are cached by program signature so elites carrying 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.fitness import fitness_fn
from engine.objectives import BinaryAUROCObjective, Objective
from engine.program import Program, crossover, mutate, random_program


def _assign_id(program: Program, generation: int, idx: int) -> None:
    program.program_id = f"g{generation}c{idx}"


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


def run_gp(
    M_train: pd.DataFrame,
    y_train: np.ndarray,
    pool: Sequence[str],
    *,
    objective: Objective | None = None,
    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,
    log_top_k: int = 8,
    on_generation: Callable[[dict], None] | None = None,
) -> tuple[list[dict], Program, float]:
    """Run the GP loop. Returns (per-generation log, best program, best fitness).

    ``on_generation`` is invoked once per generation with the just-built log
    entry; the API layer uses it to stream generation events over SSE.
    """
    obj = objective or BinaryAUROCObjective()
    py_rng = random.Random(seed)

    population: list[Program] = [
        random_program(py_rng, pool) for _ in range(population_size)
    ]
    for i, p in enumerate(population):
        _assign_id(p, 0, i)
        p.born = True
        p.parents = []

    fitness_cache: dict[tuple, float] = {}

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

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

    for gen in range(n_generations):
        fitnesses = evaluate(population)

        ranked = sorted(
            zip(population, fitnesses),
            key=lambda kv: kv[1],
            reverse=True,
        )
        top = ranked[:log_top_k]
        gen_entry = {
            "generation": gen,
            "best_fitness": float(top[0][1]),
            "median_fitness": float(np.median(fitnesses)),
            "elitism": elitism,
            "candidates": [
                {
                    "id": prog.program_id,
                    "gene_ids": prog.gene_ids,
                    "feature_sets": [list(s) for s in prog.feature_sets],
                    "program_repr": prog.program_repr(),
                    "fitness": float(fit),
                    "n_genes": prog.n_genes,
                    "parents": list(prog.parents),
                    "born": bool(prog.born),
                    "survived": rank < elitism,
                }
                for rank, (prog, fit) in enumerate(top)
            ],
        }
        log.append(gen_entry)
        if on_generation is not None:
            on_generation(gen_entry)

        if top[0][1] > best_overall_fit:
            best_overall = top[0][0]
            best_overall_fit = top[0][1]

        if gen == n_generations - 1:
            break

        elites = [prog for prog, _ in ranked[:elitism]]
        new_pop: list[Program] = []
        for ei, e in enumerate(elites):
            carried = Program(
                feature_sets=[list(s) for s in e.feature_sets],
                parents=[e.program_id],
                born=False,
            )
            _assign_id(carried, gen + 1, ei)
            new_pop.append(carried)

        i = elitism
        while len(new_pop) < population_size:
            p1 = _tournament_select(py_rng, population, fitnesses, tournament_k)
            p2 = _tournament_select(py_rng, population, fitnesses, tournament_k)
            child = crossover(py_rng, p1, p2)
            child = mutate(py_rng, child, pool, p_mut=p_mutate)
            _assign_id(child, gen + 1, i)
            child.born = True
            new_pop.append(child)
            i += 1
        population = new_pop

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