from transformers import PretrainedConfig

class MiniMindConfig(PretrainedConfig):
    model_type = "minimind"

    def __init__(
            self,
            dropout: float = 0.0,
            bos_token_id: int = 1,
            eos_token_id: int = 2,
            hidden_act: str = 'silu',
            hidden_size: int = 512,
            intermediate_size: int = None,
            max_position_embeddings: int = 32768,
            num_attention_heads: int = 8,
            num_hidden_layers: int = 8,
            num_key_value_heads: int = 2,
            vocab_size: int = 6400,
            rms_norm_eps: float = 1e-05,
            rope_theta: int = 1000000.0,
            inference_rope_scaling: bool = False,
            flash_attn: bool = True,
            ####################################################
            # Here are the specific configurations of MOE
            # When use_moe is false, the following is invalid
            ####################################################
            use_moe: bool = False,
            num_experts_per_tok: int = 2,
            n_routed_experts: int = 4,
            n_shared_experts: int = 1,
            scoring_func: str = 'softmax',
            aux_loss_alpha: float = 0.1,
            seq_aux: bool = True,
            norm_topk_prob: bool = True,
            **kwargs
    ):
        super().__init__(**kwargs)
        self.dropout = dropout
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        self.hidden_act = hidden_act
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.max_position_embeddings = max_position_embeddings
        self.num_attention_heads = num_attention_heads
        self.num_hidden_layers = num_hidden_layers
        self.num_key_value_heads = num_key_value_heads
        self.vocab_size = vocab_size
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.inference_rope_scaling = inference_rope_scaling
        # 外推长度 = factor * original_max_position_embeddings
        self.rope_scaling = {
            "beta_fast": 4,
            "beta_slow": 1,
            "factor": 4,
            "original_max_position_embeddings": 2048,
            "type": "yarn"
        } if self.inference_rope_scaling else None
        self.flash_attn = flash_attn
        ####################################################
        # Here are the specific configurations of MOE
        # When use_moe is false, the following is invalid
        ####################################################
        self.use_moe = use_moe
        self.num_experts_per_tok = num_experts_per_tok  # 每个token选择的专家数量
        self.n_routed_experts = n_routed_experts  # 总的专家数量
        self.n_shared_experts = n_shared_experts  # 共享专家
        self.scoring_func = scoring_func  # 评分函数，默认为'softmax'
        self.aux_loss_alpha = aux_loss_alpha  # 辅助损失的alpha参数
        self.seq_aux = seq_aux  # 是否在序列级别上计算辅助损失
        self.norm_topk_prob = norm_topk_prob  # 是否标准化top-k概率

# RMSNorm 类，继承自 nn.Module 以获得所有基础功能
class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        # 初始化父类
        super().__init__()
        # eps 是一个很小的数，防止计算中除以零
        self.eps = eps
        # 定义一个可学习的缩放参数 weight，并初始化为全1
        # nn.Parameter 会将这个张量注册为模型的参数，使其可以被优化器更新
        self.weight = nn.Parameter(torch.ones(dim))

    # 这是一个内部辅助函数，实现了 RMSNorm 的核心数学计算
    def _norm(self, x):
        # 1. x.pow(2): 计算输入 x 中每个元素的平方
        # 2. .mean(-1, keepdim=True): 沿着最后一个维度（特征维度）计算平方的均值
        #    keepdim=True 保持了维度，例如 (batch, seq_len, dim) -> (batch, seq_len, 1)
        # 3. + self.eps: 加上极小值 eps
        # 4. torch.rsqrt(...): 计算结果的平方根的倒数，即 1/sqrt(...)，这比直接除法效率更高
        # 5. x * ...: 将原始输入 x 乘以这个倒数，完成归一化
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    # forward 方法定义了数据如何流过这个模块
    def forward(self, x):
        # 1. self._norm(x.float()): 调用核心计算函数对输入 x 进行归一化。
        #    x.float() 是为了保证计算精度
        # 2. .type_as(x): 将计算结果的数据类型转换回原始输入 x 的类型 (例如 float16)
        # 3. self.weight * ...: 将归一化后的结果乘以可学习的缩放参数 weight
        return self.weight * self._norm(x.float()).type_as(x)

