# -*- coding: utf-8 -*-
"""
attention_from_scratch.py · 第 0003 课动手脚本（自验证式）

从零实现 scaled dot-product attention + causal mask + multi-head，
然后与两个参照实现对拍——全绿才算毕业，脚本自己就是判卷老师：
    B1: 单头 attention  vs 朴素逐位置循环实现（慢但直观）
    B2: 多头 attention  vs torch.nn.functional.scaled_dot_product_attention（官方）

最后用"玩具语义嵌入"画一张中文句子的 attention 热力图，亲眼看"谁在看谁"。

用法：
    conda activate py310_qwenpaw    # 或你自己的环境
    python scripts/attention_from_scratch.py
"""

import math
import sys
from pathlib import Path

# Windows 编码兜底：无论控制台还是重定向，强制 UTF-8 输出
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

import matplotlib
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

matplotlib.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei"]
matplotlib.rcParams["axes.unicode_minus"] = False

torch.manual_seed(0)
OUT_PNG = Path(__file__).parent / "attention_heatmap.png"

# ============================================================
# A. 从零实现
# ============================================================

def scaled_dot_product_attention(Q, K, V, mask=None):
    """Q,K,V: (batch, n_heads, seq_len, head_dim)；mask: (L, L) bool，True=可见"""
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)       # ① 相似度 + 缩放: (b,h,L,L)
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))   # ② causal mask: 前方位置置 -inf
    attn = torch.softmax(scores, dim=-1)                    # ③ 归一化: 每行和为 1
    return attn @ V, attn                                   # ④ 聚合: 对 V 加权平均


class MultiHeadAttention(nn.Module):
    """多头注意力：投影 → 切头 → 并行 attention → 拼接 → 输出投影"""

    def __init__(self, hidden: int, n_heads: int):
        super().__init__()
        assert hidden % n_heads == 0, "hidden 必须能被头数整除（GLM-4.5 的 head_dim 独立设置是例外）"
        self.h, self.d_head = n_heads, hidden // n_heads
        self.wq = nn.Linear(hidden, hidden, bias=False)
        self.wk = nn.Linear(hidden, hidden, bias=False)
        self.wv = nn.Linear(hidden, hidden, bias=False)
        self.wo = nn.Linear(hidden, hidden, bias=False)

    def forward(self, x, causal=True):
        B, L, _ = x.shape

        def split(t):  # (B, L, hidden) -> (B, h, L, d_head)
            return t.view(B, L, self.h, self.d_head).transpose(1, 2)

        q, k, v = split(self.wq(x)), split(self.wk(x)), split(self.wv(x))
        mask = torch.tril(torch.ones(L, L, dtype=torch.bool)) if causal else None
        out, attn = scaled_dot_product_attention(q, k, v, mask)
        out = out.transpose(1, 2).reshape(B, L, self.h * self.d_head)   # 拼回多头
        return self.wo(out), attn


# ============================================================
# B. 对拍验证（脚本自判卷）
# ============================================================

def naive_single_head(Q, K, V):
    """朴素逐位置循环实现：位置 i 只聚合 j<=i。慢但直观，用来验证向量化版本。"""
    B, H, L, d = Q.shape
    out = torch.zeros(B, H, L, d, dtype=Q.dtype)
    attn_ref = torch.zeros(B, H, L, L, dtype=Q.dtype)
    for b in range(B):
        for h in range(H):
            for i in range(L):
                scores = torch.tensor(
                    [float(Q[b, h, i] @ K[b, h, j]) / math.sqrt(d) for j in range(i + 1)],
                    dtype=Q.dtype,
                )
                w = torch.softmax(scores, dim=0)
                attn_ref[b, h, i, : i + 1] = w
                for t, j in enumerate(range(i + 1)):
                    out[b, h, i] += w[t] * V[b, h, j]
    return out, attn_ref


