"""
MoE模型在Snapdragon NPU上的部署优势

原理: 不是所有参数同时激活，而是路由到最相关的"专家"
效果: 30B参数模型，实际计算量可能只有8-10B
"""

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

class MoEExpert(nn.Module):
    """单个专家网络"""
    def __init__(self, d_model: int = 768, d_ff: int = 3072):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
    
    def forward(self, x):
        return self.fc2(F.gelu(self.fc1(x)))


class MoELayer(nn.Module):
    """
    Mix of Experts 层
    
    路由每个token到最合适的专家
    
    对NPU的意义:
    - Element Accelerator加速路由计算
    - 共享内存减少专家切换开销
    - 30B参数模型实际计算量~10B
    """
    def __init__(self, d_model: int = 768, num_experts: int = 8,
                 top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.experts = nn.ModuleList([
            MoEExpert(d_model) for _ in range(num_experts)
        ])
        self.gate = nn.Linear(d_model, num_experts, bias=False)
    
    def forward(self, x):
        """
        Args:
            x: (B, L, d_model)
        Returns:
            output: (B, L, d_model)
        """
        B, L, D = x.shape
        
        # 路由
        gate_logits = self.gate(x)  # (B, L, num_experts)
        gate_probs = F.softmax(gate_logits, dim=-1)
        
        # Top-K选择
        topk_vals, topk_idx = gate_probs.topk(self.top_k, dim=-1)
        topk_vals = topk_vals / topk_vals.sum(dim=-1, keepdim=True)
        
        # 专家计算
        output = torch.zeros_like(x)
        for k in range(self.top_k):
            for b in range(B):
                for l in range(L):
                    expert_idx = topk_idx[b, l, k].item()
                    weight = topk_vals[b, l, k]
                    output[b, l] += weight * self.experts[expert_idx](x[b, l])
        
        return output


# DMS专用MoE模型示例
class MoE_DMS_Model(nn.Module):
    """
    DMS专用MoE模型
    
    场景: 不同驾驶场景使用不同专家
    - 专家1: 日间正常驾驶
    - 专家2: 夜间驾驶
    - 专家3: 隧道场景
    - 专家4: 高速公路
    - 专家5: 城市道路
    - 专家6: 停车状态
    - 专家7: 恶劣天气
    - 专家8: 疲劳状态
    
    效果: 总参数8x，实际计算1x
    """
    def __init__(self):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 64, 3, 2, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, 2, 1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.moe = MoELayer(d_model=128, num_experts=8, top_k=2)
        self.head = nn.Linear(128, 10)  # 10类行为
        
    def forward(self, x):
        feat = self.backbone(x).flatten(1).unsqueeze(1)  # (B, 1, 128)
        out = self.moe(feat).squeeze(1)
        return self.head(out)


if __name__ == "__main__":
    # 测试MoE DMS模型
    model = MoE_DMS_Model()
    x = torch.randn(4, 3, 224, 224)
    out = model(x)
    
    total_params = sum(p.numel() for p in model.parameters())
    active_params = sum(
        p.numel() for name, p in model.named_parameters()
        if 'experts.0' in name or 'experts.1' in name or 
           'backbone' in name or 'head' in name or 'gate' in name
    )
    
    print(f"输出: {out.shape}")
    print(f"总参数: {total_params:,}")
    print(f"激活参数(Top-2): ~{active_params:,}")
    print(f"稀疏比: {1 - active_params/total_params:.1%}")
    print("\n✅ MoE模型测试通过")