"""GEPA Population component — the unbounded candidate pool plus its Pareto frontier.

Port of the archive half of ``GEPAState`` (``gepa/core/state.py``): the list of every accepted
candidate (``program_candidates``) and, per frontier key, the best score any candidate has reached
(``pareto_front_valset`` / ``objective_pareto_front``) together with the set of candidates holding it
(``program_at_pareto_front_*``). GEPA never prunes — dominance decides *selection*, never eviction —
so this store is uncapped and ``add`` only ever appends.

``add`` also carries GEPA's **acceptance test** (``gepa/strategies/acceptance.py``), because in
galapagos admission is the one bit of policy the store keeps: a proposal enters the pool only when it
beats the parent it was mutated from, exactly as the engine's ``_accept_reflective_proposal`` gate
does before it runs the full valset evaluation.

**Single-instance search.** A galapagos task is ONE optimization problem scored by one
``evaluate(program_path)`` call, so there is no valset the framework can subset. That is precisely
GEPA's own ``optimize_anything(dataset=None, valset=None)`` mode: the engine builds a dataset holding
a single ``_SINGLE_INSTANCE_SENTINEL`` example, sets ``reflection_minibatch_size = 1``, and the
minibatch, the acceptance test, and the "full valset" evaluation all collapse onto that one example.
The instance frontier therefore has exactly one key (the candidate's aggregate score) and the
*objective* frontier — GEPA's ``side_info["scores"]`` axis, which the upstream single-task examples
select with ``frontier_type="objective"`` — carries the task's declared metrics.
"""
from __future__ import annotations

from ...components.population import Population
from ...records import Genome

#: The one pseudo-example of single-instance mode — GEPA's ``_SINGLE_INSTANCE_SENTINEL``.
_SINGLE_INSTANCE = 0

#: ``EngineConfig.frontier_type`` (``FrontierType`` in ``gepa/core/state.py``).
FRONTIER_TYPES = ("instance", "objective", "hybrid", "cartesian")
#: ``EngineConfig.acceptance_criterion`` (``gepa/strategies/acceptance.py``).
ACCEPTANCE_CRITERIA = ("strict_improvement", "improvement_or_equal")


