"""
Expert-Aware Scheduler - Router-guided sparse expert loading for MoE.

This is the key innovation for trillion-parameter MoE models:

A 1T MoE model might have 128 experts per layer, but only 2 are activated
per token. AirLLM would load all 128 experts (~128x unnecessary I/O).
Dokodemo's scheduler:

1. Loads only the gating/router network first (tiny: ~1MB)
2. Runs the router to determine which experts are selected
3. Loads ONLY the selected experts (e.g., 2 out of 128)
4. Achieves up to 64x I/O reduction for MoE layers

The scheduler also implements speculative expert prefetching: based on
the routing patterns observed so far, it can predict which experts will
be needed for the next layer and start loading them in advance.
"""

from __future__ import annotations

import logging
from collections import Counter, defaultdict
from typing import Optional

import torch
import torch.nn.functional as F

from dokodemo_ai.engine.streaming import TensorStreamer
from dokodemo_ai.graph.compiler import LayerInfo, ModelGraph

logger = logging.getLogger(__name__)


class ExpertScheduler:
    """
    Manages expert selection and loading for MoE models.

    Responsibilities:
    1. Execute router forward pass to get expert selections
    2. Load only selected experts from disk
    3. Track expert usage statistics for prefetching
    4. Speculatively prefetch likely-needed experts
    """

    def __init__(
        self,
        model_graph: ModelGraph,
        streamer: TensorStreamer,
        device: torch.device,
        enable_speculation: bool = True,
    ):
        self.graph = model_graph
        self.streamer = streamer
        self.device = device
        self.enable_speculation = enable_speculation

        # Expert usage statistics for speculative prefetching
        # Maps (layer_index) -> Counter of expert indices
        self._expert_frequency: dict[int, Counter] = defaultdict(Counter)

        # Cross-layer expert correlation
        # Maps (layer_i, expert_j) -> Counter of (layer_i+1, expert_k)
        self._expert_transitions: dict[
            tuple[int, int], Counter
        ] = defaultdict(Counter)

        # Last selected experts per layer (for correlation tracking)
        self._last_selections: dict[int, list[int]] = {}

    def route_and_load(
        self,
        layer: LayerInfo,
        hidden_states: torch.Tensor,
        router_weights: dict[str, torch.Tensor],
    ) -> tuple[list[int], torch.Tensor, list[dict[str, torch.Tensor]]]:
        """
        Run the router and load only selected experts.

        Args:
            layer: MoE layer info.
            hidden_states: Input hidden states [batch, seq, hidden].
            router_weights: Pre-loaded router weight tensors.

        Returns:
            Tuple of:
            - expert_indices: List of selected expert indices
            - routing_weights: Tensor of routing weights for selected experts
            - expert_tensors: List of weight dicts for selected experts
        """
        # Step 1: Run the router
        expert_indices, routing_weights = self._run_router(
            layer, hidden_states, router_weights
        )

        # Step 2: Update statistics
        self._update_stats(layer.index, expert_indices)

        # Step 3: Load selected experts
        expert_tensors = []
        for idx in expert_indices:
            tensors = self.streamer.load_expert(
                layer, idx, to_device=self.device
            )
            expert_tensors.append(tensors)

        # Step 4: Speculatively prefetch for next layer
        if self.enable_speculation:
            self._speculative_prefetch(layer.index, expert_indices)

        return expert_indices, routing_weights, expert_tensors

    def _run_router(
        self,
        layer: LayerInfo,
        hidden_states: torch.Tensor,
        router_weights: dict[str, torch.Tensor],
    ) -> tuple[list[int], torch.Tensor]:
        """
        Execute the router/gating network.

        The router is a simple linear layer that maps hidden_states
        to expert selection logits. We apply top-k to select experts.
        """
        num_active = layer.num_active_experts or 2

        # Find the router weight tensor
        router_w = None
        router_b = None
        for name, tensor in router_weights.items():
            if "weight" in name:
                router_w = tensor
            elif "bias" in name:
                router_b = tensor

        if router_w is None:
            raise RuntimeError(
                f"Could not find router weight in layer {layer.index}. "
                f"Available: {list(router_weights.keys())}"
            )

        # Router forward: linear projection
        # hidden_states: [batch, seq_len, hidden_size]
        # router_w: [num_experts, hidden_size]
        # We use the last token's hidden state for routing (during generation)
        if hidden_states.dim() == 3:
            routing_input = hidden_states[:, -1, :]  # [batch, hidden]
        else:
            routing_input = hidden_states

        logits = F.linear(routing_input, router_w, router_b)  # [batch, num_experts]

        # Softmax + top-k selection
        routing_probs = F.softmax(logits, dim=-1)
        top_k_weights, top_k_indices = torch.topk(
            routing_probs, num_active, dim=-1
        )

        # Normalize top-k weights
        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)

        # For single-sequence inference, flatten to lists
        expert_indices = top_k_indices[0].tolist()
        routing_weights_out = top_k_weights[0]

        return expert_indices, routing_weights_out

    def _update_stats(
        self, layer_index: int, expert_indices: list[int]
    ) -> None:
        """Update expert usage statistics."""
        for idx in expert_indices:
            self._expert_frequency[layer_index][idx] += 1

        # Track cross-layer transitions
        prev = self._last_selections.get(layer_index - 1)
        if prev is not None:
            for prev_idx in prev:
                for curr_idx in expert_indices:
                    self._expert_transitions[
                        (layer_index - 1, prev_idx)
                    ][curr_idx] += 1

        self._last_selections[layer_index] = expert_indices

    def _speculative_prefetch(
        self, current_layer: int, current_experts: list[int]
    ) -> None:
        """
        Speculatively prefetch experts for the next MoE layer.

        Uses two signals:
        1. Expert frequency: Load the most commonly used experts
        2. Cross-layer correlation: If expert A in layer N often leads
           to expert B in layer N+1, prefetch B
        """
        # Find next MoE layer
        next_moe = None
        for layer in self.graph.layers:
            if layer.index > current_layer and layer.is_moe:
                next_moe = layer
                break

        if next_moe is None:
            return

        # Predict experts for next layer
        predicted: Counter = Counter()

        # Signal 1: Frequency-based prediction
        freq = self._expert_frequency.get(next_moe.index)
        if freq:
            for expert_idx, count in freq.most_common(4):
                predicted[expert_idx] += count

        # Signal 2: Transition-based prediction
        for curr_idx in current_experts:
            transitions = self._expert_transitions.get(
                (current_layer, curr_idx)
            )
            if transitions:
                for next_idx, count in transitions.most_common(2):
                    predicted[next_idx] += count * 2  # Weight transitions higher

        # Prefetch top predicted experts
        num_to_prefetch = min(
            next_moe.num_active_experts * 2,  # Prefetch 2x active count
            len(predicted),
        )

        if num_to_prefetch > 0:
            top_predicted = [
                idx for idx, _ in predicted.most_common(num_to_prefetch)
            ]
            self.streamer.prefetch_experts(next_moe, top_predicted)
            logger.debug(
                "Speculative prefetch: layer %d experts %s",
                next_moe.index,
                top_predicted,
            )

    def get_stats(self) -> dict:
        """Return expert usage statistics."""
        stats = {}
        for layer_idx, freq in self._expert_frequency.items():
            total = sum(freq.values())
            stats[f"layer_{layer_idx}"] = {
                "total_activations": total,
                "top_experts": freq.most_common(5),
                "unique_experts_used": len(freq),
            }
        return stats

    def reset_stats(self) -> None:
        """Reset all statistics (e.g., between prompts)."""
        self._expert_frequency.clear()
        self._expert_transitions.clear()
        self._last_selections.clear()
