"""
=============================================================================
第 1 讲 · 代码 04：注意力 —— 软的字典查找、为什么除以 √d、因果遮罩、归纳头
=============================================================================

对应讲义：lectures/lecture01_history_of_lms.md §6（Transformer）

【这个脚本要回答的问题】
    1. 注意力在算什么？为什么说它是一个"软的字典查找"？（演示 1）
    2. 公式里为什么要除以 √d？不除会怎样？（演示 2）
    3. GPT 这种"从左往右写"的模型，怎样保证一个位置看不见它右边的词？（演示 3）
    4. 两层注意力就能实现"看到 A B … A，预测 B"的抄写电路（归纳头）。为什么两层缺一不可？
       为什么需要位置信息？（演示 4）
    5. 我们训练出来的迷你 GPT，回答问题时在"看"哪里？它会抄写重复的序列吗？（演示 5）

【怎么读输出】
    注意力权重是一行加起来等于 1 的数：第 t 行第 s 列表示"位置 t 从位置 s 取了多少信息"。

【运行】python3 04_attention.py      （几秒钟；需要迷你 GPT 的缓存，没有的话会先训练约 50 秒）
【依赖】numpy、PyTorch；import 同目录的 attention_np.py、tiny_gpt.py、pretrain_world.py、console.py。
"""

import numpy as np
import torch

from attention_np import (HandBuiltInduction, attention, causal_mask, copy_accuracy,
                          repeated_sequences, softmax)
from console import pct, print_header, print_observations, print_table
from pretrain_world import RELATIONS, question_text
from tiny_gpt import get_model


# ---------------------------------------------------------------------------
def demo1_soft_lookup():
    """演示 1：注意力 = 软的字典查找。"""
    print_header("演示 1：注意力是一个'软的字典查找'")
    # 一个有 4 个条目的"字典"：key 描述条目是什么，value 是条目里存的东西
    keys = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.7, 0.7, 0.0]])
    names = ["颜色", "数字", "地点", "颜色+数字"]
    values = np.array([[10.0], [20.0], [30.0], [40.0]])
    rows = []
    for qname, q in [("找'颜色'", [1, 0, 0]), ("找'地点'", [0, 0, 1]), ("又像颜色又像数字", [1, 1, 0]),
                     ("和谁都不像", [0, 0, 0])]:
        for beta in [1.0, 10.0]:
            Q = beta * np.array([q], dtype=float)
            out, W = attention(Q, keys, values, scale=False)
            rows.append([qname, beta, "  ".join(f"{w:.2f}" for w in W[0]), f"{out[0, 0]:.1f}"])
    print(f"4 个条目的 key = {names}，value = [10, 20, 30, 40]；分数 = β · (query · key)")
    print_table(["query", "锐利度 β", "对 4 个条目的权重", "取回的 value"], rows)
    print_observations([
        "普通的字典查找是'完全匹配才返回'；注意力是'越像，取得越多'，结果是所有 value 的加权平均——所以叫**软**查找。",
        "β 越大（分数差距越大），权重越接近 one-hot，就越接近硬查找；β 小时权重平摊，取回的是一个'平均值'。",
        "和谁都不像时，权重完全平均：注意力不会说'没找到'，它总会返回点什么。",
        "在 Transformer 里，query、key、value 都是从每个位置的向量**学出来**的线性变换：模型自己学会'该找什么、该被怎样找到'。",
    ])


def demo2_scaling(seed=0):
    """演示 2：为什么除以 √d。"""
    print_header("演示 2：为什么要除以 √d——不缩放时，维度越高，softmax 越接近 one-hot")
    rng = np.random.default_rng(seed)
    T, trials = 16, 2000
    rows = []
    for d in [4, 16, 64, 256, 1024]:
        stats = {}
        for scale in [False, True]:
            mx, ent = [], []
            for _ in range(trials):
                q, K = rng.standard_normal((1, d)), rng.standard_normal((T, d))
                w = softmax(q @ K.T / (np.sqrt(d) if scale else 1.0))[0]
                mx.append(w.max())
                ent.append(-(w * np.log(w + 1e-30)).sum())
            stats[scale] = (np.mean(mx), np.mean(ent))
        rows.append([d, f"{stats[False][0]:.3f}", f"{stats[False][1]:.2f}", f"{stats[True][0]:.3f}", f"{stats[True][1]:.2f}"])
    print(f"随机的 q、k（每一维 ~ N(0,1)），{T} 个 key，各 {trials} 次；均匀分布的熵 = ln {T} = {np.log(T):.2f}")
    print_table(["维度 d", "不缩放：最大权重", "不缩放：熵", "除以 √d：最大权重", "除以 √d：熵"], rows)
    print_observations([
        "q·k 是 d 个随机数之和，标准差是 √d。不缩放时 d 越大，分数差距越大，最大权重趋近 1、熵趋近 0——"
        "每个位置只看一个位置，而且 softmax 在这种饱和状态下梯度几乎为 0，训练会很困难。",
        "除以 √d 之后，无论 d 多大，权重的分布都差不多：模型可以自己通过学习把分数拉开，而不是一开始就被维度'逼'成 one-hot。",
    ])


