"""
tool_world.py —— 第 10 讲：花园集市的账本，和一个计算器工具
=====================================================================

【这个文件做什么】
    第 7 讲的集市模型算不了 20 以上的数、算不了三步；这一讲让语言模型学会**调用计算器**。
    这里提供：
        (1) 语料：一篇篇小账本，例如
                Ann has 1 7 apples . Bob gives Ann 2 5 apples . so Ann has 4 2 apples . that is 4 2 apples .
            数字按**单个数位**切成 token（"42" → 4 2），和真实 LLM 的分词器一样（Llama、GPT-4 的分词器都把
            数字切成一位或几位）。这样任何数字都能被表示、被照抄，但"算出"一个大数需要逐位进位。
            两成的账本带一句复述（"that is 4 2 apples ."），两成在总数前带一条方括号里的算式注释
            （"so Ann has [ 1 7 + 2 5 = 4 2 ] 4 2 apples ."）。预训练会因此学会"照抄上文刚出现的数"——
            第 6 讲说过，这种复制能力是很多下游本事的基础，Toolformer 的过滤也靠它（02 号脚本）。
        (2) 计算器：模型写出 [ 1 7 + 2 5 = ，我们暂停生成，用 Python 算出 4 2，接上 ] 后让模型继续——
            工具调用的格式和预训练里的算式注释一样，区别只在于"=" 后面的数由计算器给出，而不是模型自己算。
        (3) 训练（带缓存）、带工具的生成、评估。

【LLM 联系】
    Toolformer（Schick et al. 2023）的 API 调用格式是 "[Calculator(400 / 1400) → 0.29]"（这里用 = 代替 →）；
    今天的函数调用（function calling）把调用写成 JSON，由推理框架在生成到特殊 token 时暂停、执行、把结果拼回上下文。
"""


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", "lecture06"):
    sys.path.insert(0, os.path.join(HERE, "..", _d))

from garden_lm import Vocab, clm_batches, clm_loss, load_or_train, train  # noqa: E402

NAMES = ["Ann", "Bob", "Cid", "Dee", "Eve", "Fay", "Gus", "Hal", "Ida", "Jon"]
FRUITS = ["apples", "pears", "plums", "figs", "limes", "peaches", "grapes", "mangos"]
ACTS = {"gives": 1, "buys": 1, "eats": -1, "loses": -1}
CALL, ARROW, END = "[", "=", "]"
WORDS = ["has", "gives", "eats", "buys", "loses", "so", "that", "is", ".", ","]
DIGITS = [str(d) for d in range(10)]


def build_vocab():
    return Vocab(["<bos>", "<eos>", "<pad>"] + NAMES + FRUITS + WORDS + DIGITS + [CALL, ARROW, END, "+", "-"])


VOCAB = build_vocab()
BOS, EOS, PAD = 0, 1, 2


def num(n):
    """整数 → 数位 token 列表：42 → ['4', '2']。"""
    return list(str(n))


# =====================================================================
# 1. 账本
# =====================================================================

class Ledger:
    """
    一本账：主角 a 一开始有 start 个 fruit，经过若干步 (动作, 数, 给东西的人)，最后 "so a has 总数 fruit ."。

    【方法】
        prefix()：到 "so a has" 为止的词（模型要接着说出总数）；
        words(echo, call)：完整的账本；echo=True 时末尾加一句复述；call 是可选的 (表达式词列表, 结果)，
                           插在总数之前（Toolformer 式的增强数据）；
        answer()：总数；numbers()：账本里出现的数（按顺序）。
    """

    def __init__(self, a, fruit, start, steps):
        self.a, self.fruit, self.start, self.steps = a, fruit, start, steps

    def answer(self):
        v = self.start
        for act, x, _ in self.steps:
            v += ACTS[act] * x
        return v

    def numbers(self):
        return [self.start] + [x for _, x, _ in self.steps]

    def prefix(self):
        w = [self.a, "has"] + num(self.start) + [self.fruit, "."]
        for act, x, giver in self.steps:
            if act == "gives":
                w += [giver, "gives", self.a] + num(x) + [self.fruit, "."]
            else:
                w += [self.a, act] + num(x) + [self.fruit, "."]
        return w + ["so", self.a, "has"]

    def suffix(self, echo):
        w = num(self.answer()) + [self.fruit, "."]
        if echo:
            w += ["that", "is"] + num(self.answer()) + [self.fruit, "."]
        return w

    def words(self, echo=False, call=None, note=False):
        """
        【参数】echo：末尾加复述句；note：在总数之前加一条方括号里的算式注释 "[ 1 7 + 2 5 = 4 2 ]"；
                call：(表达式词列表, 结果)——和 note 同样的格式，但表达式和结果可以是任意的（02 号脚本的候选调用）。
        """
        w = self.prefix()
        if call is not None:
            expr, res = call
            w += [CALL] + expr + [ARROW] + num(res) + [END]
        elif note:
            w += [CALL] + expr_of(self) + [ARROW] + num(self.answer()) + [END]
        return w + self.suffix(echo)


