"""Beam Search — maintain a fixed-width beam of promising programs, expand one per step.

Ported from SkyDiscover's ``BeamSearchDatabase`` (UC Berkeley Sky Computing Lab). One module per component:

  population.py        -> BeamPopulation          (fixed-width beam + diversity-aware prune + depth)
  selection_policy.py  -> BeamSelectionPolicy     (best / stochastic / round_robin / diversity_weighted)
  prompt_builder.py    -> BeamSearchPromptBuilder (the default multi-section template)
  proposer.py          -> BeamSearchProposer      (SEARCH/REPLACE diff)
  evaluator.py         -> BeamSearchEvaluator     (task-supplied)
  memory.py            -> BeamSearchMemory        (none)
  scaffold.py          -> BeamSearchScaffold      (the orchestrator that composes the six)

The two beam-specific axes map cleanly onto two components: the beam structure + diversity-aware
admission is the Population, the parent choice is the SelectionPolicy. Both share one fitness notion
(``combined_score`` with an optional depth penalty) via :meth:`BeamPopulation.score`, keeping pruning
and selection consistent.
"""
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 Beam Search scaffold method)
from .memory import BeamSearchMemory
from .population import BeamPopulation
from .prompt_builder import BeamSearchPromptBuilder
from .proposer import BeamSearchProposer
from .selection_policy import BeamSelectionPolicy


@register_scaffold("beam_search")
class BeamSearchScaffold(GalapagosScaffold):
    name = "beam_search"

    @classmethod
    def build_components(cls, config: GalapagosConfig, model: GalapagosModel | None) -> dict:
        seed = int(config.seed)
        pop = config.population
        sel = config.selection_policy
        diversity_weight = float(pop.beam_diversity_weight)
        return {
            "population": BeamPopulation(
                beam_width=int(pop.beam_width),
                diversity_weight=diversity_weight,
                depth_penalty=float(pop.beam_depth_penalty),
            ),
            "selection_policy": BeamSelectionPolicy(
                seed=seed,
                strategy=str(sel.beam_selection_strategy),
                temperature=float(sel.beam_temperature),
                diversity_weight=diversity_weight,
                num_inspirations=int(sel.num_inspirations),
            ),
            "prompt_builder": BeamSearchPromptBuilder(),
            "proposer": BeamSearchProposer(),
            "memory": BeamSearchMemory(),
        }