"""
=============================================================================
第 2 讲 · 代码 06：上下文学习从哪里来 —— 在"每段文字一套新规律"的数据上训练出归纳头
=============================================================================

对应讲义：lectures/lecture02_in_context_learning.md §8

【这个脚本要回答的问题】
    上下文学习是一种**在推理时**学习的能力。它是怎样在**训练时**长出来的？
    Olsson et al.（2022）认为，大模型里很大一部分上下文学习由"归纳头"（induction head）完成：
    一个两层注意力电路，看到当前词 A，就去前文里找 A 上一次出现的位置，把 A 后面那个词 B 抄过来（第 1 讲 §6.6）。
    这个脚本在一个干净的玩具任务上，从零训练两个小 Transformer，亲眼看这个过程：

    任务：每条训练序列来自一条**新的、随机的**马尔可夫链（3 个符号，转移概率每条序列重新抽）。
        模型不可能"背下"转移规律——每条序列的规律都不一样，只能**从这条序列的前文里**估计。
        这正是上下文学习：参数固定，规律从上下文里读出来（Edelman et al. 2024 研究过同一个任务）。

    1. 训练曲线：2 层和 1 层模型的损失，和三条"理论线"比较：不看上下文（均匀猜）、
       只数前文里每个符号出现几次（一元统计）、数前文里"当前符号后面跟什么"（二元统计 = 归纳头做的事）。（演示 1）
    2. 打开模型看注意力：哪个头在看"前一个位置"，哪个头在看"和当前符号相同的符号的下一个位置"？（演示 2）
    3. 上下文越长学得越好吗？每个位置的损失（"上下文学习曲线"）。（演示 3）
    4. 回到第 1 讲和本讲的两个迷你 GPT：它们能照着前文抄一段没见过的词序列吗？（演示 4）

【运行】python3 06_induction_heads.py     （约 60 秒：训练两个小模型，每个约 25 秒；不写缓存）
【依赖】PyTorch；import 第 1 讲的 tiny_gpt.py（模型结构和学习率表）、console.py，同目录的 icl_lab.py（演示 4）。
"""

import math
import os
import random
import sys
import time

import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import print_header, print_observations, print_table  # noqa: E402
from tiny_gpt import GPTConfig, TinyGPT, lr_at  # noqa: E402

K = 3            # 符号个数
T = 64           # 每条序列的长度（= 模型的上下文窗口）
STEPS = 1500
LOG_EVERY = 50


# ---------------------------------------------------------------------------
# 数据：每条序列一条新的马尔可夫链
# ---------------------------------------------------------------------------

def random_chain(rng):
    """
    一条随机的马尔可夫链：K × K 的转移矩阵，每一行从 Dirichlet(1, ..., 1) 里抽（也就是"在所有概率分布里均匀地抽一个"）。

    【Python 小知识】Dirichlet(1,…,1) 可以这样抽：K 个独立的 Gamma(1) 随机数，再除以它们的和。
    """
    rows = []
    for _ in range(K):
        g = [rng.gammavariate(1.0, 1.0) for _ in range(K)]
        rows.append([x / sum(g) for x in g])
    return rows


def sample_sequence(rng):
    """按一条新抽的链生成 T + 1 个符号（第一个符号均匀地抽）。"""
    P = random_chain(rng)
    s = rng.randrange(K)
    seq = [s]
    for _ in range(T):
        s = rng.choices(range(K), weights=P[s])[0]
        seq.append(s)
    return seq


def batch(rng, n):
    """n 条序列 → (输入 x, 目标 y)，y 是 x 向右错开一位。"""
    S = torch.tensor([sample_sequence(rng) for _ in range(n)])
    return S[:, :-1].contiguous(), S[:, 1:].contiguous()


