"""
toy_world.py —— 第 1 讲的"花园世界"（Garden World），之后几讲也会用
=====================================================================

【这个文件要回答的问题】
    要讲清楚"语言模型是什么"，我们需要一份语料，而且最好**确切知道**它背后的
    真实概率分布 p_true。原因有三个：
    (1) 可以精确算出任意一句话的真实概率，验证"链式法则"；
    (2) 可以算出"理论上最好的模型"能达到的困惑度（任何模型都不可能更好），
        给 n-gram 和神经网络的成绩一个参照；
    (3) 可以判断模型生成的一句话"合不合法"（真实概率是否为 0），
        这样第 1 讲比较各种解码策略时，"质量"就能被精确测量，而不靠肉眼。
    真实英语做不到这三点，所以我们自己造一门小语言。

【这门语言长什么样】
    一篇"文档"由若干句话组成，最后以 <eos>（end of sequence，序列结束）收尾。
    每句话的结构是：

        <主语>  <动词>  <限定词>  [<形容词>]  <名词>  [<尾巴短语>]  .

    例如（第 1 讲 slides 里的那句话就是一个合法句子）：
        I eat the apple from my own garden .
        she picks her ripe plum .
        the children feed the hungry goat every day .

    里面藏着几条规律（"语法"），模型要从数据里自己学会：
    - 主谓一致：she / he / my friend 后面的动词要加 s（eats），其余主语不加（eat）。
    - 物主代词一致：限定词可以是 the，也可以是"自己的"（I → my，she → her，
      the children → their……），必须和主语一致。
    - 动词决定名词的类别：eat 后面是水果或蔬菜，pick 后面是水果或花，
      feed 后面是动物……（"I feed the apple" 不合法）。
    - 形容词决定名词的类别：ripe / sweet / red 只修饰水果，hungry / little 只修饰动物。
    - 名词决定尾巴：植物可以说 "from my own garden"，动物不行。
      尾巴里的 my / her / their 又要和**句首的主语**一致——这是一条
      跨度 5～7 个词的"长程依赖"，第 1 讲用它展示固定窗口模型的局限。
    - 每个主语有自己的"爱好"（the children 更爱喂动物，she 更爱摘东西）；
      每类名词内部的频率服从 Zipf 定律（apple 很常见，quince 很少见）。

【真实概率是怎么算的】
    生成过程是一个"有限状态机"：每一步处在某个状态，按该状态下的分布选下一个词，
    然后转移到新状态。这门语言经过专门设计，**任何一句话只有唯一的生成路径**
    （各类词互不重叠），所以：
        p_true(句子) = 每一步所选词的条件概率的乘积   ← 就是链式法则
    `next_dist(state)` 给出"在当前状态下，下一个词的精确分布"；
    `logprob(tokens)` 把整句话沿状态机走一遍，把每步的 log 概率加起来。

【怎样使用】
    from toy_world import GardenWorld, standard_splits
    world = GardenWorld()
    world.logprob("I eat the apple .".split() + ["<eos>"])   # 精确的 log 概率
    splits = standard_splits()      # {"train": [...], "valid": [...], "test": [...]}
    直接运行本文件会打印几篇样例文档和语料统计：python3 toy_world.py

【LLM 联系】
    真实的 LLM 研究也常用"人造语言 / 合成数据"做受控实验（例如 Physics of Language
    Models 系列用合成的传记和语法来研究 LLM 怎样存储知识、学习语法）。好处是
    我们知道正确答案，能把"模型学到了什么"和"数据里有什么"分开来研究。
"""

import math
import random
from collections import namedtuple

# =====================================================================
# 1. 词表与"世界设定"
# =====================================================================

