"""
reasoning_world.py —— 第 7 讲：两个"标准答案已知"的小任务，用来检验模型是在推理还是在匹配模式
=====================================================================

【这个文件做什么】
    第 7 讲（Ondřej Bojar 的客座讲座）的主题是：评估一个模型"会不会推理"时，我们很容易一厢情愿。
    要把这件事讲清楚，需要能**精确控制**测试数据和训练数据之间的差别。这里造两个小世界：

    (1) 花园 → Floral 的"翻译"：
            输入若干句花园语，输出对应的 Floral（第 3 讲虚构的外语：逐词替换 + 词序改成 主-宾-谓）。
            用来做"序列长度本身就是一个领域"（Variš & Bojar 2021）的实验：只在某些长度上训练，看别的长度。
    (2) 花园集市的算术应用题（仿 GSM8K / GSM-Symbolic）：
            Ann has 7 apples . Bob gives Ann 5 apples . how many apples does Ann have ? <x0> 12
            用来做"换名字、换数字、加一句无关的话、多一步推理"的扰动评估，以及"数据泄漏"的实验。

    两个世界共用一个词表：前 100 个 token 和第 2 讲的扩展词表完全相同（<bos>=0、<eos>=1、<pad>=90、
    <x0>…<x7>），后面接 Floral 词、集市用词和数字 0～99。所以第 6 讲 garden_tasks.py 里的
    批次、微调、贪心生成函数可以直接拿来用（它们只依赖 <bos>、<eos>、<pad>、<x0> 的编号）。

【LLM 联系】
    GSM-Symbolic（Mirzadeh et al. 2025）正是这样做的：把 GSM8K 的题目改写成模板，名字和数字变成变量，
    再大量地重新实例化。我们在一个从零训练的小模型上重复它的几个实验——好处是我们完全知道训练数据里有什么。
"""

import os
import random
import sys

import torch

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

from garden_lm import BASE_ARCH, Vocab, ext_vocab  # noqa: E402
from toy_world import HELDOUT_PAIRS, GardenWorld  # noqa: E402
from web_world import FLORAL, to_floral  # noqa: E402

WORLD = GardenWorld(HELDOUT_PAIRS)

# =====================================================================
# 1. 词表
# =====================================================================
NAMES = ["Ann", "Bob", "Cid", "Dee", "Eve", "Fay", "Gus", "Hal", "Ida", "Jon"]
FRUITS = ["apples", "pears", "plums", "figs", "limes", "peaches", "grapes", "mangos"]
MARKET_WORDS = ["has", "gives", "eats", "buys", "loses", "how", "many", "does", "have", "?", "of", "them",
                "are", "red", "small", "then", "and", "more", "<tool>", "</tool>", "+", "-", "=", "<calc>"]
NUMBERS = [str(i) for i in range(100)]
ROLE_TAGS = ["T:subj", "T:verb", "T:det", "T:adj", "T:noun", "T:poss", "T:tail", "T:end"]   # 04 号脚本的"辅助标签"
FLORAL_WORDS = sorted(set(FLORAL.values()) - {"."})


def build_vocab():
    """
    【返回】Vocab：第 2 讲扩展词表（100 个）+ Floral 词 + 名字 + 水果复数 + 集市用词 + 数字 0～99 + 8 个角色标签。
    【边界】Floral 词、英文名字、数字都和花园词不重名（花园词表里没有大写名字和数字）。
    """
    base = ext_vocab().itos
    extra = []
    for w in FLORAL_WORDS + NAMES + FRUITS + MARKET_WORDS + NUMBERS + ROLE_TAGS:
        if w not in base and w not in extra:
            extra.append(w)
    return Vocab(base + extra)


VOCAB = build_vocab()
SEP = "<x0>"


def new_model(seed=0, **overrides):
    """
    一个随机初始化的因果 Transformer（第 2 讲的 TransformerLM，默认 d=64、2 层、4 头、正弦位置编码），
    词表是本文件的 VOCAB。本讲所有模型都从零训练。
    """
    from transformer import TransformerLM
    arch = dict(BASE_ARCH, **overrides)
    torch.manual_seed(seed)
    return TransformerLM(len(VOCAB), attn="causal", max_len=512, **arch)


