"""
Tests for the Expert-Aware Scheduler.

Verifies:
1. Router correctly selects top-k experts
2. Expert usage statistics are accumulated correctly
3. Speculative prefetching predicts plausible experts
4. Stats reset works correctly
"""

import pytest
import torch

from dokodemo_ai.engine.scheduler import ExpertScheduler
from dokodemo_ai.graph.compiler import ComponentInfo, LayerInfo, ModelGraph, TensorRef


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_moe_graph(
    num_layers: int = 4,
    num_experts: int = 8,
    num_active: int = 2,
    hidden_size: int = 64,
) -> ModelGraph:
    from dokodemo_ai.graph.compiler import EmbeddingInfo, HeadInfo

    # Build experts with a fixed layer index (0) — shared across all layers for
    # testing purposes. The scheduler tests don't depend on exact layer indices
    # in component names.
    experts = [
        ComponentInfo(
            name=f"model.layers.0.block_sparse_moe.experts.{j}",
            kind="expert",
            params=[
                TensorRef(
                    name=f"model.layers.0.block_sparse_moe.experts.{j}.w1",
                    file="model.safetensors",
                    dtype="F16",
                    shape=[256, hidden_size],
                )
            ],
        )
        for j in range(num_experts)
    ]

    router = ComponentInfo(
        name=f"model.layers.0.block_sparse_moe.gate",
        kind="router",
        params=[
            TensorRef(
                name=f"model.layers.0.block_sparse_moe.gate.weight",
                file="model.safetensors",
                dtype="F16",
                shape=[num_experts, hidden_size],
            )
        ],
    )

    layers = [
        LayerInfo(
            index=i,
            module_name=f"model.layers.{i}",
            is_moe=True,
            num_experts=num_experts,
            num_active_experts=num_active,
            components=([router] if i == 0 else []) + experts,
        )
        for i in range(num_layers)
    ]

    return ModelGraph(
        model_path="/fake",
        config={
            "hidden_size": hidden_size,
            "vocab_size": 1000,
            "num_local_experts": num_experts,
            "num_experts_per_tok": num_active,
        },
        embedding=EmbeddingInfo(module_name="embed_tokens"),
        layers=layers,
        head=HeadInfo(module_name="lm_head"),
        safetensors_index={},
        num_layers=num_layers,
        hidden_size=hidden_size,
        vocab_size=1000,
    )


class MockStreamer:
    """Fake streamer that records which experts were loaded."""

    def __init__(self):
        self.loaded_experts = []
        self.prefetched_experts = []

    def load_expert(self, layer, expert_index, to_device=None):
        self.loaded_experts.append((layer.index, expert_index))
        # Return fake tensors
        return {f"expert_{expert_index}.weight": torch.randn(64, 64)}

    def prefetch_experts(self, layer, indices):
        self.prefetched_experts.extend(
            (layer.index, idx) for idx in indices
        )


class TestRouterForward:
    def test_topk_selection(self):
        """Router should select exactly num_active_experts experts."""
        graph = _make_moe_graph(num_experts=8, num_active=2, hidden_size=64)
        streamer = MockStreamer()
        scheduler = ExpertScheduler(
            model_graph=graph,
            streamer=streamer,
            device=torch.device("cpu"),
            enable_speculation=False,
        )

        # Fake router weights: [num_experts, hidden_size]
        router_w = torch.randn(8, 64)
        router_tensors = {
            "model.layers.0.block_sparse_moe.gate.weight": router_w
        }
        hidden = torch.randn(1, 1, 64)
        layer = graph.layers[0]

        expert_indices, routing_weights, _ = scheduler.route_and_load(
            layer, hidden, router_tensors
        )

        assert len(expert_indices) == 2  # num_active_experts
        assert routing_weights.shape[0] == 2
        assert all(0 <= idx < 8 for idx in expert_indices)

    def test_routing_weights_sum_to_one(self):
        """Routing weights should sum to 1 after normalization."""
        graph = _make_moe_graph(num_experts=8, num_active=3, hidden_size=32)
        streamer = MockStreamer()
        scheduler = ExpertScheduler(
            model_graph=graph,
            streamer=streamer,
            device=torch.device("cpu"),
            enable_speculation=False,
        )

        router_tensors = {
            "model.layers.0.block_sparse_moe.gate.weight": torch.randn(8, 32)
        }
        hidden = torch.randn(1, 1, 32)
        layer = graph.layers[0]

        _, routing_weights, _ = scheduler.route_and_load(
            layer, hidden, router_tensors
        )

        total = routing_weights.sum().item()
        assert abs(total - 1.0) < 1e-5, f"Weights sum to {total}, expected ~1.0"

    def test_only_selected_experts_loaded(self):
        """Only the top-k selected experts should be loaded from streamer."""
        graph = _make_moe_graph(num_experts=8, num_active=2, hidden_size=32)
        streamer = MockStreamer()
        scheduler = ExpertScheduler(
            model_graph=graph,
            streamer=streamer,
            device=torch.device("cpu"),
            enable_speculation=False,
        )

        router_tensors = {
            "model.layers.0.block_sparse_moe.gate.weight": torch.randn(8, 32)
        }
        hidden = torch.randn(1, 1, 32)
        layer = graph.layers[0]

        expert_indices, _, _ = scheduler.route_and_load(
            layer, hidden, router_tensors
        )

        # Only 2 experts should be loaded (not all 8)
        loaded_indices = [idx for _, idx in streamer.loaded_experts]
        assert len(loaded_indices) == 2
        assert set(loaded_indices) == set(expert_indices)