def demo3_causal_mask(seed=1):
    """演示 3：因果遮罩。"""
    print_header("演示 3：因果遮罩——GPT 的每个位置只能看见自己和左边")
    words = "the queen wears a crown .".split()
    rng = np.random.default_rng(seed)
    X = rng.standard_normal((len(words), 8))
    _, W_full = attention(X, X, X)
    _, W_causal = attention(X, X, X, mask=causal_mask(len(words)))
    for title, W in [("不遮罩（BERT 式的双向注意力）", W_full), ("因果遮罩（GPT 式）", W_causal)]:
        print(f"\n{title}：")
        print_table(["位置 t \\ 看位置 s"] + words,
                    [[w] + [f"{x:.2f}" for x in W[i]] for i, w in enumerate(words)])
    print_observations([
        "因果遮罩把每个位置右边的分数设成 −∞，softmax 之后权重恰好是 0：第 3 个词的表示只由前 3 个词决定。",
        "这样训练时可以**一次**算出所有位置的'下一个词'预测（每个位置都不会偷看答案），这正是 Transformer 比 RNN 训练快的原因之一。",
        "生成时也一样：模型只根据已经写下的内容决定下一个词。所以对 GPT 来说，**提示的内容只能影响它后面的部分**——"
        "把指令放在要处理的文本前面还是后面，效果可能不同（第 2 讲）。",
    ])


def mean_weight_on_answer(m, data):
    """
    后一半每个位置上，归纳头的输出里"正确的下一个词"得到的权重（输出是若干 one-hot 的加权平均，所以它在 0–1 之间）。
    """
    half = data.shape[1] // 2
    tot, n = 0.0, 0
    for seq in data:
        X = m.embed(seq)
        if m.layers >= 2:
            X, _ = m.layer1_prev_token(X)
        out, _ = m.layer2_induction(X)
        for t in range(half, len(seq) - 1):
            tot += out[t, seq[t + 1]]
            n += 1
    return tot / n


def demo4_induction():
    """演示 4：手工搭建的归纳头。"""
    print_header("演示 4：两层注意力实现'看到 A B … A，预测 B'（手工设定权重的归纳头）")
    V, half = 20, 10
    data = repeated_sequences(V=V, half=half, n=300, seed=0)
    rows = []
    variants = [
        ("完整电路（β = 20）", HandBuiltInduction(V, 2 * half)),
        ("锐利度 β = 5", HandBuiltInduction(V, 2 * half, beta=5.0)),
        ("锐利度 β = 2", HandBuiltInduction(V, 2 * half, beta=2.0)),
        ("没有位置信息", HandBuiltInduction(V, 2 * half, use_pos=False)),
        ("只有第 2 层（没有前一个词头）", HandBuiltInduction(V, 2 * half, layers=1)),
    ]
    res, conf = {}, {}
    for name, m in variants:
        a1, a2 = copy_accuracy(lambda s, m=m: m.predict(s)[0], data)
        res[name] = a2
        conf[name] = mean_weight_on_answer(m, data)
        rows.append([name, pct(a1, 1), pct(a2, 1), f"{conf[name]:.2f}"])
    print(f"测试数据：300 条序列，每条是 {half} 个不同的随机词（词表 {V} 个）重复两遍；随机猜的准确率 = 1/{V} = {pct(1, V)}")
    print_table(["模型", "前一半（没法抄）", "后一半（可以抄）", "后一半：正确的词得到的权重"], rows)

    seq = data[0]
    m = HandBuiltInduction(V, 2 * half)
    X1, W1 = m.layer1_prev_token(m.embed(seq))
    _, W2 = m.layer2_induction(X1)
    t = half + 3
    print(f"\n一个例子：序列 = {seq.tolist()}")
    print(f"  位置 {t} 的词是 {seq[t]}；第 1 层让每个位置看它左边：位置 {t} 看位置 {int(W1[t].argmax())}。")
    print(f"  第 2 层：位置 {t} 去找'前一个词是 {seq[t]}'的位置 → 位置 {int(W2[t].argmax())}"
          f"（权重 {W2[t].max():.2f}），那里的词是 {seq[int(W2[t].argmax())]}，于是预测下一个词是 {seq[int(W2[t].argmax())]}"
          f"（真实的下一个词：{seq[t + 1]}）。")
    print_observations([
        "完整电路在后一半做到 100%：第 1 层把'前一个词是谁'写到每个位置上，第 2 层拿'我是谁'去匹配'谁的前一个词是我'。"
        "这就是 Olsson et al.（2022）在真实 Transformer 里找到的**归纳头**，被认为是上下文学习的重要机制之一。",
        "前一半是 0%（比随机猜 5% 还低）：没有可抄的东西时，注意力平摊在前面的词上，它总是'抄'一个已经出现过的词，"
        "而前一半的下一个词恰好总是新词。",
        f"没有位置信息就完全不行（{pct(res['没有位置信息'], 1)}）：注意力本身对顺序是'无感'的，不告诉它谁在谁左边，"
        "第 1 层就找不到'前一个词'。所以 Transformer 必须额外加入位置编码。",
        f"只有第 2 层也不行（{pct(res['只有第 2 层（没有前一个词头）'], 1)}）：一个头只能做一次'查找'，"
        "'先查前一个词、再按它去匹配'需要两层**串起来**。这是'深度'的一个具体用处。",
        f"锐利度从 20 降到 2，准确率仍是 {pct(res['锐利度 β = 2'], 1)}（正确的词仍然权重最大），但它得到的权重从 "
        f"{conf['完整电路（β = 20）']:.2f} 降到 {conf['锐利度 β = 2']:.2f}：软查找取回的是好几个位置的混合，"
        "模型'抄'得越来越没把握。准确率看不出来的退化，置信度看得出来（第 5 讲讲校准时会再遇到）。",
    ])


