from __future__ import annotations

import importlib

import numpy as np
import pytest

from zen_fronts.pareto_core.pareto_rating import rate_population


def _mc_ext_available() -> bool:
    try:
        importlib.import_module("zen_fronts.selection_core._mc_rank")
        return True
    except Exception:
        return False


@pytest.mark.skipif(
    not _mc_ext_available(),
    reason="Cython extension is not built (zen_fronts.selection_core._mc_rank)",
)
def test_mc_rank_select_sigma0_matches_single_order() -> None:
    from zen_fronts.selection_core._mc_rank import mc_rank_select  # type: ignore

    # sigma2==0 => every sample identical => counts are exact.
    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,
    )
    sigma2 = np.zeros_like(mu)

    n_samples = 25
    percentile = 0.2  # k=1 for N=5

    winners, losers, win_count, lose_count, rank_sum = mc_rank_select(
        mu, sigma2, n_samples,
        percentile=percentile,
        seed=123,
        mode_i=1,
        within_mode_i=0,
    )

    rating = rate_population(mu, mode="ranks", within_mode="crowding", use_cython=True)
    expected_winner = int(rating.order[0])
    expected_loser = int(rating.order[-1])

    assert winners.tolist() == [expected_winner]
    assert losers.tolist() == [expected_loser]

    assert int(win_count[expected_winner]) == n_samples
    assert int(lose_count[expected_loser]) == n_samples

    # Rank sum should be position * n_samples for each point.
    for pos, idx in enumerate(rating.order.tolist()):
        assert int(rank_sum[idx]) == pos * n_samples


@pytest.mark.skipif(
    not _mc_ext_available(),
    reason="Cython extension is not built (zen_fronts.selection_core._mc_rank)",
)
def test_mc_rank_counts_invariants() -> None:
    from zen_fronts.selection_core._mc_rank import mc_rank_counts  # type: ignore

    rng = np.random.default_rng(0)
    N, M = 31, 3
    mu = rng.normal(size=(N, M)).astype(np.float64)
    sigma2 = (0.1 + rng.random((N, M))).astype(np.float64)

    n_samples = 123
    percentile = 0.2
    k = int(np.ceil(percentile * N))

    win_count, lose_count, rank_sum = mc_rank_counts(
        mu, sigma2, n_samples,
        percentile=percentile,
        seed=42,
        mode_i=1,
        within_mode_i=0,
    )

    assert win_count.shape == (N,)
    assert lose_count.shape == (N,)
    assert rank_sum.shape == (N,)

    assert int(win_count.sum()) == n_samples * k
    assert int(lose_count.sum()) == n_samples * k

    # Each sample contributes sum_{pos=0..N-1} pos = N*(N-1)/2
    expected_rank_sum_total = n_samples * (N * (N - 1) // 2)
    assert int(rank_sum.sum()) == expected_rank_sum_total


@pytest.mark.skipif(
    not _mc_ext_available(),
    reason="Cython extension is not built (zen_fronts.selection_core._mc_rank)",
)
def test_mc_rank_is_deterministic_for_fixed_seed() -> None:
    from zen_fronts.selection_core._mc_rank import mc_rank_counts  # type: ignore

    rng = np.random.default_rng(123)
    mu = rng.normal(size=(40, 4)).astype(np.float64)
    sigma2 = (0.5 + rng.random((40, 4))).astype(np.float64)

    a = mc_rank_counts(mu, sigma2, 200, percentile=0.2, seed=7)
    b = mc_rank_counts(mu, sigma2, 200, percentile=0.2, seed=7)
    for x, y in zip(a, b):
        assert np.array_equal(x, y)


@pytest.mark.skipif(
    not _mc_ext_available(),
    reason="Cython extension is not built (zen_fronts.selection_core._mc_rank)",
)
def test_mc_rank_changes_with_seed_on_typical_case() -> None:
    from zen_fronts.selection_core._mc_rank import mc_rank_counts  # type: ignore

    rng = np.random.default_rng(321)
    mu = rng.normal(size=(64, 3)).astype(np.float64)
    sigma2 = (0.2 + rng.random((64, 3))).astype(np.float64)

    a = mc_rank_counts(mu, sigma2, 400, percentile=0.2, seed=1)
    b = mc_rank_counts(mu, sigma2, 400, percentile=0.2, seed=2)

    # Very unlikely to be identical across all outputs.
    assert not (np.array_equal(a[0], b[0]) and np.array_equal(a[1], b[1]) and np.array_equal(a[2], b[2]))