class TestExpertStats:
    def test_stats_accumulated(self):
        """Expert usage should be tracked across multiple calls."""
        graph = _make_moe_graph(num_experts=4, num_active=1, hidden_size=16)
        streamer = MockStreamer()
        scheduler = ExpertScheduler(
            model_graph=graph,
            streamer=streamer,
            device=torch.device("cpu"),
            enable_speculation=False,
        )

        # Route 5 times
        router_tensors = {
            "model.layers.0.block_sparse_moe.gate.weight": torch.randn(4, 16)
        }
        layer = graph.layers[0]

        for _ in range(5):
            hidden = torch.randn(1, 1, 16)
            scheduler.route_and_load(layer, hidden, router_tensors)

        stats = scheduler.get_stats()
        assert "layer_0" in stats
        layer_stats = stats["layer_0"]
        assert layer_stats["total_activations"] == 5

    def test_reset_clears_stats(self):
        """After reset, stats should be empty."""
        graph = _make_moe_graph(num_experts=4, num_active=1, hidden_size=16)
        streamer = MockStreamer()
        scheduler = ExpertScheduler(
            model_graph=graph,
            streamer=streamer,
            device=torch.device("cpu"),
            enable_speculation=False,
        )

        router_tensors = {
            "model.layers.0.block_sparse_moe.gate.weight": torch.randn(4, 16)
        }
        layer = graph.layers[0]
        scheduler.route_and_load(layer, torch.randn(1, 1, 16), router_tensors)

        scheduler.reset_stats()
        stats = scheduler.get_stats()
        assert len(stats) == 0


class TestSpeculativePrefetching:
    def test_prefetching_triggered(self):
        """With speculation enabled, experts for next layer should be prefetched."""
        graph = _make_moe_graph(num_experts=8, num_active=2, hidden_size=32)
        streamer = MockStreamer()
        scheduler = ExpertScheduler(
            model_graph=graph,
            streamer=streamer,
            device=torch.device("cpu"),
            enable_speculation=True,
        )

        router_tensors = {
            "model.layers.0.block_sparse_moe.gate.weight": torch.randn(8, 32)
        }
        layer = graph.layers[0]

        # After several calls, speculation should kick in
        for _ in range(3):
            scheduler.route_and_load(
                layer, torch.randn(1, 1, 32), router_tensors
            )

        # Prefetching may or may not trigger on the first call
        # but should build up stats
        stats = scheduler.get_stats()
        assert "layer_0" in stats

    def test_no_prefetch_without_speculation(self):
        """With speculation disabled, no prefetching should occur."""
        graph = _make_moe_graph(num_experts=8, num_active=2, hidden_size=32)
        streamer = MockStreamer()
        scheduler = ExpertScheduler(
            model_graph=graph,
            streamer=streamer,
            device=torch.device("cpu"),
            enable_speculation=False,
        )

        router_tensors = {
            "model.layers.0.block_sparse_moe.gate.weight": torch.randn(8, 32)
        }
        layer = graph.layers[0]
        scheduler.route_and_load(layer, torch.randn(1, 1, 32), router_tensors)

        assert len(streamer.prefetched_experts) == 0
