"""第 17 章：字符起步的微型 BPE、字节与嵌入查表；只用标准库。"""

from collections import Counter
from random import Random
import unicodedata


TEXTS = [
    "打印机准备好了。", "打印机正在打印。",
    "请打印报告。", "请打印照片。",
    "打印报告。", "打印照片。",
]
SPECIAL = ["<PAD>", "<BOS>", "<EOS>", "<UNK>"]


def normalize(text):
    # 保留大小写与空格，只统一换行和规范等价的 Unicode 写法。
    return unicodedata.normalize("NFC", text.replace("\r\n", "\n").replace("\r", "\n"))


def merge(tokens, pair):
    result, index = [], 0
    while index < len(tokens):
        if index+1 < len(tokens) and tuple(tokens[index:index+2]) == pair:
            result.append("".join(pair))
            index += 2
        else:
            result.append(tokens[index])
            index += 1
    return result


class TinyBPE:
    def __init__(self, texts, max_merges=6):
        sequences = [list(normalize(text)) for text in texts]
        self.base = sorted({token for tokens in sequences for token in tokens})
        self.merges, self.trace = [], []
        for _ in range(max_merges):
            counts = Counter(pair for tokens in sequences
                             for pair in zip(tokens, tokens[1:]))
            if not counts:
                break
            # 同频时按字符串排序，避免运行顺序改变结果。
            pair = min(counts, key=lambda p: (-counts[p], p))
            if counts[pair] < 2:
                break
            self.merges.append(pair)
            sequences = [merge(tokens, pair) for tokens in sequences]
            self.trace.append((pair, counts[pair], sum(map(len, sequences))))
        self.vocab = list(dict.fromkeys(SPECIAL+self.base+
                                       ["".join(p) for p in self.merges]))
        self.ids = {token: index for index, token in enumerate(self.vocab)}

    def tokenize(self, text):
        tokens = [c if c in self.base else "<UNK>" for c in normalize(text)]
        for pair in self.merges:
            tokens = merge(tokens, pair)
        return tokens

    def encode(self, text):
        tokens = ["<BOS>"]+self.tokenize(text)+["<EOS>"]
        return [self.ids[token] for token in tokens]

    def decode(self, ids):
        ignored = {"<PAD>", "<BOS>", "<EOS>"}
        return "".join(self.vocab[i] for i in ids if self.vocab[i] not in ignored)


def main():
    tokenizer = TinyBPE(TEXTS)
    for pair, count, total in tokenizer.trace:
        print(f"合并 {pair[0]} + {pair[1]}，出现 {count} 次，训练文本剩 {total} 个词元")
    text = "请打印报告。"
    ids = tokenizer.encode(text)
    print(f"字符切分：{list(text)}")
    print(f"BPE 切分：{tokenizer.tokenize(text)}")
    print(f"含起止符的 ID：{ids}，词表大小={len(tokenizer.vocab)}")
    print(f"解码结果：{tokenizer.decode(ids)}")
    print(f"未见字符：{tokenizer.decode(tokenizer.encode('打印🙂'))}")
    raw = "茶🙂".encode("utf-8")
    print(f"'茶🙂' 的 UTF-8 字节：{list(raw)}，恢复={raw.decode('utf-8')}")
    same = normalize("e\u0301") == normalize("é")
    print(f"NFC 合并等价写法：{same}")
    print(f"NFC 保留全角：{normalize('Ａ')}；NFKC 转为："
          f"{unicodedata.normalize('NFKC', 'Ａ')}")
    rng = Random(7)
    table = [[rng.uniform(-1, 1) for _ in range(3)] for _ in tokenizer.vocab]
    vectors = [table[index] for index in ids]
    print(f"嵌入查表形状：{len(vectors)}×{len(vectors[0])}（随机值，未经训练）")
    assert tokenizer.decode(ids) == normalize(text)


if __name__ == "__main__":
    main()