import random

from multi_swarm.ga.selection import elite_select, tournament_select
from multi_swarm.genome.hypothesis import HypothesisAgentGenome, ModelTier


def make(idx: int) -> HypothesisAgentGenome:
    return HypothesisAgentGenome(
        system_prompt=f"p-{idx}",
        feature_access=["close"],
        temperature=0.9,
        top_p=0.95,
        model_tier=ModelTier.C,
        lookback_window=100,
        cognitive_style="x",
    )


def test_tournament_picks_best_in_sample() -> None:
    population = [make(i) for i in range(10)]
    fitnesses = {g.id: float(i) for i, g in enumerate(population)}
    rng = random.Random(0)
    winner = tournament_select(population, fitnesses, k=5, rng=rng)
    assert isinstance(winner, HypothesisAgentGenome)
    assert fitnesses[winner.id] >= 0.0


def test_tournament_size_one_is_random() -> None:
    population = [make(i) for i in range(10)]
    fitnesses = {g.id: float(i) for i, g in enumerate(population)}
    rng = random.Random(0)
    picks = [tournament_select(population, fitnesses, k=1, rng=rng) for _ in range(50)]
    distinct = {p.id for p in picks}
    assert len(distinct) > 1


def test_elite_select_returns_top_k() -> None:
    population = [make(i) for i in range(10)]
    fitnesses = {g.id: float(i) for i, g in enumerate(population)}
    elites = elite_select(population, fitnesses, k=3)
    elite_fitnesses = sorted([fitnesses[g.id] for g in elites], reverse=True)
    assert elite_fitnesses == [9.0, 8.0, 7.0]