"""
=============================================================================
前置讲 P1 · 代码 02：字节级 BPE —— 分词器是怎样"学"出来的
=============================================================================

对应讲义：lectures/prereq1_text_bytes_tokens.md §3

【这个脚本要回答的问题】
    1. BPE 训练时到底在做什么？它最先学会的是哪些 token？（演示 1）
    2. 字节级 BPE 和第 1 讲的词级分词器比，好在哪里？为什么它"永远不会遇到不认识的字"？（演示 2）
    3. 合并次数（词表大小）越多越好吗？（演示 3：压缩率 vs 词表大小）
    4. 词表里那些"很少用到的 token"是怎么回事？（演示 4）

【怎么读输出】
    "字节 / token" 越大，说明同样的文字被切成的 token 越少——提示越"便宜"、上下文窗口装得越多。
    token 显示成 '<e6>' 这样的，是一个不完整的 UTF-8 片段（例如半个汉字）。

【运行】python3 02_bpe.py     （约 10 秒）
【依赖】只用 Python 标准库；import 同目录的 text_lab.py 和第 1 讲的 console.py、pretrain_world.py。
"""

import os
import sys
from collections import Counter

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import pct, print_header, print_observations, print_table  # noqa: E402
from pretrain_world import WordTokenizer, build_corpus, build_world, corpus_text  # noqa: E402
from text_lab import ByteBPE, bilingual_corpus, held_out_texts  # noqa: E402


def bytes_per_token(bpe, texts):
    """平均每个 token 代表几个 UTF-8 字节（越大越省）。"""
    nb = sum(len(t.encode("utf-8")) for t in texts)
    nt = sum(len(bpe.encode(t)) for t in texts)
    return nb / nt


def demo1_training(bpe):
    """演示 1：BPE 最先学会了什么。"""
    print_header("演示 1：BPE 训练——反复把最常一起出现的两个 token 合成一个")
    rows = []
    for k in list(range(20)) + [49, 99, 199, 399, 799, 1599]:
        a, b = bpe.merges[k]
        rows.append([k + 1, repr(bpe.token_str(a)), repr(bpe.token_str(b)), repr(bpe.token_str(256 + k))])
    print_table(["第几次合并", "左边", "右边", "合成的新 token"], rows)
    longest = sorted(range(256, len(bpe)), key=lambda i: (-len(bpe.vocab[i]), i))[:10]
    print("\n  最长的 10 个 token：" + ", ".join(repr(bpe.token_str(i)) for i in longest))
    zh_tokens = [i for i in range(256, len(bpe)) if any("一" <= ch <= "鿿" for ch in bpe.token_str(i))]
    print(f"  词表 {len(bpe)} 个 token 里，包含完整汉字的有 {len(zh_tokens)} 个："
          + "".join(bpe.token_str(i) for i in zh_tokens[:40]) + (" ……" if len(zh_tokens) > 40 else ""))
    print_observations([
        "最早的合并是英文里最常见的字母组合和'空格 + 字母'（' t'、'he'、' the'……）：BPE 只认频率，不认语法。",
        "空格被合进了单词的**前面**（' capital'、' of'）：同一个词在句首（没有空格）和句中是两个不同的 token。",
        f"训练语料里中文只占 10% 的字符，词表里也有 {len(zh_tokens)} 个 token 含完整汉字：常用字成了一个 token。"
        "但一个汉字 token 只顶一个字，而英文 token 往往顶好几个字母（演示 3）；常用字之外的字更是要 3 个字节 token（演示 2）。",
    ])


def demo2_roundtrip(bpe):
    """演示 2：字节级 BPE vs 第 1 讲的词级分词器。"""
    print_header("演示 2：字节级 BPE 永远不会遇到'不认识的字'")
    world = build_world()
    word_tok = WordTokenizer(corpus_text(build_corpus(world)))
    samples = [
        "the capital of volel is grorken .",
        "The Capital of Volel is Grorken.",
        "Q: what is the capital of volel? A:",
        "汤凉了，但是面包很新鲜。",
        "鲲鹏展翅九万里。",
        "Café ☕ costs 3.50 dollars 😀",
    ]
    rows = []
    for s in samples:
        ids = bpe.encode(s)
        ok = bpe.decode(ids) == s
        w = word_tok.encode(s)
        rows.append([repr(s), len(ids), "✓" if ok else "✗", len(w), sum(1 for i in w if i == word_tok.unk)])
    print_table(["文字", "BPE token 数", "BPE 往返无损", "词级 token 数", "其中 <unk>"], rows)
    print()
    for s in samples[1:]:
        print(f"  {s!r}\n    → {bpe.pieces(s)}")
    common, rare = len(bpe.encode(samples[3])), len(bpe.encode(samples[4]))
    print_observations([
        "第 1 讲的词级分词器只认识训练语料里出现过的 290 个小写词：大写、标点紧贴单词、中文、emoji 统统变成 <unk>，"
        "信息直接丢了。",
        f"字节级 BPE 对任何文字都能**无损**往返：认识的片段用一个 token，不认识的就退回单个字节。"
        f"'汤凉了，但是面包很新鲜。'里都是训练语料中的常用字，每个字一个 token（共 {common} 个）；"
        f"'鲲鹏展翅九万里。'里的'鲲''鹏''翅''万'不在训练语料的常用字里，各被切成 2–3 个字节 token，8 个字用了 {rare} 个 token。"
        "不熟悉的文字很'贵'。",
        "大小写也会改变切法：'The Capital' 和 'the capital' 是完全不同的 token 序列。模型对它们的'理解'可能相近，"
        "但并不保证完全一样——这是提示对大小写、格式敏感的一个原因（第 1 讲 §8.4）。",
    ])


