import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Tuple, List

# ==============================================================
# 1. 专家模块（Expert Module）
# ==============================================================

class Expert(nn.Module):
    """
    GMoE的专家模块。
    一个标准的FFN（前馈神经网络），包含两层线性变换和激活函数。
    """
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        dropout: float = 0.1,
        activation: str = "gelu"
    ):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)
        
        if activation == "gelu":
            self.act = nn.GELU()
        elif activation == "relu":
            self.act = nn.ReLU()
        elif activation == "silu":
            self.act = nn.SiLU()
        else:
            raise ValueError(f"Unknown activation: {activation}")
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (batch_size, seq_len, d_model)
        Returns:
            (batch_size, seq_len, d_model)
        """
        return self.w2(self.dropout(self.act(self.w1(x))))


# ==============================================================
# 2. GRU-based 全局路由器（Global Router）
# ==============================================================

class GRURouter(nn.Module):
    """
    GMoE的全局路由器，基于GRU的循环路由组件。
    核心创新：跨层传播路由对数（logits），解决路径坍缩问题。
    """
    def __init__(
        self,
        d_model: int,
        num_global_experts: int,
        gru_hidden_size: int = 128,
        num_experts_per_token: int = 2,
    ):
        super().__init__()
        self.num_global_experts = num_global_experts
        self.num_experts_per_token = num_experts_per_token
        
        # 输入投影：将输入映射到GRU隐藏空间
        self.input_proj = nn.Linear(d_model, gru_hidden_size, bias=False)
        
        # GRU单元：跨层传递路由状态
        self.gru_cell = nn.GRUCell(gru_hidden_size, gru_hidden_size)
        
        # 路由头部：为每个token生成专家选择分数
        # 输出维度为 num_global_experts
        self.routing_head = nn.Linear(gru_hidden_size, num_global_experts, bias=False)
        
        # 初始化
        self._init_weights()
    
    def _init_weights(self):
        for name, param in self.named_parameters():
            if 'weight' in name:
                nn.init.xavier_uniform_(param)
            elif 'bias' in name:
                nn.init.zeros_(param)
    
    def forward(
        self,
        x: torch.Tensor,
        prev_hidden: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (batch_size, seq_len, d_model) 当前层输入
            prev_hidden: (batch_size, seq_len, gru_hidden_size) 前一层GRU隐藏状态
        Returns:
            routing_weights: (batch_size, seq_len, num_global_experts) 路由权重
            selected_experts: (batch_size, seq_len, num_experts_per_token) 选中的专家索引
            hidden_state: (batch_size, seq_len, gru_hidden_size) 当前GRU状态，传递给下一层
        """
        batch_size, seq_len, _ = x.shape
        
        # 1. 输入投影
        projected = self.input_proj(x)  # (batch, seq, gru_hidden)
        
        # 2. GRU状态更新
        if prev_hidden is None:
            # 第一层，初始化为零状态
            prev_hidden = torch.zeros(
                batch_size, seq_len, self.gru_cell.hidden_size,
                device=x.device, dtype=x.dtype
            )
        
        # 将序列维度展平，GRUCell需要2D输入
        flat_projected = projected.view(-1, projected.size(-1))  # (batch*seq, gru_hidden)
        flat_prev = prev_hidden.view(-1, prev_hidden.size(-1))    # (batch*seq, gru_hidden)
        
        # GRU前向
        flat_hidden = self.gru_cell(flat_projected, flat_prev)    # (batch*seq, gru_hidden)
        hidden_state = flat_hidden.view(batch_size, seq_len, -1)  # (batch, seq, gru_hidden)
        
        # 3. 生成路由对数
        logits = self.routing_head(hidden_state)  # (batch, seq, num_experts)
        
        # 4. Top-K选择
        routing_weights, selected_experts = self._top_k_routing(logits)
        
        return routing_weights, selected_experts, hidden_state
    
    def _top_k_routing(
        self, logits: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Top-K路由选择。
        对每个token，选择分数最高的K个专家。
        """
        # 对logits应用softmax
        scores = F.softmax(logits, dim=-1)  # (batch, seq, num_experts)
        
        # Top-K选择
        routing_weights, selected_experts = torch.topk(
            scores, self.num_experts_per_token, dim=-1
        )
        
        # 归一化
        routing_weights = routing_weights / (
            routing_weights.sum(dim=-1, keepdim=True) + 1e-8
        )
        
        return routing_weights, selected_experts


# ==============================================================
# 3. GMoE层（GMoE Layer）
# ==============================================================

class GMoELayer(nn.Module):
    """
    GMoE的核心层。
    包含：一个局部专家（Local Expert）+ 全局专家池（Global Experts）+ 全局路由器。
    """
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_global_experts: int,
        num_experts_per_token: int = 2,
        dropout: float = 0.1,
        gru_hidden_size: int = 128,
        layer_id: int = 0,
    ):
        super().__init__()
        self.layer_id = layer_id
        self.d_model = d_model
        self.num_global_experts = num_global_experts
        self.num_experts_per_token = num_experts_per_token
        
        # 局部专家（每层唯一）
        self.local_expert = Expert(d_model, d_ff, dropout)
        
        # 注意：全局专家池在GMoEModel中统一管理，
        # 这里只存储对全局专家的引用
        self.global_experts: Optional[nn.ModuleList] = None
        
        # 全局路由器（带GRU）
        self.router = GRURouter(
            d_model=d_model,
            num_global_experts=num_global_experts,
            gru_hidden_size=gru_hidden_size,
            num_experts_per_token=num_experts_per_token,
        )
        
        # LayerNorm
        self.norm = nn.LayerNorm(d_model)
    
    def set_global_experts(self, global_experts: nn.ModuleList):
        """设置全局专家池的引用"""
        self.global_experts = global_experts
    
    def forward(
        self,
        x: torch.Tensor,
        prev_routing_state: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (batch, seq, d_model) 输入
            prev_routing_state: 前一层路由状态
        Returns:
            output: (batch, seq, d_model) 输出
            routing_state: 当前层路由状态，传递给下一层
        """
        residual = x
        x = self.norm(x)
        
        batch_size, seq_len, _ = x.shape
        
        # 1. 局部专家输出
        local_output = self.local_expert(x)  # (batch, seq, d_model)
        
        # 2. 路由选择
        routing_weights, selected_experts, routing_state = self.router(x, prev_routing_state)
        # routing_weights: (batch, seq, K), selected_experts: (batch, seq, K)
        
        # 3. 全局专家输出（稀疏激活）
        global_output = self._sparse_global_forward(x, routing_weights, selected_experts)
        
        # 4. 合并输出
        output = residual + local_output + global_output
        
        return output, routing_state
    
    def _sparse_global_forward(
        self,
        x: torch.Tensor,
        routing_weights: torch.Tensor,
        selected_experts: torch.Tensor
    ) -> torch.Tensor:
        """
        稀疏激活全局专家。
        每个token只激活K个专家，大幅降低计算量。
        """
        assert self.global_experts is not None, "Global experts not set!"
        
        batch_size, seq_len, d_model = x.shape
        K = self.num_experts_per_token
        num_experts = len(self.global_experts)
        
        # 初始化输出
        final_output = torch.zeros_like(x)
        
        # 展平批次和序列维度
        flat_x = x.view(-1, d_model)  # (batch*seq, d_model)
        flat_weights = routing_weights.view(-1, K)  # (batch*seq, K)
        flat_experts = selected_experts.view(-1, K)  # (batch*seq, K)
        
        total_tokens = flat_x.size(0)
        
        # 对每个专家，收集需要处理的token，批量计算
        for expert_idx in range(num_experts):
            # 找出哪些token选择了这个专家
            # flat_experts: (total_tokens, K)
            mask = (flat_experts == expert_idx)  # (total_tokens, K)
            
            if not mask.any():
                continue
            
            # 获取对应的token索引和权重
            token_indices, expert_positions = torch.where(mask)
            # token_indices: 选择了该专家的token在flat_x中的索引
            # expert_positions: 该专家在Top-K中的位置(0到K-1)
            
            # 提取对应的token输入
            selected_x = flat_x[token_indices]  # (selected_count, d_model)
            
            # 提取对应的权重
            selected_weights = flat_weights[token_indices, expert_positions]  # (selected_count,)
            
            # 专家前向传播
            expert_output = self.global_experts[expert_idx](selected_x)  # (selected_count, d_model)
            
            # 加权累加
            weighted_output = expert_output * selected_weights.unsqueeze(-1)  # (selected_count, d_model)
            
            # 放回最终输出
            final_output.view(-1, d_model).index_add_(
                0, token_indices, weighted_output
            )
        
        return final_output


# ==============================================================
# 4. 完整GMoE模型
# ==============================================================

class GMoEModel(nn.Module):
    """
    完整的GMoE Transformer模型。
    包含：嵌入层 + N个GMoE层 + 全局专家池 + 输出层。
    """
    def __init__(
        self,
        vocab_size: int = 50257,
        d_model: int = 768,
        d_ff: int = 3072,
        num_layers: int = 12,
        num_global_experts: int = 16,
        num_experts_per_token: int = 2,
        num_heads: int = 12,
        dropout: float = 0.1,
        max_seq_len: int = 1024,
        gru_hidden_size: int = 128,
    ):
        super().__init__()
        self.d_model = d_model
        self.num_layers = num_layers
        self.num_global_experts = num_global_experts
        
        # 词嵌入
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = PositionalEncoding(d_model, max_seq_len, dropout)
        
        # 全局专家池（所有层共享）
        self.global_experts = nn.ModuleList([
            Expert(d_model, d_ff, dropout)
            for _ in range(num_global_experts)
        ])
        
        # GMoE层
        self.layers = nn.ModuleList([
            GMoELayer(
                d_model=d_model,
                d_ff=d_ff,
                num_global_experts=num_global_experts,
                num_experts_per_token=num_experts_per_token,
                dropout=dropout,
                gru_hidden_size=gru_hidden_size,
                layer_id=i,
            )
            for i in range(num_layers)
        ])
        
        # 为每层设置全局专家引用
        for layer in self.layers:
            layer.set_global_experts(self.global_experts)
        
        # 输出层
        self.final_norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        
        # 绑定嵌入权重
        self.lm_head.weight = self.embedding.weight
        
        # 初始化参数
        self._init_weights()
    
    def _init_weights(self):
        for name, param in self.named_parameters():
            if 'weight' in name and param.dim() >= 2:
                nn.init.xavier_uniform_(param)
            elif 'bias' in name:
                nn.init.zeros_(param)
    
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        Args:
            input_ids: (batch_size, seq_len)
        Returns:
            logits: (batch_size, seq_len, vocab_size)
        """
        # 嵌入
        x = self.embedding(input_ids)  # (batch, seq, d_model)
        x = self.pos_encoding(x)
        
        # 逐层前向传播，传递路由状态
        routing_state = None
        for layer in self.layers:
            x, routing_state = layer(x, routing_state)
        
        # 输出
        x = self.final_norm(x)
        logits = self.lm_head(x)
        
        return logits
    
    def count_active_parameters(self) -> dict:
        """
        统计活跃参数数量。
        注意：GMoE中，全局专家虽然总参数量大，但每个token只激活K个。
        """
        total_params = sum(p.numel() for p in self.parameters())
        
        # 必须参数（嵌入、注意力、输出等）
        must_params = 0
        for name, p in self.named_parameters():
            if 'global_experts' not in name:
                must_params += p.numel()
        
        # 全局专家参数
        expert_params = sum(p.numel() for p in self.global_experts.parameters())
        
        # 每个token激活的专家参数
        active_expert_params = expert_params * self.num_global_experts // self.num_global_experts * 2
        
        return {
            "total_params": total_params,
            "must_params": must_params,
            "expert_params": expert_params,
            "active_params_per_token": must_params + active_expert_params,
        }


class PositionalEncoding(nn.Module):
    """正弦位置编码"""
    def __init__(self, d_model: int, max_len: int, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

# ==============================================================
# 5. 负载均衡损失
# ==============================================================

class GMoELoss(nn.Module):
    """
    GMoE的辅助损失函数。
    包含负载均衡损失（Load Balancing Loss）和路由Z损失（Router Z-Loss）。
    """
    def __init__(self, alpha: float = 0.01):
        super().__init__()
        self.alpha = alpha
    
    def forward(
        self,
        logits: torch.Tensor,           # 模型输出logits
        labels: torch.Tensor,           # 真实标签
        routing_weights_list: List[torch.Tensor],  # 每层的路由权重 (batch, seq, num_experts)
        selected_experts_list: List[torch.Tensor], # 每层选中的专家 (batch, seq, K)
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        计算总损失 = 交叉熵损失 + alpha * 负载均衡损失
        
        Returns:
            total_loss, ce_loss, balance_loss
        """
        # 1. 交叉熵损失
        vocab_size = logits.size(-1)
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        
        ce_loss = F.cross_entropy(
            shift_logits.view(-1, vocab_size),
            shift_labels.view(-1),
            ignore_index=-100,
        )
        
        # 2. 负载均衡损失
        balance_loss = 0.0
        for routing_weights, selected_experts in zip(
            routing_weights_list, selected_experts_list
        ):
            balance_loss += self._load_balancing_loss(
                routing_weights, selected_experts
            )
        balance_loss = balance_loss / len(routing_weights_list)
        
        # 3. 总损失
        total_loss = ce_loss + self.alpha * balance_loss
        
        return total_loss, ce_loss, balance_loss
    
    def _load_balancing_loss(
        self,
        routing_weights: torch.Tensor,
        selected_experts: torch.Tensor,
    ) -> torch.Tensor:
        """
        负载均衡损失。
        鼓励所有专家被均匀使用。
        
        参考：https://arxiv.org/abs/2101.03961 (Switch Transformer)
        """
        batch_size, seq_len, num_experts = routing_weights.shape
        K = selected_experts.size(-1)
        total_tokens = batch_size * seq_len
        
        # 计算每个专家被选中的频率
        # selected_experts: (batch, seq, K) -> (total_tokens, K)
        flat_experts = selected_experts.view(-1, K)  # (total_tokens, K)
        
        # 对每个专家，统计被选中的token数
        expert_freq = torch.zeros(num_experts, device=routing_weights.device)
        for k in range(K):
            expert_freq += torch.bincount(
                flat_experts[:, k].long(),
                minlength=num_experts,
            ).float()
        expert_freq = expert_freq / total_tokens  # 归一化
        
        # 计算每个专家的平均路由权重
        # routing_weights: (batch, seq, num_experts) -> (total_tokens, num_experts)
        flat_weights = routing_weights.view(-1, num_experts)
        expert_weight = flat_weights.mean(dim=0)  # (num_experts,)
        
        # 负载均衡损失 = sum(freq_i * weight_i)
        # 当所有专家被均匀使用时，freq_i = 1/N, weight_i = 1/N, 损失最小
        balance_loss = num_experts * (expert_freq * expert_weight).sum()
        
        return balance_loss

# ==============================================================
# 6. 训练循环
# ==============================================================

def train_step(
    model: GMoEModel,
    loss_fn: GMoELoss,
    optimizer: torch.optim.Optimizer,
    batch: Tuple[torch.Tensor, torch.Tensor],
    device: torch.device,
) -> dict:
    """
    单步训练。
    """
    input_ids, labels = batch
    input_ids = input_ids.to(device)
    labels = labels.to(device)
    
    optimizer.zero_grad()
    
    # 前向传播
    logits = model(input_ids)
    
    # 收集路由权重用于损失计算
    # 注意：这里需要从模型中获取每层的路由信息
    # 实际实现中，model.forward应该返回路由信息
    # 为简化示例，我们这里假设模型返回了路由信息
    routing_weights_list = []
    selected_experts_list = []
    for layer in model.layers:
        # 在完整实现中，这些信息在forward时返回
        # 这里仅为演示损失函数的使用
        dummy_weights = torch.zeros(
            input_ids.size(0), input_ids.size(1), model.num_global_experts,
            device=device
        )
        dummy_experts = torch.zeros(
            input_ids.size(0), input_ids.size(1), model.num_global_experts,
            device=device, dtype=torch.long
        )
        routing_weights_list.append(dummy_weights)
        selected_experts_list.append(dummy_experts)
    
    # 计算损失
    total_loss, ce_loss, balance_loss = loss_fn(
        logits, labels, routing_weights_list, selected_experts_list
    )
    
    # 反向传播
    total_loss.backward()
    
    # 梯度裁剪
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    
    # 优化器步进
    optimizer.step()
    
    return {
        "total_loss": total_loss.item(),
        "ce_loss": ce_loss.item(),
        "balance_loss": balance_loss.item(),
    }


# ==============================================================
# 7. 参数效率对比实验
# ==============================================================

def parameter_efficiency_comparison():
    """
    对比传统MoE和GMoE的参数效率。
    """
    print("=" * 70)
    print("GMoE vs 传统MoE 参数效率对比")
    print("=" * 70)
    
    # 配置
    d_model = 768
    d_ff = 3072
    num_layers = 12
    traditional_experts_per_layer = 64
    num_global_experts = 16
    num_experts_per_token = 2
    vocab_size = 50257
    
    # 传统MoE参数估算
    # 每个专家参数 = 2 * d_model * d_ff (w1 + w2, 无bias)
    expert_params = 2 * d_model * d_ff
    
    traditional_total_expert_params = (
        num_layers * traditional_experts_per_layer * expert_params
    )
    
    # GMoE参数估算
    gmoe_global_expert_params = num_global_experts * expert_params
    gmoe_local_expert_params = num_layers * 1 * expert_params
    gmoe_total_expert_params = gmoe_global_expert_params + gmoe_local_expert_params
    
    # 非专家参数（嵌入、注意力、LayerNorm、输出等）
    # 嵌入: vocab_size * d_model
    # 注意力: 4 * d_model * d_model (Q, K, V, O)
    # LayerNorm: 2 * d_model per layer
    # 输出: d_model * vocab_size
    non_expert_params = (
        vocab_size * d_model +                                    # 嵌入
        num_layers * 4 * d_model * d_model +                      # 注意力
        num_layers * 2 * d_model +                                # LayerNorm
        d_model * vocab_size +                                    # 输出层
        num_layers * gru_hidden_size * (d_model + gru_hidden_size + d_model + num_global_experts)  # 路由器
    )
    
    # 总参数
    traditional_total = traditional_total_expert_params + non_expert_params
    gmoe_total = gmoe_total_expert_params + non_expert_params
    
    # 活跃参数（每个token）
    # 传统MoE: 每层激活K个专家
    traditional_active_expert = num_layers * num_experts_per_token * expert_params
    # GMoE: 每层激活K个全局专家 + 1个局部专家
    gmoe_active_expert = num_layers * (num_experts_per_token + 1) * expert_params
    
    print(f"\n模型配置:")
    print(f"  d_model={d_model}, d_ff={d_ff}, num_layers={num_layers}")
    print(f"  传统MoE: 每层{traditional_experts_per_layer}个专家")
    print(f"  GMoE: {num_global_experts}个全局专家 + 每层1个局部专家")
    print(f"  每个token激活专家数: K={num_experts_per_token}")
    
    print(f"\n--- 总参数量对比 ---")
    print(f"  传统MoE总参数: {traditional_total:,}")
    print(f"  GMoE总参数:    {gmoe_total:,}")
    print(f"  参数压缩比:    {gmoe_total / traditional_total * 100:.1f}%")
    print(f"  参数减少:      {(1 - gmoe_total / traditional_total) * 100:.1f}%")
    
    print(f"\n--- 专家参数对比 ---")
    print(f"  传统MoE专家参数: {traditional_total_expert_params:,}")
    print(f"  GMoE专家参数:    {gmoe_total_expert_params:,}")
    print(f"  专家参数压缩比:  {gmoe_total_expert_params / traditional_total_expert_params * 100:.1f}%")
    
    print(f"\n--- 活跃参数对比（每个token）---")
    print(f"  传统MoE活跃参数: {traditional_active_expert + non_expert_params:,}")
    print(f"  GMoE活跃参数:    {gmoe_active_expert + non_expert_params:,}")
    
    return {
        "traditional_total": traditional_total,
        "gmoe_total": gmoe_total,
        "compression_ratio": gmoe_total / traditional_total,
    }


if __name__ == "__main__":
    # 运行参数效率对比
    results = parameter_efficiency_comparison()
    
    print("\n" + "=" * 70)
    print("创建GMoE模型实例...")
    print("=" * 70)
    
    # 创建小模型用于测试
    model = GMoEModel(
        vocab_size=50257,
        d_model=256,
        d_ff=1024,
        num_layers=4,
        num_global_experts=8,
        num_experts_per_token=2,
        num_heads=4,
        dropout=0.1,
        max_seq_len=512,
        gru_hidden_size=64,
    )
    
    # 统计参数
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"  模型总参数: {total_params:,}")
    print(f"  可训练参数: {trainable_params:,}")
    print(f"  GPU内存估算: {total_params * 4 / 1024 / 1024:.2f} MB (FP32)")
    
    # 模拟前向传播
    batch_size, seq_len = 2, 128
    dummy_input = torch.randint(0, 50257, (batch_size, seq_len))
    
    with torch.no_grad():
        output = model(dummy_input)
    
    print(f"  输入形状: {dummy_input.shape}")
    print(f"  输出形状: {output.shape}")
    print(f"  模型前向传播成功!")

