"""
garden_tasks.py —— 第 6 讲：把花园世界的句子变成"下游任务"，以及微调 / 评估的公用工具
=====================================================================

【这个文件做什么】
    第 2 讲的 Transformer 只学过"续写花园文档"。第 6 讲要让它做**新任务**，比较几条路：
    提示（不改参数）、全参数微调、参数高效微调（PEFT）、多任务指令微调。这里提供：

        (1) 任务数据：
            noun_example(sent)          抽取式问答："这句话里吃 / 摘 / 种 / 浇 / 喂的是什么？" → 一个名词
            rewrite(sent, subj)         改写：把主语换成 subj，动词和所有物主代词跟着变
                                        （I eat my apple from my own garden . → she eats her apple from her own garden .）
            sample_sentences(n, seed)   按真实分布采样不重复的句子
        (2) 统一的序列格式（"提示模板"）：
                <bos> [<x1> 指令 <x2>] 输入句子 <x0> 输出 <eos>
            <x0> 是第 2 讲扩展词表里的一个哨兵 token，这里当"答案从这里开始"的分隔符用——
            相当于真实聊天模型模板里的 "Assistant:"。
            只有 <x0> 之后的 token 计入损失（损失掩码，第 1 讲 §2.2）。
        (3) 批次、损失、贪心生成、精确匹配准确率、"当普通语言模型用时的困惑度"（测量遗忘）。

【为什么要自己造任务】
    和前几讲一样：花园世界里每个任务的**正确答案是确定的**，可以精确判对错；
    而且我们知道预训练数据里**有什么、没有什么**（从来没有 <x0>，从来没有"改写"这种格式），
    所以"模型的本事是从预训练来的，还是从微调来的"能分得很清楚。

【LLM 联系】
    真实的指令微调数据（FLAN、P3、Alpaca）也是这种"模板 + 输入 + 输出"的格式，
    训练时也只对回答部分算损失（HuggingFace TRL 的 completion-only loss / DataCollatorForCompletionOnlyLM）。
"""

import os
import random
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 (BOS, EOS, PAD, POSSESSIVES, clm_eval, ext_vocab, lr_lambda,  # noqa: E402
                       standard_splits, train_causal_lm)
from toy_world import (DOC_START, HELDOUT_PAIRS, NOUN_CAT, SUBJECTS, VERBS, GardenWorld)  # noqa: E402

VOCAB = ext_vocab()
WORLD = GardenWorld(HELDOUT_PAIRS)
SEP, INST, INST_END = "<x0>", "<x1>", "<x2>"   # 答案分隔符、指令开始 / 结束符（都是预训练从没见过的 token）
BOS_ID, EOS_ID, PAD_ID = VOCAB.stoi[BOS], VOCAB.stoi[EOS], VOCAB.stoi[PAD]
SEP_ID, INST_ID, INST_END_ID = VOCAB.stoi[SEP], VOCAB.stoi[INST], VOCAB.stoi[INST_END]
SUBJ_TOKENS = [list(toks) for toks, _, _ in SUBJECTS]   # [['I'], ['you'], ..., ['the', 'children'], ['my', 'friend']]
SUBJ_NAMES = [" ".join(t) for t in SUBJ_TOKENS]


# =====================================================================
# 1. 任务数据
# =====================================================================

def sample_sentences(n, seed, exclude=()):
    """
    按花园世界的真实分布采样 n 个**互不相同**的句子。

    【参数】n：句子数；seed：随机种子；exclude：要排除的句子集合（元素是 tuple），
            用来保证测试句子没在训练里出现过。
    【返回】list[list[str]]，每句以 "." 结尾。
    【边界】花园世界的短句很常见（最常见的一句概率约 0.1%），不去重的话测试集会和训练集大量重叠，
            "准确率"就混进了"背答案"的成分（第 3 讲 §6 的污染问题）。
    """
    rng = random.Random(seed)
    seen, out = set(exclude), []
    while len(out) < n:
        s = WORLD.sample_sentence(rng)
        if tuple(s) not in seen:
            seen.add(tuple(s))
            out.append(s)
    return out


