#注意力机制 (Attention) 详解：为什么"注意力"就是你所需要的一切

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

def scaled_dot_product_attention(
    Q: torch.Tensor, 
    K: torch.Tensor, 
    V: torch.Tensor, 
    mask: torch.Tensor = None
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    实现Transformer的核心缩放点积注意力
    
    参数形状说明：
    - Q: (batch_size, num_heads, seq_len_q, d_k)
    - K: (batch_size, num_heads, seq_len_k, d_k)
    - V: (batch_size, num_heads, seq_len_v, d_v)
      （通常seq_len_q = seq_len_k = seq_len_v，除了特殊应用场景）
    - mask: (batch_size, 1, seq_len_q, seq_len_k)，可选掩码（遮挡padding或未来信息）
    
    返回：
    - output: (batch_size, num_heads, seq_len_q, d_v)，注意力加权后的输出
    - attention_weights: (batch_size, num_heads, seq_len_q, seq_len_k)，可解释的注意力权重
    """
    # 获取Q的最后一个维度 d_k，用于缩放
    d_k = Q.size(-1)

    # -------------------------- 步骤1：计算相似度矩阵 --------------------------
    # 用Query和每个Key做点积，点积越大表示越相关
    # 再除以 sqrt(d_k) 进行缩放，防止d_k太大时点积结果过大，导致softmax梯度消失
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    # -------------------------- 步骤2：可选的掩码操作 --------------------------
    if mask is not None:
        # mask中为0的位置（比如padding、未来的词），把score设为-1e9（接近负无穷）
        # 这样softmax后这些位置的权重几乎为0，完全不会被关注
        scores = scores.masked_fill(mask == 0, -1e9)

    # -------------------------- 步骤3：归一化得到注意力权重 --------------------------
    # 对最后一个维度（seq_len_k）做softmax，把score变成[0,1]的概率分布，总和为1
    attention_weights = F.softmax(scores, dim=-1)

    # -------------------------- 步骤4：用权重加权Value得到输出 --------------------------
    output = torch.matmul(attention_weights, V)

    return output, attention_weights

import matplotlib.pyplot as plt
import seaborn as sns

def plot_attention_heatmap(
    attention_weights: torch.Tensor,
    source_tokens: list[str],
    target_tokens: list[str],
    save_path: str = "attention_heatmap.png"
) -> None:
    """
    绘制Attention的热力图
    
    参数：
    - attention_weights: (seq_len_t, seq_len_s)，单头无batch的注意力权重
    - source_tokens: 源语言的token列表
    - target_tokens: 目标语言的token列表
    - save_path: 图片保存路径
    """
    # 先把torch张量转成numpy数组，方便绘图
    weights_np = attention_weights.detach().cpu().numpy()

    # 设置绘图风格，用seaborn的暖色调，更直观
    sns.set_style("whitegrid")
    plt.figure(figsize=(10, 8))
    heatmap = sns.heatmap(
        weights_np,
        xticklabels=source_tokens,
        yticklabels=target_tokens,
        cmap="YlOrRd",  # 黄橙红渐变，颜色越深关注度越高
        annot=True,      # 在热力图上显示具体的权重数值
        fmt=".2f",       # 数值保留两位小数
        cbar_kws={"label": "Attention Weight"},  # 给颜色条加标签
    )

    # 设置图表标题和轴标签
    heatmap.set_title("机器翻译中的Attention权重热力图", fontsize=14, pad=20)
    heatmap.set_xlabel("Source Tokens（中文）", fontsize=12)
    heatmap.set_ylabel("Target Tokens（英文）", fontsize=12)

    # 调整布局，防止标签被截断
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)  # dpi设高一点，图片更清晰
    plt.show()

# -------------------------- 测试代码 --------------------------
if __name__ == "__main__":
    # 用上面的模拟注意力矩阵
    test_weights = torch.tensor([
        [0.92, 0.05, 0.03],
        [0.04, 0.91, 0.05],
        [0.02, 0.04, 0.94]
    ])
    test_source = ["我", "爱", "PythonAI"]
    test_target = ["I", "love", "PythonAI"]

    plot_attention_heatmap(test_weights, test_source, test_target)