"""
=============================================================================
CMU 17-630 · 前置讲 P1 共享模块：text_lab.py —— 字符、字节、token
=============================================================================

【这个模块是做什么的】
    你写给模型的每一个字，在到达模型之前要经过三次变形：

        你看到的字符（"中"、"é"、"😀"）
          → Unicode 码位（U+4E2D、U+00E9、U+1F600）
          → UTF-8 字节（e4 b8 ad、c3 a9、f0 9f 98 80）
          → token（分词器把字节串切成词表里的片段，每个片段一个整数 id）

    模型看到的只有最后那一串整数。上下文窗口、计费、"数不清 strawberry 里有几个 r"、
    "数字算不对"、"提示末尾多一个空格结果变差"……这些现象都发生在这三次变形里。

    这个模块提供：
        1. 检查文字的工具：逐字符看码位和字节（describe）、找隐形字符（find_invisible）、
           找混用的文字系统（mixed_scripts，用来发现"长得一样的假字母"）。
        2. 预分词（pretokenize）：先把文字切成"词块"，BPE 只在词块内部合并。
        3. 字节级 BPE 分词器（ByteBPE）：和 GPT-2 以来大多数 LLM 用的是同一类算法，
           从 256 个单字节出发，反复把最常一起出现的两个 token 合并成一个新 token。
        4. 从模型输出里稳健地取出 JSON（extract_json）。
        5. 一份双语训练语料（bilingual_corpus），以及带缓存的标准分词器（load_or_train_bpe）：一半来自第 1 讲的虚构世界（英文版 + 中文版），
           一半是"长尾杂记"（按 Zipf 分布抽取的几千个假词，模拟真实文本里大量的罕见词）；
           可以调中英文的比例，看分词器怎样"偏心"。

【运行】python3 text_lab.py     # 自检：训练一个小 BPE，往返编码几句话
【依赖】只用 Python 标准库；会 import ../lecture01/pretrain_world.py（只读）。
"""

import os
import random
import re
import sys
import unicodedata
from collections import Counter

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from pretrain_world import build_corpus, build_world, make_review  # noqa: E402  第 1 讲的虚构世界


# ---------------------------------------------------------------------------
# 1. 看清楚一段文字：码位、字节、隐形字符、混用的文字系统
# ---------------------------------------------------------------------------

def describe(s):
    """
    逐个字符（码位）列出：字符本身、码位、Unicode 名字、UTF-8 字节。

    【返回】[(字符, "U+XXXX", 名字, "e4 b8 ad"), ...]
    【边界】没有名字的码位（例如控制字符）名字记为 "<无名>"。
    【Python 小知识】
        ord("中") → 20013：字符 → 码位（整数）；chr(20013) → "中"：反过来。
        "中".encode("utf-8") → b'\\xe4\\xb8\\xad'：字符串 → 字节串（bytes）。
        f"{x:04X}"：把整数写成至少 4 位的大写十六进制。
    """
    rows = []
    for ch in s:
        rows.append((ch, f"U+{ord(ch):04X}", unicodedata.name(ch, "<无名>"),
                     " ".join(f"{b:02x}" for b in ch.encode("utf-8"))))
    return rows


# 常见的"看不见"或"看起来像别的东西"的字符
INVISIBLE = {
    "​": "ZERO WIDTH SPACE（零宽空格）",
    "‌": "ZERO WIDTH NON-JOINER（零宽不连字）",
    "‍": "ZERO WIDTH JOINER（零宽连字，emoji 组合用）",
    "⁠": "WORD JOINER（词连接符）",
    "": "BYTE ORDER MARK / ZERO WIDTH NO-BREAK SPACE（BOM）",
    " ": "NO-BREAK SPACE（不换行空格，长得和空格一样）",
    "‮": "RIGHT-TO-LEFT OVERRIDE（把后面的文字反着显示）",
    "
": "LINE SEPARATOR（另一种换行）",
}