def parse(sent):
    """
    沿真实状态机走一遍，标出每个词在句子里的"角色"。

    【返回】list[str]，和 sent 等长，取值：
            "subj"（主语）、"verb"、"det"（限定词：the 或物主代词）、"adj"、"noun"、
            "poss"（尾巴里的物主代词）、"tail"（尾巴里的其他词）、"end"（句号）。
    【边界】输入必须是合法句子（否则状态机走不下去，抛 ValueError）。
    """
    roles, st = [], DOC_START
    for w in sent:
        stage = st.stage
        if stage in ("S0", "S1", "E"):
            roles.append("subj")
        elif stage == "V":
            roles.append("verb")
        elif stage == "D":
            roles.append("det")
        elif stage == "A":
            roles.append("noun" if w in NOUN_CAT else "adj")
        elif stage == "N":
            roles.append("noun")
        elif w == ".":
            roles.append("end")
        else:
            roles.append("poss" if (stage == "R" and w in POSSESSIVES) else "tail")
        st = WORLD.step(st, w)
        if st is None:
            raise ValueError(f"不合法的句子：{' '.join(sent)}")
    return roles


def subject_index(sent):
    """句子主语在 SUBJECTS 里的下标（0..7）。【思路】句首连续的 "subj" 角色就是主语的词。"""
    roles = parse(sent)
    n = 0
    while roles[n] == "subj":
        n += 1
    return SUBJ_TOKENS.index(sent[:n])


_THIRD_TO_BASE = {third: base for base, third in VERBS.items()}


def rewrite(sent, subj):
    """
    把句子的主语换成 SUBJECTS[subj]，动词形式和物主代词（限定词位置、尾巴里）跟着变。

    【参数】sent：合法句子；subj：目标主语的下标（0..7），例如 4 = she。
    【返回】改写后的词列表（仍然是合法句子）。
    【例子】rewrite("I eat my ripe apple from my own garden .".split(), 4)
            → she eats her ripe apple from her own garden .
    【边界】限定词是 the 的保持 the（"the" 和主语无关）；只有物主代词要换。
    【为什么选这个任务】它同时需要 (a) 照抄输入里的大部分词（宾语、形容词、尾巴），
            (b) 预训练学到的语法知识（主谓一致、物主代词一致）。前者是"新本事"，后者是"老本事"——
            正好用来看各种微调方法能不能"借用"预训练的知识。
    """
    toks, third, poss = SUBJECTS[subj]
    roles = parse(sent)
    out = list(toks)
    for w, r in zip(sent, roles):
        if r == "subj":
            continue
        if r == "verb":
            base = _THIRD_TO_BASE.get(w, w)
            out.append(VERBS[base] if third else base)
        elif r in ("det", "poss") and w in POSSESSIVES:
            out.append(poss)
        else:
            out.append(w)
    return out


def noun_of(sent):
    """句子的宾语名词（抽取式问答的标准答案）。"""
    return next(w for w, r in zip(sent, parse(sent)) if r == "noun")


# =====================================================================
# 2. 序列格式与批次
# =====================================================================

def make_example(src, tgt, instruction=None):
    """
    组装一条训练 / 测试样本。

    【参数】src：输入词列表；tgt：输出词列表（不含 <eos>，这里会自动加上）；
            instruction：可选的指令词列表（夹在 <x1> 和 <x2> 之间，放在输入前面）。
    【返回】(prompt_ids, target_ids)：
            prompt_ids = <bos> [<x1> 指令 <x2>] 输入 <x0>
            target_ids = 输出 <eos>
    【边界】把"提示"和"答案"分开存，批次函数再把它们拼起来并做损失掩码；生成时只喂 prompt_ids。
    """
    prompt = [BOS_ID]
    if instruction is not None:
        prompt += [INST_ID] + VOCAB.encode(instruction) + [INST_END_ID]
    prompt += VOCAB.encode(src) + [SEP_ID]
    return prompt, VOCAB.encode(tgt) + [EOS_ID]


def task_batch(examples):
    """
    把若干条 (prompt_ids, target_ids) 拼成一个批次。

    【返回】(ids, pad, loss_mask)：都是 (B, L)。ids = prompt + target（右侧补齐）；
            pad 为 True 的是补齐位置；loss_mask[b, t] 为 True 表示"位置 t 的 token 要被预测并计入损失"，
            即 target 部分（包括 <eos>）。
    【思路】和第 2 讲的 clm_loss 一样"整体错开一位"：用 ids[:, :-1] 预测 ids[:, 1:]，
            只是多乘一个 loss_mask[:, 1:]——提示部分的 token 不算损失。
    """
    L = max(len(p) + len(t) for p, t in examples)
    B = len(examples)
    ids = torch.full((B, L), PAD_ID, dtype=torch.long)
    pad = torch.ones(B, L, dtype=torch.bool)
    lm = torch.zeros(B, L, dtype=torch.bool)
    for i, (p, t) in enumerate(examples):
        seq = p + t
        ids[i, :len(seq)] = torch.tensor(seq)
        pad[i, :len(seq)] = False
        lm[i, len(p):len(seq)] = True
    return ids, pad, lm


