"""
chat_lab.py —— 农场"聊天助手"：指令数据、对话模板、指令微调、生成与评估、人类偏好（第 11、12 讲）
=====================================================================

【这个文件要回答的问题】
    slides：语言模型被训练成"续写"——你给它半句话，它接着写。可一个聊天助手应该**按指令回答，并且知道什么时候停**。
    做法（指令微调，instruction fine-tuning）：造一个"指令 → 回答"的数据集，在上面微调预训练模型。
    这个文件把农场语言变成一个小小的指令数据集（9 种任务，每种 6 种说法），并提供训练、生成、评估的工具。

【对话格式】（模仿 slides 里 Llama-2 Chat 的 [INST] ... [/INST] 和 Alpaca 的 ### Instruction / ### Response）
    <s> <user> 指令的说法 : 输入 <assistant> 回答 </s>
    例：<s> <user> what is the subject ? : the dogs near the cat chase the cows at night . <assistant> the dogs </s>

【9 种任务】
    subject / verb / object / time   从句子里找出主语 / 动词 / 宾语 / 时间短语
    plural / singular                把主语（和动词）改成复数 / 单数
    short                            去掉所有介词短语
    repeat                           原样复述
    describe                         输入是一个主语（"the dogs"），回答一句关于它的话——开放式生成，第 12 讲的 RLHF 用它
    每种任务 6 种说法：前 4 种用于训练，后 2 种（"give me the ..."、"find the ..." 一类）训练时从不出现，用来测"换个说法还懂不懂"。

【本文件提供】
    TASKS / TEMPLATES / CHAT_WORDS     任务、说法、新增的词
    load_chat_base                     第 2 讲的农场 GPT + 扩充词表（新词的嵌入初始化为旧词嵌入的平均）
    make_example / sample_ift          造指令数据
    encode_chat / train_chat           编码（只在回答上算损失，或者全部都算）、训练
    generate / evaluate_tasks          贪心生成（遇到 </s> 停）、按任务算准确率
    human_reward                       "人类"对 describe 任务回答的偏好分数（第 11 讲 06、第 12 讲用）
"""

import copy
import math
import os
import random
import sys

import torch
import torch.nn as nn
import torch.nn.functional as F

_HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture10", "lecture03", "lecture02", "lecture01"):
    _p = os.path.join(_HERE, "..", _d)
    if _p not in sys.path:
        sys.path.insert(0, _p)
from decoding import farm_true_logprob  # noqa: E402
from kv_cache import load_farm_gpt  # noqa: E402
from peft_lab import time_phrase_start  # noqa: E402
from toy_corpus import ANIMALS, VERBS, Vocab, generate_sentence, noun_form  # noqa: E402

IGNORE = -100
USER, ASSISTANT = "<user>", "<assistant>"

# 每种任务 6 种说法：前 4 种训练用，后 2 种只在测试时出现
TEMPLATES = {
    "subject": [["what", "is", "the", "subject", "?"], ["name", "the", "subject"], ["subject", "?"],
                ["who", "does", "it", "?"], ["give", "me", "the", "subject"], ["find", "the", "subject"]],
    "verb": [["what", "is", "the", "verb", "?"], ["name", "the", "verb"], ["verb", "?"],
             ["what", "do", "they", "do", "?"], ["give", "me", "the", "verb"], ["find", "the", "verb"]],
    "object": [["what", "is", "the", "object", "?"], ["name", "the", "object"], ["object", "?"],
               ["to", "what", "?"], ["give", "me", "the", "object"], ["find", "the", "object"]],
    "time": [["what", "is", "the", "time", "?"], ["name", "the", "time"], ["when", "?"],
             ["when", "does", "it", "happen", "?"], ["give", "me", "the", "time"], ["find", "the", "time"]],
    "plural": [["make", "it", "plural"], ["plural", "?"], ["say", "it", "in", "plural"],
               ["make", "the", "subject", "plural"], ["give", "me", "the", "plural"], ["write", "it", "plural"]],
    "singular": [["make", "it", "singular"], ["singular", "?"], ["say", "it", "in", "singular"],
                 ["make", "the", "subject", "singular"], ["give", "me", "the", "singular"],
                 ["write", "it", "singular"]],
    "short": [["make", "it", "shorter"], ["remove", "the", "details"], ["shorter", "?"], ["say", "it", "shorter"],
              ["give", "me", "a", "shorter", "sentence"], ["write", "it", "shorter"]],
    "repeat": [["repeat", "it"], ["say", "it", "again"], ["repeat", "?"], ["say", "the", "sentence"],
               ["give", "me", "the", "sentence"], ["write", "it", "again"]],
    "describe": [["describe"], ["tell", "me", "about"], ["write", "a", "sentence", "about"], ["describe", "?"],
                 ["give", "me", "a", "sentence", "about"], ["say", "something", "about"]],
}
TASKS = list(TEMPLATES)
N_TRAIN_TEMPLATES = 4