# 主语：(主语的词列表, 动词是否要加 s, 对应的物主代词)
# 【Python 小知识】元组（tuple）用圆括号写，创建后不能修改。
#   ("I",) 末尾的逗号不能省：("I") 只是一个加了括号的字符串，("I",) 才是只有一个元素的元组。
SUBJECTS = [
    (("I",), False, "my"),
    (("you",), False, "your"),
    (("we",), False, "our"),
    (("they",), False, "their"),
    (("she",), True, "her"),
    (("he",), True, "his"),
    (("the", "children"), False, "their"),
    (("my", "friend"), True, "their"),
]
# 每个主语出现的相对频率（和 SUBJECTS 一一对应，加起来等于 1）
SUBJECT_WEIGHTS = [0.22, 0.12, 0.12, 0.10, 0.12, 0.12, 0.10, 0.10]

# 动词原形 -> 第三人称单数形式
VERBS = {"eat": "eats", "pick": "picks", "grow": "grows", "water": "waters", "feed": "feeds"}

# 每个主语的"爱好"：选各个动词的概率（顺序与 SUBJECTS 对应，每行加起来等于 1）
# 【Python 小知识】字典（dict）用花括号写，"键: 值"成对出现；从 Python 3.7 起，
#   字典会记住插入顺序，所以遍历顺序就是书写顺序（这对可复现性很重要）。
VERB_PREFS = [
    {"eat": 0.35, "pick": 0.20, "grow": 0.20, "water": 0.15, "feed": 0.10},  # I
    {"eat": 0.30, "pick": 0.25, "grow": 0.15, "water": 0.15, "feed": 0.15},  # you
    {"eat": 0.20, "pick": 0.10, "grow": 0.35, "water": 0.25, "feed": 0.10},  # we
    {"eat": 0.30, "pick": 0.20, "grow": 0.10, "water": 0.10, "feed": 0.30},  # they
    {"eat": 0.10, "pick": 0.35, "grow": 0.20, "water": 0.25, "feed": 0.10},  # she
    {"eat": 0.30, "pick": 0.10, "grow": 0.20, "water": 0.10, "feed": 0.30},  # he
    {"eat": 0.20, "pick": 0.30, "grow": 0.05, "water": 0.05, "feed": 0.40},  # the children
    {"eat": 0.20, "pick": 0.10, "grow": 0.30, "water": 0.30, "feed": 0.10},  # my friend
]

# 四类名词。每类内部按 Zipf 定律分配频率：排第 r 位的词，权重 ∝ 1/r。
NOUNS = {
    "fruit": ["apple", "pear", "plum", "cherry", "peach", "grape", "lemon", "mango",
              "fig", "melon", "kiwi", "lime", "orange", "apricot", "quince"],
    "veg": ["carrot", "bean", "potato", "tomato", "onion", "pumpkin", "cabbage",
            "radish", "lettuce", "leek", "beet", "turnip"],
    "flower": ["rose", "tulip", "daisy", "lily", "poppy", "iris", "violet", "lilac"],
    "animal": ["cat", "dog", "hen", "duck", "goat", "rabbit", "pony", "goose"],
}
# 每类名词专属的形容词（各类之间互不重叠，所以看到形容词就知道名词属于哪一类）
ADJECTIVES = {
    "fruit": ["ripe", "sweet", "red"],
    "veg": ["fresh", "green"],
    "flower": ["pretty", "white"],
    "animal": ["hungry", "little"],
}
# 动词 -> 宾语名词类别的分布（eat 只吃水果蔬菜，feed 只喂动物……）
VERB_CATEGORIES = {
    "eat": {"fruit": 0.7, "veg": 0.3},
    "pick": {"fruit": 0.6, "flower": 0.4},
    "grow": {"fruit": 0.3, "veg": 0.4, "flower": 0.3},
    "water": {"veg": 0.4, "flower": 0.6},
    "feed": {"animal": 1.0},
}
P_THE = 0.6        # 限定词用 the 的概率；其余 0.4 用和主语一致的物主代词
P_ADJ = 0.3        # 名词前面加形容词的概率
P_EOS = 0.3        # 每句话结束后，整篇文档结束的概率（否则接着写下一句）

