"""
interp.py —— 第 14 讲：一个"花园 BERT"，和打开它的几种工具（注意力分析、逐层探针、损失地形）
=====================================================================

【这个文件做什么】
    slides 的可解释性分三部分：BERT 里捕捉到了什么（注意力模式、探针）；预训练模型为什么泛化得好（损失地形）；
    上下文学习在做什么（04 号脚本）。前两部分需要一个"BERT"：
        garden_bert(...)    在第 1 讲的花园世界上用掩码语言建模（MLM，第 2 讲）训练的 4 层双向 Transformer，
                            训练途中存下几个检查点（"探针随训练步数的变化"要用）；
        annotate(doc)       用花园世界的状态机给每个 token 标出精确的语言学信息：词性（主语 / 动词 / 限定词 /
                            形容词 / 名词 / 尾巴 / 句号）、这句话的主语是谁、主语的单复数、宾语名词的类别，
                            以及几种"依存关系"（名词 → 它的动词、形容词 → 它修饰的名词、物主代词 → 主语）。
                            真实研究要靠人工标注的树库（Penn Treebank、OntoNotes），这里是自动、零错误的；
        layer_states(...)   每一层（嵌入层 = 第 0 层）的隐藏状态，和每一层每个头的注意力权重；
        train_probe(...)    线性探针（逻辑回归），Tenney et al. (2019) 的"边探针"思路的最简版本；
        hessian_trace(...)  Hutchinson 估计 Hessian 的迹（损失地形的"平坦度"，03 号脚本）。

【LLM 联系】
    Clark et al. (2019) 分析 BERT 的 144 个注意力头；Tenney et al. (2019) 逐层探针发现"BERT 重新发现了经典的
    NLP 流水线"；Liu et al. (2021) 在 RoBERTa 的训练过程中做探针。这里在一个知道全部真相的小世界里重做这些实验。
"""

import os
import sys

import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture01", "lecture02"):
    sys.path.insert(0, os.path.join(HERE, "..", _d))

from garden_lm import ext_vocab, lr_lambda  # noqa: E402
from objectives import mlm_batches, mlm_loss  # noqa: E402
from toy_world import (DOC_START, NOUN_CAT, SUBJECTS, GardenWorld, State, make_docs,  # noqa: E402
                       standard_splits)

VOCAB = ext_vocab()
WORLD = GardenWorld()
CKPT_DIR = os.path.join(HERE, "checkpoints")
ARCH = dict(d_model=64, n_heads=4, n_layers=4, d_ff=256)
SAVE_STEPS = (0, 100, 300, 1000, 3000, 6000)

# =====================================================================
# 1. 标注
# =====================================================================

ROLES = ["主语", "动词", "限定词", "形容词", "名词", "尾巴", "句号"]
SUBJ_NAMES = [" ".join(t) for t, _, _ in SUBJECTS]
STAGE_ROLE = {"S0": "主语", "S1": "主语", "V": "动词", "D": "限定词", "N": "名词", "R": "尾巴"}


def annotate(doc):
    """
    给一篇文档（词列表，以 <eos> 结尾）的每个 token 标注（句号单独算一类，尾巴短语末尾的句号也是"句号"）。

    【返回】list[dict]，和 doc 等长；<eos> 的那一项是 None。每项包含：
        role：ROLES 之一；subj：这句话主语的编号（SUBJECTS 的下标）；
        third：主语是不是第三人称单数（she / he / my friend，动词要加 s——主谓一致）；
        cat：这句话宾语名词的类别（fruit / veg / flower / animal）；
        head：依存关系 {关系名: 目标 token 的下标}，例如名词的 {"宾语→动词": 动词的位置}。
    【思路】沿第 1 讲的状态机走一遍：处在哪个 stage 就知道下一个词是什么角色。
            "A" 阶段的词可能是形容词也可能是名词（看词本身），"T" 阶段可能是句号也可能是尾巴的第一个词。
    """
    out = [None] * len(doc)
    state = DOC_START
    sent = []
    for i, w in enumerate(doc):
        if w == "<eos>":
            break
        st = state.stage
        if st == "E":
            st = "S0"
            state = State("S0", -1, "", "", ())
        if w == ".":
            role = "句号"
        elif st == "A":
            role = "名词" if w in NOUN_CAT else "形容词"
        elif st == "T":
            role = "尾巴"
        else:
            role = STAGE_ROLE[st]
        sent.append((i, w, role))
        prev, state = state, WORLD.step(state, w)
        if w == ".":
            _fill_sentence(sent, out, prev)            # 说出句号之前的状态里还记着主语和类别
            sent = []
    return out


