"""
=============================================================================
第 3 讲 · 代码 5：训练 Transformer —— 注意力到底带来了什么？
=============================================================================

对应讲义：lectures/lecture03_architectures_hyperparameters.md 第 8 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
第 2 讲的"固定窗口"模型只能看到最近 8 个字节。换成 Transformer 后，模型能看到整个上下文。
这到底有什么用？我们设计一个实验让差别一目了然：

    生成大量模板化的小故事，每个故事里有两个随机起名的角色（名字从约 8000 个可能中随机取），
    名字在故事中反复出现：
        "Once upon a time, there was a brave fox named Kalomi. Kalomi loved to sing ..."
    - 第一次出现的名字是随机的，谁也猜不到；
    - 但**再次出现**时，只要能"回头看"到前文，就能照抄出来。
    窗口模型看不到几十个字节之前的名字，只能瞎猜；Transformer 应该能学会"抄写"。

这个文件：
    1. 生成数据，并标记出"重复出现的名字"所在的字节；
    2. 在相同的训练预算下，训练第 2 讲的窗口模型和第 3 讲的 Transformer，分别测量
       整体损失、以及"重复名字"字节上的损失；
    3. 打开 Transformer 的注意力，看它预测名字时"在看哪里"；
    4. 用 top-p（核）采样生成新故事，看名字是否前后一致；
    5. 稳定性实验：Pre-norm vs Post-norm，在较大学习率下谁能训练起来？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 05_train_transformer.py        （CPU 上约 2 分钟）

依赖：同目录的 04_transformer_lm.py；第 2 讲的 code/lecture02/07_training_loop.py（窗口模型、get_batch）。
"""

import importlib
import math
import random
import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

HERE = Path(__file__).parent
sys.path.append(str(HERE.parent / "lecture02"))        # 让 importlib 能找到第 2 讲的文件
_lm = importlib.import_module("04_transformer_lm")
_l2 = importlib.import_module("07_training_loop")
TransformerLM = _lm.TransformerLM
WindowMLPLM = _l2.WindowMLPLM
get_batch = _l2.get_batch
lr_cosine_schedule = _l2.lr_cosine_schedule


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


# =============================================================================
# 第 1 部分：合成故事数据
# =============================================================================
SYLLABLES = ["ka", "lo", "mi", "ra", "to", "nu", "se", "vi", "da", "po",
             "ri", "ze", "ha", "mo", "lu", "ne", "ta", "fi", "go", "be"]
ANIMALS = ["cat", "dog", "bird", "fox", "rabbit", "bear", "fish", "frog", "duck", "mouse"]
ADJS = ["little", "happy", "brave", "small", "kind", "curious", "sleepy", "funny"]
PLACES = ["park", "forest", "river", "garden", "beach", "hill", "village", "market"]
OBJECTS = ["red ball", "shiny stone", "old map", "blue kite", "tiny key", "big box", "green hat", "golden coin"]
ACTIVITIES = ["play", "sing", "run", "read", "draw", "swim", "dance", "jump"]


def make_name(rng):
    """
    随机造一个名字：2~3 个音节拼起来，首字母大写，例如 "Kalomi"。

    【为什么这样造】20 个音节，2 音节有 400 种、3 音节有 8000 种组合，
    模型不可能靠"记住常见名字"来猜，只能从上下文里抄。
    """
    return "".join(rng.choice(SYLLABLES) for _ in range(rng.choice([2, 3]))).capitalize()