def _chat_words():
    words = [USER, ASSISTANT, ":"]
    for ts in TEMPLATES.values():
        for t in ts:
            for w in t:
                if w not in words:
                    words.append(w)
    return words


CHAT_WORDS = _chat_words()


# =====================================================================
# 1. 基座模型：第 2 讲的农场 GPT + 扩充词表
# =====================================================================

def expand_vocab(model, old_vocab, new_vocab):
    """
    给一个已经训练好的 GPT 加新词：嵌入表和输出层各多出几行。

    【思路】旧词的编号不变（新词表 = 旧词表 + 新词追加在末尾），所以旧的行原样复制；新词的嵌入初始化为
        所有旧词嵌入的平均（一个"不偏向任何词"的起点），输出层的新行初始化为 0（一开始模型给新词的打分是 0）。
        这就是实践中给预训练模型加 <|user|> 这类特殊记号时的做法（例如 Hugging Face 的 resize_token_embeddings）。
    【边界】只支持第 2 讲的 GPT（tok、head 两个属性，且输出层不和嵌入共享）。
    """
    if old_vocab.itos != new_vocab.itos[:len(old_vocab)]:
        raise ValueError("新词表必须以旧词表开头")
    m = copy.deepcopy(model)
    V_old, V_new, d = len(old_vocab), len(new_vocab), model.tok.weight.shape[1]
    tok = nn.Embedding(V_new, d)
    head = nn.Linear(d, V_new, bias=False)
    with torch.no_grad():
        tok.weight[:V_old] = model.tok.weight
        tok.weight[V_old:] = model.tok.weight.mean(0)
        head.weight[:V_old] = model.head.weight
        head.weight[V_old:] = 0.0
    m.tok, m.head = tok, head
    return m


def load_chat_base(log=None):
    """
    【返回】(model, vocab, base_model, base_vocab)：扩充了词表的农场 GPT 和新词表，以及原来的模型和词表（做对比用）。
    【边界】第一次运行会读取（或训练）第 2 讲的 farm_gpt。
    """
    base, base_vocab, _ = load_farm_gpt(log=log)
    vocab = Vocab(base_vocab.itos[3:] + CHAT_WORDS)
    return expand_vocab(base, base_vocab, vocab), vocab, base, base_vocab


# =====================================================================
# 2. 指令数据
# =====================================================================

def answer(task, ex):
    """
    任务的标准答案（词列表）。ex 是 toy_corpus.Example（describe 除外）。
    【边界】plural / singular：句子本来就是那个数时，答案就是原句。
    """
    t = ex.tokens
    vi, ts = ex.verb_index, time_phrase_start(ex.tokens)
    if task == "subject":
        return ["the", t[1]]
    if task == "verb":
        return [t[vi]]
    if task == "object":
        return t[vi + 1:ts]
    if task == "time":
        return t[ts:-1]
    if task in ("plural", "singular"):
        pl = task == "plural"
        out = list(t)
        out[1] = noun_form(ex.subject, pl)
        out[vi] = VERBS[ex.verb][1 if pl else 0]
        return out
    if task == "short":
        return t[:2] + t[vi:]
    if task == "repeat":
        return list(t)
    raise ValueError(task)


def make_example(rng, task, template):
    """
    造一条指令数据。
    【返回】(prompt 词列表（不含 <s>）, 回答词列表（不含 </s>）)
    【思路】普通任务的输入是一句 0–2 个介词短语的农场句子；describe 的输入是一个主语（"the dog" / "the dogs"），
        回答是按真实分布抽的一句不带介词短语、以这个主语开头的句子。
    """
    if task == "describe":
        a, pl = rng.choice(ANIMALS), rng.random() < 0.5
        ex = generate_sentence(rng, k=0, subject=a, plural=pl)
        return [USER] + template + [":", "the", noun_form(a, pl), ASSISTANT], list(ex.tokens)
    ex = generate_sentence(rng, k_weights=(0.45, 0.35, 0.20))
    return [USER] + template + [":"] + ex.tokens + [ASSISTANT], answer(task, ex)


