"""
=============================================================================
第 12 讲 · 代码 2：看着 induction head 形成——一个只有注意力的 2 层模型，在训练中途突然学会了"回去抄"
=============================================================================

对应讲义：lectures/lecture12_mechanistic_interpretability.md 第 3、4 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
第 5 讲和第 11 讲都发现，第 1 讲的 base model 不会"回去抄"：同一个随机串出现第二次，它预测得和第一次一样差。
Olsson et al. (2022) 说，"回去抄"由一种叫 induction head 的电路完成，它在训练中某个时刻**突然**形成，
而且至少需要两层注意力。这里在最干净的设定下复现这件事。

    演示 1：训练数据——随机片段重复若干遍（片段长度随机）。第一遍谁也预测不了，第二遍以后只要会"抄"就能预测。
    演示 2：训练 2 层只有注意力的模型 2000 步，每 50 步记录：重复部分的损失、每个头的"前一个 token 分数"和
            "induction 分数"——损失的骤降和两种头的出现发生在同一时刻吗？
    演示 3：对照——只有 1 层的模型：能把损失降多少？为什么降不到底？
    演示 4：训练好的 2 层模型，在一个具体的序列上，前一个 token 头和 induction 头各自看哪里。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_induction_heads.py      # 第一次约 45 秒（训练 2 层和 1 层两个模型）；之后读缓存，约 20 秒

依赖本目录的 circuit_lab.py（2 层模型缓存在本目录的 checkpoints/；代码 03 会读同一个缓存）。不需要第 1 讲的数据或模型。
"""

import math
import os
import sys

import torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))
from circuit_lab import VOCAB, eval_batch, load_or_train_induction, repeated_batch, train_attn_only  # noqa: E402
from console import pad, print_header  # noqa: E402


def demo1_data():
    print_header("演示 1：训练数据——随机片段重复若干遍")
    x, mask, lens = repeated_batch(torch.Generator().manual_seed(0), batch=3)
    for b in range(3):
        toks = x[b].tolist()
        L = lens[b]
        print(f"片段长度 L = {L}：" + " ".join(f"{t:2d}" for t in toks[:2 * L + 3]) + " …")
        print(f"{'':>15}" + " ".join("··" if i < L else "↑↑" for i in range(2 * L + 3)))
    print(f"\n词表 {VOCAB} 个 token，每个序列 64 个；'··' 是第一遍（随机，谁也预测不了，最好也要 ln 64 = {math.log(VOCAB):.2f} nats），")
    print("'↑↑' 是第二遍及以后（只要找到前面相同的 token，抄它后面的那个，就能预测）。片段长度在 6–31 之间随机，")
    print("所以'往回数固定步数'这种只看位置的捷径行不通，必须靠内容匹配。")


def best_heads(scores):
    """从 head_scores 的结果里挑出前一个 token 分数最高的头和 induction 分数最高的头。"""
    prev = max(scores.items(), key=lambda kv: kv[1][0])
    ind = max(scores.items(), key=lambda kv: kv[1][1])
    return prev, ind


def demo2_training(history):
    print_header("演示 2：训练 2 层只有注意力的模型（4 个头 × 2 层，宽度 64，2000 步）——每 50 步的体征")
    print(f"{pad('步', 6, '>')}{pad('重复部分损失', 14, '>')}{pad('第一遍损失', 12, '>')}"
          f"{pad('最强的前一个 token 头', 24, '>')}{pad('最强的 induction 头', 22, '>')}")
    for r in history:
        if r["step"] % 100 == 0 or 600 <= r["step"] <= 1100:
            (pk, pv), (ik, iv) = best_heads(r["scores"])
            print(f"{pad(r['step'], 6, '>')}{pad(format(r['rep'], '.3f'), 14, '>')}"
                  f"{pad(format(r['first'], '.3f'), 12, '>')}"
                  f"{pad(f'第 {pk[0]} 层头 {pk[1]}：{pv[0]:.2f}', 24, '>')}{pad(f'第 {ik[0]} 层头 {ik[1]}：{iv[1]:.2f}', 22, '>')}")
    final = history[-1]["scores"]
    print("\n训练结束时每个头的两个分数（前一个 token 分数 / induction 分数）：")
    for l in range(2):
        print(f"  第 {l} 层：" + "，".join(f"头 {h} {final[(l, h)][0]:.2f} / {final[(l, h)][1]:.2f}" for h in range(4)))
    print("\n观察 Observations：")
    print("  1. 前 600 步，重复部分的损失和第一遍差不多（约 3.6–4.2 nats）：模型还不会抄，只学到了一点'前面出现过的 token 更可能再出现'。")
    print("  2. 第 700–1000 步之间，重复部分的损失**骤降**（3.58 → 0.85），同一时刻：")
    print("     第 1 层出现了 induction 分数约 0.85 的头（在第二遍里，几乎把全部注意力放在'前一个 token 和我相同'的位置上），")
    print("     第 0 层出现了前一个 token 分数较高的头。这就是 Olsson et al. 说的相变：电路的形成和能力的出现发生在同一时刻。")
    print("  3. 第一遍的损失反而**上升**（第 900 步到 5.2，之后回落到 4.6）：模型开始把概率押在'前面出现过的 token'上，")
    print("     在第一遍（全是新 token）里这个赌注是错的。学会一种新能力，常常伴随着在别处付出一点代价。")
    print("  4. 训练结束时第 1 层的 4 个头**都**成了 induction 头（分数 0.84–0.85）——它们是冗余的（代码 03 的消融会看到）。")