// ==============================================================
// GMoE推理引擎 (Go语言实现)
// 专注于高性能推理，支持端侧部署
// ==============================================================

package main

import (
	"encoding/binary"
	"fmt"
	"math"
	"os"
)

// ==============================================================
// 数据类型定义
// ==============================================================

// Matrix 表示一个二维矩阵，行优先存储
type Matrix struct {
	Rows int
	Cols int
	Data []float32
}

func NewMatrix(rows, cols int) *Matrix {
	return &Matrix{
		Rows: rows,
		Cols: cols,
		Data: make([]float32, rows*cols),
	}
}

func (m *Matrix) At(r, c int) float32 {
	return m.Data[r*m.Cols+c]
}

func (m *Matrix) Set(r, c int, v float32) {
	m.Data[r*m.Cols+c] = v
}

// ==============================================================
// 基础数学运算
// ==============================================================

// MatMul 矩阵乘法: C = A * B
// A: (M, K), B: (K, N), C: (M, N)
func MatMul(A, B, C *Matrix) {
	if A.Cols != B.Rows {
		panic(fmt.Sprintf("维度不匹配: A(%d,%d) B(%d,%d)", A.Rows, A.Cols, B.Rows, B.Cols))
	}
	if C.Rows != A.Rows || C.Cols != B.Cols {
		panic("C维度不匹配")
	}

	M, K, N := A.Rows, A.Cols, B.Cols

	for i := 0; i < M; i++ {
		for j := 0; j < N; j++ {
			var sum float32
			for k := 0; k < K; k++ {
				sum += A.At(i, k) * B.At(k, j)
			}
			C.Set(i, j, sum)
		}
	}
}

