# part 1: 导入相关的 package
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from dataclasses import dataclass

import math

torch.manual_seed(1024)

# 形状是 n_export * n_top * (seq_len * batch_size)
n_export = 5
n_top = 3
seq_len = 3
batch_size = 2

export_mask =torch.tril( torch.ones((n_export, n_top, seq_len * batch_size)))
export_mask

for i in range(n_export):
    idx,top_x = torch.where(export_mask[i])
    print(idx) # 行索引 (对应top几） 
    print(top_x)  # 列索引 （对应第几个token)
    # 行索引和列索引 对应这个export_mask的意义是什么
    # idx_t：被第i个专家选中的第t个token是top几
    # top_x_t： 被第i个专家选中的第t个token的在原始序列中的token序号
    
    # 我有这两个变量有什么用
    
    break

class BasicExpert(nn.Module):
    # 一个 Expert 可以是一个最简单的， linear 层即可
    # 也可以是 MLP 层
    # 也可以是 更复杂的 MLP 层（active function 设置为 swiglu）
    def __init__(self, feature_in, feature_out):
        super().__init__()
        self.linear = nn.Linear(feature_in, feature_out)
    
    def forward(self, x):
        return self.linear(x)

# 主要参考自 mistral MOE 的实现

class MOERouter(nn.Module):
    def __init__(self, hidden_dim, expert_number, top_k):
        super().__init__()
        self.gate = nn.Linear(hidden_dim, expert_number)
        self.expert_number = expert_number
        self.top_k = top_k
    
    def forward(self, hidden_states):
        # 计算路由logits
        router_logits = self.gate(hidden_states)  # shape is (b * s, expert_number)
        
        # 计算专家经过softmax之后的概率
        routing_probs = F.softmax(router_logits, dim=-1, dtype=torch.float)
        
        # 计算topk的专家的输出
        router_weights, selected_experts = torch.topk(
            routing_probs, self.top_k, dim=-1
        )  # shape都是 (b * s, top_k)
        
        # 专家权重归一化
        router_weights = router_weights / router_weights.sum(dim=-1, keepdim=True)
        router_weights = router_weights.to(hidden_states.dtype)
        
        # 生成专家掩码
        expert_mask = F.one_hot(
            selected_experts,
            num_classes=self.expert_number
        )  # shape是 (b * s, top_k, expert_number)
        expert_mask = expert_mask.permute(2, 1, 0)  # (expert_number, top_k, b * s)
        
        return router_logits, router_weights, selected_experts, expert_mask


class MOEConfig:
    def __init__(
            self, 
            hidden_dim, 
            expert_number, 
            top_k, 
            shared_experts_number=2,
        ):
        self.hidden_dim = hidden_dim
        self.expert_number = expert_number
        self.top_k = top_k
        self.shared_experts_number = shared_experts_number