def make_story(rng):
    """
    生成一个故事，并返回其中"重复出现的名字"所在的字符区间。

    【返回】
        (text, spans)：text 是故事文本；spans 是 [(起点, 终点), ...]，
        每个区间是某个名字的**第 2 次及以后**的出现位置（第一次出现不算，因为无法预测）。

    【边界】所有内容都是 ASCII 字符，所以字符位置 = 字节位置，方便后面按字节计算损失。
    """
    a = make_name(rng)
    b = make_name(rng)
    while b == a:
        b = make_name(rng)
    # 每一段是 (文字, 类型)：类型 "first" = 名字首次出现，"rep" = 名字再次出现，None = 普通文字
    segments = [
        ("Once upon a time, there was a " + rng.choice(ADJS) + " " + rng.choice(ANIMALS) + " named ", None),
        (a, "first"), (". ", None),
        (a, "rep"), (" loved to " + rng.choice(ACTIVITIES) + " in the " + rng.choice(PLACES) + ". One day, ", None),
        (a, "rep"), (" found a " + rng.choice(OBJECTS) + ". Then ", None),
        (a, "rep"), (" met a " + rng.choice(ANIMALS) + " named ", None),
        (b, "first"), (". ", None),
        (b, "rep"), (' said, "Can I ' + rng.choice(ACTIVITIES) + " with you, ", None),
        (a, "rep"), ('?" ', None),
        (a, "rep"), (" said yes, and ", None), (a, "rep"), (" and ", None), (b, "rep"),
        (" played all day. At night, ", None), (b, "rep"), (" went home, and ", None),
        (a, "rep"), (" went to sleep.\n", None),
    ]
    text, spans = "", []
    for piece, kind in segments:
        if kind == "rep":
            spans.append((len(text), len(text) + len(piece)))
        text += piece
    return text, spans


def build_corpus(num_stories, seed):
    """
    生成 num_stories 个故事并拼接。

    【返回】
        (tokens, name_mask)：
            tokens    : np.int64 数组，字节级 token（0~255）
            name_mask : 与 tokens 等长的布尔数组，True 表示该字节属于"重复出现的名字"
    """
    rng = random.Random(seed)
    text, spans = "", []
    for _ in range(num_stories):
        t, sp = make_story(rng)
        spans += [(s + len(text), e + len(text)) for s, e in sp]
        text += t
    tokens = np.frombuffer(text.encode("ascii"), dtype=np.uint8).astype(np.int64)
    mask = np.zeros(len(tokens), dtype=bool)
    for s, e in spans:
        mask[s:e] = True
    return tokens, mask


# =============================================================================
# 第 2 部分：训练与评估
# =============================================================================
def train(model, data, steps, lr, context_length, batch_size=16, warmup=50, log_every=None, label=""):
    """
    用 AdamW + 预热余弦调度 + 梯度裁剪训练模型（和第 2 讲的训练循环相同）。

    【参数】
        model          : 输入 (B, T) token，输出 (B, T, 256) logits 的任意模型
        data           : 训练 token（np 数组）
        steps, lr      : 训练步数、最大学习率
        context_length : 每个样本的长度 T
        warmup         : 预热步数（0 表示不预热）
        log_every      : 每隔多少步打印一次训练损失（None 表示不打印）
    【返回】
        list[float]：每 50 步记录一次的训练损失
    """
    opt = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1)
    rng = np.random.default_rng(0)
    history = []
    for it in range(steps):
        for g in opt.param_groups:
            g["lr"] = lr_cosine_schedule(it, lr, lr / 10, warmup, steps) if warmup else lr
        x, y = get_batch(data, batch_size, context_length, "cpu", rng)
        loss = F.cross_entropy(model(x).reshape(-1, 256), y.reshape(-1))
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        if (it + 1) % 50 == 0:
            history.append(loss.item())
        if log_every and (it + 1) % log_every == 0:
            print(f"    [{label}] 第 {it + 1:>4} 步，训练损失 {loss.item():.3f}")
    return history