def _fill_sentence(sent, out, last_state):
    subj = last_state.subj
    third = SUBJECTS[subj][1]
    cat = last_state.cat
    pos = {r: [i for i, _, rr in sent if rr == r] for r in ROLES}
    verb = pos["动词"][0]
    noun = pos["名词"][0]
    subj_head = pos["主语"][-1]
    for i, w, role in sent:
        head = {}
        if role == "名词":
            head["宾语→动词"] = verb
        if role == "形容词":
            head["形容词→名词"] = noun
        if role == "限定词" and w != "the":
            head["物主代词→主语"] = subj_head
        if role == "尾巴" and w in {p for _, _, p in SUBJECTS}:
            head["物主代词→主语"] = subj_head
        out[i] = {"role": role, "subj": subj, "third": third, "cat": cat, "head": head, "word": w}


# =====================================================================
# 2. 花园 BERT
# =====================================================================

def new_bert(seed=0):
    from transformer import TransformerLM
    torch.manual_seed(seed)
    return TransformerLM(len(VOCAB), attn="bidirectional", **ARCH)


def _ckpt(step, seed):
    return os.path.join(CKPT_DIR, f"garden_bert_s{seed}_step{step}.pt")


def garden_bert(step=SAVE_STEPS[-1], seed=0, verbose=False):
    """
    花园 BERT：4 层、4 个头、d = 64 的双向 Transformer，在第 1 讲的 8000 篇训练文档上做 MLM（动态掩码）。
    共训练 6000 步，在 SAVE_STEPS 处存检查点。第一次调用会训练（约 1 分钟），之后读缓存。

    【参数】step：要哪个检查点（SAVE_STEPS 之一；0 = 随机初始化）。
    【返回】模型（eval 模式）。
    """
    path = _ckpt(step, seed)
    if not os.path.exists(path):
        _train_bert(seed, verbose)
    m = new_bert(seed)
    m.load_state_dict(torch.load(path))
    m.eval()
    return m


def _train_bert(seed, verbose):
    os.makedirs(CKPT_DIR, exist_ok=True)
    docs = standard_splits()["train"]
    m = new_bert(seed)
    total = SAVE_STEPS[-1]
    opt = torch.optim.Adam(m.parameters(), lr=3e-3)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda(200, total))
    g = torch.Generator().manual_seed(seed)
    step, ep = 0, 0
    torch.save(m.state_dict(), _ckpt(0, seed))
    while step < total:
        batches = mlm_batches(docs, VOCAB, 32, seed=seed * 1000 + ep)
        for bi in torch.randperm(len(batches), generator=g).tolist():
            m.train()
            s, n = mlm_loss(m, batches[bi])
            opt.zero_grad()
            (s / n).backward()
            torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
            opt.step()
            sched.step()
            step += 1
            if step in SAVE_STEPS:
                torch.save(m.state_dict(), _ckpt(step, seed))
                if verbose:
                    print(f"    第 {step} 步：MLM 损失 {float((s / n).detach()):.3f}")
            if step >= total:
                break
        ep += 1


@torch.no_grad()
def layer_states(model, docs):
    """
    每篇文档的逐层隐藏状态和注意力权重。

    【返回】list，每篇文档一项：(states, attn)。states：(层数 + 1, T, d) 张量，第 0 层是嵌入（含位置编码），
            第 l 层是第 l 个块的输出；attn：(层数, 头数, T, T)，第 l 层每个头"每个位置看各位置"的权重。
            T = 1 + 文档长度（开头是 <bos>，相当于 BERT 的 [CLS]）。
    """
    out = []
    for d in docs:
        ids = torch.tensor([[VOCAB.stoi["<bos>"]] + VOCAB.encode(d)])
        x = model.embed(ids)
        states, attn = [x[0]], []
        for blk in model.blocks:
            x, w = blk(x, mask=None)
            states.append(x[0])
            attn.append(w["self"][0])
        out.append((torch.stack(states), torch.stack(attn)))
    return out