def bayes_losses(x, y):
    """
    两种"理论上的上下文学习者"在每个位置上的平均损失（交叉熵，nats）。

    【一元统计（unigram）】只数前文里每个符号出现了几次：P(下一个 = j) = (count_j + 1) / (t + 1 + K)
        ——"最近出现得多的，下一个也更可能是它"。注意这就是第 2 讲 §5 的**多数标签偏差**的统计版本。
    【二元统计（bigram）】数前文里"当前符号 i 后面跟了 j"的次数：P(下一个 = j | 当前 = i) = (count_ij + 1) / (count_i· + K)
        ——这是知道每条序列有自己的转移规律时的贝叶斯最优预测（Dirichlet(1) 先验下的后验均值）。
        要算它，模型必须在前文里找到"和当前符号一样的符号"，再看它们**后面**是什么——这正是归纳头的两步。
    【返回】(一元损失列表, 二元损失列表)，长度都是 T（每个位置一个，对所有序列取平均）
    """
    uni, bi = [0.0] * T, [0.0] * T
    for xs, ys in zip(x.tolist(), y.tolist()):
        cu = [1] * K
        cb = [[1] * K for _ in range(K)]
        for t in range(T):
            cu[xs[t]] += 1
            if t > 0:
                cb[xs[t - 1]][xs[t]] += 1
            uni[t] -= math.log(cu[ys[t]] / sum(cu))
            bi[t] -= math.log(cb[xs[t]][ys[t]] / sum(cb[xs[t]]))
    n = x.shape[0]
    return [u / n for u in uni], [b / n for b in bi]


# ---------------------------------------------------------------------------
# 训练
# ---------------------------------------------------------------------------