def example(src_words, tgt_words):
    """(prompt_ids, target_ids)：<bos> 输入 <x0> | 输出 <eos>。格式和第 6 讲相同。"""
    return ([0] + VOCAB.encode(src_words) + [VOCAB.stoi[SEP]], VOCAB.encode(tgt_words) + [1])


# =====================================================================
# 2. 花园 → Floral 翻译
# =====================================================================

def translation_pair(n_sent, rng):
    """
    采样 n_sent 句花园语，返回 (源词列表, Floral 词列表)。

    【思路】每句独立按真实分布采样，再逐句 to_floral（逐词替换 + 动词挪到宾语之后）。
            多句时源和目标都是各句首尾相接，句子之间没有任何依赖——
            所以"翻译一篇 5 句的文档"在原理上和"翻译 5 次单句"完全一样难，难的只是"长度没见过"。
    """
    src, tgt = [], []
    for _ in range(n_sent):
        s = WORLD.sample_sentence(rng)
        src += s
        tgt += to_floral(s)
    return src, tgt


def floral_with_roles(sent):
    """
    一句花园语的 Floral 译文，以及译文里每个词在原句中的"角色"（主语、动词、限定词……，见第 6 讲 garden_tasks.parse）。

    【返回】(floral_words, roles)：两个等长的列表；roles 的取值是 subj / verb / det / adj / noun / poss / tail / end。
    【思路】to_floral 把动词挪到宾语名词之后；这里对角色列表做同样的重排，保证一一对应。
    【用途】04 号脚本的"辅助标签"实验：在译文的每个词后面插入它的角色标签，看这种"语言学信息"有没有帮助。
    """
    from garden_tasks import parse
    from web_world import _NOUNS, THIRD_TO_BASE, VERBS
    roles = parse(sent)
    words, rl = sent[:-1], roles[:-1]
    vi = next(i for i, w in enumerate(words) if w in THIRD_TO_BASE or w in VERBS)
    j = vi + 1
    while j < len(words) and words[j] not in _NOUNS:
        j += 1
    order = list(range(vi)) + list(range(vi + 1, j + 1)) + [vi] + list(range(j + 1, len(words)))
    out = [FLORAL[words[i]] for i in order] + ["."]
    assert out == to_floral(sent)
    return out, [rl[i] for i in order] + ["end"]


def translation_data(n, lengths, seed):
    """
    n 条翻译样本，每条的句子数从 lengths（列表）里均匀地选。

    【返回】[(prompt_ids, target_ids)]。
    """
    rng = random.Random(seed)
    out = []
    for _ in range(n):
        k = rng.choice(lengths)
        out.append(example(*translation_pair(k, rng)))
    return out


# =====================================================================
# 3. 花园集市的算术应用题
# =====================================================================

class Problem:
    """
    一道应用题的"符号模板 + 具体取值"。

    【字段】
        names：出场的人（第一个是主角，问题问的是主角）；fruit：问的那种水果；
        steps：[(动作, 数), ...]，动作取 "gives"（别人给主角，+）、"eats"（主角吃掉，−）、"buys"（+）、"loses"（−）；
        start：主角一开始有几个；noop：可选的"无关的一句话"（词列表）；
        noop_pos："first" 插在第一句之后，"last" 插在问句之前。
    【方法】words() 生成题目文字；answer() 给出正确答案。
    """

    def __init__(self, names, fruit, start, steps, noop=None, ask_fruit=None, noop_pos="first"):
        self.names, self.fruit, self.start, self.steps = names, fruit, start, steps
        self.noop, self.noop_pos = noop, noop_pos
        self.ask_fruit = ask_fruit or fruit

    def words(self):
        a = self.names[0]
        w = [a, "has", str(self.start), self.fruit, "."]
        if self.noop and self.noop_pos == "first":
            w += self.noop
        for i, (act, x) in enumerate(self.steps):
            if act == "gives":
                giver = self.names[1 + i % (len(self.names) - 1)]
                w += [giver, "gives", a, str(x), self.fruit, "."]
            else:
                w += [a, act, str(x), self.fruit, "."]
        if self.noop and self.noop_pos == "last":
            w += self.noop
        w += ["how", "many", self.ask_fruit, "does", a, "have", "?"]
        return w

    def answer(self):
        if self.ask_fruit != self.fruit:
            return 0                                   # 问的是一种从没提到过的水果
        v = self.start
        for act, x in self.steps:
            v += x if act in ("gives", "buys") else -x
        return v

    def numbers(self):
        """题目里出现的所有数（包括无关句里的），按出现顺序。"""
        return [int(t) for t in self.words() if t.isdigit()]