def sample_ift(n, seed, n_templates=N_TRAIN_TEMPLATES, tasks=TASKS):
    """
    n 条指令数据：任务均匀地抽，说法从每个任务的前 n_templates 种里抽。
    【参数】n_templates=1 时每个任务只用第一种说法（"说法多样性"的对照实验）
    """
    rng = random.Random(seed)
    out = []
    for _ in range(n):
        task = rng.choice(tasks)
        template = rng.choice(TEMPLATES[task][:n_templates])
        out.append((task,) + make_example(rng, task, template))
    return out


def encode_chat(vocab, prompt, response, loss_on="response"):
    """
    编码一条对话：ids = [<s>] + prompt + response + [</s>]；目标在需要算损失的位置是下一个词，其余是 IGNORE。
    【参数】loss_on："response"（只在回答和 </s> 上算损失——指令微调的常见做法）或 "all"（整条序列都算）
    【返回】(ids, targets)，等长的列表；targets[i] 是 ids[i] 之后那个词（或 IGNORE）
    """
    ids = [vocab.bos_id] + vocab.encode(prompt) + vocab.encode(response) + [vocab.eos_id]
    start = 1 + len(prompt)                       # 第一个回答词在 ids 里的下标
    targets = []
    for i in range(len(ids) - 1):
        nxt = ids[i + 1]
        targets.append(nxt if (loss_on == "all" or i + 1 >= start) else IGNORE)
    return ids[:-1], targets