def sample_ledger(rng, n_steps=None, lo=1, hi=20, max_total=99):
    """随机生成一本账：1～2 步（或指定步数），每个数在 [lo, hi]，所有中间结果在 [0, max_total]。"""
    n_steps = n_steps or rng.choice([1, 2])
    while True:
        a, giver = rng.sample(NAMES, 2)
        fruit = rng.choice(FRUITS)
        start = rng.randint(lo, hi)
        steps, v, ok = [], start, True
        for _ in range(n_steps):
            act = rng.choice(list(ACTS))
            x = rng.randint(lo, hi)
            v += ACTS[act] * x
            if not 0 <= v <= max_total:
                ok = False
                break
            steps.append((act, x, giver))
        if ok:
            return Ledger(a, fruit, start, steps)


def expr_of(ledger):
    """正确的计算器表达式（词列表）：1 7 + 2 5 - 3。"""
    w = num(ledger.start)
    for act, x, _ in ledger.steps:
        w += ["+" if ACTS[act] > 0 else "-"] + num(x)
    return w


def pretrain_ledger(rng):
    """
    预训练语料里的一本账：步数 1 / 2 / 3 的概率 0.45 / 0.45 / 0.10；85% 的账本每个数在 1～20，
    15% 在 1～49（总数最多 199）。所以大数、三步的账本模型都"见过长什么样"，只是见得少、算起来难。
    """
    n_steps = rng.choices([1, 2, 3], weights=[0.45, 0.45, 0.10])[0]
    if rng.random() < 0.85:
        return sample_ledger(rng, n_steps=n_steps)
    return sample_ledger(rng, n_steps=n_steps, lo=1, hi=49, max_total=199)


def corpus(n, seed, echo_p=0.2, note_p=0.2):
    """
    n 本账的预训练语料（每本是一篇文档）。

    【参数】echo_p：带复述句的比例（"that is 4 2 apples ."）；
            note_p：带方括号算式注释的比例——"so Ann has [ 1 7 + 2 5 = 4 2 ] 4 2 apples ."，
            就像网上的题解、代码注释、维基百科的换算 "(42 km)" 那样，把一段计算写在正文里。
            两者都让模型学会"照抄上文刚出现的数"；注释还让它见过"[ … = … ]"这种写法。
    """
    rng = random.Random(seed)
    out = []
    for _ in range(n):
        led = pretrain_ledger(rng)
        u = rng.random()
        out.append(led.words(echo=u < echo_p, note=echo_p <= u < echo_p + note_p) + ["<eos>"])
    return out


# =====================================================================
# 2. 计算器
# =====================================================================

def calculate(expr_words):
    """
    执行计算器：把 ['1','7','+','2','5'] 拼成 "17+25" 求值。

    【返回】整数结果；表达式不合法时返回 None（模型可能写出乱七八糟的调用）。
    【边界】只允许数位和 + -，不用 eval 执行任意代码。
    """
    s = "".join(expr_words)
    if not s or any(ch not in "0123456789+-" for ch in s) or s[0] in "+-" or s[-1] in "+-" or "++" in s or "--" in s \
            or "+-" in s or "-+" in s:
        return None
    total, sign, cur = 0, 1, ""
    for ch in s + "+":
        if ch.isdigit():
            cur += ch
        else:
            total += sign * int(cur)
            cur, sign = "", (1 if ch == "+" else -1)
    return total


# =====================================================================
# 3. 模型
# =====================================================================
ARCH = dict(d_model=128, n_heads=4, n_layers=4, d_ff=512)
CKPT_DIR = os.path.join(HERE, "checkpoints")


def new_model(seed=0):
    from transformer import TransformerLM
    torch.manual_seed(seed)
    return TransformerLM(len(VOCAB), attn="causal", max_len=256, **ARCH)


def base_model(notes=True, verbose=False):
    """
    基座：在 8000 本账本上从零预训练（没有真正的工具调用，计算器从没被执行过）。第一次约 50 秒，之后读缓存。

    【参数】notes：语料里有没有两成"方括号算式注释"。False 的基座用作 02 号脚本的对照：
            它没见过"[ … = … ]"，不会利用插进上下文里的计算结果。
    """
    docs = corpus(8000, seed=0, note_p=0.2 if notes else 0.0)
    batches = clm_batches(docs, VOCAB, 32)

    def fit(m):
        return train(m, batches, clm_loss, epochs=12, lr=1e-3, warmup=200, seed=0, verbose=verbose, tag="ledger")

    name = "ledger_base_notes" if notes else "ledger_base_plain"
    m, _ = load_or_train(name, {"arch": ARCH, "docs": 8000, "epochs": 12, "notes": notes, "v": 4}, new_model, fit,
                         verbose=verbose, ckpt_dir=CKPT_DIR)
    return m


