"""
=============================================================================
第 14 讲 · 工具模块：变量追踪任务、"直接回答"与"思维链"两种格式、训练与生成
=============================================================================

对应讲义：lectures/lecture14_reasoning_cot.md 第 4 节
被本目录的 01、03 号文件使用。

-----------------------------------------------------------------------------
【任务：变量追踪】
-----------------------------------------------------------------------------
题目里有两条互相交错的"赋值链"（顺序打乱），问其中一条链最后那个变量的值（所有加法都对 10 取余）：

    题目：  d = b + 1 , a = 4 , c = a + 3 , b = 7 ; c ?
    链 1：  a = 4 → c = a + 3 = 7          （要问的是这一条）
    链 2：  b = 7 → d = b + 1 = 8          （干扰项）
    答案：  7

链长 k = 每条链里赋值语句的个数（上面的例子 k = 2）。k 越大，要"顺藤摸瓜"的步数越多。

为什么要两条链？如果只有一条链，答案 = 起始值 + 所有常数之和（不用管顺序），模型可以走捷径。
有了干扰链，模型必须先弄清楚**哪些语句属于要问的那条链** —— 这只能一步步地追。

-----------------------------------------------------------------------------
【两种输出格式】
-----------------------------------------------------------------------------
    直接回答（direct）：   <s> 题目 ans 7 </s>
    思维链（cot）：        <s> 题目 c = a + 3 ; a = 4 ; a = 4 ; c = 4 + 3 = 7 ; ans 7 </s>
                                    └─ 倒着找：从被问的变量开始，──┘ └─ 再正着算：从起始值开始，──┘
                                       一条条找出它依赖的语句           一步步算出每个变量的值

思维链的每一步都只需要做一件很简单的事：找到一条语句并抄下来，或者做一次一位数加法。
直接回答则要求模型在**一次前向传播**里完成所有这些步骤（讲义第 4 节：固定深度的 Transformer
在一次前向传播里能做的"串行步骤"是有限的）。

【用法】
    import reasoning_tasks as R
    vocab = R.Vocab()
    p = R.make_problem(3, random.Random(0))
    ids, mask = R.encode(vocab, p, "cot")
    model = R.load_or_train("cot", vocab)          # 01 号文件训练并保存
    R.accuracy(model, vocab, "cot", k=3, n=200)

直接运行本文件会打印几个例子并做自我检查（约 1 秒）。
"""

import math
import random
import sys
import time
from pathlib import Path

import numpy as np
import torch

sys.path.append(str(Path(__file__).resolve().parent.parent / "lecture07"))
from transformer import TransformerLM, lm_loss, warmup_cosine  # noqa: E402

CKPT_DIR = Path(__file__).parent / "checkpoints"
VARS = list("abcdefghijklmnopqrst")
MAX_K = 4
MAX_LEN = 128
MODEL_CFG = dict(d_model=128, n_heads=4, n_layers=3, max_len=MAX_LEN, dropout=0.0)


# =============================================================================
# 第 1 部分：生成题目
# =============================================================================
def make_problem(k, rng):
    """
    生成一道链长为 k 的变量追踪题。

    【参数】k：每条链的赋值语句个数（1 ≤ k ≤ MAX_K）；rng：random.Random
    【返回】dict：
        question : 题目的词列表（含末尾的 "; 变量 ?"）
        backward : 思维链"倒着找"的部分（词列表）
        forward  : 思维链"正着算"的部分（词列表）
        answer   : 答案（字符串 "0"~"9"）
        values   : 被问的那条链上每个变量的值（按从起点到终点的顺序），03 号文件篡改思维链时用

    【边界】两条链的变量互不相同；每条链的起始值和常数都是 0~9 的随机数。
    """
    names = rng.sample(VARS, 2 * k)
    chains = [names[:k], names[k:]]
    statements, info = [], []
    for chain in chains:
        start = rng.randint(0, 9)
        consts = [rng.randint(0, 9) for _ in range(k - 1)]
        stmts = [[chain[0], "=", str(start)]]
        vals = [start]
        for i in range(1, k):
            stmts.append([chain[i], "=", chain[i - 1], "+", str(consts[i - 1])])
            vals.append((vals[-1] + consts[i - 1]) % 10)
        statements += stmts
        info.append((chain, stmts, vals, consts))
    rng.shuffle(statements)
    question = []
    for i, s in enumerate(statements):
        question += s + ([","] if i < len(statements) - 1 else [])
    chain, stmts, vals, consts = info[0]
    question += [";", chain[-1], "?"]
    backward = []
    for s in reversed(stmts):                         # 从被问的变量倒着找到起点
        backward += s + [";"]
    forward = [chain[0], "=", str(vals[0]), ";"]      # 从起点正着算
    for i in range(1, k):
        forward += [chain[i], "=", str(vals[i - 1]), "+", str(consts[i - 1]), "=", str(vals[i]), ";"]
    return {"question": question, "backward": backward, "forward": forward, "answer": str(vals[-1]),
            "values": vals, "k": k}


