"""DeepSeekMoE：细粒度专家 + 共享专家 + 无辅助损失负载均衡。

三个和普通 MoE 不一样的地方：

1. 细粒度专家（fine-grained）：把专家切得更小更多，top-k 也调大。
   同样的激活参数量，组合数从 C(8,2) 变成 C(64,8)，专家分工能更专一。

2. 共享专家（shared expert）：抽出 1 个专家让所有 token 都过。
   通用知识（语法、常识）由它承担，路由专家就不必人手一份重复学。

3. 无辅助损失负载均衡（aux-loss-free，V3 的关键改进）：
   传统做法是加一个负载均衡损失，但它和语言建模目标是打架的，会损伤效果。
   V3 改成给每个专家挂一个**不参与梯度**的偏置 b_i，只在 top-k 选择时加上：
       选择依据 = s_i + b_i，  聚合权重仍用原始 s_i
   每步训练后，谁超载就把谁的 b_i 调小，谁欠载就调大：
       b_i += γ * sign(平均负载 - 该专家负载)
   于是负载被掰平了，而梯度完全没被污染。
"""

from typing import Tuple

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

from .layers import SwiGLU


class Gate(nn.Module):
    """路由门控：决定每个 token 该交给哪几个专家。"""

    def __init__(self, cfg):
        super().__init__()
        self.dim = cfg.dim
        self.n_experts = cfg.n_routed_experts
        self.topk = cfg.n_activated_experts
        self.n_groups = cfg.n_expert_groups
        self.topk_groups = cfg.n_limited_groups
        self.score_func = cfg.score_func
        self.route_scale = cfg.route_scale
        self.aux_alpha = cfg.aux_loss_alpha

        self.weight = nn.Parameter(torch.empty(self.n_experts, self.dim))
        nn.init.normal_(self.weight, std=0.02)

        # 负载均衡偏置：只影响"选谁"，不参与反向传播
        self.register_buffer("expert_bias", torch.zeros(self.n_experts))
        # 统计窗口内每个专家接了多少 token，供 update_bias 使用
        self.register_buffer("load_count", torch.zeros(self.n_experts), persistent=False)
        self.last_imbalance = 1.0

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """x: (N, dim) -> (weights (N,topk), indices (N,topk), aux_loss 标量)"""
        logits = F.linear(x.float(), self.weight.float())
        if self.score_func == "sigmoid":
            scores = logits.sigmoid()
        else:
            scores = logits.softmax(dim=-1)
        original = scores

        # ---- 选择用的分数 = 原始分数 + 负载偏置 ----
        sel = scores + self.expert_bias

        # ---- group-limited routing：先选组再选专家，限制跨设备通信量 ----
        if self.n_groups > 1:
            g = sel.view(-1, self.n_groups, self.n_experts // self.n_groups)
            per_group = g.size(-1)
            if per_group > 1:
                group_score = g.topk(min(2, per_group), dim=-1)[0].sum(dim=-1)
            else:
                group_score = g.squeeze(-1)
            keep = group_score.topk(self.topk_groups, dim=-1)[1]            # (N, topk_groups)
            gmask = torch.zeros_like(group_score, dtype=torch.bool).scatter_(1, keep, True)
            sel = sel.masked_fill(~gmask.unsqueeze(-1).expand_as(g).reshape_as(sel),
                                  torch.finfo(sel.dtype).min)

        indices = sel.topk(self.topk, dim=-1)[1]                            # (N, topk)
        weights = original.gather(1, indices)
        if self.score_func == "sigmoid":
            weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(1e-9)
        weights = weights * self.route_scale

        # ---- 统计负载 + 兜底的序列级辅助损失（权重很小）----
        onehot = F.one_hot(indices, self.n_experts).sum(dim=1).float()      # (N, E)
        counts = onehot.sum(dim=0)
        with torch.no_grad():
            self.load_count += counts.detach()

        if self.training and self.aux_alpha > 0:
            f = counts / counts.sum().clamp_min(1.0)                        # 实际负载占比
            p = original.mean(dim=0)                                        # 平均路由概率
            aux_loss = self.aux_alpha * self.n_experts * (f * p).sum()
        else:
            aux_loss = x.new_zeros(())

        return weights.type_as(x), indices, aux_loss

    @torch.no_grad()
    def update_bias(self, speed: float):
        """训练循环每步调一次：把超载专家的偏置压低、欠载的抬高。"""
        total = self.load_count.sum()
        if total <= 0:
            return
        self.last_imbalance = (self.load_count.max() / total * self.n_experts).item()
        target = total / self.n_experts
        err = target - self.load_count
        self.expert_bias += speed * torch.sign(err)
        self.load_count.zero_()

    @torch.no_grad()
    def load_stats(self) -> dict:
        """负载不均衡度 = 最大负载 / 平均负载。1.0 完全均衡，topk/E 的倒数为最差。"""
        total = self.load_count.sum()
        if total <= 0:
            return {"imbalance": self.last_imbalance}
        frac = self.load_count / total
        return {"imbalance": (frac.max() * self.n_experts).item()}


class MoE(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.dim = cfg.dim
        self.n_experts = cfg.n_routed_experts
        self.gate = Gate(cfg)
        self.experts = nn.ModuleList([SwiGLU(cfg.dim, cfg.moe_inter_dim)
                                      for _ in range(cfg.n_routed_experts)])
        self.shared = (SwiGLU(cfg.dim, cfg.moe_inter_dim * cfg.n_shared_experts)
                       if cfg.n_shared_experts > 0 else None)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        shape = x.shape
        x = x.view(-1, self.dim)
        weights, indices, aux = self.gate(x)

        y = torch.zeros_like(x)
        for i, expert in enumerate(self.experts):
            token_idx, slot = torch.where(indices == i)
            if token_idx.numel() == 0:
                continue
            y.index_add_(0, token_idx,
                         expert(x[token_idx]) * weights[token_idx, slot].unsqueeze(-1))

        if self.shared is not None:
            y = y + self.shared(x)
        return y.view(shape), aux


class DenseFFN(nn.Module):
    """前几层用的普通 FFN。V3 的做法：靠前的层先稠密，训练更稳。"""

    def __init__(self, cfg):
        super().__init__()
        self.ffn = SwiGLU(cfg.dim, cfg.dense_inter_dim)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        return self.ffn(x), x.new_zeros(())