# 函数：预计算旋转位置编码所需的 cos 和 sin 值
def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6,
                         rope_scaling: Optional[dict] = None):
    # 1. 计算基础频率 freqs。这是RoPE的核心，不同的维度会有不同的旋转速度。
    freqs = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))

    # (这部分是高级功能，用于处理超长文本，可以暂时忽略)
    if rope_scaling is not None:
        # ... YaRN scaling logic ...
        pass

    # 2. 生成所有可能的位置序列 t，从 0 到 end-1
    t = torch.arange(end, device=freqs.device)
    # 3. 计算每个位置 t 和每个频率 freqs 的相位角
    # torch.outer(t, freqs) 的结果是一个 (end, dim/2) 的矩阵，代表了每个位置在每个频率上的角度
    freqs = torch.outer(t, freqs).float()

    # 4. 计算所有角度的 cos 值和 sin 值，并拼接起来，形成最终的预计算表
    # freqs_cos 和 freqs_sin 的维度都是 (end, dim)
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
    return freqs_cos, freqs_sin

# 函数：将RoPE应用到 Query 和 Key 向量上
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    # 内部函数：将向量的后半部分和前半部分对调，并给后半部分取反，模拟复数中的乘以 i
    def rotate_half(x):
        return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)

    # RoPE的核心公式：q_rot = q * cos + rotate_half(q) * sin
    # 这在数学上等价于用一个旋转矩阵去乘以 q 向量
    q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
    k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
    return q_embed, k_embed

class Attention(nn.Module):
    def __init__(self, args: MiniMindConfig):
        super().__init__()
        # ... 初始化参数 ...
        self.n_local_heads = args.num_attention_heads
        self.n_local_kv_heads = args.num_key_value_heads
        self.n_rep = self.n_local_heads // self.n_local_kv_heads # 每个KV头被共享的次数
        self.head_dim = args.hidden_size // args.num_attention_heads

        # 定义Q, K, V的线性投影层
        self.q_proj = nn.Linear(args.hidden_size, args.num_attention_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        # 定义输出的线性投影层
        self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False)
        # ... dropout 和 flash attention 开关 ...
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn

    def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
        bsz, seq_len, _ = x.shape
        # 1. 将输入 x 通过线性层计算出 Q, K, V
        xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)

        # 2. 将Q,K,V的形状调整为 (batch_size, seq_len, num_heads, head_dim) 以便多头并行计算
        xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)

        # 3. 应用旋转位置编码 (RoPE)
        cos, sin = position_embeddings
        xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len])

        # 4. KV 缓存机制 (用于加速推理)
        if past_key_value is not None:
            # 将过去的 K, V 和当前的 K, V 拼接起来
            xk = torch.cat([past_key_value[0], xk], dim=1)
            xv = torch.cat([past_key_value[1], xv], dim=1)
        past_kv = (xk, xv) if use_cache else None

        # 5. 为了计算，将维度调整为 (bsz, num_heads, seq_len, head_dim)
        #    并使用 repeat_kv 实现分组查询注意力 (GQA)，让多个Q头共享K,V头
        xq, xk, xv = (
            xq.transpose(1, 2),
            repeat_kv(xk, self.n_rep).transpose(1, 2),
            repeat_kv(xv, self.n_rep).transpose(1, 2)
        )

        # 6. 计算注意力分数并得到输出
        if self.flash: # 如果支持并开启了Flash Attention
            # 使用PyTorch内置的高度优化的注意力实现，速度快，省内存
            output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
        else: # 使用手动实现的注意力
            # Q 和 K 的转置相乘，得到注意力分数
            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
            # 添加因果掩码，防止看到未来的信息
            scores = scores + torch.triu(...)
            # ... (处理额外的 attention_mask) ...
            # Softmax 将分数转换为概率
            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
            # 概率与 V 相乘，得到加权和
            output = scores @ xv

        # 7. 整理输出形状并通过输出投影层
        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
        output = self.o_proj(output)
        return output, past_kv

class FeedForward(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        # ... 自动计算中间层大小 ...
        # SwiGLU结构包含三个线性层
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act] # 激活函数，例如 SiLU

    def forward(self, x):
        # SwiGLU的计算公式: (x @ W_gate * act_fn(x @ W_up)) @ W_down
        # 这里 act_fn(self.gate_proj(x)) 和 self.up_proj(x) 的顺序与公式略有不同，但本质相同
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