class Vocab:
    """词表：<pad> 0、<s> 1、</s> 2，然后是变量名、数字、符号和 "ans"。"""

    def __init__(self):
        self.itos = ["<pad>", "<s>", "</s>"] + VARS + [str(d) for d in range(10)] + ["=", "+", ",", ";", "?", "ans"]
        self.stoi = {w: i for i, w in enumerate(self.itos)}

    def __len__(self):
        return len(self.itos)

    def encode(self, words):
        return [self.stoi[w] for w in words]

    def decode(self, ids):
        return [self.itos[i] for i in ids]


def output_tokens(problem, fmt):
    """模型要生成的部分：直接回答是 "ans 7"，思维链是 "倒着找 + 正着算 + ans 7"。"""
    tail = ["ans", problem["answer"]]
    return tail if fmt == "direct" else problem["backward"] + problem["forward"] + tail


def encode(vocab, problem, fmt):
    """
    一道题 → (编号序列, 损失掩码)。序列 = <s> 题目 输出 </s>；只在"输出 </s>"上计算损失（和第 10 讲的 SFT 一样）。

    【返回】(ids, mask)：等长的列表，mask[j] = 1 表示 ids[j] 是要模型学会生成的
    """
    prompt = [vocab.stoi["<s>"]] + vocab.encode(problem["question"])
    out = vocab.encode(output_tokens(problem, fmt)) + [vocab.stoi["</s>"]]
    return prompt + out, [0] * len(prompt) + [1] * len(out)


# =============================================================================
# 第 2 部分：训练
# =============================================================================
def new_model(vocab, seed=0):
    torch.manual_seed(seed)
    return TransformerLM(len(vocab), **MODEL_CFG)


def batch(vocab, fmt, rng, size, ks=range(1, MAX_K + 1)):
    """随机生成一批题目（链长在 ks 里均匀抽），返回 (X, Y)：输入和目标（不计损失的位置目标为 0）。"""
    seqs = [encode(vocab, make_problem(rng.choice(list(ks)), rng), fmt) for _ in range(size)]
    T = max(len(s) for s, _ in seqs) - 1
    X = torch.zeros(size, T, dtype=torch.long)
    Y = torch.zeros(size, T, dtype=torch.long)
    for b, (ids, mask) in enumerate(seqs):
        X[b, :len(ids) - 1] = torch.tensor(ids[:-1])
        Y[b, :len(ids) - 1] = torch.tensor([t if m else 0 for t, m in zip(ids[1:], mask[1:])])
    return X, Y


