"""
=============================================================================
第 12 讲 · 代码 3：证明它就是那个电路——消融、激活修补，以及直接读权重（OV 复制分数、K-composition）
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
代码 02 看到 induction head 和"会抄"在同一时刻出现。这是相关性。要说"这个电路**就是**复制能力的来源"，需要因果证据；
要说"它的算法是 [A][B] … [A] → [B]"，最好能直接从权重里读出来。

    演示 1：消融：把某个头（或一组头）的输出置零，重复部分的损失变成多少？
    演示 2：激活修补：把序列里第一遍的片段换成别的随机 token（于是第二遍"无处可抄"），再把某个头在原序列上的输出
            "修补"回来——哪些头能把能力找回来？
    演示 3：读 OV 电路：第 1 层的头是不是在"复制"（注意到 token t，就提高 t 的 logit）？
    演示 4：读 QK 电路的组合：第 1 层的头用第 0 层哪个头的输出来算"键"、"查询"、"值"？和随机初始化的基线比较。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 03_ablation_and_circuits.py      # 约 5 秒（读代码 02 缓存的模型；没有缓存时先训练，约 30 秒）

依赖本目录的 circuit_lab.py。

【PyTorch 小知识】
    - model.ablate = {(层, 头)}、model.patch = {(层, 头): 张量}：circuit_lab 的前向会在这一层把对应头的输出置零或替换。
      用完一定要清空（model.ablate = set()、model.patch = {}），否则会影响后面的计算。
"""

import os
import statistics
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, AttnOnlyTransformer, composition, eval_batch, load_or_train_induction,  # noqa: E402
                         ov_copying_score, per_token_loss, repeat_loss)
from console import pad, print_header  # noqa: E402


def demo1_ablation(model, x, mask):
    print_header("演示 1：消融——把头的输出置零，重复部分的损失（不消融时见第一行；不会抄时约 4.2）")
    with torch.no_grad():
        base = repeat_loss(model, x, mask)[0]
        print(f"{pad('消融', 34)}{pad('重复部分损失', 14, '>')}")
        print(f"{pad('（不消融）', 34)}{pad(f'{base:.3f}', 14, '>')}")
        for l in range(2):
            for h in range(4):
                model.ablate = {(l, h)}
                print(f"{pad(f'第 {l} 层头 {h}', 34)}{pad(f'{repeat_loss(model, x, mask)[0]:.3f}', 14, '>')}")
        groups = [("第 1 层的 4 个头（全部 induction 头）", {(1, h) for h in range(4)}),
                  ("第 0 层头 1 和头 3（前一个 token 头）", {(0, 1), (0, 3)}),
                  ("第 0 层头 0 和头 2", {(0, 0), (0, 2)}),
                  ("第 0 层的 4 个头", {(0, h) for h in range(4)})]
        for name, grp in groups:
            model.ablate = grp
            print(f"{pad(name, 34)}{pad(f'{repeat_loss(model, x, mask)[0]:.3f}', 14, '>')}")
        model.ablate = set()
    print("\n观察 Observations：")
    print("  1. 消融任何**一个** induction 头，损失只从 0.46 升到 0.8–1.1：第 1 层的 4 个头是冗余的，互相补位。")
    print("     把 4 个一起消融，损失回到 4.24——完全不会抄了，和第一遍（随机）一样。")
    print("  2. 消融第 0 层的两个前一个 token 头（头 1、3），损失升到 3.06；消融另外两个，只到 1.17。")
    print("     induction 头本身没有被碰，但它们读不到'前一个 token 是谁'，就找不到该抄的位置——")
    print("     这是'两层组合'的因果证据：第 1 层的能力依赖第 0 层特定的头。")


def demo2_patching(model, x, mask, lens):
    print_header("演示 2：激活修补——第一遍换成别的随机 token，第二遍就无处可抄；把某个头的'干净'输出修补回来")
    g = torch.Generator().manual_seed(7)
    xc = x.clone()
    for b, L in enumerate(lens):
        xc[b, :L] = torch.randint(0, VOCAB, (L,), generator=g)          # 破坏第一遍
    second = torch.zeros_like(mask)
    for b, L in enumerate(lens):
        second[b, L:min(2 * L, x.size(1))] = True                        # 只看第二遍（第三遍还能从第二遍抄）

    def loss2(inp):
        return per_token_loss(model, inp)[second[:, 1:]].mean().item()

    with torch.no_grad():
        clean = loss2(x)
        clean_out = [o.clone() for o in model.last_head_out]             # 干净序列上每个头的输出
        corrupt = loss2(xc)
        print(f"第二遍的损失：原序列 {clean:.3f}，破坏第一遍之后 {corrupt:.3f}")
        print(f"\n{pad('修补哪个头（用原序列上的输出替换）', 36)}{pad('第二遍损失', 12, '>')}{pad('找回的比例', 12, '>')}")
        rows = [(f"第 {l} 层头 {h}", {(l, h): clean_out[l][:, h]}) for l in range(2) for h in range(4)]
        rows += [("第 1 层的 4 个头", {(1, h): clean_out[1][:, h] for h in range(4)}),
                 ("第 0 层的 4 个头", {(0, h): clean_out[0][:, h] for h in range(4)})]
        for name, patch in rows:
            model.patch = patch
            p = loss2(xc)
            print(f"{pad(name, 36)}{pad(f'{p:.3f}', 12, '>')}{pad(f'{(corrupt - p) / (corrupt - clean):+.3f}', 12, '>')}")
        model.patch = {}
    print("\n观察 Observations：")
    print("  1. '找回的比例' = (破坏后的损失 − 修补后的损失) / (破坏后的损失 − 原来的损失)：1 表示完全找回，0 表示没用。")
    print("  2. 修补任何一个 induction 头都能找回 34%–43%，4 个一起修补找回 99.9%：复制的信息完全由这 4 个头的输出携带。")
    print("  3. 修补第 0 层的头几乎没用（甚至略有反效果）：它们在原序列上写下的'前一个 token 是谁'说的是被破坏前的 token，")
    print("     而第 1 层的头从被破坏的位置读到的'值'已经是别的 token 了——光有'找哪里'的信息，没有'抄什么'的内容。")
    print("     修补告诉我们信息**在哪里**、**以什么形式**流动，这比'消融后变差'更精确。")