def sample_problem(rng, n_steps=1, lo=1, hi=20, names=NAMES, fruits=FRUITS, max_answer=99, noop=False,
                   noop_kind=None, noop_pos="first"):
    """
    随机生成一道 n_steps 步的题。

    【参数】lo / hi：每个数的取值范围（含两端）；names / fruits：可选的名字和水果；
            max_answer：答案（以及中间结果）不超过它、不小于 0；noop：是否插入一句无关的话；
            noop_kind："red"、"other" 或 None（各一半）；noop_pos：无关句的位置（见 Problem）。
    【思路】反复抽样，直到所有中间结果都在 [0, max_answer] 之内（题目必须"合理"：吃掉的不能比有的多）。
    【无关的话】两种，各一半："K of them are red ."（GSM-NoOp 式：看起来相关，其实不影响答案）
                和 "C has K pears ."（别人、别的水果）。
    """
    while True:
        ns = rng.sample(names, 3)
        fruit = rng.choice(fruits)
        start = rng.randint(lo, hi)
        steps, v, ok = [], start, True
        for _ in range(n_steps):
            act = rng.choice(["gives", "eats", "buys", "loses"])
            x = rng.randint(lo, hi)
            v = v + x if act in ("gives", "buys") else v - x
            if not 0 <= v <= max_answer:
                ok = False
                break
            steps.append((act, x))
        if not ok:
            continue
        nop = None
        if noop:
            k = rng.randint(lo, hi)
            kind = noop_kind or ("red" if rng.random() < 0.5 else "other")
            if kind == "red":
                nop = [str(k), "of", "them", "are", rng.choice(["red", "small"]), "."]
            else:
                other = rng.choice([f for f in fruits if f != fruit])
                nop = [ns[2], "has", str(k), other, "."]
        return Problem(ns, fruit, start, steps, nop, noop_pos=noop_pos)


def problem_example(p):
    """(prompt_ids, target_ids)：题目 <x0> 答案 <eos>。答案是一个数字 token。"""
    return example(p.words(), [str(p.answer())])


def reinstantiate(p, rng, change_names=True, change_numbers=True, lo=1, hi=20, max_answer=99):
    """
    GSM-Symbolic 式的重新实例化：保持模板（步数、每步的动作、是否有无关句）不变，换名字和 / 或数字。

    【返回】新的 Problem（保证答案和中间结果在合理范围内）。
    """
    while True:
        names = rng.sample(NAMES, 3) if change_names else p.names
        start = rng.randint(lo, hi) if change_numbers else p.start
        steps = [(a, rng.randint(lo, hi) if change_numbers else x) for a, x in p.steps]
        v, ok = start, True
        for a, x in steps:
            v = v + x if a in ("gives", "buys") else v - x
            ok &= 0 <= v <= max_answer
        if ok:
            noop = None
            if p.noop is not None:
                noop = list(p.noop)
                if change_numbers:
                    noop[0 if p.noop[1] == "of" else 2] = str(rng.randint(lo, hi))
                if change_names and p.noop[1] == "has":
                    noop[0] = names[2]
            return Problem(names, p.fruit, start, steps, noop, noop_pos=p.noop_pos)


# 三个"不看懂题目也能答"的基线（第 7 讲 slides："Always use trivial baselines"）
def baseline_first(p):
    """答案 = 题目里的第一个数。"""
    return p.numbers()[0]


def baseline_sum(p):
    """答案 = 题目里所有数之和（不管是给还是吃）。"""
    return sum(p.numbers())