// Softmax 对矩阵的最后一维进行Softmax归一化
func Softmax(m *Matrix) {
	for i := 0; i < m.Rows; i++ {
		// 找最大值（数值稳定性）
		maxVal := float32(math.Inf(-1))
		for j := 0; j < m.Cols; j++ {
			if m.At(i, j) > maxVal {
				maxVal = m.At(i, j)
			}
		}

		// 计算exp和sum
		var sum float32
		for j := 0; j < m.Cols; j++ {
			val := float32(math.Exp(float64(m.At(i, j) - maxVal)))
			m.Set(i, j, val)
			sum += val
		}

		// 归一化
		for j := 0; j < m.Cols; j++ {
			m.Set(i, j, m.At(i, j)/sum)
		}
	}
}

// GELU 激活函数
func GELU(x float32) float32 {
	return float32(0.5 * float64(x) * (1 + math.Erf(float64(x)/math.Sqrt2)))
}

// ==============================================================
// GMoE专家模块 (Go版本)
// ==============================================================

type ExpertWeights struct {
	W1 *Matrix // (d_model, d_ff)
	W2 *Matrix // (d_ff, d_model)
}

func NewExpertWeights(dModel, dFF int) *ExpertWeights {
	return &ExpertWeights{
		W1: NewMatrix(dModel, dFF),
		W2: NewMatrix(dFF, dModel),
	}
}

