import torch
import torch.nn as nn


# 一個(gè)簡(jiǎn)單的專家網(wǎng)絡(luò)
class SimpleExpert(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleExpert, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, output_size)
        )


    def forward(self, x):
        return self.network(x)


# 簡(jiǎn)化版的MOE模型
class SimpleMOE(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_experts=4):
        super(SimpleMOE, self).__init__()


        # 創(chuàng)建多個(gè)專家
        self.experts = nn.ModuleList([
            SimpleExpert(input_size, hidden_size, output_size) 
            for _ in range(num_experts)
        ])


        # 路由器：決定用哪個(gè)專家
        self.router = nn.Linear(input_size, num_experts)


    def forward(self, x):
        # 1. 路由器計(jì)算每個(gè)專家的權(quán)重
        router_outputs = torch.softmax(self.router(x), dim=-1)


        # 2. 獲取權(quán)重最高的兩位專家
        top_k_weights, top_k_indices = torch.topk(router_outputs, k=2, dim=-1)


        # 3. 歸一化權(quán)重，使它們的和為1
        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)


        # 4. 初始化最終輸出
        final_output = torch.zeros_like(x)


        # 5. 對(duì)于每個(gè)樣本
        for i in range(x.shape[0]):
            # 6. 計(jì)算并組合所選專家的輸出
            for j in range(2):  # 兩位專家
                expert_idx = top_k_indices[i, j]
                weight = top_k_weights[i, j]


                # 獲取專家輸出
                expert_output = self.experts[expert_idx](x[i:i+1])


                # 加權(quán)組合
                final_output[i:i+1] += weight * expert_output


        return final_output

def simple_load_balancing(router_probs, num_experts):
    """簡(jiǎn)單的負(fù)載均衡計(jì)算"""
    # 計(jì)算每個(gè)專家實(shí)際使用頻率
    expert_usage = router_probs.sum(dim=0) / router_probs.shape[0]


    # 理想情況下每個(gè)專家使用率相等
    ideal_usage = torch.ones_like(expert_usage) / num_experts


    # 計(jì)算實(shí)際與理想的差距作為損失
    balance_loss = ((expert_usage - ideal_usage) ** 2).sum()


    return balance_loss

# 簡(jiǎn)化版的DeepSeek MOE架構(gòu)
class SimpleDeepSeekMOELayer(nn.Module):
def __init__(self, hidden_size, num_experts=8, k=2):
        super(SimpleDeepSeekMOELayer, self).__init__()


# 自注意力層（保持不變）
        self.attention = nn.MultiheadAttention(hidden_size, num_heads=8)
        self.norm1 = nn.LayerNorm(hidden_size)


# MOE替代標(biāo)準(zhǔn)FFN
        self.moe = SimpleMOE(
            input_size=hidden_size,
            hidden_size=hidden_size*4,
            output_size=hidden_size,
            num_experts=num_experts
        )
        self.norm2 = nn.LayerNorm(hidden_size)


def forward(self, x):
# 自注意力部分
        residual = x
        x = self.attention(x, x, x)[0]
        x = self.norm1(x + residual)


# MOE部分
        residual = x
        x = self.moe(x)
        x = self.norm2(x + residual)


return x