from __future__ import annotations

import numpy as np

from zen_fronts.policies import DefaultParentPolicy
from zen_fronts.selection.engine import RefreshResult


def test_deterministic_parent_policy_tie_breaks_as_specified() -> None:
    # active ids: 10, 11, 12
    rr = RefreshResult(
        active_ids=(10, 11, 12),
        winners_ids=(10, 11, 12),
        losers_ids=(12,),
        win_count=np.array([5, 5, 4], dtype=np.int32),
        lose_count=np.array([0, 0, 0], dtype=np.int32),
        rank_sum=np.array([100, 50, 60], dtype=np.int64),
        n_samples=5,
        percentile=0.34,
        seed_used=0,
    )

    pol = DefaultParentPolicy(stochastic=False)
    rng = np.random.default_rng(0)
    # win_count ties between 10 and 11; rank_sum breaks => 11 (50 < 100)
    assert pol.choose_parent(999, refresh_result=rr, rng=rng) == 11


def test_stochastic_policy_is_seeded_and_prefers_higher_win_count() -> None:
    rr = RefreshResult(
        active_ids=(1, 2, 3),
        winners_ids=(1, 2, 3),
        losers_ids=(3,),
        win_count=np.array([10, 1, 1], dtype=np.int32),
        lose_count=np.array([0, 0, 0], dtype=np.int32),
        rank_sum=np.array([0, 1, 2], dtype=np.int64),
        n_samples=10,
        percentile=0.34,
        seed_used=0,
    )
    pol = DefaultParentPolicy(stochastic=True)

    a = pol.choose_parent(0, refresh_result=rr, rng=np.random.default_rng(123))
    b = pol.choose_parent(0, refresh_result=rr, rng=np.random.default_rng(123))
    assert a == b
    # With these weights, it is very likely to pick 1 for many seeds.
    assert a in (1, 2, 3)
