import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    """多头注意力：自注意力的核心实现"""
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_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)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # Step 1: 生成 Q、K、V
        Q = self.W_Q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_K(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_V(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # Step 2: 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # Step 3: Softmax + 加权求和
        attention = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention, V)

        # 合并多头并投影
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        return self.W_O(output)

class TransformerBlock(nn.Module):
    """一个 Transformer 层 = 自注意力 + 前馈网络"""
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.attention = MultiHeadAttention(d_model, n_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 自注意力 + 残差连接 + 层归一化
        attn_output = self.attention(x, mask)
        x = self.norm1(x + self.dropout(attn_output))

        # 前馈网络 + 残差连接 + 层归一化
        ff_output = self.ff(x)
        x = self.norm2(x + self.dropout(ff_output))
        return x

class MiniGPT(nn.Module):
    """一个极简的 GPT 风格 Transformer"""
    def __init__(self, vocab_size=10000, d_model=512, n_heads=8,
                 n_layers=6, max_len=1024):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(max_len, d_model)

        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, n_heads) for _ in range(n_layers)
        ])
        self.ln_final = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size)  # 预测下一个词

    def forward(self, token_ids):
        batch_size, seq_len = token_ids.shape

        # Token 嵌入 + 位置嵌入
        positions = torch.arange(seq_len, device=token_ids.device).unsqueeze(0)
        x = self.token_embedding(token_ids) + self.position_embedding(positions)

        # GPT 的因果掩码（每个位置只能看到之前的 Token）
        mask = torch.tril(torch.ones(seq_len, seq_len)).view(1, 1, seq_len, seq_len)

        for block in self.blocks:
            x = block(x, mask)

        x = self.ln_final(x)
        return self.lm_head(x)  # [batch, seq_len, vocab_size]

# 使用示例
model = MiniGPT(vocab_size=10000)
dummy_input = torch.randint(0, 10000, (2, 128))  # batch=2, seq_len=128
output = model(dummy_input)
print(f"输入形状: {dummy_input.shape}")   # [2, 128]
print(f"输出形状: {output.shape}")        # [2, 128, 10000]
print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")