def demo3_compression(bpe, held_en, held_zh):
    """演示 3：压缩率随合并次数的变化。"""
    print_header("演示 3：词表越大越好吗？——压缩率 vs 合并次数")
    rows, stats = [], {}
    for k in [0, 50, 100, 200, 400, 800, 1600]:
        t = bpe.truncated(k)
        en, zh = bytes_per_token(t, held_en), bytes_per_token(t, held_zh)
        stats[k] = (en, zh)
        rows.append([k, 256 + k, f"{en:.2f}", f"{zh:.2f}",
                     f"{sum(len(x) for x in held_en) / sum(len(t.encode(x)) for x in held_en):.2f}"])
    print("在没见过的测试文本上（英文约 2.5 万字符、中文约 1.6 万字符）")
    print_table(["合并次数", "词表大小", "英文：字节 / token", "中文：字节 / token", "英文：字符 / token"], rows)
    print_observations([
        f"0 次合并就是纯字节：每个 token 1 个字节。英文到 1600 次合并时每个 token 平均 {stats[1600][0]:.2f} 个字节，"
        f"序列短了 {stats[1600][0]:.1f} 倍——同样的上下文窗口能装下几倍的文字，同样的文字也便宜几倍。",
        f"收益递减：前 200 次合并带来的压缩远大于后面的 1400 次。真实分词器的词表通常是 3 万到 20 万，"
        "在'序列更短'和'嵌入矩阵更大、罕见 token 训练不足'之间折中。",
        f"按'字节 / token'算，中文（{stats[1600][1]:.2f}）低于英文（{stats[1600][0]:.2f}）：训练语料里中文只有 10%，"
        "一个 token 很少能超过一个汉字。但这个比较对中文不公平——一个汉字占 3 个字节、承载的意思也比一个英文字母多得多。"
        "比较两种语言的'成本'，要用**意思相同**的文本去数 token（03 号脚本演示 1），结论会不一样。",
    ])


def demo4_rare_tokens(bpe, train_texts, held_en, held_zh):
    """演示 4：很少被用到的 token。"""
    print_header("演示 4：词表里很少被用到的 token")
    use, use_train = Counter(), Counter()
    for t in held_en + held_zh:
        use.update(bpe.encode(t))
    for t in train_texts:
        use_train.update(bpe.encode(t))
    learned = range(256, len(bpe))
    never = [i for i in learned if use[i] == 0]
    absorbed = [i for i in never if use_train[i] == 0]
    specific = [i for i in never if use_train[i] > 0]
    rare = [i for i in learned if 0 < use[i] < 3]
    print(f"1600 个学出来的 token 在测试文本里：从没出现 {len(never)} 个（{pct(len(never), len(learned))}），"
          f"只出现 1–2 次 {len(rare)} 个（{pct(len(rare), len(learned))}）。从没出现的又分两种：")
    print(f"  (a) 连在**训练**语料里编码时都用不到（被后来的合并'吞掉'的中间产物）：{len(absorbed)} 个，例如 "
          + ", ".join(repr(bpe.token_str(i)) for i in absorbed[:8]))
    print(f"  (b) 训练语料里用得到、测试文本里用不到（训练语料特有的字符串）：{len(specific)} 个，例如 "
          + ", ".join(repr(bpe.token_str(i)) for i in specific[:8]))
    print_observations([
        "(a) 类是 BPE 的副作用：'<80 82>' 只是为了合成 '。' 的中间一步，'。' 出现之后它就再也不会被单独用到，"
        "却永远占着词表里的一个位置。",
        "(b) 类是 BPE 为训练语料里的高频字符串专门造的 token（这里多是长尾假词的片段）。换一份文本，它们可能很少出现。",
        "在真实的大模型里，这种'训练分词器时很常见、训练模型时却几乎没见过'的 token 会得到训练不足的向量，"
        "让模型在遇到它们时行为古怪——2023 年被发现的 ' SolidGoldMagikarp' 就是这样一个'故障 token'"
        "（glitch token）：它来自分词器训练语料里一个高频的用户名，模型却几乎没见过它。",
        "分词器和模型是分开训练的两个东西；提示里出现罕见的 token 组合（奇怪的用户名、乱码、特殊符号），"
        "是模型行为不稳定的一个来源。",
    ])


def main():
    en, zh = bilingual_corpus(en_share=0.9)
    print(f"训练语料：英文 {sum(map(len, en)):,} 字符、中文 {sum(map(len, zh)):,} 字符（中文约占 10%）；训练 1600 次合并……")
    bpe = ByteBPE(digits="group3").train(en + zh, n_merges=1600)
    held_en, held_zh = held_out_texts()
    demo1_training(bpe)
    demo2_roundtrip(bpe)
    demo3_compression(bpe, held_en, held_zh)
    demo4_rare_tokens(bpe, en + zh, held_en, held_zh)


if __name__ == "__main__":
    main()