def find_invisible(s):
    """
    找出字符串里所有"看不见 / 长得像空格"的字符，以及其他控制字符（Unicode 类别 Cc / Cf，换行和制表符除外）。

    【返回】[(下标, 码位, 名字), ...]
    【提示工程联系】从网页、PDF、聊天记录里复制来的文字常常带着这些字符：
        它们会让"看起来一样"的两个字符串不相等（精确匹配评估失败、缓存不命中），
        也会被拿来在提示里藏东西（第 12 讲的间接提示注入）。
    """
    out = []
    for i, ch in enumerate(s):
        if ch in INVISIBLE:
            out.append((i, f"U+{ord(ch):04X}", INVISIBLE[ch]))
        elif unicodedata.category(ch) in ("Cc", "Cf") and ch not in "\n\t\r":
            out.append((i, f"U+{ord(ch):04X}", unicodedata.name(ch, "<控制字符>")))
    return out


def script_of(ch):
    """
    粗略判断一个字母属于哪种文字：取 Unicode 名字的第一个词（LATIN / CYRILLIC / GREEK / CJK ……）。

    【边界】不是字母的字符（数字、标点、空格）返回 None。
    """
    if not ch.isalpha():
        return None
    return unicodedata.name(ch, "UNKNOWN").split()[0]


def mixed_scripts(word):
    """
    一个"词"里混用了几种文字。正常的英文单词只有 LATIN；"pаypal" 里如果有一个西里尔字母 а，
    就会返回 {"LATIN", "CYRILLIC"}——这是同形字（homoglyph）欺骗的典型特征。

    【返回】文字系统名字的集合
    """
    return {s for s in (script_of(ch) for ch in word) if s}


# ---------------------------------------------------------------------------
# 2. 预分词：先切成"词块"
# ---------------------------------------------------------------------------

# 数字的三种切法（见 03 号脚本演示 3）：
#   "single"：每个数字单独一块，"2026" → 2 0 2 6（Llama、Gemma 系列的做法之一）
#   "group3"：最多 3 个数字一块，"123456" → 123 456（GPT-4 的 cl100k 分词器用的是 1–3 位一块）
#   "any"   ：一串数字整个是一块，让 BPE 自己学（GPT-2 的做法，会学出不规则的切法）
_DIGIT_RULES = {"single": r"[0-9]", "group3": r"[0-9]{1,3}", "any": r"[0-9]+"}


def make_pretoken_re(digits="group3"):
    """
    构造预分词用的正则表达式。

    【思路】一段文字被切成以下几类词块（从左到右，先匹配上的先算）：
        " ?[A-Za-z]+"          英文单词，**连同它前面的一个空格**（所以 " cat" 和 "cat" 是两个不同的块）
        数字规则                见 _DIGIT_RULES
        "[\\u4e00-\\u9fff]"      一个汉字一块（真实分词器会把常见的词合成一个 token，这里为了简单只到"字"）
        " ?[^\\sA-Za-z0-9\\u4e00-\\u9fff]+"  标点、符号、emoji（连同前面的一个空格）
        "\\s+"                  空白（空格、换行、制表符）
        "."                    兜底：其他任何字符
    【为什么要预分词】如果不预先切开，BPE 可能学出 "e." "d the" 这种跨词的 token；
        预分词让每个 token 都落在一个词、一个数字或一个标点块之内，更规整，也更容易解释。
    【Python 小知识】re.compile 先把正则"编译"好，之后反复用 .findall 更快。
        正则里 "|" 表示"或"，按顺序尝试。
    """
    return re.compile(r" ?[A-Za-z]+|" + _DIGIT_RULES[digits]
                      + r"|[一-鿿]| ?[^\sA-Za-z0-9一-鿿]+|\s+|.", re.S)


def pretokenize(text, digits="group3"):
    """把文字切成词块列表。所有词块拼起来**正好**等于原文（不丢任何字符）。"""
    return make_pretoken_re(digits).findall(text)


