"""
pipeline.py
===========
Torrent Expert-Aware Multi-Batch Pipeline 主调度器（论文 §5）

核心流程（单 batch-group 迭代）：
  ┌─────────────────────────────────────────────────────────────┐
  │  Batch i:   [Attention]                                     │
  │                ↕ (io_stream 并行预取 gate + K 热专家)       │
  ├─────────────────────────────────────────────────────────────┤
  │  …（重复 n 次）                                             │
  ├─────────────────────────────────────────────────────────────┤
  │  MoE 层：gate → 冷专家预取 → 热优先执行 → 就绪顺序执行    │
  └─────────────────────────────────────────────────────────────┘
"""

from __future__ import annotations

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

import torch
import torch.nn as nn
import torch.nn.functional as F

from srep_moe.prefetch.correlation_table import CorrelationTable
from srep_moe.prefetch.async_loader import AsyncLoader
from srep_moe.runtime.expert_executor import ExpertExecutor, LayerExecutionStats
from srep_moe.runtime.memory_manager import MemoryManager, StorageTier
from srep_moe.metrics.collector import MetricsCollector, StepMetrics

logger = logging.getLogger(__name__)


# ===========================================================================
# 玩具 MoE 层（用于本地验证，不依赖真实模型权重）
# ===========================================================================

class ToyAttention(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, H = x.shape
        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v)
        return self.o_proj(out.transpose(1, 2).reshape(B, T, H))


class ToyGate(nn.Module):
    def __init__(self, hidden_size: int, num_experts: int):
        super().__init__()
        self.gate = nn.Linear(hidden_size, num_experts, bias=False)

    def forward(self, x: torch.Tensor, top_k: int) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = self.gate(x)
        weights = F.softmax(logits, dim=-1)
        top_weights, top_indices = torch.topk(weights, top_k, dim=-1)
        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True)
        return top_weights, top_indices