@torch.no_grad()
def evaluate(model, tokens, mask, context_length):
    """
    在验证数据上计算每个字节的损失，分"整体 / 重复名字 / 其他"三类报告（单位：比特/字节）。

    【思路】
        用长度为 T 的窗口滑过验证数据，每次前进 T/2，只统计窗口后一半的损失，
        这样每个字节被预测时都至少有 T/2 个字节的上文（开头第一个窗口除外）。
    【返回】
        dict：{"整体": x, "重复名字": y, "其他": z}
    """
    model.eval()
    losses = np.zeros(len(tokens) - 1)
    half = context_length // 2
    for i in range(0, len(tokens) - 1, half):
        s = max(0, i - half)                                # 窗口起点：比要统计的位置早 half 个字节
        x = torch.from_numpy(tokens[s:s + context_length])[None]
        y = torch.from_numpy(tokens[s + 1:s + context_length + 1])
        n = len(y)
        per_byte = F.cross_entropy(model(x[:, :n])[0], y, reduction="none").numpy()
        off = i - s
        losses[i:i + n - off] = per_byte[off:n]
    model.train()
    m = mask[1:]                                            # 第 j 个损失是在预测 tokens[j+1]
    to_bits = 1 / math.log(2)
    return {"整体": losses.mean() * to_bits, "重复名字": losses[m].mean() * to_bits, "其他": losses[~m].mean() * to_bits}


CTX = 256            # 上下文长度：一个故事约 300 字节，256 足够看到前文中的名字


def demo_compare(train_data, val_data, val_mask):
    """演示 1：窗口模型 vs Transformer。返回训练好的 Transformer。"""
    print_header("演示 1：同样的训练预算，窗口模型 vs Transformer")

    torch.manual_seed(0)
    mlp = WindowMLPLM()
    t0 = time.perf_counter()
    train(mlp, train_data, steps=600, lr=3e-3, context_length=CTX)
    t_mlp = time.perf_counter() - t0
    r_mlp = evaluate(mlp, val_data, val_mask, CTX)

    torch.manual_seed(0)
    tf = TransformerLM(vocab_size=256, context_length=CTX, d_model=128, num_layers=4, num_heads=4, d_ff=384)
    print("  训练 Transformer（4 层，d=128，4 头）……")
    t0 = time.perf_counter()
    train(tf, train_data, steps=600, lr=2e-3, context_length=CTX, log_every=200, label="Transformer")
    t_tf = time.perf_counter() - t0
    r_tf = evaluate(tf, val_data, val_mask, CTX)

    n_mlp = sum(p.numel() for p in mlp.parameters())
    n_tf = sum(p.numel() for p in tf.parameters())
    print(f"\n  {'模型':<24}{'参数量':>10}{'训练用时':>10}{'整体':>8}{'重复名字':>10}{'其他':>8}   （单位：比特/字节）")
    for name, n, t, r in [("窗口模型（看 8 字节）", n_mlp, t_mlp, r_mlp), ("Transformer（看 256 字节）", n_tf, t_tf, r_tf)]:
        print(f"  {name:<24}{n:>10,}{t:>9.0f}s{r['整体']:>8.3f}{r['重复名字']:>10.3f}{r['其他']:>8.3f}")
    print("""
  - 在"重复名字"的字节上，窗口模型几乎只能靠猜（名字在几十个字节之前，它看不到），
    Transformer 则接近 0 比特：它学会了从上文中把名字"抄"过来。
  - 在"其他"字节上两者差距小得多：模板文字大多靠局部规律就能预测。
  - 公平起见要注意：Transformer 参数更多、计算更贵。但即使把窗口模型做得再大，
    它也看不到 8 个字节之外的内容 —— 能力的差别来自**结构**，而不是参数量。""")
    return tf


# =============================================================================
# 第 3 部分：看看注意力在看哪里
# =============================================================================
@torch.no_grad()
def attention_weights(model, ids, layer_idx):
    """
    重新计算第 layer_idx 层的注意力权重（和 MultiHeadSelfAttention.forward 中的计算完全一致）。

    【形状】
        ids : (1, T) → 返回 (num_heads, T, T)，[h, i, j] = 第 h 个头中位置 i 对位置 j 的注意力权重
    """
    positions = torch.arange(ids.shape[1])
    x = model.token_embeddings(ids)
    for layer in model.layers[:layer_idx]:
        x = layer(x, positions)
    blk = model.layers[layer_idx]
    attn = blk.attn
    h_in = blk.ln1(x)
    q = attn.q_proj(h_in).view(1, -1, attn.num_heads, attn.d_head).transpose(1, 2)
    k = attn.k_proj(h_in).view(1, -1, attn.num_heads, attn.d_head).transpose(1, 2)
    q, k = attn.rope(q, positions), attn.rope(k, positions)
    scores = q @ k.transpose(-1, -2) / math.sqrt(attn.d_head)
    T = ids.shape[1]
    scores = scores.masked_fill(~torch.tril(torch.ones(T, T, dtype=torch.bool)), float("-inf"))
    return torch.softmax(scores, dim=-1)[0]