def task_loss(model, batch):
    """
    只在答案部分计算的交叉熵。

    【参数】model：任何 forward(ids, pad=...) → logits 的模型（基础模型、PEFT 包装后的模型都行）。
    【返回】(NLL 之和, 计入损失的 token 数)。
    """
    ids, pad, lm = batch
    logits = model(ids[:, :-1], pad=pad[:, :-1])
    tgt, keep = ids[:, 1:], lm[:, 1:]
    nll = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), tgt.reshape(-1), reduction="none")
    return (nll * keep.reshape(-1)).sum(), keep.sum()


def finetune(model, params, train_ex, steps=400, lr=1e-3, batch_size=32, warmup=30, seed=0,
             dev_ex=None, eval_every=100, eval_fn=None, extra_loss_batches=None, clip=1.0, log=None):
    """
    微调循环：只更新 params 里的参数，固定步数，按开发集挑最好的检查点。

    【参数】
        model：要训练的模型（可能是 PEFT 包装后的）；params：要训练的参数列表（其余参数不动）。
        train_ex：[(prompt_ids, target_ids)]；steps：总步数（不是遍数——数据量不同时步数相同，比较才公平）；
        lr / warmup：峰值学习率、预热步数（之后余弦衰减到 0.1 倍，和第 2 讲相同）；batch_size；seed。
        dev_ex：开发集；给了就每 eval_every 步用 eval_fn(model, dev_ex)（越大越好）评估，最后恢复最好的那一次。
        extra_loss_batches：可选的函数 step → 额外的 (损失之和, token 数)，用来在微调时混入预训练数据（"复习"）。
        log：可选的列表；给了就把 (step, 训练损失) 追加进去（画学习曲线用）。
    【返回】dict：{"best_step": 最好的步数, "best_dev": 最好的开发集分数, "curve": [(step, dev 分数)]}。
    【边界】数据少于 batch_size 时每步用全部数据；每"遍"重新打乱一次（固定种子，可复现）。
    """
    for p in model.parameters():
        p.requires_grad_(False)
    for p in params:
        p.requires_grad_(True)
    opt = torch.optim.Adam(params, lr=lr)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda(warmup, steps))
    rng = random.Random(seed)
    order, pos = [], 0
    best, best_state, best_step, curve = -1.0, None, 0, []
    for step in range(1, steps + 1):
        model.train()
        if pos + batch_size > len(order):
            order = list(range(len(train_ex)))
            rng.shuffle(order)
            pos = 0
        idx = order[pos:pos + batch_size]
        pos += batch_size
        s, n = task_loss(model, task_batch([train_ex[i] for i in idx]))
        loss = s / n
        if extra_loss_batches is not None:
            s2, n2 = extra_loss_batches(step)
            loss = loss + s2 / n2
        opt.zero_grad()
        loss.backward()
        if clip:
            torch.nn.utils.clip_grad_norm_(params, clip)
        opt.step()
        sched.step()
        if log is not None:
            log.append((step, (s / n).item()))
        if dev_ex is not None and (step % eval_every == 0 or step == steps):
            model.eval()
            score = eval_fn(model, dev_ex)
            curve.append((step, score))
            if score > best:
                best, best_step = score, step
                best_state = {k: v.detach().clone() for k, v in model.state_dict().items()}
    if best_state is not None:
        model.load_state_dict(best_state)
    model.eval()
    for p in model.parameters():
        p.requires_grad_(False)
    return {"best_step": best_step, "best_dev": best, "curve": curve}


# =====================================================================
# 3. 生成与评估
# =====================================================================

@torch.no_grad()
def greedy_generate(model, prompts, max_new=16, batch=256):
    """
    批量贪心解码：每条提示一直生成到 <eos> 或 max_new 个 token。

    【参数】prompts：[prompt_ids]；model：forward(ids, pad=...) → logits。
    【返回】list[list[int]]：每条生成的 token（不含 <eos>）。
    【思路】左侧补齐（padding on the left）：让所有提示的"最后一个 token"对齐在同一列，
            这样每一步只要看最后一列的 logits。补齐位置用 pad 掩码屏蔽（因果 + 补齐掩码）。
            没有 KV cache，每步重新算整段——序列只有二三十个 token，简单比快重要。
    【边界】左侧补齐会让真实 token 的绝对位置整体右移几格；第 2 讲的正弦位置编码只关心相对位置的话影响很小，
            但不是零。为了不引入这个偏差，这里**按提示长度分组**，同一组长度相同、不需要补齐。
    """
    model.eval()
    out = [None] * len(prompts)
    groups = {}
    for i, p in enumerate(prompts):
        groups.setdefault(len(p), []).append(i)
    for _, idxs in sorted(groups.items()):
        for s in range(0, len(idxs), batch):
            chunk = idxs[s:s + batch]
            ids = torch.tensor([prompts[i] for i in chunk])
            done = torch.zeros(len(chunk), dtype=torch.bool)
            gen = [[] for _ in chunk]
            for _ in range(max_new):
                nxt = model(ids)[:, -1].argmax(-1)
                for j, t in enumerate(nxt.tolist()):
                    if not done[j]:
                        if t == EOS_ID:
                            done[j] = True
                        else:
                            gen[j].append(t)
                if done.all():
                    break
                ids = torch.cat([ids, nxt[:, None]], 1)
            for j, i in enumerate(chunk):
                out[i] = gen[j]
    return out