func (e *ExpertWeights) Forward(input *Matrix) *Matrix {
	// input: (batch, d_model)
	// hidden = GELU(input * W1): (batch, d_ff)
	hidden := NewMatrix(input.Rows, e.W1.Cols)
	MatMul(input, e.W1, hidden)

	// 应用GELU
	for i := 0; i < hidden.Rows; i++ {
		for j := 0; j < hidden.Cols; j++ {
			hidden.Set(i, j, GELU(hidden.At(i, j)))
		}
	}

	// output = hidden * W2: (batch, d_model)
	output := NewMatrix(hidden.Rows, e.W2.Cols)
	MatMul(hidden, e.W2, output)

	return output
}

// ==============================================================
// GRU路由器 (Go版本)
// ==============================================================

type GRURouterWeights struct {
	InputProj *Matrix // (d_model, gru_hidden)

	// GRU单元参数
	Wz *Matrix // (gru_hidden, gru_hidden) 更新门
	Wr *Matrix // (gru_hidden, gru_hidden) 重置门
	Wh *Matrix // (gru_hidden, gru_hidden) 候选隐藏状态

	Uz *Matrix // (gru_hidden, gru_hidden)
	Ur *Matrix // (gru_hidden, gru_hidden)
	Uh *Matrix // (gru_hidden, gru_hidden)

	Bz []float32 // (gru_hidden,)
	Br []float32 // (gru_hidden,)
	Bh []float32 // (gru_hidden,)

	RoutingHead *Matrix // (gru_hidden, num_experts)
}

