#Transformer完整架构详解：从自注意力机制到编码器解码器设计及PyTorch实现

import torch
import torch.nn as nn
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    缩放点积注意力的核心实现
    Q/K/V: (batch_size, seq_len, d_k/d_v)
    """
    d_k = Q.size(-1)
    # 1. 计算Q和K的相似度（注意力分数）
    # 除以sqrt(d_k)是为了防止分数太大导致softmax梯度消失
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 2. 应用掩码（比如解码器里不能看未来的词）
    if mask is not None:
        scores.masked_fill_(mask == 0, -1e9)
    
    # 3. 用softmax把分数转成0-1的权重，总和为1
    attention_weights = torch.softmax(scores, dim=-1)
    
    # 4. 加权求和V得到最终输出
    output = torch.matmul(attention_weights, V)
    
    return output, attention_weights

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # 四个线性变换：Q/K/V的生成，以及多头结果的合并
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
        self.dropout = nn.Dropout(dropout)
        
    def split_heads(self, x, batch_size):
        """把大的向量拆成多个小的子向量，并行计算"""
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(1, 2)  # 变成(batch_size, num_heads, seq_len, d_k)
    
    def combine_heads(self, x, batch_size):
        """把多个子向量拼回原来的形状"""
        x = x.transpose(1, 2).contiguous()
        return x.view(batch_size, -1, self.d_model)
    
    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)
        
        # 1. 线性变换生成初始Q/K/V
        Q = self.W_q(Q)
        K = self.W_k(K)
        V = self.W_v(V)
        
        # 2. 拆成多头
        Q = self.split_heads(Q, batch_size)
        K = self.split_heads(K, batch_size)
        V = self.split_heads(V, batch_size)
        
        # 3. 计算缩放点积注意力
        scaled_attn, attn_weights = scaled_dot_product_attention(Q, K, V, mask)
        
        # 4. 合并多头
        concat_attn = self.combine_heads(scaled_attn, batch_size)
        
        # 5. 最终线性变换
        output = self.W_o(concat_attn)
        
        return output, attn_weights

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        # 创建一个足够大的位置编码矩阵
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        
        # 偶数维度用sin，奇数维度用cos
        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)
        
        # 注册为buffer，不会被优化器更新
        self.register_buffer('pe', pe.unsqueeze(0))
        
    def forward(self, x):
        """把位置编码加到词嵌入上"""
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.mha = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask=None):
        # 残差+层归一化的顺序是：LayerNorm(x + SubLayer(x))
        # 先过自注意力
        attn_out, _ = self.mha(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))
        
        # 再过前馈网络
        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_out))
        
        return x

class Encoder(nn.Module):
    def __init__(self, num_layers, d_model, num_heads, d_ff, vocab_size, max_len, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
        self.enc_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
        
    def forward(self, x, mask=None):
        # 论文里提到要把词嵌入乘以sqrt(d_model)，防止位置编码的影响被淹没
        x = self.embedding(x) * math.sqrt(self.d_model)
        x = self.pos_encoding(x)
        
        for layer in self.enc_layers:
            x = layer(x, mask)
            
        return x