class ToyExpertFFN(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

    def get_cpu_weights(self) -> Dict[str, torch.Tensor]:
        return {
            "gate_proj": self.gate_proj.weight.data.cpu(),
            "up_proj": self.up_proj.weight.data.cpu(),
            "down_proj": self.down_proj.weight.data.cpu(),
        }


# ===========================================================================
# Torrent Pipeline
# ===========================================================================

@dataclass
class TorrentConfig:
    """Torrent 流水线配置。"""
    hidden_size: int = 512
    intermediate_size: int = 1024
    num_heads: int = 8
    num_layers: int = 4
    num_experts: int = 8
    top_k: int = 2
    n: int = 4                    # batch-group 大小（由 ConstraintSolver 决定）
    dtype: torch.dtype = torch.float32
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    prefetch_top_k: int = 2
    enable_correlation: bool = True


class TorrentPipeline:
    """
    Torrent Expert-Aware Multi-Batch Pipeline。

    Args:
        config  : TorrentConfig
        metrics : MetricsCollector（可选）
    """

    def __init__(
        self,
        config: TorrentConfig,
        metrics: Optional[MetricsCollector] = None,
    ):
        self.cfg = config
        self.device = torch.device(config.device)
        self.metrics = metrics or MetricsCollector()

        self._build_model()

        self.corr_table = CorrelationTable(
            num_layers=config.num_layers,
            num_experts=config.num_experts,
            top_k=config.top_k,
        )

        if self.device.type == "cuda":
            self.io_stream = torch.cuda.Stream(device=self.device)
            self.compute_stream = torch.cuda.current_stream(self.device)
        else:
            self.io_stream = None
            self.compute_stream = None

        self.loader = AsyncLoader(
            device=self.device,
            io_stream=self.io_stream,
            compute_stream=self.compute_stream,
        )
        self.executor = ExpertExecutor(
            loader=self.loader,
            device=self.device,
            top_k=config.top_k,
        )
        self.mem_mgr = MemoryManager(device=self.device)

        self._register_experts()
        self._prev_routing: Optional[torch.Tensor] = None

    # ------------------------------------------------------------------
    # 模型构建
    # ------------------------------------------------------------------

    def _build_model(self) -> None:
        cfg = self.cfg
        dtype = cfg.dtype

        self.attention_layers = nn.ModuleList([
            ToyAttention(cfg.hidden_size, cfg.num_heads).to(dtype=dtype)
            for _ in range(cfg.num_layers)
        ])
        self.gate_layers = nn.ModuleList([
            ToyGate(cfg.hidden_size, cfg.num_experts).to(dtype=dtype)
            for _ in range(cfg.num_layers)
        ])
        self.expert_layers = nn.ModuleList([
            nn.ModuleList([
                ToyExpertFFN(cfg.hidden_size, cfg.intermediate_size).to(dtype=dtype)
                for _ in range(cfg.num_experts)
            ])
            for _ in range(cfg.num_layers)
        ])

        # attention + gate → GPU，专家权重留 CPU（offloading）
        for attn in self.attention_layers:
            attn.to(self.device)
        for gate in self.gate_layers:
            gate.to(self.device)

    def _register_experts(self) -> None:
        for layer_idx, experts in enumerate(self.expert_layers):
            for expert_id, expert in enumerate(experts):
                cpu_w = expert.get_cpu_weights()
                key = f"layer{layer_idx}_expert{expert_id}"
                self.loader.register_expert(expert_id, cpu_w)
                self.mem_mgr.register_tensor(key, expert.gate_proj.weight.data, StorageTier.CPU)

    # ------------------------------------------------------------------
    # 核心推理接口
    # ------------------------------------------------------------------

    def forward_batch_group(
        self,
        batch_inputs: List[torch.Tensor],
    ) -> List[torch.Tensor]:
        """
        对一个 batch-group（n 个 batch）执行 Torrent 流水线前向传播。

        Args:
            batch_inputs: list of [seq_len, hidden_size] 张量，长度 = n

        Returns:
            list of [seq_len, hidden_size] 输出张量
        """
        n = len(batch_inputs)
        cfg = self.cfg

        step_start = time.perf_counter()
        outputs = [None] * n

        for layer_idx in range(cfg.num_layers):
            layer_layer_stats: List[LayerExecutionStats] = []
            all_hidden = []
            all_routing_w = []
            all_routing_idx = []

            # 阶段 1：多 batch Attention + 热专家预取
            for batch_i, x in enumerate(batch_inputs):
                if x.dim() == 2:
                    x = x.unsqueeze(0)
                x = x.to(self.device, dtype=cfg.dtype)
                attn_out = self.attention_layers[layer_idx](x)

                if batch_i == 0:
                    if self._prev_routing is not None and cfg.enable_correlation:
                        hot_ids = self.corr_table.predict_hot_experts(
                            layer_idx, self._prev_routing, cfg.prefetch_top_k
                        )
                    else:
                        hot_ids = list(range(cfg.prefetch_top_k))

                    self.loader.reset_layer()
                    self.loader.prefetch_hot_experts(hot_ids)

                all_hidden.append(attn_out.squeeze(0))

            # 阶段 2：gate + 路由 → 冷专家预取
            all_rw_list = []
            all_ri_list = []
            for hidden in all_hidden:
                rw, ri = self.gate_layers[layer_idx](hidden, cfg.top_k)
                all_rw_list.append(rw)
                all_ri_list.append(ri)

            all_indices_cat = torch.cat(all_ri_list, dim=0)
            active_set = set(all_indices_cat.reshape(-1).tolist())
            cold_ids = [int(e) for e in active_set if e not in set(hot_ids)]
            self.loader.prefetch_cold_experts(cold_ids)

            if self._prev_routing is not None and cfg.enable_correlation:
                self.corr_table.update_batch(layer_idx, self._prev_routing, all_indices_cat)
            self._prev_routing = all_indices_cat.detach()

            # 阶段 3：MoE 执行（专家重排）
            for batch_i in range(n):
                hidden = all_hidden[batch_i]
                rw = all_rw_list[batch_i]
                ri = all_ri_list[batch_i]
                is_last_batch = (batch_i == n - 1)

                moe_out, stats = self.executor.execute_layer(
                    layer_idx=layer_idx,
                    hidden_states=hidden,
                    routing_weights=rw,
                    routing_indices=ri,
                    hot_expert_ids=hot_ids,
                    cold_expert_ids=cold_ids,
                    evict_after_use=is_last_batch,
                )
                layer_layer_stats.append(stats)
                all_hidden[batch_i] = hidden + moe_out

            if n > 1:
                self.loader.evict_experts([int(e) for e in active_set])

            outputs = all_hidden

            total_wait = sum(s.total_wait_s for s in layer_layer_stats)
            self.metrics.record_layer(
                layer_idx=layer_idx,
                intra_bubble_s=total_wait,
                num_expert_used=layer_layer_stats[0].num_experts_used if layer_layer_stats else 0,
                hot_hit_count=len([e for e in hot_ids if e in active_set]),
                hot_total=len(hot_ids),
            )

        step_time = time.perf_counter() - step_start
        total_tokens = sum(h.shape[0] for h in outputs)
        self.metrics.record_step(step_time_s=step_time, tokens_generated=total_tokens, n=n)

        return outputs

    def generate_step(self, input_ids_batch: List[torch.Tensor]) -> List[torch.Tensor]:
        """
        单步 decode 的 Torrent 调用入口（玩具版：随机 embedding）。
        """
        cfg = self.cfg
        batch_inputs = [
            torch.randn(ids.shape[-1], cfg.hidden_size, dtype=cfg.dtype, device="cpu")
            for ids in input_ids_batch
        ]
        return self.forward_batch_group(batch_inputs)

    def get_summary(self) -> dict:
        return {
            "pipeline_metrics": self.metrics.get_summary(),
            "memory_stats": self.mem_mgr.get_stats(),
            "correlation_coverage": self.corr_table.coverage_stats(),
        }
