GEPA/gepa

"""GEPA SelectionPolicy component — Pareto-frontier candidate sampling (Algorithm 2 in the paper).

Port of ``ParetoCandidateSelector`` (``gepa/strategies/candidate_selector.py``) and the three
utilities it stands on (``gepa/gepa_utils.py``: ``is_dominated`` → ``remove_dominated_programs`` →
``select_program_candidate_from_pareto_front``), with candidates keyed by Genome id instead of by
position in ``GEPAState.program_candidates``.

The mechanism, unchanged: take the set of candidates that hold at least one frontier key, prune the
ones that are *dominated* (every key they hold is also held by someone else — the "illumination" step
the paper credits to Mouret & Clune 2015), then sample proportionally to how many keys each surviving
candidate holds. A candidate that is best on more of the search's axes is proposed from more often,
and a candidate that is best on exactly one obscure axis is still never starved. This is what keeps
GEPA off the single-lineage hill-climb that greedy ``SelectBestCandidate`` collapses into (paper
Table 3: +12.44% over greedy, +7.33% over beam search).

**No inspirations.** GEPA's reflective mutation shows the LM one candidate and that candidate's own
execution feedback — never a gallery of rival programs. The selection is a parent and nothing else.
"""
from __future__ import annotations

import random

from ...components.selection import SelectionPolicy
from ...records import LoopContext, Selection


def _is_dominated(candidate, others: set, mapping: dict) -> bool:
    """``gepa_utils.is_dominated``: true when every frontier key ``candidate`` holds is also held by
    one of ``others`` — i.e. dropping it removes nothing from the frontier."""
    for key, front in mapping.items():
        if candidate not in front:
            continue
        if not any(other in front for other in others):
            return False
    return True


def remove_dominated(mapping: dict, scores: dict) -> dict:
    """``gepa_utils.remove_dominated_programs``: repeatedly drop a dominated candidate until none is
    left. Candidates are tried lowest-aggregate-score first, so when two candidates cover each other
    the weaker one is the one that goes."""
    frequency: dict = {}
    for front in mapping.values():
        for candidate in front:
            frequency[candidate] = frequency.get(candidate, 0) + 1

    dominated: set = set()
    candidates = sorted(frequency, key=lambda c: scores.get(c, float("-inf")))
    found_to_remove = True
    while found_to_remove:
        found_to_remove = False
        for candidate in candidates:
            if candidate in dominated:
                continue
            others = set(candidates) - {candidate} - dominated
            if _is_dominated(candidate, others, mapping):
                dominated.add(candidate)
                found_to_remove = True
                break

    survivors = {c for c in candidates if c not in dominated}
    return {key: {c for c in front if c in survivors} for key, front in mapping.items()}


def select_from_pareto_front(mapping: dict, scores: dict, rng: random.Random):
    """``gepa_utils.select_program_candidate_from_pareto_front``: prune the dominated candidates,
    then draw one with probability proportional to the number of frontier keys it holds.

    Returns ``None`` when no candidate holds a key. Upstream asserts a non-empty sampling list here
    and leaves the fallback commented out with a "TODO: Determine if we need this fallback"; a
    galapagos run must not die on it, so the caller falls back to the best aggregate score — the
    exact branch upstream sketched.
    """
    pruned = remove_dominated(mapping, scores)
    frequency: dict = {}
    for front in pruned.values():
        for candidate in front:
            frequency[candidate] = frequency.get(candidate, 0) + 1
    sampling_list = [candidate for candidate, freq in frequency.items() for _ in range(freq)]
    if not sampling_list:
        return None
    return rng.choice(sampling_list)


class ParetoCandidateSelector(SelectionPolicy):
    """Sample the parent from the pruned Pareto frontier of the candidate pool.

    Reads the frontier the :class:`~galapagos.scaffolds.gepa.population.GepaCandidatePool` maintains
    (``frontier_mapping``) and its aggregate ranking (``aggregate_score``), the same two reads the
    upstream selector makes against ``GEPAState``. All randomness flows through ``self.rng``, which
    the base class seeds and checkpoints, so a resumed run keeps sampling where it left off.
    """

    def select(self, population, ctx: LoopContext | None = None) -> Selection:
        members = population.all()
        if not members:
            raise RuntimeError("cannot select from an empty population")
        scores = {genome.id: population.aggregate_score(genome) for genome in members}
        mapping = population.frontier_mapping()
        chosen = select_from_pareto_front(mapping, scores, self.rng)
        mode = "pareto_frequency"
        if chosen is None or chosen not in scores:
            chosen = max(scores, key=lambda gid: scores[gid])
            mode = "aggregate_best"
        parent = next(genome for genome in members if genome.id == chosen)
        frontier_candidates = {c for front in mapping.values() for c in front}
        return Selection(
            parent=parent,
            inspirations=[],   # reflective mutation reflects on ONE candidate's own feedback
            pool=members,
            details={
                "selection_strategy": "pareto_frontier",
                "selection_mode": mode,
                "frontier_type": population.frontier_type,
                "frontier_key_count": len(mapping),
                "frontier_candidate_count": len(frontier_candidates),
                "parent_frontier_keys": sum(1 for front in mapping.values() if parent.id in front),
            },
        )