# 尾巴短语：名字 -> (词模板, 植物类的概率, 动物类的概率)
# 模板里的 "POSS" 会被替换成和主语一致的物主代词。
TAILS = {
    "none": ((".",), 0.45, 0.50),
    "garden": (("from", "POSS", "own", "garden", "."), 0.25, 0.00),
    "morning": (("in", "the", "morning", "."), 0.10, 0.20),
    "daily": (("every", "day", "."), 0.10, 0.20),
    "friends": (("with", "POSS", "friends", "."), 0.10, 0.10),
}

BOS, EOS, UNK = "<bos>", "<eos>", "<unk>"
SPECIALS = [BOS, EOS, UNK]


def zipf_weights(n, s=1.0):
    """
    返回长度为 n 的 Zipf 分布（第 r 个元素的概率 ∝ 1 / r^s），加起来等于 1。

    【参数】n：元素个数（>= 1）；s：Zipf 指数，越大越"头重脚轻"。
    【返回】list[float]，从大到小排列。
    【边界】n = 1 时返回 [1.0]。
    【LLM 联系】真实语料的词频大致服从 Zipf 定律：少数词极其常见，大量词极其罕见。
              这条"长尾"正是 n-gram 稀疏性和解码"长尾问题"的根源。
    """
    raw = [1.0 / (r ** s) for r in range(1, n + 1)]
    total = math.fsum(raw)               # math.fsum 是精确求和，比 sum 更不怕舍入误差
    return [x / total for x in raw]


# 预先算好每类名词的分布：{"fruit": {"apple": 0.30, "pear": 0.15, ...}, ...}
# 【Python 小知识】这是"字典推导式"：{键: 值 for ... in ...}，一行建出整个字典。
#   zip(a, b) 把两个列表按位置配对：zip(["x","y"], [1,2]) -> ("x",1), ("y",2)。
NOUN_DIST = {cat: dict(zip(words, zipf_weights(len(words)))) for cat, words in NOUNS.items()}
ADJ_DIST = {cat: {a: 1.0 / len(adjs) for a in adjs} for cat, adjs in ADJECTIVES.items()}
NOUN_CAT = {w: cat for cat, words in NOUNS.items() for w in words}   # 名词 -> 类别
ADJ_CAT = {a: cat for cat, adjs in ADJECTIVES.items() for a in adjs}  # 形容词 -> 类别
VERB_LEMMA = {}                                                      # 动词任意形式 -> 原形
for _base, _third in VERBS.items():
    VERB_LEMMA[_base] = _base
    VERB_LEMMA[_third] = _base


def all_words():
    """
    返回这门语言里出现的全部词（不含特殊符号），按固定顺序排列。

    【返回】list[str]，没有重复。
    【思路】把主语、动词、限定词、形容词、名词、尾巴里的词依次收集起来，
            用一个 set 去重，但输出顺序按"第一次出现"的顺序（保证每次运行都一样）。
    """
    words = []
    seen = set()

    def add(w):
        if w not in seen and w != "POSS":
            seen.add(w)
            words.append(w)

    for toks, _, poss in SUBJECTS:
        for w in toks:
            add(w)
    for base, third in VERBS.items():
        add(base)
        add(third)
    add("the")
    for _, _, poss in SUBJECTS:
        add(poss)
    for cat in NOUNS:
        for a in ADJECTIVES[cat]:
            add(a)
        for n in NOUNS[cat]:
            add(n)
    for toks, _, _ in TAILS.values():
        for w in toks:
            add(w)
    return words


# =====================================================================
# 2. 有限状态机：状态、下一个词的分布、状态转移
# =====================================================================