def train(model, vocab, fmt, steps=6000, batch_size=64, lr=1e-3, seed=0, verbose=True, log_every=1000):
    """
    训练：每一步现场生成一批新题目（数据无限，不存在过拟合到某一批题的问题）。

    【参数】fmt："direct" 或 "cot"；steps：训练步数；其余同前几讲
    【返回】model（eval 模式）

    【边界】两种格式训练**同样多的题目、同样多的步数**。但思维链格式每道题有几十个词要预测，
        直接回答只有 2~3 个 —— 思维链额外提供了大量"中间步骤"的监督信号。这一点在比较时要记住（讲义第 4.3 节）。
    """
    rng = random.Random(seed)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    model.train()
    t0 = time.perf_counter()
    for step in range(steps):
        X, Y = batch(vocab, fmt, rng, batch_size)
        opt.param_groups[0]["lr"] = warmup_cosine(step, steps, min(200, steps // 10 + 1), lr)
        logits, _ = model(X)
        loss = lm_loss(logits, Y)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        if verbose and (step + 1) % log_every == 0:
            print(f"    [{fmt}] 第 {step + 1} 步：损失 {loss.item():.3f}，已用 {time.perf_counter() - t0:.0f} 秒", flush=True)
    model.eval()
    return model


def load_or_train(fmt, vocab, quick=False, verbose=True):
    """
    checkpoints/{fmt}.pt 存在就加载，否则训练并保存（quick 时只训练几十步、不保存）。

    【返回】TransformerLM（eval 模式）
    """
    path = CKPT_DIR / f"{fmt}.pt"
    model = new_model(vocab)
    if path.exists() and not quick:
        model.load_state_dict(torch.load(path))
        model.eval()
        if verbose:
            print(f"    （从 checkpoints/{fmt}.pt 加载了模型）")
        return model
    if verbose:
        print(f"    （训练 {fmt} 模型{'（冒烟测试：只训练很少几步）' if quick else ''}）")
    train(model, vocab, fmt, steps=40 if quick else 6000, verbose=verbose, log_every=20 if quick else 1000)
    if not quick:
        CKPT_DIR.mkdir(exist_ok=True)
        torch.save(model.state_dict(), path)
    return model


# =============================================================================
# 第 3 部分：生成与评估
# =============================================================================
@torch.no_grad()
def generate(model, vocab, prefixes, temperature=0.0, max_new=80, seed=0):
    """
    批量生成：从每个前缀出发逐词生成，直到 </s> 或 max_new 个词。

    【参数】
        prefixes    : list[list[int]]，每个前缀的编号
        temperature : 0 = 贪心（每步取概率最大的词）；> 0 = 按温度采样
        max_new     : 最多生成多少个词

    【返回】list[list[str]]：每个前缀**新生成**的词（不含 </s>）

    【思路】和第 9 讲 lm_utils.sample 一样：每一步把所有序列重新送进模型（简单但有重复计算）。
    """
    gen = torch.Generator().manual_seed(seed)
    seqs = [list(p) for p in prefixes]
    new = [[] for _ in prefixes]
    alive = list(range(len(prefixes)))
    eos = vocab.stoi["</s>"]
    for _ in range(max_new):
        if not alive:
            break
        cur = [seqs[i] for i in alive]
        if max(len(s) for s in cur) >= MAX_LEN:
            break
        T = max(len(s) for s in cur)
        X = torch.zeros(len(cur), T, dtype=torch.long)
        for b, s in enumerate(cur):
            X[b, :len(s)] = torch.tensor(s)
        logits, _ = model(X)
        z = logits[torch.arange(len(cur)), torch.tensor([len(s) - 1 for s in cur])]
        z[:, :2] = -float("inf")                                    # 不生成 <pad>、<s>
        if temperature > 0:
            nxt = torch.multinomial(torch.softmax(z / temperature, -1), 1, generator=gen).squeeze(1).tolist()
        else:
            nxt = z.argmax(-1).tolist()
        still = []
        for i, t in zip(alive, nxt):
            if t == eos:
                continue
            seqs[i].append(t)
            new[i].append(vocab.itos[t])
            still.append(i)
        alive = still
    return new


def parse_answer(tokens):
    """从生成的词里取出 "ans" 后面的那个数字；格式不对时返回 None。"""
    for i in range(len(tokens) - 1, 0, -1):
        if tokens[i - 1] == "ans" and tokens[i].isdigit():
            return tokens[i]
    return None


def prompt_ids(vocab, problem):
    """题目的前缀：<s> + 题目。"""
    return [vocab.stoi["<s>"]] + vocab.encode(problem["question"])


def accuracy(model, vocab, fmt, k, n=300, seed=123, temperature=0.0):
    """
    链长为 k 的 n 道新题上，模型（贪心解码）的准确率。

    【返回】(准确率, 题目列表, 生成结果列表)
    【边界】fmt 只用来说明模型是按哪种格式训练的；生成时模型自己决定写不写思维链。
    """
    rng = random.Random(seed + k)
    probs = [make_problem(k, rng) for _ in range(n)]
    outs = generate(model, vocab, [prompt_ids(vocab, p) for p in probs], temperature=temperature,
                    max_new=8 if fmt == "direct" else 80, seed=seed)
    acc = float(np.mean([parse_answer(o) == p["answer"] for o, p in zip(outs, probs)]))
    return acc, probs, outs


if __name__ == "__main__":
    vocab = Vocab()
    rng = random.Random(0)
    for k in [1, 2, 3]:
        p = make_problem(k, rng)
        print(f"k = {k}")
        print(f"  题目：  {' '.join(p['question'])}")
        print(f"  思维链：{' '.join(output_tokens(p, 'cot'))}")
    # 自我检查：思维链"正着算"的最后一个值就是答案；序列长度不超过 MAX_LEN
    for k in range(1, MAX_K + 1):
        for _ in range(200):
            p = make_problem(k, rng)
            assert p["forward"][-2] == p["answer"]
            ids, mask = encode(vocab, p, "cot")
            assert len(ids) <= MAX_LEN, len(ids)
    longest = max(len(encode(vocab, make_problem(MAX_K, rng), "cot")[0]) for _ in range(200))
    print(f"\n词表 {len(vocab)} 个；k = {MAX_K} 时思维链格式的序列最长约 {longest} 个词元（上限 {MAX_LEN}）")
    print("自我检查通过 ✓")