def train_chat(model, vocab, data, steps, lr=1e-3, batch=32, loss_on="response", seed=0, params=None, log=None):
    """
    指令微调：在 (任务, prompt, 回答) 数据上训练若干步（数据用完就重新打乱再来）。
    【参数】params：只训练这些参数（LoRA 用）；None 表示训练所有 requires_grad 的参数
    【返回】最后 100 步的平均损失
    """
    rng = random.Random(seed)
    enc = [encode_chat(vocab, p, r, loss_on) for _, p, r in data]
    params = [p for p in model.parameters() if p.requires_grad] if params is None else list(params)
    opt = torch.optim.Adam(params, lr=lr)
    order, losses = [], []
    model.train()
    for step in range(1, steps + 1):
        if len(order) < batch:
            order = list(range(len(enc)))
            rng.shuffle(order)
        idx, order = order[:batch], order[batch:]
        T = max(len(enc[i][0]) for i in idx)
        x = torch.full((batch, T), vocab.pad_id, dtype=torch.long)
        y = torch.full((batch, T), IGNORE, dtype=torch.long)
        for b, i in enumerate(idx):
            ids, tg = enc[i]
            x[b, :len(ids)] = torch.tensor(ids)
            y[b, :len(tg)] = torch.tensor(tg)
        logits, _ = model(x)
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1), ignore_index=IGNORE)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(params, 1.0)
        opt.step()
        losses.append(loss.item())
        if log and step % max(steps // 4, 1) == 0:
            log(f"    step {step:5d}  平均损失 {sum(losses[-100:]) / len(losses[-100:]):.4f}")
    model.eval()
    return sum(losses[-100:]) / len(losses[-100:])


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

@torch.no_grad()
def generate(model, vocab, prompts, max_new=20, greedy=True, gen=None, temperature=1.0):
    """
    从每个 prompt（词列表，不含 <s>）开始生成，遇到 </s> 停。
    【返回】回答的词列表的列表（不含 </s>）；如果 max_new 步内没停，最后一个元素是 "…"（表示"没停下来"）
    【思路】按 prompt 长度分组（GPT 用绝对位置，不能左边补齐），每组一起逐词生成。每一步都重新算整段前向
        ——序列很短，比写 KV 缓存简单（第 3 讲讲过怎样用缓存加速）。
    """
    out = [None] * len(prompts)
    groups = {}
    for i, p in enumerate(prompts):
        groups.setdefault(len(p), []).append(i)
    for _, idx in groups.items():
        x = torch.tensor([[vocab.bos_id] + vocab.encode(prompts[i]) for i in idx])
        done = [False] * len(idx)
        resp = [[] for _ in idx]
        for _ in range(max_new):
            if x.size(1) >= model.max_len:
                break
            logits, _ = model(x)
            z = logits[:, -1].clone()
            z[:, vocab.pad_id] = float("-inf")
            z[:, vocab.bos_id] = float("-inf")
            if greedy:
                nxt = z.argmax(-1)
            else:
                nxt = torch.multinomial(F.softmax(z / temperature, -1), 1, generator=gen).squeeze(1)
            for b, t in enumerate(nxt.tolist()):
                if not done[b]:
                    if t == vocab.eos_id:
                        done[b] = True
                    else:
                        resp[b].append(vocab.itos[t])
            if all(done):
                break
            x = torch.cat([x, nxt.unsqueeze(1)], dim=1)
        for b, i in enumerate(idx):
            out[i] = resp[b] if done[b] else resp[b] + ["…"]
    return out


def describe_ok(prompt, response):
    """describe 任务的回答算不算对：是一句合语法的农场句子（真实概率 > 0），而且主语就是提示里那个。"""
    subj = prompt[-2]                              # prompt 的结尾是 ... "the" 主语 <assistant>
    return len(response) >= 2 and response[:2] == ["the", subj] and \
        farm_true_logprob(response, given_subject=True) > -math.inf


def evaluate_tasks(model, vocab, n_per_task, seed, which="train", tasks=TASKS):
    """
    每种任务造 n_per_task 道新题（新句子），贪心生成，算准确率。
    【参数】which："train"（用训练时的 4 种说法）、"heldout"（只用没见过的 2 种说法），或者一个下标列表（如 [0]）
    【返回】({任务: 准确率}, 例子字典 {任务: (prompt, 回答, 标准答案)})
    【边界】除 describe 外都要求和标准答案逐词相同（包括在正确的地方停下）。
    """
    rng = random.Random(seed)
    acc, examples = {}, {}
    for task in tasks:
        if which == "train":
            temps = TEMPLATES[task][:N_TRAIN_TEMPLATES]
        elif which == "heldout":
            temps = TEMPLATES[task][N_TRAIN_TEMPLATES:]
        else:
            temps = [TEMPLATES[task][i] for i in which]
        items = [make_example(rng, task, rng.choice(temps)) for _ in range(n_per_task)]
        outs = generate(model, vocab, [p for p, _ in items])
        if task == "describe":
            ok = [describe_ok(p, o) for (p, _), o in zip(items, outs)]
        else:
            ok = [o == r for (_, r), o in zip(items, outs)]
        acc[task] = sum(ok) / n_per_task
        examples[task] = [(p, o, r) for (p, r), o in zip(items, outs)]
    return acc, examples


# =====================================================================
# 4. "人类"偏好（第 11 讲 06、第 12 讲）
# =====================================================================

def human_reward(prompt_subject, response):
    """
    describe 任务里，"人类标注员"对一个回答的真实打分 r*(x, y)（标注员自己说不清，但他的偏好就是这样）：
        - 不是农场里可能发生的事（语法错，或者不符合这种动物的习性，比如狗在夜里追狐狸），或者主语不对：−2
          （答非所问、说错话）
        - 否则从 0 起：讲吃东西（动词 eat）+1；讲"叫声"（宾语 noise）−1；时间是 in the morning +0.5
          ——这位标注员喜欢"吃早饭"的故事，讨厌吵闹
    【参数】prompt_subject：提示里的主语形式（"dog" 或 "dogs"）；response：回答的词列表
    【边界】回答里带介词短语也算合语法。
    """
    if len(response) < 2 or response[:2] != ["the", prompt_subject] or \
            farm_true_logprob(response, given_subject=True) == -math.inf:
        return -2.0
    r = 0.0
    if any(w in ("eat", "eats") for w in response):
        r += 1.0
    if "noise" in response:
        r -= 1.0
    if response[-4:-1] == ["in", "the", "morning"]:
        r += 0.5
    return r


# =====================================================================
# 5. 缓存的指令微调模型（05 号脚本训练；06 号脚本和第 12 讲读取）
# =====================================================================

SFT_CONFIG = {"n_data": 4000, "data_seed": 1, "steps": 1500, "lr": 1e-3, "batch": 32,
              "n_templates": N_TRAIN_TEMPLATES, "loss_on": "response"}
SFT_CKPT = os.path.join(_HERE, "checkpoints", "chat_sft.pt")


def train_sft(log=None, config=SFT_CONFIG):
    """从扩充词表的农场 GPT 出发，按 config 做一次指令微调。【返回】(model, vocab)"""
    model, vocab, _, _ = load_chat_base()
    data = sample_ift(config["n_data"], config["data_seed"], n_templates=config["n_templates"])
    train_chat(model, vocab, data, config["steps"], lr=config["lr"], batch=config["batch"],
               loss_on=config["loss_on"], log=log)
    return model, vocab


def load_chat_sft(log=print):
    """
    读取缓存的指令微调模型 checkpoints/chat_sft.pt；没有（或配置不符）就训练一个并保存（约 15 秒）。
    【返回】(model, vocab)
    """
    _, vocab, _, _ = load_chat_base()
    if os.path.exists(SFT_CKPT):
        blob = torch.load(SFT_CKPT, weights_only=True)
        if blob["config"] == SFT_CONFIG:
            model, _, _, _ = load_chat_base()
            model.load_state_dict(blob["state_dict"])
            model.eval()
            if log:
                log("  （读取缓存的指令微调模型：lecture11/checkpoints/chat_sft.pt）")
            return model, vocab
    if log:
        log("  （没有缓存的指令微调模型，训练一个，约 15 秒……）")
    model, vocab = train_sft()
    os.makedirs(os.path.dirname(SFT_CKPT), exist_ok=True)
    torch.save({"config": SFT_CONFIG, "state_dict": model.state_dict()}, SFT_CKPT)
    if log:
        log("  （已保存到 lecture11/checkpoints/chat_sft.pt）")
    return model, vocab


# =====================================================================
# 6. 奖励模型（06 号脚本训练；第 12 讲读取）
# =====================================================================

class RewardModel(nn.Module):
    """
    奖励模型 r_θ(x, y)：slides——"奖励模型是第 1 步那个语言模型的一份拷贝，只是把输出层的 softmax（对词表的分布）
    换成了一个标量"。
    【思路】把 [<s>] + 提示 + 回答 + [</s>] 送进 GPT 的主干，取**最后一个位置**（</s>）的隐藏状态，经过一个
        线性层得到一个数。因果注意力下，只有最后一个位置"看过"整段对话。
    【参数】lm：指令微调后的 GPT（会被复制，原模型不变）
    """

    def __init__(self, lm):
        super().__init__()
        self.lm = copy.deepcopy(lm)
        d = self.lm.tok.weight.shape[1]
        self.lm.head = nn.Identity()               # 去掉"对词表的分布"那一层
        self.value = nn.Linear(d, 1)
        nn.init.zeros_(self.value.weight)
        nn.init.zeros_(self.value.bias)

    def forward(self, ids, lengths):
        """ids：(B, T) 右边补 <pad>；lengths：每条序列的真实长度。【返回】(B,) 的奖励"""
        h, _ = self.lm(ids)                        # 输出层换成了恒等映射，所以这里是最后一层 LayerNorm 之后的隐藏状态
        last = h[torch.arange(len(ids)), lengths - 1]
        return self.value(last).squeeze(-1)


def batch_dialogues(vocab, prompts, responses):
    """把 (提示, 回答) 编成 [<s>] + 提示 + 回答 + [</s>]，右边补齐。【返回】(ids, lengths)"""
    seqs = [[vocab.bos_id] + vocab.encode(p) + vocab.encode([w for w in r if w != "…"]) + [vocab.eos_id]
            for p, r in zip(prompts, responses)]
    T = max(len(s) for s in seqs)
    ids = torch.full((len(seqs), T), vocab.pad_id, dtype=torch.long)
    for i, s in enumerate(seqs):
        ids[i, :len(s)] = torch.tensor(s)
    return ids, torch.tensor([len(s) for s in seqs])


@torch.no_grad()
def score(rm, vocab, prompts, responses, batch=256):
    """奖励模型给一批 (提示, 回答) 打的分（列表）。"""
    rm.eval()
    out = []
    for i in range(0, len(prompts), batch):
        ids, lengths = batch_dialogues(vocab, prompts[i:i + batch], responses[i:i + batch])
        out += rm(ids, lengths).tolist()
    return out


RM_CKPT = os.path.join(_HERE, "checkpoints", "reward_model.pt")


def load_reward_model(log=print):
    """读取 06 号脚本缓存的奖励模型（lecture11/checkpoints/reward_model.pt）。【返回】(rm, vocab)；没有缓存时抛出 FileNotFoundError。"""
    sft, vocab = load_chat_sft(log=None)
    if not os.path.exists(RM_CKPT):
        raise FileNotFoundError("找不到 lecture11/checkpoints/reward_model.pt，请先运行 ../lecture11/06_reward_model.py")
    rm = RewardModel(sft)
    rm.load_state_dict(torch.load(RM_CKPT, weights_only=True)["state_dict"])
    rm.eval()
    if log:
        log("  （读取缓存的奖励模型：lecture11/checkpoints/reward_model.pt）")
    return rm, vocab