# 状态用一个"具名元组"表示，字段含义：
#   stage：处在句子的哪个位置（见下面 next_dist 里的说明）
#   subj ：主语的编号（SUBJECTS 的下标），还没选主语时是 -1
#   verb ：动词原形，还没选时是 ""
#   cat  ：名词类别（看到形容词或名词后才确定），否则是 ""
#   rest ：尾巴短语里还没说完的词（元组）
# 【Python 小知识】namedtuple 生成一个"带字段名的元组"类：
#   s = State("V", 0, "", "", ())，之后可以写 s.stage、s.subj，比 s[0]、s[1] 好读。
#   它和普通元组一样不可修改、可以当字典的键（我们用它做缓存）。
#   要"修改"某个字段，用 s._replace(stage="D") 得到一个新的状态。
State = namedtuple("State", ["stage", "subj", "verb", "cat", "rest"])

DOC_START = State("S0", -1, "", "", ())


class GardenWorld:
    """
    花园世界的"真实语言模型"：给出任意前缀之后下一个词的**精确**概率。

    【思路】
        stage 的取值和含义：
          "S0"   文档开头，下一个词是某个主语的第一个词
          "S1"   多词主语（the children / my friend）说了一半，下一个词确定
          "V"    下一个词是动词（形式由主语决定）
          "D"    下一个词是限定词：the 或物主代词
          "A"    下一个词是形容词或名词
          "N"    刚说完形容词，下一个词是同类名词
          "T"    刚说完名词，下一个词是句号或某个尾巴短语的第一个词
          "R"    尾巴短语说了一半，下一个词确定（rest[0]）
          "E"    刚说完句号：要么 <eos> 结束文档，要么开始下一句
          "DONE" 已经 <eos>，不会再有词
    【参数】
        heldout_pairs：一组 (动词原形, 名词) 组合。它们**只影响** `sample_doc(..., exclude_heldout=True)`：
                       生成训练语料时把含这些组合的句子扔掉，模拟"训练数据恰好没见过这种搭配"。
                       next_dist / logprob 永远按完整的真实分布计算。
    """

    def __init__(self, heldout_pairs=()):
        self.heldout_pairs = set(heldout_pairs)
        self._cache = {}                 # 状态 -> 下一个词的分布（同一状态反复出现，缓存能大大提速）

    # ---------------- 下一个词的分布 ----------------
    def next_dist(self, state):
        """
        给定当前状态，返回下一个词的精确分布。

        【参数】state：State 具名元组。
        【返回】dict：{词: 概率}，概率都 > 0，加起来等于 1（浮点误差内）。
                已经结束（stage == "DONE"）时返回空字典 {}。
        【边界】返回的字典是缓存里的同一个对象，调用者**不要修改**它。
        """
        if state in self._cache:
            return self._cache[state]
        dist = self._compute_dist(state)
        self._cache[state] = dist
        return dist

    def _subject_first_words(self, scale):
        """所有主语的第一个词及其概率（乘上 scale）。the children 和 my friend 的第一个词分别是 the 和 my。"""
        out = {}
        for (toks, _, _), w in zip(SUBJECTS, SUBJECT_WEIGHTS):
            out[toks[0]] = out.get(toks[0], 0.0) + scale * w
        return out

    def _compute_dist(self, s):
        if s.stage == "S0":
            return self._subject_first_words(1.0)
        if s.stage == "E":
            dist = {EOS: P_EOS}
            dist.update(self._subject_first_words(1.0 - P_EOS))
            return dist
        if s.stage == "S1":
            return {SUBJECTS[s.subj][0][1]: 1.0}
        if s.stage == "V":
            _, third_person, _ = SUBJECTS[s.subj]
            return {(VERBS[v] if third_person else v): p for v, p in VERB_PREFS[s.subj].items()}
        if s.stage == "D":
            poss = SUBJECTS[s.subj][2]
            return {"the": P_THE, poss: 1.0 - P_THE}
        if s.stage == "A":
            # 形容词：先按动词选类别 c，再在 c 的形容词里均匀选；名词同理。
            # 各类的形容词 / 名词互不重叠，所以同一个词只会从一个类别里来，直接相乘即可。
            dist = {}
            for cat, pc in VERB_CATEGORIES[s.verb].items():
                for a, pa in ADJ_DIST[cat].items():
                    dist[a] = P_ADJ * pc * pa
            for cat, pc in VERB_CATEGORIES[s.verb].items():
                for n, pn in NOUN_DIST[cat].items():
                    dist[n] = (1.0 - P_ADJ) * pc * pn
            return dist
        if s.stage == "N":
            return dict(NOUN_DIST[s.cat])
        if s.stage == "T":
            dist = {}
            for toks, p_plant, p_animal in TAILS.values():
                p = p_animal if s.cat == "animal" else p_plant   # 动物类用第 3 列的概率
                if p > 0:
                    dist[toks[0]] = dist.get(toks[0], 0.0) + p
            return dist
        if s.stage == "R":
            return {s.rest[0]: 1.0}
        if s.stage == "DONE":
            return {}
        raise ValueError(f"未知的状态 {s}")

    # ---------------- 状态转移 ----------------
    def step(self, s, word):
        """
        在状态 s 下说出词 word 之后，转移到哪个新状态。

        【参数】s：当前状态；word：刚说出的词。
        【返回】新的 State；如果 word 在状态 s 下概率为 0（不合法），返回 None。
        【思路】先查 next_dist 确认合法，再按 stage 分情况更新字段。
        """
        if word not in self.next_dist(s):
            return None
        st = s.stage
        if st in ("S0", "E"):
            if word == EOS:
                return State("DONE", -1, "", "", ())
            # 找到以 word 开头的主语（I / you / ... / the children / my friend）
            for i, (toks, _, _) in enumerate(SUBJECTS):
                if toks[0] == word:
                    if len(toks) == 1:
                        return State("V", i, "", "", ())
                    return State("S1", i, "", "", ())
        if st == "S1":
            return s._replace(stage="V")
        if st == "V":
            return s._replace(stage="D", verb=VERB_LEMMA[word])
        if st == "D":
            return s._replace(stage="A")
        if st == "A":
            if word in ADJ_CAT:
                return s._replace(stage="N", cat=ADJ_CAT[word])
            return s._replace(stage="T", cat=NOUN_CAT[word])
        if st == "N":
            return s._replace(stage="T")
        if st in ("T", "R"):
            if st == "T":
                # 找到以 word 开头的尾巴，把剩下的词放进 rest（POSS 换成主语的物主代词）
                poss = SUBJECTS[s.subj][2]
                for toks, _, _ in TAILS.values():
                    if toks[0] == word:
                        rest = tuple(poss if w == "POSS" else w for w in toks[1:])
                        break
            else:
                rest = s.rest[1:]
            if word == ".":
                return State("E", -1, "", "", ())
            return s._replace(stage="R", rest=rest)
        raise ValueError(f"无法从 {s} 转移")

    # ---------------- 打分：一串词的精确 log 概率 ----------------
    def logprob(self, tokens, state=DOC_START, return_steps=False):
        """
        沿状态机走一遍，返回 log p_true(tokens | 起始状态)。

        【参数】
            tokens：词列表，例如 "I eat the apple .".split() + ["<eos>"]。
                    可以不以 <eos> 结尾——那就是"一个前缀"的概率。
            state：起始状态，默认是文档开头。
            return_steps：为 True 时额外返回每一步的条件概率列表。
        【返回】float（自然对数）。只要有一个词不合法就返回 -inf（概率为 0）。
                return_steps=True 时返回 (logp, [每一步的 p])。
        【边界】空列表的 log 概率是 0（概率 1：什么都没说，当然"一定发生"）。
        【LLM 联系】这就是链式法则：log p(x_1..x_T) = Σ_t log p(x_t | x_<t)。
                   LLM 的训练损失就是它的相反数除以 T。
        """
        total = 0.0
        steps = []
        for w in tokens:
            p = self.next_dist(state).get(w, 0.0)
            steps.append(p)
            if p == 0.0:
                total = -math.inf
                break
            total += math.log(p)
            state = self.step(state, w)
        if return_steps:
            return total, steps
        return total

    def state_after(self, tokens, state=DOC_START):
        """走完 tokens 后所处的状态；中途不合法则返回 None。"""
        for w in tokens:
            state = self.step(state, w)
            if state is None:
                return None
        return state

    def is_valid_sentence(self, tokens):
        """一句话（从主语到句号，不含 <eos>）是否完全合法。"""
        st = self.state_after(tokens, DOC_START)
        return st is not None and st.stage == "E"

    # ---------------- 采样 ----------------
    def sample_sentence(self, rng, state=DOC_START):
        """
        从状态 state（默认文档开头）出发，按真实分布采样一句话，直到句号为止。

        【参数】rng：random.Random 实例（显式传入，保证可复现）。
        【返回】list[str]，以 "." 结尾。
        """
        out = []
        while True:
            dist = self.next_dist(state)
            words = list(dist)
            # 【Python 小知识】rng.choices(候选, weights=权重) 按权重随机抽一个（返回列表，取 [0]）。
            w = rng.choices(words, weights=[dist[x] for x in words])[0]
            out.append(w)
            state = self.step(state, w)
            if w == ".":
                return out

    def sample_doc(self, rng, exclude_heldout=False):
        """
        按真实分布采样一篇完整文档（若干句话 + <eos>）。

        【参数】
            rng：random.Random 实例。
            exclude_heldout：为 True 时，含 heldout_pairs 里"动词-名词"组合的句子会被重新采样
                             （用于生成"恰好没见过某些搭配"的训练语料）。
        【返回】list[str]，以 "<eos>" 结尾。
        【边界】exclude_heldout=True 时，得到的语料分布和真实分布略有不同（那几种搭配的概率变成 0），
                这正是我们想模拟的"训练数据覆盖不全"。
        """
        doc = []
        state = DOC_START
        while True:
            if state.stage == "E":
                # 决定文档是否结束
                if rng.random() < P_EOS:
                    doc.append(EOS)
                    return doc
            start = state if state.stage == "S0" else State("S0", -1, "", "", ())
            while True:
                sent = self.sample_sentence(rng, start)
                if not (exclude_heldout and self._has_heldout(sent)):
                    break
            doc.extend(sent)
            state = State("E", -1, "", "", ())

    def _has_heldout(self, sent):
        """句子里是否出现了某个 heldout 的 (动词原形, 名词) 组合。"""
        verb = next((VERB_LEMMA[w] for w in sent if w in VERB_LEMMA), None)
        noun = next((w for w in sent if w in NOUN_CAT), None)
        return (verb, noun) in self.heldout_pairs

    # ---------------- 枚举全部句子 ----------------
    def enumerate_sentences(self):
        """
        枚举这门语言里**所有**合法的句子（从文档开头到第一个句号）及其精确概率。

        【返回】生成器，每次产出 (词元组, 概率)。
        【思路】深度优先搜索：从 DOC_START 出发，对 next_dist 里的每个词往下走，
                走到句号（状态变成 "E"）就产出一个句子。
        【边界】这门语言是有限的（约 3 万句），所以能全部枚举；真实语言是无限的，做不到。
        【Python 小知识】函数里用 yield 就成了"生成器"：调用它不会立刻算完，
                而是每次 for 循环要下一个元素时才算一个，省内存。
        """
        stack = [((), DOC_START, 1.0)]
        while stack:
            prefix, state, p = stack.pop()
            if state.stage == "E":
                yield prefix, p
                continue
            # reversed：让输出顺序和 next_dist 的书写顺序一致（栈是后进先出）
            for w, pw in reversed(list(self.next_dist(state).items())):
                stack.append((prefix + (w,), self.step(state, w), p * pw))


