AlgoTune/algotune_agent

"""AlgoTune Agent's mutable working tree plus independently retained best snapshot."""
from __future__ import annotations

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


class AlgoTuneSelectionPolicy(SelectionPolicy):
    """Continue editing the current workspace; `revert` is handled by the proposer."""

    def __init__(self, seed: int = 0):
        super().__init__(seed)
        self.current_id: str | None = None

    def set_current(self, genome: Genome | str | None) -> None:
        self.current_id = genome if isinstance(genome, str) else (genome.id if genome else None)

    def select(self, population, ctx: LoopContext | None = None) -> Selection:
        members = population.all()
        if not members:
            raise RuntimeError("cannot select from an empty AlgoTune workspace")
        by_id = {member.id: member for member in members}
        parent = by_id.get(self.current_id) or population.best() or members[-1]
        self.current_id = parent.id
        best = population.best()
        inspirations = [best] if best is not None and best.id != parent.id else []
        return Selection(
            parent=parent,
            inspirations=inspirations,
            pool=members,
            details={
                "selection_strategy": "working_candidate",
                "selection_mode": "continue_workspace",
                "best_candidate_id": best.id if best is not None else None,
            },
        )

    def observe(self, genome: Genome, ctx: LoopContext | None = None) -> None:
        self.current_id = genome.id

    def state_dict(self) -> dict:
        return {**super().state_dict(), "current_id": self.current_id}

    def load_state_dict(self, state: dict) -> None:
        super().load_state_dict(state)
        if isinstance(state, dict):
            value = state.get("current_id")
            self.current_id = str(value) if value else None