class GepaCandidatePool(Population):
    """GEPA's candidate pool: keep-every-accepted-candidate, plus the per-key Pareto envelope.

    ``objective_directions`` names which metrics form the objective axis and which way each one
    points; :class:`~galapagos.scaffolds.gepa.scaffold.GepaScaffold` fills it from the task card's
    ``metrics:`` block. GEPA's objective contract is *higher is better* (``side_info["scores"]``), so
    a metric the card declares ``minimize`` enters the frontier negated — the front for it is then
    held by the candidate with the lowest value, which is what the card actually asked for.
    """

    def __init__(self, frontier_type: str = "hybrid",
                 acceptance_criterion: str = "strict_improvement"):
        if frontier_type not in FRONTIER_TYPES:
            raise ValueError(f"unknown selection_policy.frontier_type {frontier_type!r}; "
                             f"expected one of {list(FRONTIER_TYPES)}")
        if acceptance_criterion not in ACCEPTANCE_CRITERIA:
            raise ValueError(f"unknown population.acceptance_criterion {acceptance_criterion!r}; "
                             f"expected one of {list(ACCEPTANCE_CRITERIA)}")
        self.frontier_type = frontier_type
        self.acceptance_criterion = acceptance_criterion
        # Every task scores combined_score, so the objective axis is never empty even for a task that
        # declares no metrics block (upstream RAISES when an objective frontier finds no objectives).
        self.objective_directions: dict[str, str] = {"combined_score": "maximize"}
        self._members: list[Genome] = []
        self._by_id: dict[str, Genome] = {}
        self._front: dict = {}
        self._front_holders: dict = {}

    def set_objective_directions(self, directions: dict[str, str]) -> None:
        """Declare the objective axis from the task card (metric name → maximize | minimize)."""
        if directions:
            self.objective_directions = dict(directions)

    # ---- admission = GEPA's minibatch acceptance test -----------------------------------------
    def add(self, genome: Genome) -> bool:
        if genome.metadata.get("valid") is False:
            genome.metadata.update(admitted=False, eval_failed=True)
            return False
        if not self._accepts(genome):
            genome.metadata["admitted"] = False
            return False
        self._members.append(genome)
        self._by_id[genome.id] = genome
        self._update_frontier(genome)
        genome.metadata.pop("eval_failed", None)
        genome.metadata["admitted"] = True
        return True

    def _accepts(self, genome: Genome) -> bool:
        """``sum(scores_after) > sum(scores_before)`` on the minibatch the parent was scored on.

        The parent is already in the pool with its stored verdict, so the comparison reuses it rather
        than re-running the evaluator — which is what upstream's ``cache_evaluation=True`` does for
        the identical ``(candidate, example)`` pair it is about to re-score. The seed (and any genome
        whose parent is not in the pool) has nothing to beat and is admitted.
        """
        parent = self._by_id.get(genome.parent_id or "")
        if parent is None:
            return True
        before = self._minibatch_score(parent)
        after = self._minibatch_score(genome)
        if self.acceptance_criterion == "improvement_or_equal":
            return after >= before
        return after > before

    # ---- Pareto frontier bookkeeping -----------------------------------------------------------
    def _instance_scores(self, genome: Genome) -> dict:
        """Per-example scores. Single-instance mode: the one sentinel example's score."""
        return {_SINGLE_INSTANCE: genome.fitness}

    def _objective_scores(self, genome: Genome) -> dict:
        """The declared metrics this candidate scored, oriented higher-is-better."""
        out: dict[str, float] = {}
        for name, direction in self.objective_directions.items():
            value = genome.scores.get(name)
            if not isinstance(value, (int, float)) or isinstance(value, bool):
                continue
            out[name] = -float(value) if direction == "minimize" else float(value)
        return out

    def _frontier_scores(self, genome: Genome) -> dict:
        """This candidate's score on every frontier key — ``GEPAState._get_pareto_front_mapping``'s
        four key shapes, which is why the hybrid keys stay tagged rather than merged."""
        if self.frontier_type == "instance":
            return self._instance_scores(genome)
        if self.frontier_type == "objective":
            return self._objective_scores(genome)
        if self.frontier_type == "hybrid":
            keys: dict = {("val_id", k): v for k, v in self._instance_scores(genome).items()}
            keys.update({("objective", k): v for k, v in self._objective_scores(genome).items()})
            return keys
        return {("cartesian", val_id, name): value
                for val_id in self._instance_scores(genome)
                for name, value in self._objective_scores(genome).items()}

    def _minibatch_score(self, genome: Genome) -> float:
        """``sum(scores)`` over the minibatch — GEPA compares sums, not means, for acceptance."""
        return sum(self._instance_scores(genome).values())

    def _update_frontier(self, genome: Genome) -> None:
        """``_update_pareto_front_for_val_id``: a strictly better score takes the key alone, an equal
        score joins the incumbents on it."""
        for key, score in self._frontier_scores(genome).items():
            previous = self._front.get(key, float("-inf"))
            if score > previous:
                self._front[key] = score
                self._front_holders[key] = {genome.id}
            elif score == previous:
                self._front_holders.setdefault(key, set()).add(genome.id)

    def frontier_mapping(self) -> dict:
        """frontier key → ids of the candidates currently holding it (the selector's input)."""
        return {key: set(holders) for key, holders in self._front_holders.items()}

    def aggregate_score(self, genome: Genome) -> float:
        """The tracked score GEPA ranks candidates by — the mean over the evaluated examples, which
        in single-instance mode is the candidate's own fitness."""
        return genome.fitness

    # ---- reads --------------------------------------------------------------------------------
    def query(self, spec: dict | None = None) -> list[Genome]:
        spec = spec or {}
        members = sorted(self._members, key=lambda g: g.fitness, reverse=True)
        top = spec.get("top")
        return members[:top] if top else members

    def all(self) -> list[Genome]:
        # Admission order, so a checkpoint replays parents before their children and the acceptance
        # test reaches the same verdict it did live.
        return list(self._members)

    # ---- checkpoint / resume -------------------------------------------------------------------
    def state_dict(self) -> dict:
        """The objective axis. It comes from the task card at ``setup`` time, and ``_resume`` binds
        the card only *after* replaying population.jsonl — so without restoring it here the fronts
        would be rebuilt against the fallback axis instead of the run's own."""
        return {"objective_directions": dict(self.objective_directions)}

    def pre_load_state_dict(self, state: dict) -> None:
        """Topology before content: run before the checkpointed genomes are re-added, so every
        replayed candidate updates the same frontier keys the live run used."""
        if isinstance(state, dict):
            directions = state.get("objective_directions")
            if isinstance(directions, dict):
                self.set_objective_directions({str(k): str(v) for k, v in directions.items()})