class SparseMOE(nn.Module):
    # 稀疏 MOE 模型，这里每一个 token 都会过 topk 个专家，得到对应token 的 hidden_embeddings
    def __init__(self, config):
        super().__init__()

        self.hidden_dim = config.hidden_dim

        self.expert_number = config.expert_number
        self.top_k = config.top_k

        self.experts = nn.ModuleList(
            [
                BasicExpert(self.hidden_dim, self.hidden_dim) for _ in range(self.expert_number)
            ]
        )

        self.router = MOERouter(self.hidden_dim, self.expert_number, self.top_k)
    
    def forward(self, x):
        # x shape is (b, s, hidden_dim)
        batch_size, seq_len, hidden_dim = x.size()

        # 合并前两个维度，因为不是 Sample 维度了，而是 token 维度
        hidden_states = x.view(-1, hidden_dim) # shape is(b * s, hidden_dim)

        router_logits, router_weights, selected_experts_indices, expert_mask = self.router(hidden_states)
        # 其中 selected_experts_indices shape 是 (b * s, top_k)
        # 其中 expert_mask shape 是 (expert_number, top_k, b * s)
        
        final_hidden_states = torch.zeros(
            (batch_size * seq_len, hidden_dim),
            dtype=hidden_states.dtype,
            device=hidden_states.device
        )

        for expert_idx in range(self.expert_number):
            expert_layer = self.experts[expert_idx]
            # expert_mask[expert_idx] shape 是 (top_k, b * s)
            idx, top_x = torch.where(expert_mask[expert_idx]) 
            # idx 和 top_x 都是一维 tensor
            # idx 的值是 0 或 1, 表示这个 token 是作为当前专家的 top1 还是 top2
            # top_x 的值是 token 在 batch*seq_len 中的位置索引
            # 例如对于 batch_size=2, seq_len=4 的输入:
            # top_x 的值范围是 0-7, 表示在展平后的 8 个 token 中的位置
            # idx 的值是 0/1, 表示这个 token 把当前专家作为其 top1/top2 专家

            # hidden_states 的 shape 是 (b * s, hidden_dim)
            # 需要取到 top_x 对应的 hidden_states
            current_state = hidden_states.unsqueeze(
                0
            )[:, top_x, :].reshape(-1, hidden_dim) # （selected_token_number, hidden_dim）

            # router_weight 的 shape 是 (b * s, top_k)
            current_hidden_states = expert_layer(
                current_state
            ) * router_weights[top_x, idx].unsqueeze(-1)  # （selected_token_number, 1） 这里有广播

            # 把当前专家的输出加到 final_hidden_states 中
            # 方式1 的写法性能更好，并且方式1容易出现
            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
            # 方式2
            # final_hidden_states[top_x] += current_hidden_states.to(hidden_states.dtype)
            # 方式1 的写法性能更差，并且方式1容易出现错误，+= 操作在处理重复索引时需要多次读写内存，可能会导致竞争条件

        # 把 final_hidden_states 还原到原来的 shape
        final_hidden_states = final_hidden_states.reshape(batch_size, seq_len, hidden_dim)

        return final_hidden_states, router_logits # shape 是 (b * s, expert_number)


def test_token_level_moe():
    x = torch.rand(2, 4, 16)
    config = MOEConfig(16, 2, 2)
    token_level_moe = SparseMOE(config)
    out = token_level_moe(x)
    print(out[0].shape, out[1].shape)


test_token_level_moe()

torch.zeros(1,5).unsqueeze(-1)

test_x = torch.randn((2,3,4))
print(test_x.shape)
test_x.view(-1,4).shape

test_x = torch.tensor([[1,0,1],[1,1,0]])
print(test_x.size())
a,b=test_x.size()
torch.where(test_x,)

class SingleExpert(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = nn.Linear(input_dim, output_dim)
        
    def forward(self, x):
        return self.linear(x)
    

class ExportRouter(nn.Module):
    def __init__(self, n_export, n_top):
        super().__init__()
        self.router = nn.Linear(hidden_dim, n_export)
        self.n_top = n_top
        self.n_export = n_export
        
    def forward(self, x):
        
        export_logits = self.router(x)  # (seq_len * batch_size, n_export)
        
        export_probs = F.softmax(export_logits, dim=-1)  # (seq_len * batch_size, n_export)
        
        router_weight, export_idx = torch.topk(export_probs, self.n_top, dim=-1)  # (seq_len * batch_size, n_top)
        
        masked_export = one_hot(export_idx, self.n_export)  # (seq_len * batch_size, n_top, n_export)
        
        masked_export = masked_export.permute(2,1,0)  # (n_export, n_top, seq_len * batch_size)
        return export_logits, router_weight, export_idx, masked_export
    
    
        

class SparseMOE(nn.Module):
    def __init__(self, config):
        
        self.exports = nn.ModuleList([SingleExpert(config.hidden_dim, config.hidden_dim) for _ in range(config.n_export)])
        
        self.router = ExportRouter(config.n_export, config.n_top)
        
    def forward(self, x):
        
        seq_len, batch_size, hidden_dim = x.shape
        
        x = x.view(seq_len * batch_size, hidden_dim)  # (seq_len * batch_size, hidden_dim)
        
        final_hidden_state = torch.zeros(seq_len * batch_size, hidden_dim)
        
        export_logits, router_weight, export_idx, masked_export = self.router(x)
        export_x = self.exports[export_idx](x)  
        for export_idx in range(config.n_export):
            
            idx, top_x = torch.where(masked_export[export_idx])
            
            
            
            current_state = router_weight[top_x, idx] * export_x[top_x]
            
            final_hidden_state.add_index_(0,top_x, router_weight_i)
            
        
        return final_hidden_state.view(seq_len, batch_size, hidden_dim)