# =====================================================================
# 3. 标准语料
# =====================================================================

# 训练语料里刻意"没见过"的 5 种动词-名词搭配（但它们在真实语言里完全合法）。
HELDOUT_PAIRS = (("pick", "plum"), ("grow", "tulip"), ("feed", "goat"),
                 ("eat", "onion"), ("water", "daisy"))


def make_docs(n_docs, seed, world=None, exclude_heldout=False):
    """
    生成 n_docs 篇文档。

    【参数】n_docs：文档数；seed：随机种子；world：GardenWorld（默认新建一个，带 HELDOUT_PAIRS）；
            exclude_heldout：是否剔除含 heldout 搭配的句子（训练集用 True，验证 / 测试集用 False）。
    【返回】list[list[str]]，每篇文档以 "<eos>" 结尾。
    """
    world = world or GardenWorld(HELDOUT_PAIRS)
    rng = random.Random(seed)
    return [world.sample_doc(rng, exclude_heldout=exclude_heldout) for _ in range(n_docs)]


def standard_splits(n_train=8000, n_valid=500, n_test=1000):
    """
    第 1 讲统一使用的训练 / 验证 / 测试集。

    【返回】dict：{"train": 文档列表, "valid": ..., "test": ...}。
            训练集剔除了 HELDOUT_PAIRS 里的搭配；验证集和测试集按完整的真实分布采样。
    【边界】三个集合用不同的种子，互不重叠地独立采样（但短句难免和训练集里的句子"撞车"，
            这和真实语料一样正常）。
    """
    return {
        "train": make_docs(n_train, seed=1, exclude_heldout=True),
        "valid": make_docs(n_valid, seed=2),
        "test": make_docs(n_test, seed=3),
    }