def run_tests():
    print("=" * 62)
    print("B1. 单头 from-scratch  vs  朴素循环实现")
    print("=" * 62)
    B, H, L, d = 2, 3, 5, 8
    Q = torch.randn(B, H, L, d, dtype=torch.float64)
    K = torch.randn(B, H, L, d, dtype=torch.float64)
    V = torch.randn(B, H, L, d, dtype=torch.float64)
    mask = torch.tril(torch.ones(L, L, dtype=torch.bool))

    out_ours, attn_ours = scaled_dot_product_attention(Q, K, V, mask)
    out_ref, attn_ref = naive_single_head(Q, K, V)

    d_out = (out_ours - out_ref).abs().max().item()
    d_attn = (attn_ours - attn_ref).abs().max().item()
    print(f"  输出最大误差: {d_out:.2e}  |  注意力矩阵最大误差: {d_attn:.2e}")
    ok1 = torch.allclose(out_ours, out_ref) and torch.allclose(attn_ours, attn_ref)
    print(f"  B1 结果: {'PASS ✓' if ok1 else 'FAIL ✗'}")

    print()
    print("=" * 62)
    print("B2. 多头 from-scratch  vs  torch F.scaled_dot_product_attention")
    print("=" * 62)
    mha = MultiHeadAttention(hidden=32, n_heads=4).to(torch.float64)
    x = torch.randn(2, 6, 32, dtype=torch.float64)
    out_ours, attn_ours = mha(x, causal=True)

    q = mha.wq(x).view(2, 6, 4, 8).transpose(1, 2)
    k = mha.wk(x).view(2, 6, 4, 8).transpose(1, 2)
    v = mha.wv(x).view(2, 6, 4, 8).transpose(1, 2)
    out_official = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    out_official = out_official.transpose(1, 2).reshape(2, 6, 32)
    out_official = mha.wo(out_official)

    d2 = (out_ours - out_official).abs().max().item()
    print(f"  输出最大误差: {d2:.2e}")
    ok2 = torch.allclose(out_ours, out_official)
    print(f"  B2 结果: {'PASS ✓' if ok2 else 'FAIL ✗'}")
    return ok1 and ok2


# ============================================================
# C. 热力图：谁在看谁
# ============================================================

def run_heatmap():
    tokens = ["小猫", "追", "球", "因为", "它", "想", "玩"]
    # 手工设计的玩具嵌入（4 维）：语义相近的词向量相近
    E = {
        "小猫": [1.0, 0.9, 0.0, 0.0],
        "追":   [0.0, 0.2, 1.0, 0.0],
        "球":   [0.8, 0.1, 0.0, 1.0],
        "因为": [0.0, 0.0, 0.5, 0.3],
        "它":   [0.9, 0.8, 0.0, 0.3],
        "想":   [0.0, 0.1, 0.8, 0.4],
        "玩":   [0.7, 0.0, 0.9, 0.8],
    }
    X = torch.tensor([E[t] for t in tokens], dtype=torch.float64)
    X4 = X.unsqueeze(0).unsqueeze(0)                    # -> (1, 1, L, 4)
    mask = torch.tril(torch.ones(len(tokens), len(tokens), dtype=torch.bool))
    _, attn_toy = scaled_dot_product_attention(X4, X4, X4, mask)   # 恒等投影: Q=K=V=X
    attn_toy = attn_toy[0, 0].detach().numpy()

    # 对照组：随机初始化的多头注意力——均匀混沌，结构要靠训练长出来
    mha = MultiHeadAttention(hidden=32, n_heads=4)
    x_rand = torch.randn(1, len(tokens), 32)
    _, attn_rand = mha(x_rand, causal=True)
    attn_rand = attn_rand[0, 0].detach().numpy()

    fig, axes = plt.subplots(1, 2, figsize=(13, 5.6))
    for ax, attn, title in [
        (axes[0], attn_toy, "玩具语义嵌入 + 恒等投影：结构清晰（找找『它』在看谁）"),
        (axes[1], attn_rand, "随机初始化 MHA：均匀混沌——结构要靠训练长出来"),
    ]:
        im = ax.imshow(attn, cmap="viridis")
        ax.set_xticks(range(len(tokens)), tokens)
        ax.set_yticks(range(len(tokens)), tokens)
        ax.set_xlabel("被看的 token（Key）")
        ax.set_ylabel("正在看的 token（Query）")
        ax.set_title(title, fontsize=11)
        for i in range(len(tokens)):
            for j in range(len(tokens)):
                ax.text(j, i, f"{attn[i, j]:.2f}", ha="center", va="center",
                        color="white" if attn[i, j] < attn.max() * 0.6 else "black",
                        fontsize=8)
        fig.colorbar(im, ax=ax, fraction=0.046)
    fig.suptitle("Attention 热力图（causal mask：每行只看自己及之前）", fontsize=13)
    fig.tight_layout()
    fig.savefig(OUT_PNG, dpi=130)
    print(f"\n热力图已保存: {OUT_PNG}")
    print("观察任务：左图『它』这一行，权重最大的是哪几个 token？接近 0 的是哪些？")


if __name__ == "__main__":
    ok = run_tests()
    print()
    run_heatmap()
    print("\n" + ("全部对拍通过，attention 毕业 ✓" if ok else "存在 FAIL，请回头检查实现 ✗"))