GEPA/gepa

"""GEPA — a port of Genetic-Pareto reflective evolution (Agrawal et al., ICLR 2026).

Every component is the GEPA-specific class from :mod:`.population` … :mod:`.proposer`:

* Population   → :class:`~galapagos.scaffolds.gepa.population.GepaCandidatePool` — the unbounded
  candidate pool, its per-key Pareto envelope, and the minibatch acceptance test that gates
  admission (``GEPAState`` + ``strategies/acceptance.py``).
* Selection    → :class:`~galapagos.scaffolds.gepa.selection_policy.ParetoCandidateSelector` —
  prune dominated candidates, then sample proportionally to frontier coverage (Algorithm 2).
* PromptBuilder → :class:`~galapagos.scaffolds.gepa.prompt_builder.GepaReflectionPromptBuilder` —
  the ``optimize_anything`` reflection template over the parent's own evaluation feedback.
* Proposer     → :class:`~galapagos.scaffolds.gepa.proposer.GepaReflectiveMutationProposer` — one
  reflection call, one complete replacement candidate.
* Memory       → none.

That mapping puts GEPA's whole loop on the framework's own select → prompt → propose → evaluate →
admit cycle: the engine's two-stage "minibatch acceptance test, then full valset evaluation" is a
single evaluation here, because a galapagos task is one problem scored by one evaluator call — GEPA's
own ``optimize_anything(dataset=None, valset=None)`` single-instance mode, where the minibatch *is*
the valset. The acceptance test therefore compares the child against its parent's stored verdict,
which is what upstream's evaluation cache serves for that same pair.

Deliberately not ported, because each is structurally inert on a single-component candidate:
``MergeProposer`` (system-aware merge needs two candidates that changed *different* components of the
same ancestor — with one component, ``does_triplet_have_desirable_predictors`` is never satisfied and
the proposer reports "no merge candidates found" forever; upstream also leaves ``GEPAConfig.merge``
disabled by default) and the round-robin ``ReflectionComponentSelector`` (one component to pick).
"""
from __future__ import annotations

from ...config import GalapagosConfig
from ...models import GalapagosModel
from ..base_scaffold import GalapagosScaffold
from ..registry import register_scaffold
# one module per component (the GEPA scaffold method)
from .memory import GepaMemory
from .population import GepaCandidatePool
from .prompt_builder import GepaReflectionPromptBuilder
from .proposer import GepaReflectiveMutationProposer
from .selection_policy import ParetoCandidateSelector


def objective_directions(task) -> dict[str, str]:
    """The objective frontier axis, read off the task card's ``metrics:`` block.

    GEPA's objective scores are whatever the user's evaluator publishes under ``side_info["scores"]``
    — an explicit, higher-is-better declaration. galapagos's equivalent declaration is the card's
    metric list, so that is what the frontier uses: a metric the card never declares carries no
    direction and is informational only, and a ``minimize`` metric is oriented by the Population.
    """
    directions: dict[str, str] = {}
    for metric in getattr(task, "metrics", None) or []:
        name = getattr(metric, "name", None)
        if not name:
            continue
        direction = (getattr(metric, "metric_direction", None)
                     or getattr(metric, "direction", None) or "maximize")
        directions[str(name)] = "minimize" if str(direction) == "minimize" else "maximize"
    return directions


@register_scaffold("gepa")
class GepaScaffold(GalapagosScaffold):
    name = "gepa"

    @classmethod
    def build_components(cls, config: GalapagosConfig, model: GalapagosModel | None) -> dict:
        return {
            "population": GepaCandidatePool(
                frontier_type=str(config.selection_policy.frontier_type),
                acceptance_criterion=str(config.population.acceptance_criterion),
            ),
            "selection_policy": ParetoCandidateSelector(seed=int(config.seed)),
            "prompt_builder": GepaReflectionPromptBuilder(),
            "proposer": GepaReflectiveMutationProposer(),
            "memory": GepaMemory(),
        }

    def setup(self, task) -> None:
        # Before super(): setup scores the seed and admits it, and that admission is what initializes
        # the Pareto frontier — so the objective axis has to be known first.
        if isinstance(self.population, GepaCandidatePool):
            self.population.set_objective_directions(objective_directions(task))
        super().setup(task)