def demo_attention(model):
    """演示 2：预测一个重复出现的名字时，各个注意力头在看哪里。"""
    print_header("演示 2：打开注意力 —— 模型在'抄'哪里？")
    story, spans = make_story(random.Random(2024))
    story = story[:model.context_length]         # 【边界】故事可能比上下文长，只取模型能看到的部分
    ids = torch.tensor([list(story.encode("ascii"))])
    start, end = spans[2]                        # 第 3 次"重复出现"的名字
    name = story[start:end]
    pos = start + 1                              # 站在名字第 2 个字母的位置上，预测第 3 个字母
    print(f"  故事开头：{story[:150]!r}...")
    print(f"  现在模型读到 {story[pos - 12:pos + 1]!r}，要预测下一个字母 {story[pos + 1]!r}（名字是 {name}）")
    target = story[pos + 1]
    last = len(model.layers) - 1
    for layer_idx in [0, last]:
        w = attention_weights(model, ids, layer_idx)
        print(f"  第 {layer_idx} 层：")
        for h in range(w.shape[0]):
            j = int(w[h, pos].argmax())
            # 用方括号标出被关注的那个字符，前后各带几个字符作为上下文
            ctx = (story[max(0, j - 6):j] + "[" + story[j] + "]" + story[j + 1:j + 3]).replace("\n", " ")
            flag = "  ← 正是要预测的字母！" if story[j] == target and j < start else ""
            print(f"    头 {h}：最关注位置 {j:>3}（权重 {w[h, pos, j].item():.2f}），那里是 …{ctx}…{flag}")
    print("""
  第 0 层的注意力比较分散，看的是一些局部的、模板性的位置。
  最后一层则通常能找到头把大部分注意力集中在**前文中同一个名字里、下一个要预测的那个字母**上：
  模型在做的事情是"上次 Nu 后面跟的是 g，所以这次也是 g"，然后把它抄过来。
  这种机制叫"归纳头"（induction head），被认为是 Transformer 做上下文学习（in-context learning）的基础之一。
  （具体是哪个头、权重多大，每次训练都可能不同。）""")


# =============================================================================
# 第 4 部分：生成（温度 + top-p 采样）
# =============================================================================
@torch.no_grad()
def generate(model, prompt, max_new_tokens=300, temperature=0.8, top_p=0.9, seed=0, stop="\n"):
    """
    自回归生成：温度缩放 + top-p（核）采样（CS336 作业 1 的解码要求）。

    【top-p 采样】
        把下一个 token 的概率从大到小排序，只保留"累积概率刚好达到 p"的最小集合，
        在这个集合里重新归一化后采样。
        - 模型很确定时（一个 token 概率 0.95），集合里可能只有 1 个 token → 接近贪心；
        - 模型不确定时，集合会变大 → 保留多样性。
        它比固定的 top-k 更能适应不同位置的"确定程度"，同时砍掉了概率极低的"胡言乱语"长尾。

    【参数】
        temperature : > 0；top_p : (0, 1]；stop : 生成出这个字符串就停止（None 表示不停止）
    【边界】上下文超过 model.context_length 时，只保留最后 context_length 个 token。
    """
    if temperature <= 0 or not (0 < top_p <= 1):
        raise ValueError("temperature 必须 > 0，top_p 必须在 (0, 1] 之间")
    g = torch.Generator().manual_seed(seed)
    ids = list(prompt.encode("utf-8"))
    for _ in range(max_new_tokens):
        x = torch.tensor([ids[-model.context_length:]])
        probs = torch.softmax(model(x)[0, -1] / temperature, dim=-1)
        sorted_p, sorted_idx = probs.sort(descending=True)
        keep = sorted_p.cumsum(0) - sorted_p < top_p       # 保留"加上自己之前累积概率还不到 p"的那些
        sorted_p = sorted_p * keep
        choice = torch.multinomial(sorted_p / sorted_p.sum(), 1, generator=g).item()
        ids.append(sorted_idx[choice].item())
        if stop and bytes(ids).decode("utf-8", errors="replace").endswith(stop):
            break
    return bytes(ids).decode("utf-8", errors="replace")