# =====================================================================
# 3. 探针
# =====================================================================

def collect(states_list, ann_list, label_fn, layer):
    """
    从逐层状态里取出某一层、满足条件的 token 的 (表示, 标签)。

    【参数】label_fn(ann, i) → 标签（int）或 None（不参加这个任务）；ann 是 annotate 的结果，
            i 是 token 在文档里的下标（隐藏状态里对应 i + 1，因为开头有 <bos>）。
    【返回】(X, y)：X (n, d) 张量，y (n,) long 张量。
    """
    X, y = [], []
    for (states, _), ann in zip(states_list, ann_list):
        for i, a in enumerate(ann):
            if a is None:
                continue
            lab = label_fn(ann, i)
            if lab is not None:
                X.append(states[layer, i + 1])
                y.append(lab)
    return torch.stack(X), torch.tensor(y)


def train_probe(Xtr, ytr, Xte, yte, n_cls, epochs=300, lr=0.05, wd=1e-4):
    """
    线性探针：在冻结的表示上训练一个逻辑回归（全批量 Adam），返回测试准确率。

    【思路】探针越简单越好——如果用一个强大的分类器，"能预测出标签"可能是探针自己学会的，
            而不是表示里本来就有（Hewitt & Liang 2019 的控制任务就是为了检验这一点）。
    【边界】先按训练集的均值和标准差标准化特征。固定随机种子，结果可复现。
    """
    mu, sd = Xtr.mean(0), Xtr.std(0) + 1e-6
    Xtr, Xte = (Xtr - mu) / sd, (Xte - mu) / sd
    torch.manual_seed(0)
    W = torch.zeros(Xtr.shape[1], n_cls, requires_grad=True)
    b = torch.zeros(n_cls, requires_grad=True)
    opt = torch.optim.Adam([W, b], lr=lr, weight_decay=wd)
    for _ in range(epochs):
        loss = F.cross_entropy(Xtr @ W + b, ytr)
        opt.zero_grad()
        loss.backward()
        opt.step()
    with torch.no_grad():
        return float(((Xte @ W + b).argmax(1) == yte).float().mean())


# =====================================================================
# 4. 损失地形
# =====================================================================

def hessian_trace(loss_fn, params, n_samples=20, seed=0):
    """
    Hutchinson 估计：tr(H) ≈ E[vᵀ H v]，v 的每个分量独立地取 ±1。

    【参数】loss_fn()：返回标量损失的函数；params：参数列表。
    【返回】估计的迹（float）。迹越大，损失在参数空间里平均越"陡"（越尖锐）。
    【思路】Hv 用两次反向传播算（Pearlmutter 技巧）：先求梯度 g，再对 g·v 求梯度。不需要显式地写出 H。
    """
    gen = torch.Generator().manual_seed(seed)
    est = 0.0
    for _ in range(n_samples):
        vs = [(torch.randint(0, 2, p.shape, generator=gen) * 2 - 1).float() for p in params]
        loss = loss_fn()
        grads = torch.autograd.grad(loss, params, create_graph=True)
        gv = sum((g * v).sum() for g, v in zip(grads, vs))
        hv = torch.autograd.grad(gv, params)
        est += float(sum((h * v).sum() for h, v in zip(hv, vs)))
    return est / n_samples


def filter_normalized_direction(params, seed):
    """
    Li et al. (2018) 的"按滤波器归一化"的随机方向：每个参数矩阵的每一行（一个"滤波器"）的随机方向，
    缩放到和对应的那一行参数一样长。这样不同大小、不同尺度的模型之间的地形图才能比较。
    一维参数（偏置、LayerNorm）方向设为 0。
    """
    gen = torch.Generator().manual_seed(seed)
    out = []
    for p in params:
        d = torch.randn(p.shape, generator=gen)
        if p.dim() <= 1:
            d.zero_()
        else:
            d = d * (p.norm(dim=-1, keepdim=True) / (d.norm(dim=-1, keepdim=True) + 1e-10))
        out.append(d)
    return out


__all__ = ["VOCAB", "WORLD", "ROLES", "SUBJ_NAMES", "annotate", "garden_bert", "new_bert", "layer_states",
           "collect", "train_probe", "hessian_trace", "filter_normalized_direction", "SAVE_STEPS", "make_docs",
           "standard_splits", "ARCH", "SUBJECTS"]