from __future__ import annotations

import numpy as np

from zen_fronts.crits import CriteriaStore
from zen_fronts.pareto_core.pareto_rating import rate_population
from zen_fronts.selection import SelectionConfig, SelectionEngine


def _make_store_sigma0_example() -> CriteriaStore:
    store = CriteriaStore({"x": "max", "y": "max"})
    for pid in range(5):
        store.register_point(pid, active=True)

    mu = np.array(
        [
            [1.0, 1.0],
            [0.8, 0.8],
            [1.0, 0.0],
            [0.0, 1.0],
            [0.2, 0.2],
        ],
        dtype=np.float64,
    )
    for pid in range(5):
        store.update_crits(
            pid,
            {
                "x": {"mu": float(mu[pid, 0]), "trend": 0.0, "sigma2": 0.0},
                "y": {"mu": float(mu[pid, 1]), "trend": 0.0, "sigma2": 0.0},
            },
            t=1.0,
        )
    return store


def _within_rank_from_rating(order: np.ndarray, front_id: np.ndarray) -> np.ndarray:
    order = np.asarray(order, dtype=np.int64)
    front_id = np.asarray(front_id, dtype=np.int64)
    within = np.empty_like(front_id)
    max_front = int(front_id.max()) if front_id.size else -1
    counts = np.zeros(max_front + 1, dtype=np.int64)
    for idx in order.tolist():
        fid = int(front_id[idx])
        within[idx] = counts[fid]
        counts[fid] += 1
    return within


def test_refresh_is_deterministic_for_fixed_seed() -> None:
    store = _make_store_sigma0_example()
    eng = SelectionEngine(store, config=SelectionConfig(n_samples=50, percentile=0.2, seed=123))

    a = eng.refresh(now=1.0)
    b = eng.refresh(now=1.0)
    assert a == b


def test_sigma2_zero_matches_single_pareto_order_and_stats() -> None:
    store = _make_store_sigma0_example()
    eng = SelectionEngine(store, config=SelectionConfig(n_samples=25, percentile=0.2, seed=7))

    losers = eng.refresh(now=1.0)
    rr = eng.last_result
    assert rr is not None

    # When sigma2==0, each MC sample is identical.
    mu, _ = store.build_matrices(rr.active_ids, now=1.0)
    rating = rate_population(mu, mode="ranks", within_mode="crowding", use_cython=True)

    order = np.asarray(rating.order, dtype=np.int64)
    front = np.asarray(rating.front_id, dtype=np.int64)

    expected_winner = int(rr.active_ids[int(order[0])])
    expected_loser = int(rr.active_ids[int(order[-1])])

    assert rr.winners_ids == (expected_winner,)
    assert rr.losers_ids == (expected_loser,)
    assert losers == (expected_loser,)

    # win/lose counts are exact when every sample identical.
    wi = rr.index_of(expected_winner)
    li = rr.index_of(expected_loser)
    assert int(rr.win_count[wi]) == rr.n_samples
    assert int(rr.lose_count[li]) == rr.n_samples

    # Per-point "place" stats should be exact and have zero spread.
    pos = np.empty_like(order)
    pos[order] = np.arange(order.size, dtype=np.int64)
    within = _within_rank_from_rating(order, front)

    # Composite quality score should be deterministic when sigma2==0.
    assert rr.quality_score is not None
    # With default weights and sign: quality = -(place + front + within).
    expected_quality = -(pos.astype(np.float64) + front.astype(np.float64) + within.astype(np.float64))
    assert np.allclose(rr.quality_score, expected_quality)
    # For this small problem, auto mode should choose exact quantiles by default budget.
    assert rr.quantiles_mode_used in (0, None)

    assert rr.place_mean is not None
    assert rr.place_median is not None
    assert rr.place_std is not None
    assert rr.front_mean is not None
    assert rr.within_mean is not None

    for i in range(len(rr.active_ids)):
        assert float(rr.place_mean[i]) == float(pos[i])
        assert float(rr.place_median[i]) == float(pos[i])
        assert float(rr.place_std[i]) == 0.0

        assert float(rr.front_mean[i]) == float(front[i])
        assert float(rr.front_median[i]) == float(front[i])
        assert float(rr.front_std[i]) == 0.0

        assert float(rr.within_mean[i]) == float(within[i])
        assert float(rr.within_median[i]) == float(within[i])
        assert float(rr.within_std[i]) == 0.0