def demo_generate(model):
    """演示 3：让 Transformer 写新故事。"""
    print_header("演示 3：用 top-p 采样生成新故事")
    for seed, prompt in [(0, "Once upon a time, there was a sleepy duck named Zebeta. "),
                         (1, "Once upon a time, there was a kind bear named Mofi. ")]:
        text = generate(model, prompt, seed=seed)
        print(f"\n  {text.strip()}")
    print("""
  名字"Zebeta""Mofi"是我们在提示里现编的，训练数据里几乎不可能出现过，
  但模型在后文中基本能一致地使用它们 —— 这正是注意力让模型具备的能力。
  （仔细看，模型偶尔也会抄错一个字母，这是小模型、短训练的正常现象。）""")


# =============================================================================
# 第 5 部分：Pre-norm vs Post-norm 的稳定性
# =============================================================================
def demo_norm_stability(train_data):
    """演示 4：较大学习率、不预热时，8 层的 Post-norm 模型训练不起来。"""
    print_header("演示 4：Pre-norm vs Post-norm —— 谁能承受更大的学习率？")
    print("  8 层 Transformer（d=128），不做学习率预热，训练 200 步：")
    for norm, lr in [("pre", 3e-3), ("post", 3e-3), ("post", 1e-3)]:
        torch.manual_seed(0)
        m = TransformerLM(256, 64, 128, 8, 4, 384, norm_position=norm)
        hist = train(m, train_data, steps=200, lr=lr, context_length=64, warmup=0)
        print(f"    {norm:>4}-norm，学习率 {lr:.0e}：每 50 步的训练损失 {[round(v, 2) for v in hist]}")
    print("""
  - 学习率 3e-3 时，Post-norm 卡在约 3.1（大致只学会了字母频率），Pre-norm 正常下降；
  - 把 Post-norm 的学习率降到 1e-3，它也能正常训练。
  这与讲义第 3.1 节的结论一致：Pre-norm 让残差"主干道"保持畅通，训练更稳定、能用更大的学习率；
  Post-norm 需要更小心地调学习率（原始 Transformer 就依赖学习率预热）。
  【边界】这是一个小规模实验，具体在哪个学习率"坏掉"取决于层数、宽度、初始化等；
  我们在 6 层、d=96 的配置下观察到 Post-norm 只是开始慢一些、之后能追上。""")


def main():
    torch.manual_seed(0)
    torch.use_deterministic_algorithms(True)   # 让结果可复现（第 2 讲 10.5 节）
    print_header("准备数据：合成故事")
    train_data, _ = build_corpus(4000, seed=0)
    val_data, val_mask = build_corpus(200, seed=1)
    example, _ = make_story(random.Random(7))
    print(f"  训练集 {len(train_data):,} 字节（4000 个故事），验证集 {len(val_data):,} 字节（200 个故事）")
    print(f"  验证集中属于'重复名字'的字节占 {val_mask.mean():.1%}")
    print(f"  一个例子：\n    {example.strip()}")

    model = demo_compare(train_data, val_data, val_mask)
    demo_attention(model)
    demo_generate(model)
    demo_norm_stability(train_data)
    print("\n全部演示运行完毕。")


if __name__ == "__main__":
    main()