def exact_match(model, examples, max_new=16):
    """贪心生成的输出和标准答案（去掉 <eos>）完全相同的比例。"""
    gens = greedy_generate(model, [p for p, _ in examples], max_new=max_new)
    return sum(g == t[:-1] for g, (_, t) in zip(gens, examples)) / len(examples)


@torch.no_grad()
def first_token_accuracy(model, examples):
    """
    单 token 答案的准确率：看 <x0> 之后第一个位置的 argmax 是否等于标准答案的第一个 token。

    【用途】抽取名词这种"答案只有一个词"的任务，不用逐步生成，一次前向就够。
    """
    hit = 0
    for s in range(0, len(examples), 256):
        chunk = examples[s:s + 256]
        ids, pad, lm = task_batch(chunk)
        logits = model(ids, pad=pad)
        for i, (p, t) in enumerate(chunk):
            hit += int(logits[i, len(p) - 1].argmax().item() == t[0])
    return hit / len(examples)


def lm_ppl(model, docs=None):
    """
    把（可能被微调过的）模型当普通语言模型用，在花园测试文档上的困惑度——用来测量"遗忘"。

    【参数】docs：默认用标准测试集的前 300 篇（全部 1000 篇也行，只是慢一些）。
    """
    docs = docs if docs is not None else SPLITS["test"][:300]
    return clm_eval(model, docs, VOCAB)[1]


# =====================================================================
# 4. 两个"基座模型"
# =====================================================================

SPLITS = standard_splits()
CKPT_DIR = os.path.join(HERE, "checkpoints")
_BASE = {}


def base_model(seed=0):
    """
    基座 A：第 2 讲训练好的标准因果 Transformer（d=64、2 层、约 10.7 万参数），直接读第 2 讲的缓存。

    【返回】一个**新的**模型对象（深拷贝），调用者可以放心地就地修改、微调它。
    【它会什么、不会什么】花园文档里的句子彼此独立，**从来不需要照抄前文**——
            所以它学会了语法和搭配，但没学会"复制"（第 5 讲 §3 也看到了：同一句话重复 12 遍，
            它的困惑度反而更高）。
    """
    import copy
    if seed not in _BASE:
        name = "tf_base" if seed == 0 else f"tf_base_seed{seed}"
        m, _ = train_causal_lm(name, SPLITS, VOCAB, seed=seed, verbose=False)
        _BASE[seed] = m
    return copy.deepcopy(_BASE[seed])


def echo_docs(docs, p=0.2, seed=0):
    """
    "带复读的花园语料"：每句话说完后，以概率 p 原样再说一遍。

    【参数】docs：文档列表；p：每句话被重复的概率；seed：随机种子。
    【返回】新的文档列表（总 token 数约为原来的 1 + p 倍）。
    【为什么】真实网页里到处是重复：引用、转述、页眉页脚、代码里的变量名、列表……
            在这样的数据上预训练，模型会学会"照抄上文"（Olsson et al. 2022 的归纳头 induction head），
            而这正是很多下游任务（抽取、改写、上下文学习）的基础。花园文档没有这一点，我们把它补上。
    """
    rng = random.Random(seed)
    out = []
    for d in docs:
        new, cur = [], []
        for w in d[:-1]:
            cur.append(w)
            if w == ".":
                new += cur
                if rng.random() < p:
                    new += cur
                cur = []
        out.append(new + [EOS])
    return out