def finetune_lm(model, docs, epochs=3, lr=5e-4, seed=0, verbose=False):
    """在一组文档上继续做语言模型训练（02 号脚本：在插入了工具调用的语料上微调）。"""
    batches = clm_batches(docs, VOCAB, 32)
    train(model, batches, clm_loss, epochs=epochs, lr=lr, warmup=50, seed=seed, verbose=verbose, tag="ft")
    model.eval()
    return model


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

@torch.no_grad()
def answer_with_tools(model, ledger, use_tool=True, tamper=None, max_new=24):
    """
    给模型账本的前缀（到 "so A has"），让它一直生成到水果名为止，读出它说的总数。

    【参数】use_tool：模型写出 [CALC … → 时，是否真的调用计算器（False 时让模型自己往下编结果）；
            tamper：可选函数 结果 → 篡改后的结果（模拟工具被攻击或出错）。
    【返回】dict：answer（模型说的数，读不出时为 None）、called（是否调用了工具）、expr（调用的表达式）、
            tool_result（计算器给的结果）、text（生成的全部词）。
    【思路】逐 token 贪心生成。每生成一个 token：如果是 → 并且前面有未闭合的 [CALC，就执行计算器，
            把结果的数位和 ] 直接接到序列里（这几个 token 不是模型生成的，而是工具"写"进来的）。
    """
    ids = [BOS] + VOCAB.encode(ledger.prefix())
    out, info = [], {"answer": None, "called": False, "expr": None, "tool_result": None}
    call_start = None
    for _ in range(max_new):
        t = int(model(torch.tensor([ids]))[0, -1].argmax())
        w = VOCAB.itos[t]
        ids.append(t)
        out.append(w)
        if w == CALL:
            call_start = len(out)
            info["called"] = True
        elif w == ARROW and call_start is not None and use_tool:
            expr = out[call_start:-1]
            res = calculate(expr)
            info["expr"] = expr
            if res is not None:
                if tamper is not None:
                    res = tamper(res)
                info["tool_result"] = res
                inject = num(res) + [END]
            else:
                inject = [END]
            ids += VOCAB.encode(inject)
            out += inject
            call_start = None
        elif w in FRUITS or w in ("<eos>", "."):
            break
    # 读出总数：水果名前面紧挨着的那一串数位（无论前面有没有算式、有没有工具调用）
    stop = next((i for i, w in enumerate(out) if w in FRUITS), len(out))
    digits = []
    for w in reversed(out[:stop]):
        if w in DIGITS:
            digits.insert(0, w)
        else:
            break
    info["answer"] = int("".join(digits)) if digits else None
    info["text"] = out
    return info


def accuracy(model, ledgers, **kw):
    """模型说出的总数和正确答案一致的比例，以及调用工具的比例。【返回】(准确率, 调用率, 每条的 info)。"""
    infos = [answer_with_tools(model, l, **kw) for l in ledgers]
    acc = sum(i["answer"] == l.answer() for i, l in zip(infos, ledgers)) / len(ledgers)
    called = sum(i["called"] for i in infos) / len(ledgers)
    return acc, called, infos


@torch.no_grad()
def target_loss(model, context_words, target_words):
    """在 context 之后生成 target 的总负对数似然（自然对数）。Toolformer 过滤用。"""
    ids = torch.tensor([[BOS] + VOCAB.encode(context_words + target_words)])
    logp = F.log_softmax(model(ids[:, :-1]), -1)[0]
    n_ctx = 1 + len(context_words)
    tgt = ids[0, n_ctx:]
    return -float(logp[n_ctx - 1:n_ctx - 1 + len(tgt)].gather(1, tgt[:, None]).sum())