def demo3_one_layer():
    print_header("演示 3：对照——只有 1 层注意力的模型（同样的数据、同样的 2000 步）")
    model, history = train_attn_only(n_layers=1, steps=2000, lr=1e-3, log_every=500)
    for r in history:
        (pk, pv), (ik, iv) = best_heads(r["scores"])
        print(f"第 {r['step']:4d} 步：重复部分损失 {r['rep']:.3f}，第一遍 {r['first']:.3f}，最高 induction 分数 {iv[1]:.2f}")
    x, mask, lens = eval_batch()
    avg_distinct = sum(len(set(x[b, :L].tolist())) for b, L in enumerate(lens)) / len(lens)
    print(f"\n对照：如果只知道'答案是前面出现过的某个 token'，均匀地猜片段里的 token，损失约 ln(片段里不同 token 的个数) ≈ "
          f"ln {avg_distinct:.1f} = {math.log(avg_distinct):.2f} nats")
    print("\n观察 Observations：")
    print("  1. 1 层模型也把重复部分的损失降到了 2.9 左右，但降不下去，induction 分数一直接近 0。")
    print("  2. 它学到的是一个更弱的策略：把注意力分散在前面出现过的 token 上，提高它们的概率——")
    print("     相当于'答案是前面某个 token，但不知道是哪个'。损失（2.92）接近 ln(片段里不同 token 的个数) = 2.75。")
    print("  3. 要知道'是哪个'，需要两步：先知道每个位置的**前一个** token 是什么（第 0 层的前一个 token 头），")
    print("     再用当前 token 去匹配这个信息（第 1 层的 induction 头）。一层注意力做不了这种组合——")
    print("     Olsson et al. 的说法是：induction head 需要不同层的注意力头'组合'（composition），所以 1 层模型里不会出现。")


def demo4_patterns(model):
    print_header("演示 4：一个具体的序列（片段长度 7），两个头各自看哪里（每个位置注意力最大的那个位置）")
    gen = torch.Generator().manual_seed(5)
    A = torch.randint(0, VOCAB, (7,), generator=gen)
    x = A.repeat(10)[:24][None]
    with torch.no_grad():
        model(x)
    prev_head = max(range(4), key=lambda h: torch.diagonal(model.last_attn[0][0, h], offset=-1).mean().item())
    toks = x[0].tolist()
    print(f"序列：{' '.join(f'{t:2d}' for t in toks)}")
    print(f"\n第 0 层头 {prev_head}（前一个 token 头）与第 1 层头 0（induction 头），位置 q → 注意力最大的位置 k（那里的 token）：")
    print(f"{pad('q', 4, '>')}{pad('x[q]', 6, '>')}{pad('第 0 层：k', 12, '>')}{pad('第 1 层：k', 12, '>')}{pad('x[k]（第 1 层抄到的）', 22, '>')}{pad('真正的下一个', 14, '>')}")
    for q in range(5, 20):
        k0 = int(model.last_attn[0][0, prev_head, q].argmax())
        k1 = int(model.last_attn[1][0, 0, q].argmax())
        nxt = toks[q + 1] if q + 1 < len(toks) else "—"
        print(f"{pad(q, 4, '>')}{pad(toks[q], 6, '>')}{pad(f'{k0}（{toks[k0]}）', 12, '>')}{pad(k1, 12, '>')}{pad(toks[k1], 22, '>')}{pad(nxt, 14, '>')}")
    print("\n观察 Observations：")
    print("  1. 第 0 层的前一个 token 头：位置 q 主要看 q − 1。它把'我前面是谁'写进位置 q 的残差流。")
    print("  2. 第 1 层的 induction 头：从第二遍开始（q ≥ 7），位置 q 去看前面那个'前一个 token 等于 x[q]'的位置 k，")
    print("     也就是上一次出现 x[q] 之后的那个位置；它把 x[k] 抄到输出——这里 13 个位置里有 11 个正好是真正的下一个 token。")
    print("     q = 8、10 这两个位置头 0 看错了：第 1 层的另外 3 个 induction 头会补上（代码 03 的消融）。")
    print("     算法：[A][B] … [A] → [B]。第一遍里（q < 7）没有可抄的地方，它的注意力没有意义。")


if __name__ == "__main__":
    demo1_data()
    model, history = load_or_train_induction()
    demo2_training(history)
    demo3_one_layer()
    demo4_patterns(model)