def echo_base_model(seed=0, p=0.2, verbose=False):
    """
    基座 B：结构、超参数、训练步骤和基座 A 完全相同，只是预训练语料换成 echo_docs(训练集, p)。

    【返回】新的模型对象（深拷贝）。第一次调用时训练（约 10～15 秒），之后读本讲 checkpoints/ 里的缓存。
    【边界】标准测试集上的困惑度会比基座 A 略高（约 3.94 vs 3.87）：它把一点概率留给了"这句话再说一遍"，
            而标准测试文档里从不复读。
    """
    import copy
    from garden_lm import BASE_ARCH, BASE_TRAIN, clm_batches, clm_loss, load_or_train, train
    from transformer import TransformerLM
    key = (seed, p)
    if key not in _BASE:
        cfg = {"arch": BASE_ARCH, "train": dict(BASE_TRAIN, seed=seed), "data": f"garden_echo_p{p}_v1"}
        batches = clm_batches(echo_docs(SPLITS["train"], p, seed=0), VOCAB, BASE_TRAIN["batch"])

        def build():
            torch.manual_seed(seed)
            return TransformerLM(len(VOCAB), attn="causal", **BASE_ARCH)

        def fit(m):
            return train(m, batches, clm_loss, epochs=BASE_TRAIN["epochs"], lr=BASE_TRAIN["lr"],
                         warmup=BASE_TRAIN["warmup"], seed=seed, tag="echo_base", verbose=verbose,
                         valid_fn=lambda mm: clm_eval(mm, SPLITS["valid"], VOCAB)[1])

        _BASE[key], _ = load_or_train(f"echo_base_p{p}_s{seed}", cfg, build, fit, verbose=verbose,
                                      ckpt_dir=CKPT_DIR)
    return copy.deepcopy(_BASE[key])


def fresh_model(seed=0):
    """同样结构、随机初始化的模型（"不预训练、直接在任务数据上从零训练"的对照组）。"""
    from transformer import TransformerLM
    from garden_lm import BASE_ARCH
    torch.manual_seed(seed)
    return TransformerLM(len(VOCAB), attn="causal", **BASE_ARCH)


# =====================================================================
# 5. 标准的任务数据划分
# =====================================================================

def task_split(seed=0):
    """
    第 6 讲统一使用的句子划分：先采 3000 个互不相同的句子，打乱后切成 测试 300 / 开发 200 / 训练 2500。

    【返回】(test, dev, train)：三个句子列表，互不重叠。
    【为什么不分别采样】如果先采测试集、再"排除测试集"采训练集，后采的集合会系统性地偏向罕见（更长）的句子——
            三个集合的分布就不一样了。先采一个大池子再随机切分，三者同分布。
    """
    pool = sample_sentences(3000, seed)
    random.Random(seed).shuffle(pool)
    return pool[:300], pool[300:500], pool[500:]


def noun_examples(sents, instruction=None):
    """抽取任务：输入句子，输出它的宾语名词（一个 token）。"""
    return [make_example(s, [noun_of(s)], instruction) for s in sents]


def rewrite_examples(sents, subj=4, instruction=None):
    """改写任务：把主语换成 SUBJECTS[subj]（默认 4 = she）。"""
    return [make_example(s, rewrite(s, subj), instruction) for s in sents]


CATEGORY_WORDS = {"fruit": "ripe", "veg": "fresh", "flower": "pretty", "animal": "hungry"}


def category_examples(sents, instruction=None):
    """
    分类任务：输入句子，输出宾语名词的类别。

    【标签词（verbalizer）】类别用该类专属的一个形容词表示：水果 → ripe，蔬菜 → fresh，花 → pretty，动物 → hungry。
            标签必须是词表里的 token；选"和类别有关的词"还是"随便一个 token"，第 9 讲会专门比较。
    """
    return [make_example(s, [CATEGORY_WORDS[NOUN_CAT[noun_of(s)]]], instruction) for s in sents]


def count_params(params):
    return sum(p.numel() for p in params)


def nll_of(model, examples):
    """答案部分的平均 NLL（每 token）。"""
    s_tot, n_tot = 0.0, 0
    with torch.no_grad():
        for s in range(0, len(examples), 256):
            a, b = task_loss(model, task_batch(examples[s:s + 256]))
            s_tot += a.item()
            n_tot += int(b)
    return s_tot / n_tot


__all__ = ["VOCAB", "WORLD", "SPLITS", "SEP", "INST", "INST_END", "SEP_ID", "INST_ID", "INST_END_ID",
           "BOS_ID", "EOS_ID", "PAD_ID", "SUBJ_TOKENS", "SUBJ_NAMES", "CATEGORY_WORDS",
           "sample_sentences", "parse", "subject_index", "rewrite", "noun_of",
           "make_example", "task_batch", "task_loss", "finetune", "greedy_generate", "exact_match",
           "first_token_accuracy", "lm_ppl", "base_model", "echo_docs", "echo_base_model", "fresh_model",
           "task_split", "noun_examples", "rewrite_examples", "category_examples", "count_params", "nll_of"]