# ---------------------------------------------------------------------------
# 3. 字节级 BPE
# ---------------------------------------------------------------------------

class ByteBPE:
    """
    字节级 BPE（byte-pair encoding）分词器。

    【思路】
        起点：词表 = 256 个单字节（id 0–255）。任何文字都能写成 UTF-8 字节，所以任何文字都能编码，
              **永远不会出现 <unk>**（这正是"字节级"的意义；第 1 讲的词级分词器做不到）。
        训练：把语料预分词，统计每个词块出现几次；每个词块写成一串 id（一开始是它的字节）。
              重复 n_merges 次：找出所有词块里"相邻两个 id"出现次数最多的一对 (a, b)，
              给它分配一个新 id，把所有词块里的 a b 都换成这个新 id。
        编码：先预分词，再对每个词块按训练时**学到合并的先后顺序**依次合并。
        解码：把每个 id 对应的字节拼起来，再按 UTF-8 解成文字。
    【用法】
        bpe = ByteBPE(digits="group3").train(texts, n_merges=800)
        ids = bpe.encode("The capital of volel is grorken.")
        bpe.decode(ids) == 原文
    """

    def __init__(self, digits="group3"):
        self.digits = digits
        self.merges = []                                  # [(a, b), ...] 按学到的先后顺序
        self.vocab = {i: bytes([i]) for i in range(256)}  # id → 字节串
        self.rank = {}                                    # (a, b) → 新 id

    def train(self, texts, n_merges=800, verbose=False):
        """
        在一组文字上学 n_merges 次合并。

        【参数】texts：字符串列表；n_merges：合并次数（词表大小 = 256 + n_merges）
        【返回】self
        【边界】如果所有词块都已经合并成单个 token、找不到可合并的对，会提前停止。
            出现次数并列时，按 (a, b) 的大小取最小的那一对，保证每次训练结果完全相同。
        【复杂度】每次合并都重新数一遍所有不同词块里的相邻对。这份语料的不同词块只有几千个，
            800 次合并在几秒内完成；真实分词器在几十 GB 文本上训练，需要更聪明的增量计数。
        """
        word_freq = Counter()
        for t in texts:
            word_freq.update(pretokenize(t, self.digits))
        words = {w: list(w.encode("utf-8")) for w in word_freq}
        for step in range(n_merges):
            pairs = Counter()
            for w, ids in words.items():
                f = word_freq[w]
                for a, b in zip(ids, ids[1:]):
                    pairs[(a, b)] += f
            if not pairs:
                break
            best = max(pairs, key=lambda p: (pairs[p], -p[0], -p[1]))
            new_id = 256 + len(self.merges)
            self.merges.append(best)
            self.rank[best] = new_id
            self.vocab[new_id] = self.vocab[best[0]] + self.vocab[best[1]]
            for w, ids in words.items():
                if len(ids) > 1:
                    words[w] = self._merge(ids, best, new_id)
            if verbose and (step + 1) % 200 == 0:
                print(f"    第 {step + 1} 次合并：{self.token_str(new_id)!r}（出现 {pairs[best]} 次）")
        return self

    @staticmethod
    def _merge(ids, pair, new_id):
        """把 ids 里所有相邻的 pair 换成 new_id（从左到右，不重叠）。"""
        out, i = [], 0
        while i < len(ids):
            if i + 1 < len(ids) and ids[i] == pair[0] and ids[i + 1] == pair[1]:
                out.append(new_id)
                i += 2
            else:
                out.append(ids[i])
                i += 1
        return out

    def _encode_chunk(self, chunk):
        """
        编码一个词块：反复找"学得最早"的那一对合并，直到没有可以合并的对。

        【为什么按学到的先后顺序】训练时 "t"+"h" 先合并成 "th"，之后才有 "th"+"e" → "the"。
            编码时如果先合并别的对，可能得到训练时从没出现过的切法。
        """
        ids = list(chunk.encode("utf-8"))
        while len(ids) > 1:
            cand = [(self.rank[p], p) for p in zip(ids, ids[1:]) if p in self.rank]
            if not cand:
                break
            r, p = min(cand)
            ids = self._merge(ids, p, r)
        return ids

    def encode(self, text):
        """文字 → token id 列表。"""
        out = []
        for chunk in pretokenize(text, self.digits):
            out.extend(self._encode_chunk(chunk))
        return out

    def decode(self, ids, errors="replace"):
        """
        token id 列表 → 文字。

        【边界】一段 id 如果在一个汉字的 3 个字节中间断开（例如流式输出时只收到了一半），
            按 UTF-8 解码会出错；errors="replace" 把坏掉的字节显示成 "�"，errors="strict" 则直接报错。
        """
        return b"".join(self.vocab[i] for i in ids).decode("utf-8", errors=errors)

    def token_str(self, i):
        """一个 token 的"可读形式"：能解成文字就显示文字，否则显示它的十六进制字节（例如半个汉字）。"""
        b = self.vocab[i]
        try:
            return b.decode("utf-8")
        except UnicodeDecodeError:
            return "<" + " ".join(f"{x:02x}" for x in b) + ">"

    def pieces(self, text):
        """把文字切成 token 的可读片段列表，例如 ['The', ' capital', ' of', ...]。"""
        return [self.token_str(i) for i in self.encode(text)]

    def truncated(self, k):
        """
        只保留前 k 次合并的副本。

        【为什么这样就等于"只训练 k 次"】BPE 的训练是贪心的：第 k+1 次合并只依赖前 k 次的结果，
            所以"训练 1600 次再截取前 k 次"和"直接训练 k 次"得到的分词器完全一样，省去重复训练。
        """
        t = ByteBPE(self.digits)
        t.merges = self.merges[:k]
        t.rank = {p: 256 + i for i, p in enumerate(t.merges)}
        t.vocab = {i: self.vocab[i] for i in range(256 + len(t.merges))}
        return t

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