class Vocab:
    """
    词表：在"词"和"编号（id）"之间来回转换。

    【思路】itos（index to string）是一个列表，下标就是 id；stoi（string to index）是反过来的字典。
    【参数】words：所有词的列表（不能有重复），顺序决定 id。
    【边界】encode 遇到词表外的词会换成 <unk>（词表里必须有 <unk>）。
    【LLM 联系】真实 LLM 的词表有 3 万到 25 万个 token，原理完全一样：
               模型内部只认识 id，文字和 id 之间的转换由分词器 + 词表完成。
    """

    def __init__(self, words):
        self.itos = list(words)
        self.stoi = {w: i for i, w in enumerate(self.itos)}
        if len(self.stoi) != len(self.itos):
            raise ValueError("词表里有重复的词")

    def __len__(self):
        # 【Python 小知识】定义了 __len__，就可以对这个对象用 len(vocab)。
        return len(self.itos)

    def encode(self, tokens):
        unk = self.stoi.get(UNK)
        return [self.stoi.get(t, unk) for t in tokens]

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


def garden_vocab():
    """花园世界的词表：3 个特殊符号 + 87 个词 = 90 个 token。<bos> 的 id 是 0，<eos> 是 1。"""
    return Vocab(SPECIALS + all_words())


def detokenize(tokens):
    """把词列表拼回一段文字（句号前不加空格，<eos> 换成换行）。"""
    text = " ".join(t for t in tokens if t != EOS)
    return text.replace(" .", ".")


# =====================================================================
# 4. 直接运行：打印样例与统计
# =====================================================================
if __name__ == "__main__":
    world = GardenWorld(HELDOUT_PAIRS)
    rng = random.Random(0)
    print("=== 花园世界的 5 篇样例文档 ===")
    for _ in range(5):
        doc = world.sample_doc(rng)
        print(f"  [{len(doc):2d} 个词] {' '.join(doc)}")

    splits = standard_splits()
    for name, docs in splits.items():
        n_tok = sum(len(d) for d in docs)
        n_sent = sum(d.count(".") for d in docs)
        print(f"{name:5s}: {len(docs):5d} 篇文档, {n_sent:6d} 句, {n_tok:7d} 个词")
    print(f"词表（不含特殊符号）共 {len(all_words())} 个词")