def demo3_ov(model):
    print_header("演示 3：读 OV 电路——'注意到 token t，往输出里写什么'（V × V 矩阵 W_E W_V W_O W_U 的特征值）")
    print(f"{pad('头', 14)}{pad('复制分数 Σλ/Σ|λ|', 20, '>')}{pad('M[t] 最大的正好是 t 的比例', 30, '>')}")
    for l in range(2):
        for h in range(4):
            sc, diag = ov_copying_score(model, l, h)
            print(f"{pad(f'第 {l} 层头 {h}', 14)}{pad(f'{sc:+.3f}', 20, '>')}{pad(f'{diag:.3f}', 30, '>')}")
    torch.manual_seed(5)
    r = AttnOnlyTransformer(2)
    rand = [ov_copying_score(r, l, h)[0] for l in range(2) for h in range(4)]
    print(f"（对照：随机初始化的模型，复制分数在 {min(rand):+.3f} 到 {max(rand):+.3f} 之间；'正好是 t'的随机水平约 1/64 = 0.016）")
    print("\n观察 Observations：")
    print("  1. 第 1 层 4 个头的复制分数都接近 +1（0.997–0.998）：OV 电路几乎是一个'恒等映射'——")
    print("     注意到哪个 token，就提高哪个 token 的 logit。这正是 induction 算法的'抄'那一半，直接写在权重里。")
    print("     （对超过一半的 token，它们自己的 logit 就是这一行里最大的。）")
    print("  2. 第 0 层的头复制分数是负的：它们不是用来直接影响输出的，而是往残差流里写'前一个 token 是谁'，给第 1 层读。")


def demo4_composition(model):
    print_header("演示 4：读 QK 电路的组合——第 1 层的头在多大程度上'读'第 0 层各头写下的东西")
    vals = {k: [] for k in "KQV"}
    for seed in range(5):
        torch.manual_seed(100 + seed)
        r = AttnOnlyTransformer(2)
        for k in "KQV":
            vals[k] += [composition(r, a, b, k) for a in range(4) for b in range(4)]
    for kind, label in [("K", "K-composition（用来算键）"), ("Q", "Q-composition（用来算查询）"), ("V", "V-composition（用来算值）")]:
        print(f"\n{label}，行 = 第 0 层的头，列 = 第 1 层的头：")
        for h0 in range(4):
            print(f"  第 0 层头 {h0}：" + "  ".join(f"{composition(model, h0, h1, kind):.3f}" for h1 in range(4)))
        print(f"  （5 个随机初始化的模型：平均 {statistics.mean(vals[kind]):.3f}，最大 {max(vals[kind]):.3f}）")
    print("\n观察 Observations：")
    print("  1. 第 0 层的两个前一个 token 头（头 1、3）和第 1 层所有 induction 头之间的 K-composition 是 0.16–0.17，")
    print("     高于随机模型的最大值（0.144）；它们的 Q-composition 和 V-composition 反而**低于**随机水平（0.06–0.09）。")
    print("  2. 这正是 induction 电路的签名：第 1 层的头在算**键**的时候读'前一个 token 是谁'（K-composition），")
    print("     查询这一侧只用当前 token 本身，值这一侧只搬运 token 本身——前一个 token 的信息在这两处都被'绕开'了。")
    print("  3. 差距不算大（0.17 vs 随机 0.12）：分数是对整个 64 维空间的平均，而电路只用到其中几个方向。")
    print("     这也说明了直接读权重的局限：在更大的模型里，信号会更淡，需要配合消融、修补这些因果方法一起用。")


if __name__ == "__main__":
    model, _ = load_or_train_induction()
    x, mask, lens = eval_batch()
    demo1_ablation(model, x, mask)
    demo2_patching(model, x, mask, lens)
    demo3_ov(model)
    demo4_composition(model)