# ---------------------------------------------------------------------------
# 3.5 从模型的输出里稳健地取出 JSON
# ---------------------------------------------------------------------------

def extract_json(text):
    """
    从一段"模型写的文字"里找出第一个完整的 JSON 对象（{...}）或数组（[...]）。

    【返回】(解析出的 Python 对象, 状态说明)；找不到时返回 (None, 失败原因)
    【思路】模型常常不只输出 JSON：前面一句"好的，结果如下："，外面包一层 ```json 代码块，
        后面再补一句解释。所以不能直接 json.loads(整段文字)。做法是：
        从左到右找每一个 '{' 或 '['，用 json.JSONDecoder().raw_decode 尝试"从这里开始解析一个 JSON 值"，
        第一个成功的就是答案。raw_decode 会告诉我们 JSON 在哪里结束，后面多余的文字不影响。
    【边界】
        - 不修补格式错误（单引号、末尾多余的逗号、注释、Python 的 True/None……）：这些情况返回失败原因，
          让调用方决定是重试、还是让模型改正（第 8 讲）。悄悄"修好"可能把错误的内容当成正确的。
        - 文字里有好几个 JSON 时只返回第一个；如果第一个 '{' 只是文字里的一个括号（例如 "{注意}"），
          会跳过它继续往后找。
    【Python 小知识】json.JSONDecoder().raw_decode(s, idx) 从 s[idx] 开始解析，返回 (对象, 结束位置)；
        解析失败时抛出 json.JSONDecodeError，它的 .msg、.lineno、.colno 说明了错在哪里。
    """
    import json
    dec = json.JSONDecoder()
    first_error = None
    for i, ch in enumerate(text):
        if ch in "{[":
            try:
                obj, _ = dec.raw_decode(text, i)
                return obj, "ok" if i == 0 else f"ok（跳过了前面 {i} 个字符）"
            except json.JSONDecodeError as e:
                if first_error is None:
                    first_error = f"第 {e.lineno} 行第 {e.colno} 列：{e.msg}"
    return None, ("没有找到 JSON" if first_error is None else f"解析失败（{first_error}）")


