from __future__ import annotations

import numpy as np

from zen_fronts.selection_core import mc_rank_select_detailed


def test_auto_quantiles_budget_can_force_streaming_mode() -> None:
    rng = np.random.default_rng(0)
    N, M = 12, 3
    mu = rng.normal(size=(N, M)).astype(np.float64)
    sigma2 = (0.2 + rng.random((N, M))).astype(np.float64)

    res = mc_rank_select_detailed(
        mu,
        sigma2,
        10,
        percentile=0.2,
        seed=1,
        quantiles_mode_i=2,
        quantiles_budget=0,
    )
    place_stats = res[5]
    front_stats = res[6]
    within_stats = res[7]
    quality_score = res[8]

    assert int(place_stats["quantiles_mode_used"]) == 1
    assert int(front_stats["quantiles_mode_used"]) == 1
    assert int(within_stats["quantiles_mode_used"]) == 1
    assert quality_score.shape == (N,)


def test_auto_quantiles_budget_chooses_exact_for_small_problems() -> None:
    rng = np.random.default_rng(1)
    N, M = 12, 3
    mu = rng.normal(size=(N, M)).astype(np.float64)
    sigma2 = (0.2 + rng.random((N, M))).astype(np.float64)

    res = mc_rank_select_detailed(
        mu,
        sigma2,
        10,
        percentile=0.2,
        seed=1,
        quantiles_mode_i=2,
        quantiles_budget=10**9,
    )
    place_stats = res[5]
    front_stats = res[6]
    within_stats = res[7]
    quality_score = res[8]

    assert int(place_stats["quantiles_mode_used"]) == 0
    assert int(front_stats["quantiles_mode_used"]) == 0
    assert int(within_stats["quantiles_mode_used"]) == 0
    assert quality_score.shape == (N,)
