"""
=============================================================================
第 1 讲 · 代码 2：三种最简单的分词器 —— 字符级、字节级、词级
=============================================================================

对应讲义：lectures/lecture01_overview_tokenization.md 第 5、7、8、9 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
分词器（tokenizer）要做两件事：
    encode（编码）：字符串 "Hello"  →  整数列表 [15496]
    decode（解码）：整数列表 [15496] →  字符串 "Hello"

最容易想到的三种做法是：
    1. 字符级（character）：一个 Unicode 字符 = 一个 token
    2. 字节级（byte）     ：一个 UTF-8 字节  = 一个 token
    3. 词级（word）       ：一个"单词"       = 一个 token

这个文件把三种都实现出来，并用同一把尺子——**压缩率**——来比较它们。
你会亲眼看到每一种的致命缺陷，从而理解为什么最终大家都用 BPE（下一个文件）。

-----------------------------------------------------------------------------
【整体思路】
-----------------------------------------------------------------------------
    第 1 部分：定义"分词器"的统一接口（抽象基类 Tokenizer）和评价指标（压缩率）
    第 2 部分：字符级分词器 —— 词表巨大
    第 3 部分：字节级分词器 —— 词表完美（256），但序列太长
    第 4 部分：词级分词器   —— 序列短，但词表无界、会遇到未知词（UNK）
    第 5 部分：把三者放在一起比较，并算一算"序列长度"对计算量意味着什么

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 02_simple_tokenizers.py

只用到 Python 标准库。需要 data/sample_corpus.txt（本目录下已提供）。

-----------------------------------------------------------------------------
【注释约定】见 01_unicode_and_utf8.py 开头，全课通用：
    【思路】【参数】【返回】【边界】【Python 小知识】【大模型联系】
"""

# =============================================================================
# 第 0 部分：导入与工具函数
# =============================================================================
# 【Python 小知识】
#   re          ：正则表达式（regular expression）库，用"模式"来查找/切分文本。Python 自带。
#   Counter     ：collections 库里的"计数器"，是一种特殊的字典，专门用来数东西出现了几次。
#                 Counter(["a", "b", "a"]) → Counter({"a": 2, "b": 1})
#   Path        ：pathlib 库里表示"文件路径"的类，比直接拼字符串更安全、跨平台。
import re
from collections import Counter
from pathlib import Path

# 语料文件的位置。
# 【Python 小知识】
#   __file__ 是一个特殊变量，表示"当前这个 .py 文件的路径"。
#   Path(__file__).parent 就是"这个文件所在的文件夹"。
#   用 / 运算符可以拼接路径：Path("a") / "b" / "c.txt" → a/b/c.txt。
#   这样写的好处是：无论你在哪个目录下运行这个脚本，都能找到数据文件。
CORPUS_PATH = Path(__file__).parent / "data" / "sample_corpus.txt"

# 语料里用来分隔不同文档（故事）的特殊标记，和 GPT-2 / TinyStories 的约定一致。
END_OF_TEXT = "<|endoftext|>"