@torch.no_grad()
def demo5_tiny_gpt():
    """演示 5：迷你 GPT 回答问题时在看哪里；它会不会抄写。"""
    print_header("演示 5：迷你 GPT 回答问题时在'看'哪里？它会抄写重复的序列吗？")
    model, tok, info = get_model("base")
    world = info["world"]
    L, H = model.cfg.n_layer, model.cfg.n_head
    to_entity = np.zeros((L, H))
    to_relation = np.zeros((L, H))
    rel_words = {"capital", "language", "speak", "known", "leads", "of", "for", "in"}
    n = 0
    for c in world.countries:
        for rel in RELATIONS:
            prompt = f"<doc> faq : q : {question_text(c.name, rel)} a :"
            ids = tok.encode(prompt)
            _, _, atts = model(torch.tensor([ids]), return_attn=True)
            pos = ids.index(tok.stoi[c.name])
            rel_pos = [i for i, x in enumerate(ids) if tok.itos[x] in rel_words]
            for l in range(L):
                w = atts[l][0, :, -1]                      # 最后一个位置（"a :" 的 ":"）对所有位置的权重
                to_entity[l] += w[:, pos].numpy()
                to_relation[l] += w[:, rel_pos].sum(dim=1).numpy()
            n += 1
    rows = [[f"第 {l} 层 头 {h}", f"{to_entity[l, h] / n:.2f}", f"{to_relation[l, h] / n:.2f}"]
            for l in range(L) for h in range(H)]
    print(f"提示：'<doc> faq : q : <问题> a :'，{n} 道题（24 国 × 4 种关系）。看最后一个位置（马上要写答案）的注意力：")
    print_table(["注意力头", "放在国名上的平均权重", "放在关系词上的平均权重"], rows)
    best = np.unravel_index(np.argmax(to_entity), to_entity.shape)

    # 抄写测试：用模型词表里的首都名、人名拼随机序列，重复两遍
    rng = np.random.default_rng(0)
    pool = [tok.stoi[c.capital] for c in world.countries] + [tok.stoi[c.leader] for c in world.countries]
    hits = total = 0
    for _ in range(100):
        seq = list(rng.choice(pool, size=12, replace=False))
        ids = [tok.stoi["<doc>"]] + seq + seq
        logits = model(torch.tensor([ids]))[0][0]
        pred = logits.argmax(-1).tolist()
        for t in range(13, len(ids) - 1):        # 第二遍里的位置
            hits += pred[t] == ids[t + 1]
            total += 1
    print(f"\n抄写测试：把 12 个首都名 / 人名随机排成一串，重复两遍，看模型能不能在第二遍预测出下一个词：{pct(hits, total)}（{total} 个位置）")
    print_observations([
        f"第 {best[0]} 层头 {best[1]} 在要写答案时把约 {to_entity[best] / n:.0%} 的注意力放在国名上：它在'取回问的是哪个国家'；"
        "另有一个头主要看关系词（capital、of、known……），取回'问的是什么'。",
        "这和演示 1 的软查找是同一回事：答案位置的 query 去找'国名'这个 key。注意力权重是理解模型的一扇窗，"
        "但它只告诉你'从哪里取了信息'，不告诉你'取回来之后怎么用'（那部分在 MLP 里）。",
        f"可这个迷你 GPT **完全不会抄写**（{pct(hits, total)}）：它的训练语料里从来没有'重复出现的随机序列'，"
        "所以它没有长出归纳头。真实语料里名字、短语反复出现，归纳头在训练早期就会出现。"
        "能力从哪里来？——来自**训练数据奖励了什么**。第 2 讲会让一个迷你模型在合适的数据上长出它。",
    ])


def main():
    demo1_soft_lookup()
    demo2_scaling()
    demo3_causal_mask()
    demo4_induction()
    demo5_tiny_gpt()


if __name__ == "__main__":
    torch.set_num_threads(1)
    main()