"""
expert_executor.py
==================
Torrent 专家层内重排执行器（论文 §5 核心创新）

执行顺序原则：
  1. 相同专家的 token 连续计算（提高权重复用）
  2. 热专家优先执行（为冷专家传输争取最多重叠时间）
  3. 其余专家按"传输完成顺序"执行（is_ready() 轮询，减少 bubble）
  4. 每个专家计算完成后立即回收 GPU 显存

专家 FFN 类型：SwiGLU（Mixtral / Qwen3MoE 风格）
  out = (silu(gate_proj(x)) * up_proj(x)) @ down_proj
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn.functional as F

from srep_moe.prefetch.async_loader import AsyncLoader, ExpertState

logger = logging.getLogger(__name__)


@dataclass
class ExpertResult:
    expert_id: int
    token_indices: List[int]
    output: Optional[torch.Tensor]
    compute_time_s: float = 0.0
    wait_time_s: float = 0.0


@dataclass
class LayerExecutionStats:
    layer_idx: int
    total_tokens: int
    num_experts_used: int
    hot_expert_ids: List[int]
    cold_expert_ids: List[int]
    total_compute_s: float = 0.0
    total_wait_s: float = 0.0
    peak_gpu_experts: int = 0


class ExpertExecutor:
    """
    MoE 层专家重排执行器。

    Args:
        loader          : AsyncLoader 实例
        device          : 目标设备
        top_k           : 每 token 激活专家数
        poll_interval_us: 冷专家就绪轮询间隔（微秒）
    """

    def __init__(
        self,
        loader: AsyncLoader,
        device: torch.device,
        top_k: int = 2,
        poll_interval_us: float = 50.0,
    ):
        self.loader = loader
        self.device = device
        self.top_k = top_k
        self.poll_interval_s = poll_interval_us * 1e-6

    # ------------------------------------------------------------------
    # 主接口
    # ------------------------------------------------------------------

    def execute_layer(
        self,
        layer_idx: int,
        hidden_states: torch.Tensor,
        routing_weights: torch.Tensor,
        routing_indices: torch.Tensor,
        hot_expert_ids: List[int],
        cold_expert_ids: List[int],
        evict_after_use: bool = True,
    ) -> Tuple[torch.Tensor, LayerExecutionStats]:
        num_tokens, hidden_size = hidden_states.shape
        output = torch.zeros_like(hidden_states)

        expert_to_tokens = self._build_expert_token_map(routing_indices, routing_weights, num_tokens)

        active_experts = set(expert_to_tokens.keys())
        hot_active = [e for e in hot_expert_ids if e in active_experts]
        cold_active = [e for e in cold_expert_ids if e in active_experts]
        unpredicted = [e for e in active_experts if e not in set(hot_active) | set(cold_active)]
        cold_active.extend(unpredicted)

        stats = LayerExecutionStats(
            layer_idx=layer_idx,
            total_tokens=num_tokens,
            num_experts_used=len(active_experts),
            hot_expert_ids=hot_active,
            cold_expert_ids=cold_active,
        )

        # 阶段 1：执行热专家
        for expert_id in hot_active:
            result = self._run_single_expert(expert_id, hidden_states, expert_to_tokens)
            self._accumulate_output(output, result, expert_to_tokens[expert_id])
            stats.total_compute_s += result.compute_time_s
            stats.total_wait_s += result.wait_time_s
            if evict_after_use:
                self.loader.evict_expert(expert_id)

        # 阶段 2：冷专家按就绪顺序执行
        remaining = list(cold_active)
        while remaining:
            executed_any = False
            for expert_id in remaining:
                if self.loader.is_ready(expert_id):
                    result = self._run_single_expert(expert_id, hidden_states, expert_to_tokens)
                    self._accumulate_output(output, result, expert_to_tokens[expert_id])
                    stats.total_compute_s += result.compute_time_s
                    remaining.remove(expert_id)
                    if evict_after_use:
                        self.loader.evict_expert(expert_id)
                    executed_any = True
                    break

            if not executed_any:
                wait_start = time.perf_counter()
                expert_id = remaining[0]
                weights = self.loader.wait_expert(expert_id)
                stats.total_wait_s += time.perf_counter() - wait_start
                result = self._run_single_expert(
                    expert_id, hidden_states, expert_to_tokens, weights=weights
                )
                self._accumulate_output(output, result, expert_to_tokens[expert_id])
                stats.total_compute_s += result.compute_time_s
                remaining.remove(expert_id)
                if evict_after_use:
                    self.loader.evict_expert(expert_id)

        return output, stats

    # ------------------------------------------------------------------
    # 内部辅助
    # ------------------------------------------------------------------

    def _build_expert_token_map(
        self,
        routing_indices: torch.Tensor,
        routing_weights: torch.Tensor,
        num_tokens: int,
    ) -> Dict[int, Tuple[List[int], torch.Tensor]]:
        expert_map: Dict[int, Tuple[List[int], List[float]]] = {}
        for token_idx in range(num_tokens):
            for k in range(self.top_k):
                eid = int(routing_indices[token_idx, k].item())
                w = float(routing_weights[token_idx, k].item())
                if eid not in expert_map:
                    expert_map[eid] = ([], [])
                expert_map[eid][0].append(token_idx)
                expert_map[eid][1].append(w)

        result = {}
        for eid, (tids, ws) in expert_map.items():
            result[eid] = (tids, torch.tensor(ws, dtype=torch.float32, device=self.device))
        return result

    def _run_single_expert(
        self,
        expert_id: int,
        hidden_states: torch.Tensor,
        expert_to_tokens: Dict,
        weights: Optional[Dict[str, torch.Tensor]] = None,
    ) -> ExpertResult:
        wait_start = time.perf_counter()
        if weights is None:
            weights = self.loader.wait_expert(expert_id)
        wait_time = time.perf_counter() - wait_start

        token_indices, routing_ws = expert_to_tokens[expert_id]
        if not token_indices:
            return ExpertResult(
                expert_id=expert_id, token_indices=[], output=None, wait_time_s=wait_time
            )

        x = hidden_states[token_indices]
        compute_start = time.perf_counter()
        gate_out = F.silu(x @ weights["gate_proj"].T)
        up_out = x @ weights["up_proj"].T
        out = (gate_out * up_out) @ weights["down_proj"].T

        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        compute_time = time.perf_counter() - compute_start

        out = out * routing_ws.unsqueeze(-1)

        return ExpertResult(
            expert_id=expert_id,
            token_indices=token_indices,
            output=out,
            compute_time_s=compute_time,
            wait_time_s=wait_time,
        )

    def _accumulate_output(
        self,
        output: torch.Tensor,
        result: ExpertResult,
        expert_token_info: Tuple[List[int], torch.Tensor],
    ) -> None:
        if result.output is None or not result.token_indices:
            return
        idx_tensor = torch.tensor(result.token_indices, device=self.device)
        output.index_add_(0, idx_tensor, result.output)