def print_header(title):
    """打印一个带分隔线的标题（同 01 号文件）。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


def load_corpus_documents(path=CORPUS_PATH):
    """
    读取语料文件，按 <|endoftext|> 切成一篇篇文档。

    【参数】
        path : Path 或 str，语料文件路径。默认是本目录下的 data/sample_corpus.txt。

    【返回】
        list[str]，每个元素是一篇文档（去掉了首尾空白，并丢弃空文档）。

    【边界】
        - 文件不存在：抛出 FileNotFoundError，并提示正确的位置。
        - 文件必须是 UTF-8 编码。我们显式写 encoding="utf-8"，因为在某些 Windows 电脑上，
          Python 默认编码不是 UTF-8，不写的话读中文会乱码或报错。

    【Python 小知识】
        - path.read_text(encoding="utf-8") 一次性把整个文件读成一个字符串。
          对大文件（几个 GB）不能这么做，会撑爆内存；我们的小语料没问题。
        - text.split(sep) 按分隔符把字符串切成列表。
        - [d.strip() for d in docs if d.strip()] 叫"列表推导式"：
          对 docs 里每个 d，如果 d.strip() 非空，就把 d.strip() 放进新列表。
          strip() 会去掉字符串首尾的空格和换行。
    """
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"找不到语料文件 {path}。请确认 data/sample_corpus.txt 与本脚本在同一目录结构下。")
    text = path.read_text(encoding="utf-8")
    docs = text.split(END_OF_TEXT)
    return [d.strip() for d in docs if d.strip()]


# =============================================================================
# 第 1 部分：分词器的统一接口 + 评价指标
# =============================================================================
class Tokenizer:
    """
    所有分词器的"抽象基类"（abstract base class）。

    【思路】
        不管分词器内部怎么实现，对外都只提供两个方法：encode 和 decode。
        先把"接口"定下来，后面每种分词器只需要"填空"。
        这样，比较函数（比如 get_compression_ratio）就可以对任何分词器通用。

    【Python 小知识】类的继承
        class ByteTokenizer(Tokenizer): 表示 ByteTokenizer "继承"自 Tokenizer，
        它自动拥有 Tokenizer 的所有方法，并可以"重写（override）"其中的方法。
        基类里的 raise NotImplementedError 意思是"子类必须自己实现这个方法"，
        如果子类忘了实现，调用时就会报错提醒你。
    """

    name = "抽象分词器"

    def encode(self, text):
        """把字符串变成 token id 列表。子类必须实现。"""
        raise NotImplementedError

    def decode(self, ids):
        """把 token id 列表变回字符串。子类必须实现。"""
        raise NotImplementedError

    def vocab_size(self):
        """词表大小（token id 的取值个数）。子类必须实现。"""
        raise NotImplementedError


def get_compression_ratio(text, ids):
    """
    计算压缩率 = UTF-8 字节数 / token 数。

    【思路】
        压缩率表示"平均每个 token 装了多少字节的原文"。越大越好：
        同样长度的 token 序列能装下更多文字，模型能"看"得更远，计算也更省。
        以字节（而不是字符）为分子，是为了让不同语言的文本可以公平比较。

    【参数】
        text : str，原始文本。
        ids  : list[int]，这段文本编码后的 token 序列。

    【返回】
        float，压缩率。字节级分词器恒等于 1.0；GPT 系列在英文上大约是 4 左右。

    【边界】
        - ids 为空（比如 text 是空字符串）：压缩率没有意义（除以 0），抛出 ValueError。
    """
    if len(ids) == 0:
        raise ValueError("token 序列为空，无法计算压缩率（会除以 0）")
    num_bytes = len(text.encode("utf-8"))
    return num_bytes / len(ids)


def check_round_trip(tokenizer, text):
    """
    检查 decode(encode(text)) 是否等于 text（"往返一致性"）。

    【思路】
        一个好的分词器应当是**无损**的：编码后再解码，必须原样还原。
        否则模型训练时"看到的"就和原文不一样，生成时也可能输出错误的文字。

    【返回】
        bool，True 表示无损。
    """
    return tokenizer.decode(tokenizer.encode(text)) == text


# =============================================================================
# 第 2 部分：字符级分词器
# =============================================================================
class CharacterTokenizer(Tokenizer):
    """
    字符级分词器：token id 就是字符的 Unicode 码位。

    【思路】
        最直接的想法：Unicode 已经给每个字符编好了号，直接拿来用。
            encode("Hi你") = [ord('H'), ord('i'), ord('你')] = [72, 105, 20320]

    【优点】
        - 实现只有一行；
        - 无损（任何字符串都能编码再还原）。

    【缺点】（讲义第 7.3 节）
        - 词表巨大：码位范围有 1,114,112 个，已分配字符约 15 万个。
          模型的输入嵌入矩阵和输出层都要为每个 token 准备一行参数，
          其中绝大多数字符在训练数据里几乎不出现 → 参数浪费、罕见字符学不好。
        - 压缩率一般：英文每个 token 只有 1 个字节。

    【边界】
        - decode 时如果 id 不是合法码位（< 0 或 > 0x10FFFF），chr() 会抛 ValueError。
          模型的输出层如果真的有 111 万个输出，这种情况不会发生；我们仍然检查并给出中文报错。
    """

    name = "字符级"

    def encode(self, text):
        # 【Python 小知识】列表推导式：对 text 中每个字符 c，计算 ord(c)，组成新列表
        return [ord(c) for c in text]

    def decode(self, ids):
        try:
            return "".join(chr(i) for i in ids)
        except ValueError as e:
            raise ValueError(f"存在非法码位，无法解码：{e}") from e

    def vocab_size(self):
        # 所有可能的码位个数。实际系统可能只保留"训练数据里见过的字符"，但原理相同。
        return 0x110000


# =============================================================================
# 第 3 部分：字节级分词器
# =============================================================================
class ByteTokenizer(Tokenizer):
    """
    字节级分词器：先把字符串编码成 UTF-8 字节，每个字节就是一个 token。

    【思路】
        encode("Hi你") = list("Hi你".encode("utf-8")) = [72, 105, 228, 189, 160]

    【优点】
        - 词表极小且固定：只有 256 个 token（0~255）；
        - 永远不会遇到"未知 token"：任何文本（任何语言、emoji、代码、乱码）都能编码。

    【缺点】（讲义第 8.3 节）
        - 压缩率恒为 1.0：序列非常长。中文每个字要 3 个 token，emoji 要 4 个。
        - Transformer 的注意力计算量随序列长度 **平方** 增长，序列长 4 倍，
          注意力的计算量就大约变成 16 倍。

    【边界】
        - decode 时，如果 id 不在 0~255 之间，bytes() 会抛 ValueError。
        - 如果 ids 不是合法的 UTF-8（比如模型生成了"半个汉字"），
          我们用 errors="replace" 把非法部分替换成 '�'，而不是让程序崩溃。
          这意味着对"非法的 id 序列"，decode 不是无损的——但对 encode 产生的
          合法序列，它总是无损的。
    """

    name = "字节级"

    def encode(self, text):
        # 【Python 小知识】对 bytes 调用 list()，得到每个字节的整数值组成的列表
        return list(text.encode("utf-8"))

    def decode(self, ids):
        try:
            data = bytes(ids)
        except ValueError as e:
            raise ValueError(f"字节级 token id 必须在 0~255 之间：{e}") from e
        return data.decode("utf-8", errors="replace")

    def vocab_size(self):
        return 256


# =============================================================================
# 第 4 部分：词级分词器
# =============================================================================
# 【思路】
#   用正则表达式把文本切成"词"和"符号"，然后给每个不同的词分配一个编号。
#
#   我们用的模式是 r"\w+|."（CS336 课上的例子也是这个）：
#       \w+  ：一个或多个"单词字符"（字母、数字、下划线；在 Python 3 里也包括中文等）
#       |    ：或者
#       .    ：任意一个其他字符（标点、空格……）
#
#   【边界：一个隐蔽的坑】
#   在正则里，"." 默认**不匹配换行符 \n**！如果不加 re.DOTALL 标志，
#   re.findall 会悄悄跳过所有换行符，导致分词结果丢失信息、解码无法还原。
#   我们在 demo 里专门演示了这个坑。
WORD_PATTERN = re.compile(r"\w+|.", flags=re.DOTALL)
WORD_PATTERN_BUGGY = re.compile(r"\w+|.")   # 故意不加 DOTALL，用于演示 bug


class WordTokenizer(Tokenizer):
    """
    词级分词器：一个"词"（或一个标点/空格）= 一个 token。

    【思路】
        1. 训练（train）：在语料上把所有词切出来，数一数每个词出现几次，
           给最常见的若干个词分配编号，构成词表。
        2. 编码（encode）：把新文本切成词，查表得到编号；
           查不到的词（训练时没见过）统一映射为特殊 token <unk>（unknown）。
        3. 解码（decode）：查表把编号变回词，再拼起来。

    【优点】
        - 英文上压缩率高：一个常见单词（平均约 4~5 个字母）只占 1 个 token。

    【缺点】（讲义第 9.3 节）
        - 词表无界：新词、人名、拼写错误、数字组合……永远有没见过的词。
        - 未知词变成 <unk>，信息**永久丢失**，decode 无法还原 → 有损。
        - 很多词只出现一两次，模型学不好它们。
        - 中文没有空格，\\w+ 会把一整句中文当成"一个词"——完全不可用。
          （这里写成两个反斜杠，是因为在普通字符串里单个反斜杠有特殊含义，
           Python 3.12+ 看到 "\\w" 这种不认识的转义会给出 SyntaxWarning。）

    【边界】
        - 必须先调用 train() 才能 encode；否则词表里只有 <unk>，所有词都会变成 <unk>。
        - max_vocab_size 包含 <unk> 本身，所以至少为 1。
    """

    name = "词级"
    UNK = "<unk>"

    def __init__(self, pattern=WORD_PATTERN):
        """
        【参数】
            pattern : 编译好的正则表达式对象，用来切词。默认 WORD_PATTERN。
        """
        self.pattern = pattern
        # id 0 固定留给 <unk>
        self.id_to_word = [self.UNK]
        self.word_to_id = {self.UNK: 0}

    def split(self, text):
        """
        把文本切成"词"的列表（尚未转成 id）。

        【返回】
            list[str]。例如 "I'll go!" → ['I', "'", 'll', ' ', 'go', '!']
        """
        # 【Python 小知识】re.findall(模式, 文本) 返回所有匹配到的子串组成的列表
        return self.pattern.findall(text)

    def train(self, texts, max_vocab_size=None, min_count=1):
        """
        在一组文本上"训练"词表。

        【参数】
            texts          : list[str]，训练文本（比如若干篇文档）。
            max_vocab_size : int 或 None。词表最大容量（含 <unk>）。None 表示不限制。
            min_count      : int，至少出现这么多次的词才收进词表（>= 1）。

        【返回】
            self（返回自己，方便链式调用：tok = WordTokenizer().train(texts)）

        【边界】
            - max_vocab_size < 1：不合法，抛 ValueError。
            - 多次调用 train 会**重置**词表，而不是累加。
        """
        if max_vocab_size is not None and max_vocab_size < 1:
            raise ValueError("max_vocab_size 至少为 1（要给 <unk> 留位置）")

        counts = Counter()
        for t in texts:
            counts.update(self.split(t))   # Counter.update 会把列表里每个元素的计数 +1

        # 【Python 小知识】counts.most_common() 返回 [(词, 次数), ...]，按次数从大到小排序
        self.id_to_word = [self.UNK]
        for word, c in counts.most_common():
            if c < min_count:
                break
            if max_vocab_size is not None and len(self.id_to_word) >= max_vocab_size:
                break
            self.id_to_word.append(word)
        # 反向映射：词 → id。enumerate 同时给出下标和元素
        self.word_to_id = {w: i for i, w in enumerate(self.id_to_word)}
        self.train_counts = counts      # 保存下来，演示 Zipf 定律时用
        return self

    def encode(self, text):
        unk_id = self.word_to_id[self.UNK]
        # dict.get(key, default)：查不到 key 时返回 default，而不是报错
        return [self.word_to_id.get(w, unk_id) for w in self.split(text)]

    def decode(self, ids):
        try:
            return "".join(self.id_to_word[i] for i in ids)
        except IndexError as e:
            raise ValueError(f"token id 超出词表范围（词表大小 {len(self.id_to_word)}）") from e

    def vocab_size(self):
        return len(self.id_to_word)


# =============================================================================
# 第 5 部分：演示
# =============================================================================
COURSE_EXAMPLE = "Hello, 🌍! 你好!"   # CS336 第 1 讲使用的例句


def show_encoding(tokenizer, text, max_show=30):
    """
    打印一个分词器对一段文本的编码结果（辅助函数）。

    【参数】
        tokenizer : 任何 Tokenizer 子类的实例
        text      : str
        max_show  : int，最多显示多少个 id（太长就截断，后面用 ... 表示）
    """
    ids = tokenizer.encode(text)
    shown = ids[:max_show]
    more = " ..." if len(ids) > max_show else ""
    print(f"  文本：{text!r}")
    print(f"  ids ：{shown}{more}")
    print(f"  token 数 = {len(ids)}，字节数 = {len(text.encode('utf-8'))}，"
          f"压缩率 = {get_compression_ratio(text, ids):.2f}，"
          f"往返一致 = {check_round_trip(tokenizer, text)}")


def demo_character_tokenizer(docs):
    """演示 1：字符级分词器。"""
    print_header("演示 1：字符级分词器（token = Unicode 码位）")
    tok = CharacterTokenizer()
    show_encoding(tok, COURSE_EXAMPLE)
    print(f"  词表大小（所有码位）= {tok.vocab_size():,}")

    # 看看语料里实际出现了多少种字符、它们的频率分布如何
    char_counts = Counter("".join(docs))
    total = sum(char_counts.values())
    print(f"\n  我们的小语料共 {total:,} 个字符，其中不同的字符有 {len(char_counts)} 种。")
    rare = [c for c, n in char_counts.items() if n == 1]
    print(f"  只出现过 1 次的字符有 {len(rare)} 种，例如：{rare[:15]}")
    print("  → 哪怕在这个很小的语料里，也有大量'只见过一次'的字符。")
    print("    要是把 111 万个码位都放进词表，绝大多数 token 永远学不到。")


def demo_byte_tokenizer():
    """演示 2：字节级分词器。"""
    print_header("演示 2：字节级分词器（token = UTF-8 字节）")
    tok = ByteTokenizer()
    show_encoding(tok, COURSE_EXAMPLE)
    print(f"  词表大小 = {tok.vocab_size()}")

    zh = "从前，有一个叫小明的男孩。"
    print()
    show_encoding(tok, zh)
    print(f"  这句中文只有 {len(zh)} 个字符，却变成了 {len(tok.encode(zh))} 个 token（每个汉字 3 个）。")

    # 边界：decode 一段"不完整"的字节
    half = tok.encode("你")[:2]
    print(f"\n  【边界】decode({half})（'你' 的前 2 个字节）= {tok.decode(half)!r}  ← 非法部分变成了 �")


def demo_word_tokenizer(docs):
    """演示 3：词级分词器。"""
    print_header("演示 3：词级分词器（token = 单词或符号）")

    # --- 先演示 re.DOTALL 的坑 ---
    text_with_newline = "Hi!\nBye."
    good = WORD_PATTERN.findall(text_with_newline)
    bad = WORD_PATTERN_BUGGY.findall(text_with_newline)
    print(f"  【边界】切分 {text_with_newline!r}：")
    print(f"    带 re.DOTALL  → {good}")
    print(f"    不带 re.DOTALL → {bad}   ← 换行符 '\\n' 被悄悄丢掉了！")
    print(f"    拼回去是否还原：带 DOTALL {''.join(good) == text_with_newline}，"
          f"不带 {''.join(bad) == text_with_newline}")

    # --- 训练 ---
    tok = WordTokenizer().train(docs)
    print(f"\n  在 {len(docs)} 篇文档上训练后，词表大小 = {tok.vocab_size()}（含 <unk>）")

    # --- Zipf 定律：少数词非常常见，大量词非常罕见 ---
    counts = tok.train_counts
    print("  最常见的 10 个'词'：", counts.most_common(10))
    hapax = sum(1 for n in counts.values() if n == 1)
    print(f"  只出现 1 次的'词'有 {hapax} 个，占词表的 {hapax / len(counts):.0%}")
    print("  → 这就是 Zipf 定律（长尾分布）：词表的大部分位置被罕见词占据。")

    # --- 编码课上的例句：遇到没见过的词 ---
    print()
    example = "I'll say supercalifragilisticexpialidocious!"
    print(f"  切分结果：{tok.split(example)}")
    show_encoding(tok, example)
    ids = tok.encode(example)
    print(f"  解码结果：{tok.decode(ids)!r}")
    print("  → 没见过的词变成了 <unk>，原文信息永久丢失（有损）。")

    # --- 中文：没有空格，整句被当成一个"词" ---
    zh = "我爱自然语言处理"
    print(f"\n  中文 {zh!r} 的切分结果：{tok.split(zh)}")
    print("  → \\w+ 在 Python 3 中也匹配汉字，于是没有空格的一整串汉字被当成'一个词'。")
    print("    中文需要专门的分词工具（如 jieba），而且切法本身有歧义，")
    print("    例如'南京市长江大桥'可以切成'南京市/长江/大桥'或'南京/市长/江大桥'。")


def demo_compare_all(docs):
    """演示 4：三种分词器在同一组文本上比较。"""
    print_header("演示 4：三种分词器横向比较")

    word_tok = WordTokenizer().train(docs[:-2])   # 留出最后 2 篇文档不参与训练，模拟"新文本"
    tokenizers = [CharacterTokenizer(), ByteTokenizer(), word_tok]

    test_texts = {
        "课程例句": COURSE_EXAMPLE,
        "英文句子": "Once upon a time, there was a little dog named Max.",
        "英文新词": "The astronaut photographed Jupiter's moons.",
        "中文句子": "小青蛙和小蜗牛一起看彩虹。",
    }
    header = f"{'文本':<8}" + "".join(f"{t.name + '(token数/压缩率/无损)':>26}" for t in tokenizers)
    print(header)
    for name, text in test_texts.items():
        row = f"{name:<8}"
        for t in tokenizers:
            ids = t.encode(text)
            ratio = get_compression_ratio(text, ids)
            ok = "是" if check_round_trip(t, text) else "否"
            row += f"{len(ids):>14} / {ratio:>4.2f} / {ok}"
        print(row)
    print(f"\n词表大小：字符级 {tokenizers[0].vocab_size():,}，字节级 {tokenizers[1].vocab_size()}，"
          f"词级 {tokenizers[2].vocab_size()}（只在很小的语料上训练，真实语料会大得多）")
    print("""
