SkyDiscover/gepa_native

"""SkyDiscover GEPA Native population: full archive, elite pool, and rejection history."""
from __future__ import annotations

import collections
from typing import Any

from ...checkpoint import genome_from_dict, genome_to_dict
from ...components.population import Population
from ...records import Genome


def score_metrics(metrics: dict[str, Any] | None) -> float:
    """Mirror SkyDiscover ``get_score`` exactly for native algorithm decisions."""
    if not metrics:
        return 0.0
    if "combined_score" in metrics:
        try:
            return float(metrics["combined_score"])
        except (TypeError, ValueError):
            pass
    numeric = [
        value
        for value in metrics.values()
        if isinstance(value, (int, float)) and not isinstance(value, bool)
    ]
    return float(sum(numeric) / len(numeric)) if numeric else 0.0


def genome_score(genome: Genome | None) -> float:
    return score_metrics(genome.scores if genome is not None else None)


class GepaNativePopulation(Population):
    """Accepted full archive plus a fixed-size sampling pool.

    Normal children are admitted only when they strictly improve over their stored parent (unless the
    gate is disabled).  Rejected normal mutations go only to ``rejection_history``.  Merge children
    use the separate upstream criterion ``score >= max(parent_a, parent_b)`` and merge rejections are
    deliberately not reflective-history entries.
    """

    def __init__(
        self,
        population_size: int = 40,
        max_rejection_history: int = 20,
        acceptance_gating: bool = True,
    ):
        self.population_size = int(population_size)
        self.acceptance_gating = bool(acceptance_gating)
        self.max_rejection_history = int(max_rejection_history)
        self._programs: dict[str, Genome] = {}
        self.elite_pool: list[str] = []
        self.rejection_history: collections.deque[Genome] = collections.deque(
            maxlen=self.max_rejection_history
        )
        self.metric_best: dict[str, tuple[str, float]] = {}
        self.program_at_metric_front: dict[str, set[str]] = {}
        self.initial_program_id: str | None = None
        self.best_program_id: str | None = None

    @property
    def programs(self) -> dict[str, Genome]:
        """Read-only-by-convention view matching the upstream database name."""
        return self._programs

    def get(self, program_id: str | None) -> Genome | None:
        return self._programs.get(program_id) if program_id else None

    def _merge_is_acceptable(self, genome: Genome) -> bool:
        score_a = genome.metadata.get("merge_score_a")
        score_b = genome.metadata.get("merge_score_b")
        if not isinstance(score_a, (int, float)) or isinstance(score_a, bool):
            score_a = genome_score(self.get(genome.parent_id))
        if not isinstance(score_b, (int, float)) or isinstance(score_b, bool):
            context_ids = genome.metadata.get("other_context_ids") or []
            score_b = genome_score(self.get(context_ids[0] if context_ids else None))
        return genome_score(genome) >= max(float(score_a), float(score_b))

    def add(self, genome: Genome) -> bool:
        # Galapagos records invalid normal and merge attempts in the evolutionary trajectory, but
        # neither is an archive member. This is intentionally stricter than the source merge path's
        # score-only acceptance check and keeps the contract uniform across registered scaffolds.
        if genome.metadata.get("valid") is False:
            genome.metadata["gepa_native_admitted"] = False
            genome.metadata.update(admitted=False, eval_failed=True)
            return False

        is_merge = bool(genome.metadata.get("gepa_native_merge"))
        if genome.parent_id is not None:
            if is_merge:
                if not self._merge_is_acceptable(genome):
                    genome.metadata["gepa_native_admitted"] = False
                    genome.metadata["admitted"] = False
                    return False
            else:
                if self.acceptance_gating:
                    parent = self.get(genome.parent_id)
                    parent_score = genome_score(parent) if parent is not None else 0.0
                    if genome_score(genome) <= parent_score:
                        self.add_rejected(genome)
                        return False

        self._add_to_archive(genome)
        genome.metadata.pop("eval_failed", None)
        genome.metadata["gepa_native_admitted"] = True
        genome.metadata["admitted"] = True
        return True

    def _add_to_archive(self, genome: Genome) -> None:
        if not self._programs:
            self.initial_program_id = genome.id
        self._programs[genome.id] = genome

        if genome.id not in self.elite_pool:
            self.elite_pool.append(genome.id)
        self.elite_pool.sort(
            key=lambda program_id: genome_score(self._programs.get(program_id)),
            reverse=True,
        )

        # Preserve SkyDiscover's pinning loop exactly.  Because pinned members do not consume the
        # ``len(keep) < population_size`` allowance consistently, the sampling pool can occasionally
        # exceed the nominal bound; tightening it would change the source algorithm.
        if len(self.elite_pool) > self.population_size:
            pinned = {
                self.best_program_id,
                self.initial_program_id,
                genome.id,
            } - {None}
            keep: list[str] = []
            for program_id in self.elite_pool:
                if program_id in pinned or len(keep) < self.population_size:
                    keep.append(program_id)
            self.elite_pool = keep

        for metric_name, value in (genome.scores or {}).items():
            # Upstream uses this raw isinstance check, so bool-valued metrics participate here even
            # though they are excluded from scalar score fallback.
            if not isinstance(value, (int, float)):
                continue
            current = self.metric_best.get(metric_name)
            if current is None or value > current[1]:
                self.metric_best[metric_name] = (genome.id, value)
                self.program_at_metric_front[metric_name] = {genome.id}
            elif value == current[1]:
                self.program_at_metric_front.setdefault(metric_name, set()).add(genome.id)

        if (
            self.best_program_id is None
            or genome_score(genome) > genome_score(self._programs.get(self.best_program_id))
        ):
            self.best_program_id = genome.id

    def add_rejected(self, genome: Genome) -> None:
        genome.metadata["gepa_native_admitted"] = False
        genome.metadata["admitted"] = False
        self.rejection_history.append(genome)

    def get_rejection_history(self, limit: int | None = None) -> list[Genome]:
        rejected = list(self.rejection_history)
        return rejected[-limit:] if limit is not None else rejected

    def elite_programs(self) -> list[Genome]:
        return [
            self._programs[program_id]
            for program_id in self.elite_pool
            if program_id in self._programs
        ]

    def query(self, spec: dict | None = None) -> list[Genome]:
        spec = spec or {}
        members = self.elite_programs() if spec.get("elite") else list(self._programs.values())
        if spec.get("sorted") or "top" in spec:
            members.sort(key=genome_score, reverse=True)
        top = spec.get("top")
        return members[:top] if top is not None else members

    def all(self) -> list[Genome]:
        """Return the accepted full archive; the elite pool is sampling-only."""
        return list(self._programs.values())

    def best(self) -> Genome | None:
        if self.best_program_id in self._programs:
            return self._programs[self.best_program_id]
        return max(self._programs.values(), key=genome_score) if self._programs else None

    def state_dict(self) -> dict:
        return {
            "elite_pool": list(self.elite_pool),
            "initial_program_id": self.initial_program_id,
            "metric_best": {
                metric: [program_id, value]
                for metric, (program_id, value) in self.metric_best.items()
            },
            "program_at_metric_front": {
                metric: list(program_ids)
                for metric, program_ids in self.program_at_metric_front.items()
            },
            "rejection_history": [
                genome_to_dict(genome) for genome in self.rejection_history
            ],
        }

    def load_state_dict(self, state: dict) -> None:
        if not isinstance(state, dict):
            return
        if not state:
            self._rebuild_native_indexes()
            return

        self.elite_pool = [
            program_id
            for program_id in state.get("elite_pool", [])
            if program_id in self._programs
        ]
        initial = state.get("initial_program_id")
        self.initial_program_id = initial if initial in self._programs else self.initial_program_id

        self.metric_best.clear()
        for metric, value in (state.get("metric_best") or {}).items():
            if (
                isinstance(value, (list, tuple))
                and len(value) >= 2
                and value[0] in self._programs
            ):
                self.metric_best[str(metric)] = (value[0], value[1])

        self.program_at_metric_front = {
            str(metric): {
                program_id for program_id in (program_ids or [])
                if program_id in self._programs
            }
            for metric, program_ids in (state.get("program_at_metric_front") or {}).items()
        }

        self.rejection_history.clear()
        for raw in state.get("rejection_history", []):
            if not isinstance(raw, dict):
                continue
            try:
                self.rejection_history.append(genome_from_dict(raw))
            except Exception:
                continue

        if not self.elite_pool and self._programs:
            self._rebuild_native_indexes()

    def _rebuild_native_indexes(self) -> None:
        """SkyDiscover legacy-checkpoint fallback: derive native indexes from the archive."""
        self.elite_pool = sorted(
            self._programs,
            key=lambda program_id: genome_score(self._programs[program_id]),
            reverse=True,
        )[: self.population_size]
        if self._programs:
            self.initial_program_id = min(
                self._programs,
                key=lambda program_id: int(
                    self._programs[program_id].metadata.get("iteration", 0) or 0
                ),
            )
            self.best_program_id = max(
                self._programs,
                key=lambda program_id: genome_score(self._programs[program_id]),
            )

        self.metric_best.clear()
        self.program_at_metric_front.clear()
        for program_id, genome in self._programs.items():
            for metric_name, value in (genome.scores or {}).items():
                if not isinstance(value, (int, float)):
                    continue
                current = self.metric_best.get(metric_name)
                if current is None or value > current[1]:
                    self.metric_best[metric_name] = (program_id, value)
                    self.program_at_metric_front[metric_name] = {program_id}
                elif value == current[1]:
                    self.program_at_metric_front.setdefault(metric_name, set()).add(program_id)