SkyDiscover/gepa_native

"""SkyDiscover GEPA Native epsilon-greedy / best / metric-front parent selection."""
from __future__ import annotations

from ...components.selection import SelectionPolicy
from ...records import LoopContext, Selection
from .pareto_utils import select_program_candidate_from_pareto_front
from .population import GepaNativePopulation, genome_score


class GepaNativeSelectionPolicy(SelectionPolicy):
    """Use one seeded RNG for parent selection and merge-pair selection, as upstream does."""

    def __init__(
        self,
        seed: int = 42,
        candidate_selection_strategy: str = "epsilon_greedy",
        epsilon: float = 0.1,
        num_inspirations: int = 4,
    ):
        # SkyDiscover uses ``random_seed or 42``.
        super().__init__(int(seed or 42))
        self.candidate_selection_strategy = str(candidate_selection_strategy)
        self.epsilon = float(epsilon)
        self.num_inspirations = int(num_inspirations) or 4

    def _parent(self, population: GepaNativePopulation):
        if self.candidate_selection_strategy == "best" or not population.elite_pool:
            return population.best()

        if self.candidate_selection_strategy == "pareto":
            if not population.program_at_metric_front or len(population.programs) < 2:
                return population.best()
            scores = {
                program_id: genome_score(genome)
                for program_id, genome in population.programs.items()
            }
            try:
                program_id = select_program_candidate_from_pareto_front(
                    population.program_at_metric_front,
                    scores,
                    self.rng,
                )
            except AssertionError:
                return population.best()
            return population.programs[program_id]

        if self.rng.random() < self.epsilon and len(population.elite_pool) > 1:
            return population.programs[self.rng.choice(population.elite_pool)]
        return population.best()

    def _context(self, population: GepaNativePopulation, parent_id: str) -> list:
        seen = {parent_id}
        context = []
        for program_id in population.elite_pool:
            if program_id in seen or program_id not in population.programs:
                continue
            context.append(population.programs[program_id])
            seen.add(program_id)
            if len(context) >= self.num_inspirations:
                break

        for _metric, (program_id, _value) in population.metric_best.items():
            if program_id in seen or program_id not in population.programs:
                continue
            context.append(population.programs[program_id])
            seen.add(program_id)
        return context[: self.num_inspirations]

    def select(
        self,
        population: GepaNativePopulation,
        ctx: LoopContext | None = None,
    ) -> Selection:
        if not population.programs:
            raise RuntimeError("cannot select from an empty population")
        parent = self._parent(population)
        inspirations = self._context(population, parent.id)
        return Selection(
            parent=parent,
            inspirations=inspirations,
            pool=population.all(),
            details={
                "selection_strategy": self.candidate_selection_strategy,
                "selection_mode": self.candidate_selection_strategy,
                "epsilon": self.epsilon,
                "elite_pool_size": len(population.elite_pool),
                "context_limit": self.num_inspirations,
            },
        )

    def merge_candidates(self, population: GepaNativePopulation):
        """Select complementary metric leaders, then best + random top-five fallback."""
        if len(population.elite_pool) < 2:
            best = population.best()
            return best, best

        leaders: dict[str, str] = {}
        for metric_name, (program_id, _score) in population.metric_best.items():
            if program_id in population.programs and program_id in population.elite_pool:
                leaders[metric_name] = program_id
        unique_leaders = sorted(set(leaders.values()))
        if len(unique_leaders) >= 2:
            first, second = self.rng.sample(unique_leaders, 2)
            return population.programs[first], population.programs[second]

        best = population.best()
        top_five = [
            program_id
            for program_id in population.elite_pool[:5]
            if program_id != best.id
        ]
        if top_five:
            other = self.rng.choice(top_five)
            return best, population.programs[other]
        return best, best