# ---------------------------------------------------------------------------
# 4. 双语训练语料：第 1 讲的虚构世界，英文版 + 中文版
# ---------------------------------------------------------------------------

def detokenize(text):
    """
    把第 1 讲那种"每个词之间都有空格、全小写"的玩具文本，改写成正常的英文：
    标点前不留空格、句首大写。

    【例子】"the capital of volel is grorken . people in volel speak driltese ."
          → "The capital of volel is grorken. People in volel speak driltese."
    """
    t = re.sub(r" ([.,?!:;'])", r"\1", text)
    t = t.replace(" 's", "'s")
    return re.sub(r"(^|[.?!] )([a-z])", lambda m: m.group(1) + m.group(2).upper(), t)


# 中文版的词表（自己写的模板，不来自任何外部文本）
_ZH_SYL = list("阿巴达尔福格克拉里马尼萨斯塔图维亚兰洛恩索加卡莫贝多")
ZH_REGIONS = {"north": "北", "south": "南", "east": "东", "west": "西"}
ZH_EXPORTS = {"copper": "铜", "wheat": "小麦", "salt": "盐", "tea": "茶", "wool": "羊毛", "iron": "铁",
              "silk": "丝绸", "rice": "大米", "coal": "煤", "timber": "木材", "olives": "橄榄", "glass": "玻璃"}
ZH_ASPECTS = ["汤", "面包", "咖啡", "服务", "房间", "床", "价格", "音乐", "风景", "服务员", "菜", "早餐"]
ZH_POS = ["很好", "很新鲜", "很热情", "很干净", "很可爱", "很快", "很好吃", "很便宜", "很安静", "很舒服"]
ZH_NEG = ["凉了", "不新鲜", "很粗鲁", "太慢了", "很脏", "很糟糕", "太吵了", "没有味道", "太贵了", "软塌塌的"]
ZH_PEOPLE = [("国王", "王后"), ("王子", "公主"), ("男人", "女人"), ("男孩", "女孩"), ("父亲", "母亲"), ("儿子", "女儿")]
ZH_CONTEXTS = ["住在宫殿里", "戴着王冠", "在田里干活", "去学校上学", "在市场上散步", "唱了一首歌", "吃了一个苹果", "很善良"]


def _zh_name(rng, used):
    """用音译常用字随机拼一个 2–3 个字的中文名字（不重复）。"""
    while True:
        n = "".join(rng.choice(_ZH_SYL) for _ in range(rng.choice([2, 3])))
        if n not in used:
            used.add(n)
            return n