func NewGRURouterWeights(dModel, gruHidden, numExperts int) *GRURouterWeights {
	return &GRURouterWeights{
		InputProj:   NewMatrix(dModel, gruHidden),
		Wz:          NewMatrix(gruHidden, gruHidden),
		Wr:          NewMatrix(gruHidden, gruHidden),
		Wh:          NewMatrix(gruHidden, gruHidden),
		Uz:          NewMatrix(gruHidden, gruHidden),
		Ur:          NewMatrix(gruHidden, gruHidden),
		Uh:          NewMatrix(gruHidden, gruHidden),
		Bz:          make([]float32, gruHidden),
		Br:          make([]float32, gruHidden),
		Bh:          make([]float32, gruHidden),
		RoutingHead: NewMatrix(gruHidden, numExperts),
	}
}

func (r *GRURouterWeights) Forward(
	x *Matrix,
	prevHidden *Matrix,
) (*Matrix, *Matrix) {
	// x: (batch, d_model)
	// prevHidden: (batch, gru_hidden)
	// 返回: routingWeights (batch, num_experts), newHidden (batch, gru_hidden)

	batch := x.Rows
	gruHidden := r.Wz.Rows
	numExperts := r.RoutingHead.Cols

	// 1. 输入投影
	projected := NewMatrix(batch, gruHidden)
	MatMul(x, r.InputProj, projected)

	// 2. GRU前向
	// z_t = sigmoid(W_z * x_t + U_z * h_{t-1} + b_z)
	// r_t = sigmoid(W_r * x_t + U_r * h_{t-1} + b_r)
	// h_t' = tanh(W_h * x_t + U_h * (r_t * h_{t-1}) + b_h)
	// h_t = (1 - z_t) * h_{t-1} + z_t * h_t'

	newHidden := NewMatrix(batch, gruHidden)

	for i := 0; i < batch; i++ {
		for j := 0; j < gruHidden; j++ {
			var zSum, rSum, hSum float32

			// W_z * x_t + b_z
			for k := 0; k < projected.Cols; k++ {
				zSum += projected.At(i, k) * r.Wz.At(k, j)
				rSum += projected.At(i, k) * r.Wr.At(k, j)
				hSum += projected.At(i, k) * r.Wh.At(k, j)
			}
			zSum += r.Bz[j]
			rSum += r.Br[j]
			hSum += r.Bh[j]

			// U_z * h_{t-1}
			if prevHidden != nil {
				for k := 0; k < gruHidden; k++ {
					zSum += prevHidden.At(i, k) * r.Uz.At(k, j)
					rSum += prevHidden.At(i, k) * r.Ur.At(k, j)
				}
			}

			z := sigmoid(zSum)
			r := sigmoid(rSum)

			// U_h * (r * h_{t-1})
			if prevHidden != nil {
				for k := 0; k < gruHidden; k++ {
					hSum += (r * prevHidden.At(i, k)) * r.Uh.At(k, j)
				}
			}

			hCandidate := float32(math.Tanh(float64(hSum)))
			newHidden.Set(i, j, (1-z)*hCandidate+z*0) // 简化：假设初始h_0=0
		}
	}

	// 3. 路由计算
	logits := NewMatrix(batch, numExperts)
	MatMul(newHidden, r.RoutingHead, logits)

	// 4. Softmax归一化
	Softmax(logits)

	return logits, newHidden
}

func sigmoid(x float32) float32 {
	return 1.0 / (1.0 + float32(math.Exp(float64(-x))))
}

// ==============================================================
// GMoE层 (Go版本)
// ==============================================================

type GMoELayerWeights struct {
	LocalExpert *ExpertWeights
	Router      *GRURouterWeights
	LayerNorm   *LayerNormWeights
}

type LayerNormWeights struct {
	Gamma []float32 // (d_model,)
	Beta  []float32 // (d_model,)
}

