import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import copy

def attention(Q, K, V, mask=None, dropout=None):
    """缩放点积注意力函数"""
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    attn_weights = F.softmax(scores, dim=-1)
    if dropout is not None:
        attn_weights = dropout(attn_weights)
    output = torch.matmul(attn_weights, V)
    return output, attn_weights


def clones(module, N):
    """克隆模块N次"""
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class MultiHeadedAttention(nn.Module):
    """多头注意力机制 - 支持自注意力/因果注意力/交叉注意力"""
    def __init__(self, h, d_model, dropout=0.1):
        """
        h: 注意力头数 (如8)
        d_model: 模型维度 (如512)
        """
        super().__init__()
        assert d_model % h == 0, "d_model必须能被头数h整除"
        
        self.d_k = d_model // h  # 每个头的维度
        self.h = h               # 头数
        self.d_model = d_model
        
        # 4个全连接层: W_Q, W_K, W_V, W_O
        # 用一个大的全连接层代替h个小全连接层
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None  # 保存注意力得分, 用于可视化
        self.dropout = nn.Dropout(p=dropout)
    
    def forward(self, query, key, value, mask=None):
        """
        query, key, value: [batch, seq_len, d_model]
        mask: [batch, 1, seq_len] 或 None
        
        支持三种模式:
        - 自注意力: query=key=value
        - 因果注意力: query=key=value + causal mask
        - 交叉注意力: query来自解码器, key=value来自编码器
        """
        if mask is not None:
            # mask增加头维度: [batch, 1, 1, seq_len]
            mask = mask.unsqueeze(1)
        
        batch_size = query.size(0)
        
        # 1. 线性映射并切分多头
        # [batch, seq_len, d_model] → [batch, seq_len, h, d_k] → [batch, h, seq_len, d_k]
        query, key, value = [
            l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
            for l, x in zip(self.linears, (query, key, value))
        ]
        
        # 2. 注意力计算
        x, self.attn = attention(
            query, key, value, mask=mask, dropout=self.dropout
        )
        
        # 3. 拼接多头
        # [batch, h, seq_len, d_k] → [batch, seq_len, h, d_k] → [batch, seq_len, d_model]
        x = x.transpose(1, 2).contiguous().view(
            batch_size, -1, self.h * self.d_k
        )
        
        # 4. 最终线性映射
        return self.linears[-1](x)


# 使用示例
if __name__ == "__main__":
    d_model = 512
    h = 8
    batch_size = 2
    seq_len = 10
    
    mha = MultiHeadedAttention(h, d_model)
    
    # 1. 自注意力 (编码器中)
    x = torch.randn(batch_size, seq_len, d_model)
    output = mha(x, x, x, mask=None)
    print(f"自注意力输出: {output.shape}")  # [2, 10, 512]
    
    # 2. 交叉注意力 (解码器中)
    dec_input = torch.randn(batch_size, 5, d_model)   # 解码器输入
    enc_output = torch.randn(batch_size, 10, d_model)  # 编码器输出
    output = mha(dec_input, enc_output, enc_output)
    print(f"交叉注意力输出: {output.shape}")  # [2, 5, 512]