def chinese_docs(world, n_docs, seed=0):
    """
    生成 n_docs 篇中文文档：百科词条、评论、小故事三种，内容和第 1 讲的世界对应。

    【返回】字符串列表
    【边界】国名等名字是随机音译的，和英文名字没有一一对应的读音关系（够用了：我们只关心分词）。
        评论里会混入价格（"价格是 128 元"），故事里偶尔有年份，给数字分词的实验用。
    """
    rng = random.Random(9000 + seed)
    used = set()
    zh = {c.name: (_zh_name(rng, used), _zh_name(rng, used), _zh_name(rng, used) + "语", _zh_name(rng, used))
          for c in world.countries}
    docs = []
    for i in range(n_docs):
        kind = i % 3
        if kind == 0:
            c = rng.choice(world.countries)
            name, cap, lang, leader = zh[c.name]
            parts = [f"{name}是一个位于{ZH_REGIONS[c.region]}方的国家。", f"{name}的首都是{cap}。",
                     f"{name}的人们说{lang}。", f"{name}以{ZH_EXPORTS[c.export]}闻名。", f"{name}的领导人是{leader}。"]
            docs.append("百科：" + "".join(parts[:1] + rng.sample(parts[1:], 4)))
        elif kind == 1:
            items = []
            for _ in range(rng.randint(1, 3)):
                pol = rng.choice([1, -1])
                a1, a2 = rng.sample(ZH_ASPECTS, 2)
                adj = lambda p: rng.choice(ZH_POS if p > 0 else ZH_NEG)
                if rng.random() < 0.3:
                    body = f"{a1}{adj(-pol)}，但是{a2}{adj(pol)}"
                else:
                    body = f"{a1}{adj(pol)}，而且{a2}{adj(pol)}"
                price = rng.choice([18, 35, 68, 128, 256, 1024, 2999])
                items.append(f"评论：{body}，价格是 {price} 元。评分：{'好评' if pol > 0 else '差评'}。")
            docs.append("".join(items))
        else:
            sents = []
            for _ in range(rng.randint(2, 3)):
                m, f = rng.choice(ZH_PEOPLE)
                who, pron = (m, "他") if rng.random() < 0.5 else (f, "她")
                sents.append(f"{who}{rng.choice(ZH_CONTEXTS)}。{pron}{rng.choice(ZH_CONTEXTS)}。")
            if rng.random() < 0.3:
                sents.append(f"那是 {rng.choice([1998, 2003, 2019, 2024, 2026])} 年的事。")
            docs.append("故事：" + "".join(sents))
    return docs


def english_docs(world, n_docs, seed=0):
    """
    取第 1 讲语料里的 n_docs 篇文档（百科、评论、故事三种），改写成正常的英文；评论里加上价格。

    【返回】字符串列表
    """
    rng = random.Random(8000 + seed)
    pool = [d for d in build_corpus(world, seed=seed, scale=0.5) if d.kind in ("article", "reviews", "story")]
    rng.shuffle(pool)
    out = []
    for d in pool[:n_docs]:
        t = detokenize(d.text)
        if d.kind == "reviews":
            t = t.replace(" Rating:", f" It cost {rng.choice([18, 35, 68, 128, 256, 1024, 2999])} dollars. Rating:")
        out.append(t)
    return out


# 约 500 个常用汉字（按常见程度粗略排列），给"长尾"的中文假词用
_ZH_COMMON = ("的一是不了人我在有他这中大来上国个到说们为子和你地出道也时年得就那要下以生会自着去之过家学对可她里后小么"
              "心多天而能好都然没日于起还发成事只作当想看文无开手十用主行方又如前所本见经头面公同三已老从动两长知民样现分将外"
              "但身些与高意进把法此实回二理美点月明其种声全工己话儿者向情部正名定女问力机给等几很业最间新什打便位因重被走电四"
              "第门相次东政海口使教西再平真听世气信北少关并内加化由却代军产入先山五太水万市眼体别处总才场师书比住员九笑性通目"
              "华报立马命张活难神数件安表原车白应路期叫死常提感金何更反合放做系计或司利受光王果亲界及今京务制解任至清物台象记"
              "边共风战干接它许八特觉望直服毛林题建南度统色字请交爱让认算论百吃义科怎元社术结六功指思非流每青管夫连远资队跟带"
              "花快条院变联言权往展该领传近留红治决周保达办运武半候七必城父强步完革深区")
_EN_ONSET = ["b", "c", "d", "f", "g", "h", "j", "k", "l", "m", "n", "p", "r", "s", "t", "v", "w", "z",
             "br", "ch", "cl", "dr", "fl", "gr", "pl", "pr", "sh", "sl", "st", "th", "tr"]
_EN_VOWEL = ["a", "e", "i", "o", "u", "ai", "ea", "ee", "oo", "ou"]