class MOEFeedForward(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        # 1. 创建多个独立的 FeedForward 网络作为“专家”
        self.experts = nn.ModuleList([FeedForward(config) for _ in range(config.n_routed_experts)])
        # 2. 创建一个门控网络，用于决定使用哪些专家
        self.gate = MoEGate(config)
        # ... (共享专家逻辑) ...

    def forward(self, x):
        # ...
        # 3. 使用门控网络为每个词元选择 top-k 个专家，并计算权重
        topk_idx, topk_weight, aux_loss = self.gate(x)
        # ...
        # 4. 将输入 x 发送给被选中的专家进行处理
        #    (训练和推理时的实现细节不同，但目标一致)
        #    例如训练时：
        y = torch.empty_like(x_repeated)
        for i, expert in enumerate(self.experts):
            # 找到所有应该由第 i 个专家处理的 token
            # expert(...) 对这些 token 进行 FFN 计算
            y[flat_topk_idx == i] = expert(x[flat_topk_idx == i])

        # 5. 将多个专家的输出根据门控权重进行加权求和
        y = (y.view(...) * topk_weight.unsqueeze(-1)).sum(dim=1)
        # ...
        return y

class MiniMindBlock(nn.Module):
    def __init__(self, layer_id: int, config: MiniMindConfig):
        super().__init__()
        # 初始化一个 Attention 模块
        self.self_attn = Attention(config)
        # 初始化输入和后注意力层的归一化模块
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # 根据配置，选择使用普通的 FFN 还是 MoE 模块
        self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)

    def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
        # 1. 第一个残差连接 (Residual Connection)
        residual = hidden_states

        # 2. Pre-Norm: 先归一化，再进 Attention
        hidden_states, present_key_value = self.self_attn(
            self.input_layernorm(hidden_states),
            position_embeddings,
            past_key_value, use_cache, attention_mask
        )

        # 3. 将 Attention 的输出与原始输入相加
        hidden_states += residual

        # 4. 第二个残差连接
        residual = hidden_states

        # 5. Pre-Norm: 先归一化，再进 MLP (FFN 或 MoE)
        mlp_output = self.mlp(self.post_attention_layernorm(hidden_states))

        # 6. 将 MLP 的输出与 Attention 之后的输入相加
        hidden_states = residual + mlp_output

        return hidden_states, present_key_value

class MiniMindModel(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        # 1. 词嵌入层：将 token ID 转换为向量
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        # 2. 创建多个 MiniMindBlock 层（config中的8层）
        self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(config.num_hidden_layers)])
        # 3. 最后的归一化层
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # 4. 预计算 RoPE 的 cos 和 sin 值，并注册为 buffer
        freqs_cos, freqs_sin = precompute_freqs_cis(...)
        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
        self.register_buffer("freqs_sin", freqs_sin, persistent=False)

    def forward(self, input_ids, ...):
        # 5. 将输入 ID 转换为词嵌入向量
        hidden_states = self.embed_tokens(input_ids)

        # 6. 准备当前序列需要的 RoPE 位置编码
        position_embeddings = (
            self.freqs_cos[start_pos:start_pos + seq_length],
            self.freqs_sin[start_pos:start_pos + seq_length]
        )

        # 7. 循环，让数据依次流过每一个 MiniMindBlock 层
        for layer in self.layers:
            hidden_states, present = layer(
                hidden_states,
                position_embeddings,
                ...
            )

        # 8. 对最后一层的输出进行最终的归一化
        hidden_states = self.norm(hidden_states)
        return hidden_states, ...

class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = MiniMindConfig

    def __init__(self, config: MiniMindConfig):
        super().__init__(config)
        # 1. 包含一个完整的 MiniMindModel 作为其骨干网络
        self.model = MiniMindModel(config)
        # 2. 创建一个线性层作为“语言模型头”，用于预测下一个词
        #    它的输出维度是词汇表大小 vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # 3. 权重绑定：让输入嵌入层和输出预测头的权重共享
        self.model.embed_tokens.weight = self.lm_head.weight

    def forward(self, input_ids, ...):
        # 4. 首先，调用骨干网络 self.model 进行所有核心计算
        h, past_kvs, aux_loss = self.model(
            input_ids=input_ids,
            ...
        )

        # 5. 将骨干网络的输出 h 通过 lm_head，得到最终的 logits
        #    logits 的维度是 (batch_size, seq_len, vocab_size)
        logits = self.lm_head(h)

        # 6. 将 logits 和其他信息 (如kv缓存) 封装并返回
        return CausalLMOutputWithPast(logits=logits, past_key_values=past_kvs, ...)