func NewGMoELayerWeights(dModel, dFF, gruHidden, numExperts int) *GMoELayerWeights {
	return &GMoELayerWeights{
		LocalExpert: NewExpertWeights(dModel, dFF),
		Router:      NewGRURouterWeights(dModel, gruHidden, numExperts),
		LayerNorm: &LayerNormWeights{
			Gamma: make([]float32, dModel),
			Beta:  make([]float32, dModel),
		},
	}
}

// ==============================================================
// GMoE推理引擎 (Go版本)
// ==============================================================

type GMoEInferenceEngine struct {
	// 模型配置
	Config ModelConfig

	// 全局专家池
	GlobalExperts []*ExpertWeights

	// 各层参数
	Layers []*GMoELayerWeights

	// 嵌入层
	Embedding *Matrix // (vocab_size, d_model)

	// 输出层
	LMHead *Matrix // (d_model, vocab_size)

	// 最终LayerNorm
	FinalNorm *LayerNormWeights
}

type ModelConfig struct {
	VocabSize         int
	DModel            int
	DFF               int
	NumLayers         int
	NumGlobalExperts  int
	NumExpertsPerToken int
	GRUHiddenSize     int
	MaxSeqLen         int
}

func NewGMoEInferenceEngine(config ModelConfig) *GMoEInferenceEngine {
	engine := &GMoEInferenceEngine{
		Config:        config,
		GlobalExperts: make([]*ExpertWeights, config.NumGlobalExperts),
		Layers:        make([]*GMoELayerWeights, config.NumLayers),
		Embedding:     NewMatrix(config.VocabSize, config.DModel),
		LMHead:        NewMatrix(config.DModel, config.VocabSize),
		FinalNorm: &LayerNormWeights{
			Gamma: make([]float32, config.DModel),
			Beta:  make([]float32, config.DModel),
		},
	}

	for i := 0; i < config.NumGlobalExperts; i++ {
		engine.GlobalExperts[i] = NewExpertWeights(config.DModel, config.DFF)
	}

	for i := 0; i < config.NumLayers; i++ {
		engine.Layers[i] = NewGMoELayerWeights(
			config.DModel,
			config.DFF,
			config.GRUHiddenSize,
			config.NumGlobalExperts,
		)
	}

	return engine
}

// Forward 执行推理前向传播
func (e *GMoEInferenceEngine) Forward(inputIDs []int) []float32 {
	// inputIDs: (seq_len,)
	seqLen := len(inputIDs)
	batchSize := 1

	// 1. 嵌入层
	hidden := NewMatrix(batchSize, e.Config.DModel)
	for i := 0; i < seqLen; i++ {
		tokenID := inputIDs[i]
		if tokenID >= e.Config.VocabSize {
			tokenID = 0 // UNK
		}
		// 复制嵌入向量
		for j := 0; j < e.Config.DModel; j++ {
			hidden.Set(i, j, e.Embedding.At(tokenID, j))
		}
	}

	// 2. 逐层前向传播
	var routingState *Matrix
	for layerIdx := 0; layerIdx < e.Config.NumLayers; layerIdx++ {
		layer := e.Layers[layerIdx]

		// LayerNorm
		hidden = e.applyLayerNorm(hidden, layer.LayerNorm)

		// 局部专家
		localOutput := layer.LocalExpert.Forward(hidden)

		// 路由
		routingWeights, routingState := layer.Router.Forward(hidden, routingState)

		// 全局专家（稀疏激活）
		globalOutput := e.sparseGlobalForward(hidden, routingWeights)

		// 合并输出
		// output = input + local_output + global_output
		for i := 0; i < hidden.Rows; i++ {
			for j := 0; j < hidden.Cols; j++ {
				hidden.Set(i, j,
					hidden.At(i, j)+localOutput.At(i, j)+globalOutput.At(i, j))
			}
		}
	}

	// 3. 最终LayerNorm
	hidden = e.applyLayerNorm(hidden, e.FinalNorm)

	// 4. 输出层
	logits := NewMatrix(batchSize, e.Config.VocabSize)
	MatMul(hidden, e.LMHead, logits)

	// 返回最后一个token的logits
	result := make([]float32, e.Config.VocabSize)
	for i := 0; i < e.Config.VocabSize; i++ {
		result[i] = logits.At(seqLen-1, i)
	}

	return result
}

func (e *GMoEInferenceEngine) applyLayerNorm(
	x *Matrix, norm *LayerNormWeights,
) *Matrix {
	output := NewMatrix(x.Rows, x.Cols)

	for i := 0; i < x.Rows; i++ {
		// 计算均值和方差
		var mean, variance float32
		for j := 0; j < x.Cols; j++ {
			mean += x.At(i, j)
		}
		mean /= float32(x.Cols)

		for j := 0; j < x.Cols; j++ {
			diff := x.At(i, j) - mean
			variance += diff * diff
		}
		variance /= float32(x.Cols)

		std := float32(math.Sqrt(float64(variance + 1e-5)))

		// 归一化
		for j := 0; j < x.Cols; j++ {
			normalized := (x.At(i, j) - mean) / std
			output.Set(i, j, normalized*norm.Gamma[j]+norm.Beta[j])
		}
	}

	return output
}

func (e *GMoEInferenceEngine) sparseGlobalForward(
	x *Matrix, routingWeights *Matrix,
) *Matrix {
	// x: (batch, d_model)
	// routingWeights: (batch, num_experts)
	// 选择Top-K专家

	numExperts := e.Config.NumGlobalExperts
	K := e.Config.NumExpertsPerToken
	batchSize := x.Rows

	output := NewMatrix(batchSize, e.Config.DModel)

	for b := 0; b < batchSize; b++ {
		// 找Top-K专家
		type expertScore struct {
			idx   int
			score float32
		}
		scores := make([]expertScore, numExperts)
		for i := 0; i < numExperts; i++ {
			scores[i] = expertScore{idx: i, score: routingWeights.At(b, i)}
		}

		// 简单选择排序找Top-K
		for i := 0; i < K; i++ {
			maxIdx := i
			for j := i + 1; j < numExperts; j++ {
				if scores[j].score > scores[maxIdx].score {
					maxIdx = j
				}
			}
			scores[i], scores[maxIdx] = scores[maxIdx], scores[i]
		}

		// 归一化Top-K权重
		var weightSum float32
		for i := 0; i < K; i++ {
			weightSum += scores[i].score
		}

		// 稀疏激活专家
		for i := 0; i < K; i++ {
			expertIdx := scores[i].idx
			weight := scores[i].score / weightSum

			// 提取当前token的输入
			tokenInput := NewMatrix(1, e.Config.DModel)
			for j := 0; j < e.Config.DModel; j++ {
				tokenInput.Set(0, j, x.At(b, j))
			}

			// 专家前向
			expertOutput := e.GlobalExperts[expertIdx].Forward(tokenInput)

			// 加权累加
			for j := 0; j < e.Config.DModel; j++ {
				output.Set(b, j,
					output.At(b, j)+weight*expertOutput.At(0, j))
			}
		}
	}

	return output
}