def _lexicon(seed=0, n_en=3000, n_zh=1500):
    """
    一份"长尾词典"：3000 个英文假词（2–4 个音节）和 1500 个中文假词（2–3 个常用字）。

    【为什么需要它】只用第 1 讲的模板语料，不同的词只有几百个，BPE 合并几百次之后
        每个词都成了一个 token，再也无词可合（第一版就是这样提前停止的）。真实文本有一条
        长长的"长尾"：少数词很常见，大量的词很罕见。词表的大小总是不够覆盖所有的词，
        罕见词只能被切成子词——这才是真实分词器所处的局面。
    """
    rng = random.Random(7000 + seed)
    en, zh, seen = [], [], set()
    while len(en) < n_en:
        w = "".join(rng.choice(_EN_ONSET) + rng.choice(_EN_VOWEL) for _ in range(rng.randint(2, 4)))
        if rng.random() < 0.5:
            w += rng.choice(["n", "r", "s", "t", "l", "m", "ck", "nd", "st"])
        if w not in seen:
            seen.add(w)
            en.append(w)
    while len(zh) < n_zh:
        w = "".join(rng.choice(_ZH_COMMON) for _ in range(rng.randint(2, 3)))
        if w not in seen:
            seen.add(w)
            zh.append(w)
    return en, zh


def _zipf_pick(rng, words, a=1.1):
    """
    按 Zipf 分布抽一个词：排名第 r 的词被抽中的概率正比于 1 / r^a（少数词很常见，多数词很罕见）。

    【Python 小知识】rng.choices(词表, cum_weights=累积权重) 用二分查找抽样，比每次传 weights= 快得多
        （传 weights 时，每抽一次都要重新累加一遍几千个权重）。累积权重只算一次，存在函数属性 cache 里。
    """
    cw = _zipf_pick.cache.get((len(words), a))
    if cw is None:
        cw, total = [], 0.0
        for r in range(len(words)):
            total += 1.0 / (r + 1) ** a
            cw.append(total)
        _zipf_pick.cache[(len(words), a)] = cw
    return rng.choices(words, cum_weights=cw)[0]


_zipf_pick.cache = {}


def longtail_docs(lang, n_docs, seed=0):
    """
    用长尾词典写的"杂记"：每篇 3–6 句，句子里的人名、地名、物品名按 Zipf 分布从词典里抽。

    【参数】lang："en" 或 "zh"；n_docs：篇数；seed：抽词用的随机种子（词典本身固定）
    【返回】字符串列表
    """
    return list(_longtail_iter(lang, seed, n_docs))


def _longtail_iter(lang, seed=0, n_docs=None):
    """longtail_docs 的"生成器"版本：一篇一篇地产生，要多少取多少（n_docs=None 表示无限）。

    【Python 小知识】函数体里有 yield 的函数是**生成器**：调用它不会立刻算出全部结果，
        而是每次被要下一个元素时才往下执行一段。配合 _take 使用，就不会白白生成用不到的文档。
    """
    en_lex, zh_lex = _lexicon()
    rng = random.Random(6000 + seed + (0 if lang == "en" else 1))
    i = 0
    while n_docs is None or i < n_docs:
        i += 1
        sents = []
        for _ in range(rng.randint(3, 6)):
            x, y, z = (_zipf_pick(rng, en_lex if lang == "en" else zh_lex) for _ in range(3))
            year = rng.randint(1800, 2026)
            if lang == "en":
                t = rng.choice(["{X} visited {Y} in {n}.", "The {z} of {X} was sold for {m} coins.",
                                "{X} and {Y} wrote about the {z}.", "In {n}, {X} moved to {Y}."])
                sents.append(t.format(X=x.capitalize(), Y=y.capitalize(), z=z, n=year, m=rng.randint(2, 9999)))
            else:
                t = rng.choice(["{X}在{n}年访问了{Y}。", "{X}的{z}卖了{m}个金币。", "{X}和{Y}写过关于{z}的文章。",
                                "{n}年，{X}搬到了{Y}。"])
                sents.append(t.format(X=x, Y=y, z=z, n=year, m=rng.randint(2, 9999)))
        yield (" " if lang == "en" else "").join(sents)


