"""Unit tests for the EPLB rebalancing policy: structural self-consistency and determinism of rebalance_experts_global.

The original checks run on CPU; the new placement cases cover CPU/CUDA and skip CUDA when no GPU is available.
"""

from __future__ import annotations

import pytest
import torch

from deepinfer.model_executor.eplb.policy import _batch_greedy_indices, rebalance_experts_global

# (num_steps, num_moe_layers, num_logical_experts, num_physical_experts, num_local_physical_experts)
# Coverage: multi-step / multi-layer, local=1 (placement skipped), redundant=0, and redundant > logical (multiple replicas of one expert).
CASES = [
    (1, 1, 16, 24, 3),
    (3, 2, 64, 72, 9),
    (2, 4, 128, 160, 20),
    (1, 1, 8, 8, 1),
    (2, 1, 4, 16, 1),
    (5, 3, 4, 32, 4),
    (4, 2, 64, 128, 16),
    (3, 3, 16, 64, 8),
]


def _check_structure(
    phy2log: torch.Tensor,
    log2phy: torch.Tensor,
    log_count: torch.Tensor,
    num_physical_experts: int,
) -> None:
    """Structural self-consistency: phy2log and log2phy are inverses, log_count matches the valid slots, and the total replica count equals the physical expert count."""
    num_layers = phy2log.size(0)
    # log_count == the number of valid slots in log2phy
    assert log_count.equal((log2phy >= 0).sum(-1))
    # every physical slot can be found in its logical expert's log2phy
    for layer in range(num_layers):
        for physical_index, logical_index in enumerate(phy2log[layer].tolist()):
            assert physical_index in log2phy[layer][logical_index].tolist()
    # total replica count == physical expert count
    assert log_count.sum(-1).equal(torch.full((num_layers,), num_physical_experts, dtype=torch.int))


@pytest.mark.parametrize("num_steps,num_layers,num_logical,num_physical,num_local", CASES)
def test_rebalance_experts_global_structure(num_steps: int, num_layers: int, num_logical: int, num_physical: int, num_local: int) -> None:
    torch.manual_seed(num_steps * 100 + num_layers * 10 + num_physical)
    tokens = torch.randint(0, 5000, (num_steps, num_layers, num_logical))
    phy2log, log2phy, log_count = rebalance_experts_global(tokens, num_physical, num_local)
    _check_structure(phy2log, log2phy, log_count, num_physical)


@pytest.mark.parametrize("num_steps,num_layers,num_logical,num_physical,num_local", CASES)
def test_rebalance_experts_global_deterministic(
    num_steps: int, num_layers: int, num_logical: int, num_physical: int, num_local: int
) -> None:
    torch.manual_seed(num_steps * 100 + num_layers * 10 + num_physical)
    tokens = torch.randint(0, 5000, (num_steps, num_layers, num_logical))
    first = rebalance_experts_global(tokens.clone(), num_physical, num_local)
    second = rebalance_experts_global(tokens.clone(), num_physical, num_local)
    for x, y in zip(first, second, strict=True):
        assert torch.equal(x, y)


@pytest.mark.parametrize("distribution", ["uniform", "skewed", "zeros", "spike"])
def test_rebalance_experts_global_edge_distributions(distribution: str) -> None:
    """Boundary distributions: all zeros (all replicas score the same → many ties), a single spike, etc."""
    num_steps, num_layers, num_logical, num_physical, num_local = 2, 3, 16, 24, 3
    torch.manual_seed(0)
    if distribution == "uniform":
        tokens = torch.randint(1, 5000, (num_steps, num_layers, num_logical))
    elif distribution == "skewed":
        tokens = (torch.randn(num_steps, num_layers, num_logical).exp() * 1000).long()
    elif distribution == "zeros":
        tokens = torch.zeros(num_steps, num_layers, num_logical, dtype=torch.int64)
    else:  # spike
        tokens = torch.nn.functional.one_hot(torch.randint(0, num_logical, (num_steps, num_layers)), num_logical) * 99999
    phy2log, log2phy, log_count = rebalance_experts_global(tokens, num_physical, num_local)
    _check_structure(phy2log, log2phy, log_count, num_physical)


@pytest.fixture(
    params=["cpu", pytest.param("cuda", marks=[pytest.mark.gpu, pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")])]
)
def policy_device(request: pytest.FixtureRequest) -> torch.device:
    return torch.device(request.param)


def test_rebalance_does_not_modify_history(policy_device: torch.device) -> None:
    # the FP32 single-step input stays contiguous after permute, preventing in-place division from modifying the input when the working tensor shares its storage.
    tokens = torch.arange(8, dtype=torch.float32, device=policy_device).reshape(1, 1, 8)
    original = tokens.clone()
    rebalance_experts_global(tokens, 12, 3)
    assert torch.equal(tokens, original)


@pytest.mark.parametrize(
    "history,physical,expected",
    [
        ([[0, 0, 0, 0, 0, 0]], [0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4, 5]),
        # after the expert numbering is shuffled, an even load still preserves the original physical slot order.
        ([[0, 0, 0]], [2, 0, 1, 0, 2, 1], [0, 1, 2, 3, 4, 5]),
        # when placing expert 1 both GPUs score the same, and rank 0 is still preferred even though its accumulated load is higher.
        ([[10, 0, 0, 0, 0, 0], [0, 6, 4, 0, 0, 0]], [0, 1, 2, 3, 4, 5], [0, 1, 3, 2, 4, 5]),
        ([[1, 1, 1, 1, 1, 1]], [0, 1, 2, 3, 4, 5], [0, 2, 4, 1, 3, 5]),
        ([[10, 0, 6, 0, 4, 0], [0, 10, 0, 6, 0, 4]], [0, 1, 2, 3, 4, 5], [0, 1, 5, 2, 3, 4]),
    ],
)
def test_batch_greedy_objective_and_ties(
    history: list[list[int]], physical: list[int], expected: list[int], policy_device: torch.device
) -> None:
    """A hand-built layout locks the original slot order, tie-breaks to the smallest rank ID, and selects for load complementarity across batches."""
    tokens = torch.tensor(history, device=policy_device).unsqueeze(1)
    counts = torch.bincount(torch.tensor(physical), minlength=tokens.size(-1)).unsqueeze(0).to(policy_device)
    physical_map = torch.tensor([physical], dtype=torch.int32, device=policy_device)
    indices = _batch_greedy_indices(tokens, physical_map, counts, 3)
    assert indices.tolist() == [expected]