// ==============================================================
// 模型导出与加载
// ==============================================================

// SaveWeights 将模型权重保存到文件
func (e *GMoEInferenceEngine) SaveWeights(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()

	// 写入配置
	binary.Write(f, binary.LittleEndian, int32(e.Config.VocabSize))
	binary.Write(f, binary.LittleEndian, int32(e.Config.DModel))
	binary.Write(f, binary.LittleEndian, int32(e.Config.DFF))
	binary.Write(f, binary.LittleEndian, int32(e.Config.NumLayers))
	binary.Write(f, binary.LittleEndian, int32(e.Config.NumGlobalExperts))
	binary.Write(f, binary.LittleEndian, int32(e.Config.NumExpertsPerToken))
	binary.Write(f, binary.LittleEndian, int32(e.Config.GRUHiddenSize))
	binary.Write(f, binary.LittleEndian, int32(e.Config.MaxSeqLen))

	// 写入权重数据...
	// 实际实现中需要遍历所有矩阵并写入

	fmt.Printf("模型权重已保存到: %s\n", path)
	return nil
}

// ==============================================================
// 主函数：演示推理
// ==============================================================

func main() {
	fmt.Println("=" + strings.Repeat("=", 69))
	fmt.Println("  GMoE推理引擎 (Go语言实现)")
	fmt.Println("=" + strings.Repeat("=", 69))

	// 配置小模型
	config := ModelConfig{
		VocabSize:         50257,
		DModel:            256,
		DFF:               1024,
		NumLayers:         4,
		NumGlobalExperts:  8,
		NumExpertsPerToken: 2,
		GRUHiddenSize:     64,
		MaxSeqLen:         512,
	}

	// 创建引擎
	engine := NewGMoEInferenceEngine(config)

	fmt.Printf("\n模型配置:\n")
	fmt.Printf("  VocabSize:        %d\n", config.VocabSize)
	fmt.Printf("  DModel:           %d\n", config.DModel)
	fmt.Printf("  DFF:              %d\n", config.DFF)
	fmt.Printf("  NumLayers:        %d\n", config.NumLayers)
	fmt.Printf("  NumGlobalExperts: %d\n", config.NumGlobalExperts)
	fmt.Printf("  NumExpertsPerToken: %d\n", config.NumExpertsPerToken)

	// 计算总参数量
	expertParams := 2 * config.DModel * config.DFF
	globalExpertParams := config.NumGlobalExperts * expertParams
	localExpertParams := config.NumLayers * expertParams
	routerParams := config.DModel*config.GRUHiddenSize + // InputProj
		6*config.GRUHiddenSize*config.GRUHiddenSize + // GRU gates
		3*config.GRUHiddenSize + // biases
		config.GRUHiddenSize*config.NumGlobalExperts // RoutingHead
	embedParams := config.VocabSize * config.DModel
	outputParams := config.DModel * config.VocabSize
	attentionParams := config.NumLayers * 4 * config.DModel * config.DModel

	totalParams := globalExpertParams + localExpertParams +
		routerParams + embedParams + outputParams + attentionParams

	fmt.Printf("\n参数量统计:\n")
	fmt.Printf("  全局专家参数:      %d\n", globalExpertParams)
	fmt.Printf("  局部专家参数:      %d\n", localExpertParams)
	fmt.Printf("  路由器参数:        %d\n", routerParams)
	fmt.Printf("  嵌入层参数:        %d\n", embedParams)
	fmt.Printf("  总参数量:          %d (%.2fM)\n",
		totalParams, float64(totalParams)/1e6)

	// 模拟推理
	fmt.Printf("\n模拟推理...\n")
	dummyInput := []int{101, 202, 303, 404, 505}
	logits := engine.Forward(dummyInput)
	fmt.Printf("  输入序列长度: %d\n", len(dummyInput))
	fmt.Printf("  输出logits维度: %d\n", len(logits))
	fmt.Printf("  Top-5预测: ")
	// 找Top-5
	type pred struct {
		id    int
		score float32
	}
	top5 := make([]pred, 5)
	for i := 0; i < 5; i++ {
		top5[i] = pred{id: -1, score: float32(math.Inf(-1))}
	}
	for i, score := range logits {
		if score > top5[4].score {
			top5[4] = pred{id: i, score: score}
			// 重新排序
			for j := 4; j > 0; j-- {
				if top5[j].score > top5[j-1].score {
					top5[j], top5[j-1] = top5[j-1], top5[j]
				}
			}
		}
	}
	for i, p := range top5 {
		fmt.Printf("    %d. token_id=%d, score=%.4f\n", i+1, p.id, p.score)
	}

	fmt.Printf("\n✅ GMoE推理引擎运行成功!\n")
}

// 需要导入strings包
import "strings"

// 注意：Go代码中import需要在文件顶部
// 这里为了代码完整性放在此处说明

# ==============================================================
# GMoE INT8量化部署
# ==============================================================

import torch
import torch.nn as nn
import numpy as np