def bilingual_corpus(en_share=0.9, total_chars=200_000, seed=0):
    """
    中英文混合的训练语料，按**字符数**控制两种语言的比例。

    【参数】en_share：英文占总字符数的比例；total_chars：总字符数（大约）；seed：随机种子
    【返回】(英文文档列表, 中文文档列表)
    【提示工程联系】真实分词器的训练语料里英文通常占大头，这会让其他语言"更贵"（03 号脚本演示 1）。
    """
    world = build_world()
    en_target, zh_target = en_share * total_chars, (1 - en_share) * total_chars
    en = _take(english_docs(world, 6000, seed), en_target / 2) + _take(_longtail_iter("en", seed), en_target / 2)
    zh = _take(chinese_docs(world, 20000, seed), zh_target / 2) + _take(_longtail_iter("zh", seed), zh_target / 2)
    return en, zh


def _take(docs, n_chars):
    """从文档列表（或生成器）开头取文档，直到总字符数达到 n_chars。"""
    out, total = [], 0
    for d in docs:
        if total >= n_chars:
            break
        out.append(d)
        total += len(d)
    return out


def held_out_texts(seed=1):
    """一份和训练语料不同 seed 的测试文本：(英文文档列表, 中文文档列表)。一半是模板文档，一半是长尾杂记。"""
    world = build_world()
    en = english_docs(world, 75, seed) + longtail_docs("en", 60, seed)
    zh = chinese_docs(world, 150, seed) + longtail_docs("zh", 130, seed)
    return en, zh


CACHE_DIR = os.path.join(HERE, "cache")


def load_or_train_bpe(en_share=0.9, n_merges=1600, digits="group3"):
    """
    取一个"标准"分词器：缓存里有就直接读，没有就训练（约 4 秒）并存进 code/prereq1/cache/。

    【为什么要缓存】前置讲 P2 和后面几讲要用它来数 token、算费用，每次 import 都重新训练太慢。
    【返回】ByteBPE
    【边界】缓存文件名里含训练配置；改了语料生成代码之后请删掉 cache/ 目录（缓存不会自动失效——
        这里为了简单没有像第 1 讲那样把语料全文放进哈希）。
    【Python 小知识】json 不能直接存元组和 bytes，所以只存 merges（整数对的列表），
        读回来之后按顺序重建 vocab 和 rank——这和训练时的构造方式完全一样。
    """
    import json
    path = os.path.join(CACHE_DIR, f"bpe-en{int(en_share * 100)}-m{n_merges}-{digits}.json")
    bpe = ByteBPE(digits=digits)
    if os.path.exists(path):
        with open(path, encoding="utf-8") as f:
            merges = [tuple(p) for p in json.load(f)["merges"]]
        for i, pair in enumerate(merges):
            bpe.merges.append(pair)
            bpe.rank[pair] = 256 + i
            bpe.vocab[256 + i] = bpe.vocab[pair[0]] + bpe.vocab[pair[1]]
        return bpe
    en, zh = bilingual_corpus(en_share=en_share)
    bpe.train(en + zh, n_merges=n_merges)
    os.makedirs(CACHE_DIR, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump({"en_share": en_share, "n_merges": n_merges, "digits": digits, "merges": bpe.merges}, f)
    return bpe


if __name__ == "__main__":
    en, zh = bilingual_corpus(0.9)
    print(f"训练语料：英文 {len(en)} 篇 / {sum(map(len, en)):,} 字符，中文 {len(zh)} 篇 / {sum(map(len, zh)):,} 字符")
    print("  英文例子：", en[0][:120])
    print("  中文例子：", zh[1][:60])
    bpe = ByteBPE().train(en + zh, n_merges=400)
    for s in ["The capital of volel is grorken.", "汤凉了，但是面包很新鲜。", "Emoji 😀 and café, 123456"]:
        ids = bpe.encode(s)
        assert bpe.decode(ids) == s
        print(f"  {s!r} → {len(ids)} 个 token：{bpe.pieces(s)}")
    print("  往返编码全部无损 ✓")