小结：
  - 字符级：无损，但词表巨大、大部分 token 罕见。
  - 字节级：无损、词表小，但压缩率恒为 1，序列太长。
  - 词级  ：英文压缩率高，但遇到新词就变 <unk>（有损），对中文几乎不可用。
  → 我们想要一个"折中"：常见片段合成一个 token，罕见片段退化成字节。这就是 BPE。""")


def demo_sequence_length_cost():
    """演示 5：压缩率对"能看多远"和"计算量"的影响。"""
    print_header("演示 5：为什么压缩率这么重要？")

    context_len = 1024   # 假设模型的上下文窗口是 1024 个 token
    print(f"假设上下文窗口 = {context_len} 个 token：")
    for name, ratio in [("字节级", 1.0), ("字符级(英文)", 1.0), ("字符级(中文)", 3.0), ("BPE(英文，GPT 系列)", 4.4)]:
        print(f"  {name:<20} 压缩率 {ratio:>3}  → 一次能看到约 {context_len * ratio:,.0f} 字节的原文")

    print("\n反过来：同一段 4,400 字节的英文，用不同分词器需要多少 token？注意力计算量（∝ n²）差多少？")
    text_bytes = 4400
    base_tokens = text_bytes / 4.4
    for name, ratio in [("BPE", 4.4), ("字节级", 1.0)]:
        n = text_bytes / ratio
        print(f"  {name:<6} n = {n:>6,.0f} 个 token，注意力计算量 ∝ n² = {n * n:>12,.0f}"
              f"（是 BPE 的 {(n / base_tokens) ** 2:.1f} 倍）")


def main():
    docs = load_corpus_documents()
    demo_character_tokenizer(docs)
    demo_byte_tokenizer()
    demo_word_tokenizer(docs)
    demo_compare_all(docs)
    demo_sequence_length_cost()
    print("\n全部演示运行完毕。下一个文件 03_bpe_from_scratch.py 会实现 BPE。")


if __name__ == "__main__":
    main()