class GMoEInt8Quantizer:
    """
    GMoE模型的INT8量化器。
    利用全局专家池的共享特性，实现高效量化。
    """
    
    def __init__(self, model: GMoEModel, calibration_data: torch.Tensor):
        self.model = model
        self.calibration_data = calibration_data
        self.scales = {}
        self.zero_points = {}
    
    def calibrate(self):
        """
        校准量化参数。使用校准数据集统计每层激活值的范围。
        """
        self.model.eval()
        
        # 收集每层和每个专家的激活值范围
        with torch.no_grad():
            x = self.model.embedding(self.calibration_data)
            x = self.model.pos_encoding(x)
            
            routing_state = None
            for layer_idx, layer in enumerate(self.model.layers):
                x = layer.norm(x)
                
                # 收集局部专家激活值
                local_out = layer.local_expert(x)
                self._update_scale(f"layer_{layer_idx}_local", local_out)
                
                # 收集路由权重
                weights, experts, routing_state = layer.router(x, routing_state)
                self._update_scale(f"layer_{layer_idx}_router", weights)
                
                # 收集全局专家激活值（每个专家单独统计）
                for expert_idx in range(len(self.model.global_experts)):
                    mask = (experts == expert_idx)
                    if mask.any():
                        selected_x = x[mask]
                        expert_out = self.model.global_experts[expert_idx](selected_x)
                        self._update_scale(
                            f"global_expert_{expert_idx}", expert_out
                        )
                
                x = x + local_out
                x = x + self._sparse_forward_demo(x, weights, experts)
        
        print(f"量化校准完成，共收集 {len(self.scales)} 个scale值")
    
    def _update_scale(self, name: str, tensor: torch.Tensor):
        """更新指定张量的量化scale"""
        if name not in self.scales:
            self.scales[name] = tensor.abs().max().item()
            self.zero_points[name] = 0
        else:
            self.scales[name] = max(
                self.scales[name], tensor.abs().max().item()
            )
    
    def _sparse_forward_demo(self, x, weights, experts):
        """演示用的稀疏前向，实际实现见前面章节"""
        K = weights.size(-1)
        batch_size, seq_len, d_model = x.shape
        output = torch.zeros_like(x)
        
        for k in range(K):
            expert_indices = experts[:, :, k]  # (batch, seq)
            expert_weights = weights[:, :, k]  # (batch, seq)
            
            for expert_idx in range(len(self.model.global_experts)):
                mask = (expert_indices == expert_idx)
                if not mask.any():
                    continue
                
                selected_x = x[mask]
                selected_w = expert_weights[mask].unsqueeze(-1)
                expert_out = self.model.global_experts[expert_idx](selected_x)
                output[mask] += selected_w * expert_out
        
        return output
    
    def quantize_weights(self) -> dict:
        """
        量化模型权重为INT8。
        返回量化后的参数字典。
        """
        quantized = {}
        
        for name, param in self.model.named_parameters():
            if 'global_experts' in name or 'local_expert' in name:
                # 专家权重使用per-channel量化
                scale = param.abs().max(dim=-1, keepdim=True)[0] / 127.0
                quantized_data = torch.round(param / scale).to(torch.int8)
                quantized[f"{name}_quant"] = quantized_data
                quantized[f"{name}_scale"] = scale
            else:
                # 非专家权重使用per-tensor量化
                scale = param.abs().max().item() / 127.0
                quantized_data = torch.round(param / scale).to(torch.int8)
                quantized[f"{name}_quant"] = quantized_data
                quantized[f"{name}_scale"] = torch.tensor(scale)
        
        # 计算量化后的模型大小
        total_bytes = sum(
            p.numel() for p in quantized.values() if p.dtype == torch.int8
        )
        print(f"量化后模型大小: {total_bytes / 1024 / 1024:.2f} MB")
        
        return quantized
    
    def save_for_deployment(self, path: str):
        """保存量化模型用于端侧部署"""
        quantized = self.quantize_weights()
        torch.save(quantized, path)
        print(f"量化模型已保存到: {path}")


# 端侧部署示例
def edge_deployment_example():
    """
    端侧部署完整流程示例。
    """
    print("=" * 70)
    print("GMoE端侧部署流程")
    print("=" * 70)
    
    # 1. 创建模型
    model = GMoEModel(
        vocab_size=50257,
        d_model=256,
        d_ff=1024,
        num_layers=4,
        num_global_experts=8,
        num_experts_per_token=2,
        num_heads=4,
        dropout=0.0,  # 推理时关闭dropout
        max_seq_len=512,
        gru_hidden_size=64,
    )
    model.eval()
    
    # 2. 统计部署资源
    total_params = sum(p.numel() for p in model.parameters())
    fp32_size = total_params * 4 / 1024 / 1024  # MB
    int8_size = total_params / 1024 / 1024  # MB
    
    # 计算活跃参数（每个token实际参与计算的参数）
    # 每层: 1个局部专家 + K个全局专家
    d_model, d_ff = 256, 1024
    expert_params = 2 * d_model * d_ff  # w1 + w2
    active_params_per_layer = (2 + 1) * expert_params  # K=2 + 1 local
    active_params = active_params_per_layer * 4  # 4 layers
    
    print(f"\n部署资源需求:")
    print(f"  总参数量: {total_params:,}")
    print(f"  FP32模型大小: {fp32_size:.2f} MB")
    print(f"  INT8模型大小: {int8_size:.2f} MB")
    print(f"  每token活跃参数: ~{active_params:,}")
    print(f"  推荐部署设备: 手机/平板/边缘网关")
    
    # 3. 推理延迟估算
    # 假设设备推理速度为 10 GFLOPS (手机端)
    # 每token计算量 ≈ 2 * d_model * d_ff * (K+1) * num_layers
    flops_per_token = 2 * d_model * d_ff * (2 + 1) * 4
    estimated_latency = flops_per_token / (10 * 1e9) * 1000  # ms
    print(f"  每token计算量: ~{flops_per_token/1e6:.1f} MFLOPs")
    print(f"  估算延迟: ~{estimated_latency:.2f} ms/token")
    print(f"  生成速度: ~{1000/estimated_latency:.0f} tokens/s")
    
    print(f"\n✅ 端侧部署方案可行!")


# 运行示例
if __name__ == "__main__":
    edge_deployment_example()