def train(n_layer, eval_x, eval_y, seed=0):
    """
    训练一个 n_layer 层的 TinyGPT（4 个头、d_model = 64），每 LOG_EVERY 步记录一次验证损失。

    【返回】(model, 历史 [(步数, 验证损失)], 用时秒数)
    【边界】数据是无穷的（每一步都抽新序列），所以不存在"过拟合训练集"；验证集是固定的 512 条序列。
    """
    torch.manual_seed(seed)
    rng = random.Random(seed)
    model = TinyGPT(GPTConfig(vocab_size=K, block_size=T, n_layer=n_layer, n_head=4, d_model=64))
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0)
    hist, t0 = [], time.time()
    for step in range(STEPS):
        for g in opt.param_groups:
            g["lr"] = lr_at(step, STEPS, 1e-3, 50)
        x, y = batch(rng, 32)
        _, loss, _ = model(x, y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        if (step + 1) % LOG_EVERY == 0 or step == 0:
            with torch.no_grad():
                _, vl, _ = model(eval_x, eval_y)
            hist.append((step + 1, vl.item()))
    model.eval()
    return model, hist, time.time() - t0


def demo1(eval_x, eval_y, uni, bi):
    """演示 1：两个模型的训练曲线 vs 三条理论线。"""
    print_header(f"演示 1：每条序列一条新的马尔可夫链（{K} 个符号、长 {T}）——训练曲线")
    u_avg, b_avg = sum(uni) / T, sum(bi) / T
    print(f"  理论线：均匀猜 {math.log(K):.4f}；前文一元统计 {u_avg:.4f}；前文二元统计（贝叶斯最优）{b_avg:.4f}")
    models, hists = {}, {}
    for L in (2, 1):
        m, h, secs = train(L, eval_x, eval_y)
        models[L], hists[L] = m, h
        print(f"  训练 {L} 层模型：{m.n_params():,} 个参数，{STEPS} 步")
    rows = []
    marks = [1, 50, 100, 150, 200, 250, 300, 400, 500, 750, 1000, 1500]
    for step in marks:
        cells = [step]
        for L in (2, 1):
            v = dict(hists[L])[step]
            where = "≈ 均匀" if v > (math.log(K) + u_avg) / 2 else ("≈ 一元" if v > (u_avg + b_avg) / 2 else "≈ 二元")
            cells.append(f"{v:.4f}（{where}）")
        rows.append(cells)
    print_table(["训练步数", "2 层模型的验证损失", "1 层模型的验证损失"], rows)
    h2 = hists[2]
    first_uni = next(s for s, v in h2 if v < (math.log(K) + u_avg) / 2)
    first_bi = next(s for s, v in h2 if v < (u_avg + b_avg) / 2)
    gap1 = hists[1][-1][1] - b_avg
    print_observations([
        f"2 层模型先走到'一元统计'那条线附近（第 {first_uni} 步左右），停一会儿，再在第 {first_bi} 步左右跌到二元统计的水平；"
        f"训练结束时 {h2[-1][1]:.4f}，离贝叶斯最优 {b_avg:.4f} 只差 {h2[-1][1] - b_avg:.4f}。",
        f"1 层模型一直停在一元统计附近（最后 {hists[1][-1][1]:.4f}，比二元统计高 {gap1:.4f}）：只有一层注意力，"
        "它能数'前文里哪个符号多'，却做不到'找到和当前符号相同的位置，再看它的下一个'——后者需要两层注意力串起来（演示 2）。",
        "'先学会数频率、再学会看前后关系'的阶段性，Edelman et al.（2024）在同一个任务上报告过；Olsson et al.（2022）"
        "在真实语言模型里观察到的'相变'——归纳头形成的那一小段训练里，上下文学习能力突然变强——是同一类现象。",
        "注意一元统计那个阶段：模型的策略是'前文里哪个多，就猜哪个'。放到少样本提示里，这就是**多数标签偏差**；"
        "二元阶段的'照着前文抄'，放到提示里就是从示例里复制格式和标签。上下文学习的能力和它的偏差，来自同一套机制。",
    ])
    return models


def attention_scores(model, x):
    """
    每层每个头的两种"注意力得分"（在 x 的所有序列、所有位置上平均）：

        前一位得分：位置 t 放在 t − 1 上的注意力权重
        归纳得分  ：位置 t 放在 {s < t : x[s − 1] == x[t]} 上的注意力权重之和
                    （"前一个符号和我一样"的那些位置——它们存着"我这个符号后面跟了什么"）
    【返回】({(层, 头): (前一位得分, 归纳得分)}, (均匀注意力时的前一位得分, 均匀注意力时的归纳得分))
    【对照】如果注意力在 0…t 上均匀分布，前一位得分 = 1/(t+1) 的平均，归纳得分 = "符合条件的位置所占比例"的平均——
        两个对照值也在同一批序列上算出来，一起返回。
    """
    with torch.no_grad():
        _, _, atts = model(x, return_attn=True)
    xs = x.tolist()
    out = {}
    up = ui = 0.0
    nb = 0
    for b in range(len(xs)):
        for t in range(2, len(xs[b])):
            up += 1 / (t + 1)
            ui += sum(1 for s in range(1, t + 1) if xs[b][s - 1] == xs[b][t]) / (t + 1)
            nb += 1
    for li, att in enumerate(atts):
        B, H, Tn, _ = att.shape
        for h in range(H):
            prev = ind = 0.0
            cnt = 0
            for b in range(B):
                a = att[b, h]
                for t in range(2, Tn):
                    prev += a[t, t - 1].item()
                    ind += sum(a[t, s].item() for s in range(1, t + 1) if xs[b][s - 1] == xs[b][t])
                    cnt += 1
            out[(li + 1, h + 1)] = (prev / cnt, ind / cnt)
    return out, (up / nb, ui / nb)


def demo2(models, eval_x):
    """演示 2：打开模型看注意力。"""
    print_header("演示 2：哪些头在做什么？（64 条验证序列上的平均注意力得分）")
    x = eval_x[:64]
    rows = []
    best = {}
    for L, m in [(2, models[2]), (1, models[1])]:
        sc, (up, ui) = attention_scores(m, x)
        for (layer, head), (p, ind) in sc.items():
            rows.append([f"{L} 层模型", f"第 {layer} 层 · 头 {head}", f"{p:.3f}", f"{ind:.3f}"])
        best[L] = sc
    rows.append(["（对照）", "注意力在前文上均匀分布", f"{up:.3f}", f"{ui:.3f}"])
    print_table(["模型", "头", "前一位得分", "归纳得分"], rows)
    l1_prev = max(v[0] for (layer, _), v in best[2].items() if layer == 1)
    l2_ind = max(v[1] for (layer, _), v in best[2].items() if layer == 2)
    one_ind = max(v[1] for v in best[1].values())
    print_observations([
        f"2 层模型的第 1 层里有'前一位头'（前一位得分最高 {l1_prev:.3f}，均匀注意力时 {up:.3f}）：每个位置把'前一个符号是什么'"
        "抄到自己身上。",
        f"第 2 层里有'归纳头'（四个头的归纳得分 {min(v[1] for (layer, _), v in best[2].items() if layer == 2):.3f}–{l2_ind:.3f}，"
        f"均匀注意力时 {ui:.3f}）：它们找的是'前一个符号和当前符号相同'的位置——"
        "那些位置在第 1 层已经记下了'我前面是谁'，于是第 2 层能读出'当前符号上次后面跟了什么'。",
        f"1 层模型的头归纳得分最高只有 {one_ind:.3f}，和均匀注意力差不多：没有第 1 层先把'前一个符号'抄过去，"
        "单独一层找不到'和我相同的符号的下一个'。"
        "这就是第 1 讲 §6.6 手工搭的那个两层电路——这一次它是梯度下降自己找到的。",
    ])


def demo3(models, eval_x, eval_y, uni, bi):
    """演示 3：上下文学习曲线——每个位置的损失。"""
    print_header("演示 3：上下文越长，预测越准吗？——每个位置的平均损失")
    with torch.no_grad():
        per_pos = {}
        for L, m in models.items():
            logits = m(eval_x)[0]
            loss = F.cross_entropy(logits.reshape(-1, K), eval_y.reshape(-1), reduction="none").view(eval_x.shape)
            per_pos[L] = loss.mean(0).tolist()
    rows = []
    for t in [0, 1, 2, 4, 8, 16, 32, 48, 63]:
        rows.append([t + 1, f"{per_pos[2][t]:.3f}", f"{per_pos[1][t]:.3f}", f"{uni[t]:.3f}", f"{bi[t]:.3f}"])
    print_table(["已经看到的符号数", "2 层模型", "1 层模型", "一元统计", "二元统计"], rows)
    print_observations([
        f"位置 1（只看到一个符号）两个模型都在均匀猜附近（{per_pos[2][0]:.3f}、{per_pos[1][0]:.3f}；ln 3 = {math.log(3):.3f}）："
        "还没有上下文可以学。",
        f"2 层模型的损失随位置一路下降，到第 64 个位置是 {per_pos[2][63]:.3f}，紧贴二元统计那一列（{bi[63]:.3f}）：上下文越长，"
        "它对这条链的转移规律估计得越准。这就是'上下文学习曲线'——示例越多越好的最简单的版本。",
        f"1 层模型下降得少得多（第 64 个位置 {per_pos[1][63]:.3f}），贴着一元统计那一列：更多的上下文只帮它把频率数得更准。",
    ])


def demo4():
    """演示 4：第 1 讲 / 本讲的两个迷你 GPT 能照着前文抄吗？"""
    print_header("演示 4：两个迷你 GPT 能照着前文抄一段没见过的词序列吗？")
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from icl_lab import get_base_model, get_bursty_model
    base, tok = get_base_model()
    bursty, _ = get_bursty_model()
    rng = random.Random(0)
    words = [w for w in tok.itos if w not in ("<unk>", "<doc>", ".", ":", "?", "q", "a")]
    seqs = [[rng.choice(words) for _ in range(12)] for _ in range(200)]
    rows = []
    res = {}
    for name, m in [("base", base), ("bursty", bursty)]:
        first = second = 0
        n = 0
        with torch.no_grad():
            for s in seqs:
                ids = torch.tensor([[tok.stoi["<doc>"]] + [tok.stoi[w] for w in s + s]])
                pred = m(ids)[0][0].argmax(-1).tolist()
                target = ids[0, 1:].tolist()
                for t in range(1, 12):           # 第一遍：预测第 t+1 个词（随机，猜不到）
                    first += pred[t] == target[t]
                for t in range(13, 24):          # 第二遍：前面已经出现过一模一样的序列
                    second += pred[t] == target[t]
                n += 11
        res[name] = (first / n, second / n)
        rows.append([name, f"{100 * first / n:.1f}%", f"{100 * second / n:.1f}%"])
    print_table(["模型", "第一遍（没见过）", "第二遍（前文刚出现过）"], rows)
    print("  做法：从词表里随机抽 12 个词，连写两遍，看模型在第二遍里能不能预测出'下一个词'。")
    print_observations([
        f"第二遍的准确率：base {100 * res['base'][1]:.1f}%、bursty {100 * res['bursty'][1]:.1f}%（第一遍分别是 "
        f"{100 * res['base'][0]:.1f}%、{100 * res['bursty'][0]:.1f}%）。",
        "两个模型都几乎不会'照着前文抄'一段随机词：它们的训练语料是模板化的句子，没有'同一段随机内容重复出现'的结构，"
        "归纳头没有理由长出来。这和 04 号脚本的发现一致：它们的少样本能力靠的是认格式，不是从示例里复制。",
        "真实的网页文本里重复无处不在（名字、代码里的变量、表格的列）。Olsson et al. 认为这正是大模型普遍长出归纳头的原因——"
        "**上下文学习能力由预训练数据的结构决定**，这是本讲两个实验（bursty 语料、马尔可夫链）共同的结论。",
    ])


def main():
    torch.set_num_threads(1)
    eval_x, eval_y = batch(random.Random(123), 512)
    uni, bi = bayes_losses(eval_x, eval_y)
    models = demo1(eval_x, eval_y, uni, bi)
    demo2(models, eval_x)
    demo3(models, eval_x, eval_y, uni, bi)
    demo4()


if __name__ == "__main__":
    main()