@torch.no_grad()
def sample_calls(model, ledger, n=3, temperature=1.0, rng=None, max_len=14):
    """
    Toolformer 的第一步：让模型自己提出候选调用。在账本前缀后面接上 "["，按温度采样，直到写出 "=" 为止。

    【返回】去重后的表达式列表（词列表）；写不出 "=" 的样本丢掉。
    【思路】基座在预训练里见过"[ 算式 = 结果 ]"的注释，所以它会写出像样的算式——但算式对不对、写在哪里，
            要靠下一步的过滤来判断。
    """
    gen = torch.Generator().manual_seed(rng.randrange(10 ** 9) if rng else 0)
    prefix = [BOS] + VOCAB.encode(ledger.prefix() + [CALL])
    out = []
    for _ in range(n):
        ids, expr = list(prefix), []
        for _ in range(max_len):
            p = torch.softmax(model(torch.tensor([ids]))[0, -1] / temperature, -1)
            t = int(torch.multinomial(p, 1, generator=gen))
            w = VOCAB.itos[t]
            if w == ARROW:
                if expr and expr not in out:
                    out.append(expr)
                break
            expr.append(w)
            ids.append(t)
    return out


# =====================================================================
# 5. Toolformer：自监督地给语料加上工具调用
# =====================================================================

def annotate(base, n=4000, seed=2, n_samples=3, tau=1.0, echo_p=0.2):
    """
    Toolformer 的三步（Schick et al. 2023），对 n 本新的账本逐本处理：
        ① 采样：在总数的位置，让基座自己提出 n_samples 个候选调用（sample_calls）；
        ② 执行：用计算器算出每个候选的结果；
        ③ 过滤：比较基座预测"后面的总数"的损失——
              L⁺ = 插入"[ 表达式 = 结果 ]"之后的损失；
              L⁻ = min(不插入任何调用的损失, 插入"[ 表达式 = ]"（有调用没结果）的损失)；
           只有 L⁻ − L⁺ ≥ tau（结果确实让后面的词更好预测）时才保留，多个都满足就留 L⁺ 最小的那个。
    【返回】(docs, stats)：docs 是增强后的语料（保留了调用的账本插入调用，其余原样）；
            stats 记录候选数、表达式正确的候选数、保留下来的调用里结果正确 / 错误的个数。
    【为什么用 L⁻ 的两种情况取最小】既要排除"本来就能预测对、调用多此一举"的情况，
            也要排除"光是写出调用的格式就让预测变好、和结果无关"的情况。
    """
    rng, r_call = random.Random(seed), random.Random(seed + 1)
    docs = []
    stats = {"账本": 0, "候选": 0, "候选表达式正确": 0, "保留（结果正确）": 0, "保留（结果错误）": 0,
             "加了调用的账本": 0}
    for _ in range(n):
        led = pretrain_ledger(rng)
        echo = rng.random() < echo_p
        pre, ans = led.prefix(), num(led.answer())
        l_none = target_loss(base, pre, ans)
        best = None
        for e in sample_calls(base, led, n=n_samples, rng=r_call):
            stats["候选"] += 1
            stats["候选表达式正确"] += e == expr_of(led)
            res = calculate(e)
            if res is None or res < 0:
                continue
            l_empty = target_loss(base, pre + [CALL] + e + [ARROW, END], ans)
            l_res = target_loss(base, pre + [CALL] + e + [ARROW] + num(res) + [END], ans)
            if min(l_none, l_empty) - l_res >= tau:
                stats["保留（结果正确）" if res == led.answer() else "保留（结果错误）"] += 1
                if best is None or l_res < best[0]:
                    best = (l_res, e, res)
        stats["账本"] += 1
        if best is not None:
            stats["加了调用的账本"] += 1
            docs.append(led.words(echo=echo, call=(best[1], best[2])) + ["<eos>"])
        else:
            docs.append(led.words(echo=echo) + ["<eos>"])
    return docs, stats


def tool_model(docs=None, epochs=3, verbose=False):
    """
    会用计算器的模型：从基座（有注释的那个）出发，在 annotate() 产生的增强语料上继续做语言模型训练。

    【参数】docs：增强语料；不给就现场调用 annotate()（约 50 秒）。有缓存时直接读缓存，docs 被忽略。
    【返回】模型。缓存名 ledger_tool。
    """
    import copy

    def build():
        return copy.deepcopy(base_model(notes=True))

    def fit(m):
        d = docs if docs is not None else annotate(m)[0]
        finetune_lm(m, d, epochs=epochs, verbose=verbose)
        return {}

    m, _ = load_or_train("ledger_tool", {"arch": ARCH, "annotate": "n4000_s2_k3_tau1", "epochs": epochs, "v": 4},
                         build, fit, verbose=verbose, ckpt_dir=CKPT_DIR)
    return m


__all__ = ["VOCAB", "NAMES", "FRUITS", "ACTS", "CALL", "ARROW", "END", "DIGITS", "num", "Ledger", "sample_ledger",
           "expr_of", "corpus", "calculate", "new_model", "base_model", "finetune_lm", "answer_with_tools",
           "accuracy", "target_loss", "CKPT_DIR", "ARCH", "pretrain_ledger", "sample_calls", "annotate",
           "tool_model"]