def baseline_keyword(p):
    """答案 = 按关键词加减：has / gives / buys 加，eats / loses 减；无关句里的数也照算（看见数就用）。"""
    w = p.words()
    v = 0
    for i, t in enumerate(w):
        if t.isdigit():
            prev = w[i - 1]
            v += -int(t) if prev in ("eats", "loses") else int(t)
    return v


@torch.no_grad()
def predict_numbers(model, problems, batch=256):
    """
    模型对每道题的答案：<x0> 之后第一个位置上，在 0～99 这 100 个数字 token 里取概率最大的。

    【返回】list[int]。
    【为什么只在数字里取 argmax】答案永远是一个数；限制在数字 token 里，避免"格式错误"和"算错"混在一起。
    """
    num_ids = torch.tensor([VOCAB.stoi[str(i)] for i in range(100)])
    out = []
    from garden_tasks import task_batch
    for s in range(0, len(problems), batch):
        chunk = [problem_example(p) for p in problems[s:s + batch]]
        ids, pad, _ = task_batch(chunk)
        logits = model(ids, pad=pad)
        for i, (pr, _) in enumerate(chunk):
            out.append(int(logits[i, len(pr) - 1, num_ids].argmax()))
    return out


def accuracy(preds, problems):
    return sum(a == p.answer() for a, p in zip(preds, problems)) / len(problems)


# =====================================================================
# 4. 集市题的标准数据与模型
# =====================================================================

MARKET_ARCH = dict(d_model=128, n_heads=4, n_layers=4, d_ff=512)    # 比第 2 讲的模型大：算术比花园语法难
CKPT_DIR = os.path.join(HERE, "checkpoints")


def market_split():
    """
    第 7 讲统一使用的集市题：训练 6000、开发 200、测试 300（"原题"），都是 1 步或 2 步、数字 1～20。

    【返回】(train, dev, test)：Problem 列表。同一个种子，每次调用结果相同。
    【边界】训练集里从来没有：无关的句子、3 步的题、大于 20 的运算数、问一种没提到的水果。
    """
    rng = random.Random(0)
    train = [sample_problem(rng, n_steps=rng.choice([1, 2])) for _ in range(6000)]
    dev = [sample_problem(rng, n_steps=rng.choice([1, 2])) for _ in range(200)]
    test = [sample_problem(rng, n_steps=rng.choice([1, 2])) for _ in range(300)]
    return train, dev, test


def train_market(name, problems, dev, steps=4000, lr=1e-3, seed=0, verbose=False):
    """
    在一组集市题上从零训练（或读缓存）一个 MARKET_ARCH 模型。

    【参数】name：缓存名，必须能唯一地说明训练数据是怎样构造的（数据本身不进哈希）；
            problems：训练题；dev：挑检查点用的开发题；steps / lr / seed：训练超参数。
    【返回】模型（eval 模式）。第一次约 70 秒，之后读本讲 checkpoints/ 里的缓存。
    【边界】改了数据的构造方式但没改 name 时，缓存不会失效——要手动删除 checkpoints/ 里对应的文件。
    """
    from garden_lm import load_or_train
    from garden_tasks import finetune
    cfg = {"arch": MARKET_ARCH, "steps": steps, "lr": lr, "seed": seed, "data": name, "n": len(problems),
           "vocab": len(VOCAB)}

    def build():
        return new_model(seed, **MARKET_ARCH)

    def fit(m):
        return finetune(m, list(m.parameters()), [problem_example(p) for p in problems], steps=steps, lr=lr,
                        warmup=100, seed=seed, dev_ex=dev,
                        eval_fn=lambda mm, ex: accuracy(predict_numbers(mm, ex), ex), eval_every=500)

    model, _ = load_or_train(name, cfg, build, fit, verbose=verbose, ckpt_dir=CKPT_DIR)
    return model


__all__ = ["VOCAB", "WORLD", "NAMES", "FRUITS", "SEP", "new_model", "example", "translation_pair",
           "translation_data", "Problem", "sample_problem", "problem_example", "reinstantiate",
           "baseline_first", "baseline_sum", "baseline_keyword", "predict_numbers", "accuracy", "FLORAL",
           "ROLE_TAGS", "MARKET_ARCH", "market_split", "train_market", "floral_with_roles"]