"""
bpe.py —— 第 3 讲 §6：BPE 分词器的训练（教学版 + 增量计数的快速版）
=====================================================================

【这个文件要回答的问题】
    第 1 讲的 toy_tokenizers.BPETokenizer 每做一次合并，都把整个语料里所有相邻对重新数一遍。
    slides 指出这个朴素做法很慢（"Step of finding most frequent pairs has a naive runtime of O(n²)"），
    解决办法是"缓存出现次数"。这里实现：

        char_bpe_merges(word_counts, k)   slides 上 hug / pug / pun / bun / hugs 的字符级例子，一步步打印
        FastBPE                          字节级 BPE，增量维护"每对相邻 token 出现多少次"和
                                         "这一对出现在哪些词里"，每次合并只更新受影响的词

    FastBPE 和第 1 讲的朴素实现在同样的输入上学出**完全相同**的合并规则（06 号脚本逐条核对），只是快得多。

【增量计数的思路】
    合并 (a, b) → ab 时，只有"含有 a b 相邻"的那些词会改变。对每个这样的词：
        1. 先把它旧的所有相邻对从总计数里减掉（乘以词频）；
        2. 在它内部把 a b 换成 ab；
        3. 再把新的所有相邻对加回总计数。
    同时维护 where[pair] = 含有这一对的词的集合，于是"找受影响的词"不用扫描整个语料。
    （工业实现，如 HuggingFace tokenizers，还会用堆来找最大的一对；这里每次扫一遍计数表，已经够快。）

【LLM 联系】
    GPT-2 起的字节级 BPE：从 256 个字节出发，任何文字都能编码，不会有 <unk>；
    Llama 3 的词表 128K、GPT-4o 的约 200K 都是用类似的算法在 TB 级文本上训练出来的。
"""

import os
import sys
from collections import Counter, defaultdict

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))
from toy_tokenizers import pretokenize, show  # noqa: E402


def char_bpe_merges(word_counts, num_merges, verbose=True):
    """
    字符级 BPE，直接在"词 → 频率"表上做（slides 上的高效实现思路）。

    【参数】word_counts：{词: 频率}；num_merges：合并次数。
    【返回】(合并列表 [(a, b), ...], 最终每个词的切分 {词: [token, ...]})。
    【例子】slides：("hug",10) ("pug",5) ("pun",12) ("bun",4) ("hugs",5)
            第 1 次合并 u+g（出现 10+5+5 = 20 次）……
    【边界】出现次数并列时选字典序最小的一对，保证结果确定。
    """
    seqs = {w: list(w) for w in word_counts}
    merges = []
    for it in range(num_merges):
        pairs = Counter()
        for w, c in word_counts.items():
            s = seqs[w]
            for a, b in zip(s, s[1:]):
                pairs[(a, b)] += c
        if not pairs:
            break
        best = min(pairs, key=lambda p: (-pairs[p], p))
        merges.append(best)
        for w in seqs:
            s, out, i = seqs[w], [], 0
            while i < len(s):
                if i + 1 < len(s) and (s[i], s[i + 1]) == best:
                    out.append(s[i] + s[i + 1])
                    i += 2
                else:
                    out.append(s[i])
                    i += 1
            seqs[w] = out
        if verbose:
            vocab = sorted({t for s in seqs.values() for t in s} | {a + b for a, b in merges})
            seg = "，".join(f"({' '.join(seqs[w])}, {c})" for w, c in word_counts.items())
            print(f"    第 {it + 1} 次：合并 '{best[0]}' + '{best[1]}'（出现 {pairs[best]} 次）→ {seg}")
    return merges, seqs


class FastBPE:
    """
    字节级 BPE，增量计数版本。接口与第 1 讲的 BPETokenizer 相同：train / encode / decode / tokens / vocab_size。

    【训练的思路】见文件开头。
    【编码的思路】和第 1 讲一样：对每个预分词的小块，按合并规则的先后顺序反复合并；
            同一个小块（例如常见词 " the"）的结果会被缓存，长文本编码快很多。
    【边界】
        - 出现次数并列时，选 (第一个 id, 第二个 id) 最小的那一对——和第 1 讲的实现一致，
          所以两者学出的合并规则逐条相同。
        - 合并从不跨越预分词的边界。
    """

    def __init__(self):
        self.merges, self.ranks = [], {}
        self.token_bytes = {i: bytes([i]) for i in range(256)}
        self._cache = {}

    def train(self, text=None, num_merges=100, word_counts=None):
        """
        【参数】text：训练文字（或直接给 word_counts：{预分词小块: 次数}）；num_merges：合并次数。
        【返回】self。
        """
        if word_counts is None:
            word_counts = Counter(pretokenize(text))
        words = [list(w.encode("utf-8")) for w in word_counts]
        freqs = list(word_counts.values())
        pair_counts = Counter()
        where = defaultdict(set)
        for wi, (s, f) in enumerate(zip(words, freqs)):
            for p in zip(s, s[1:]):
                pair_counts[p] += f
                where[p].add(wi)
        for m in range(num_merges):
            pair_counts = +pair_counts                    # 去掉计数为 0 的项
            if not pair_counts:
                break
            best = max(pair_counts, key=lambda p: (pair_counts[p], -p[0], -p[1]))
            new_id = 256 + m
            self.token_bytes[new_id] = self.token_bytes[best[0]] + self.token_bytes[best[1]]
            self.ranks[best] = m
            self.merges.append(best)
            for wi in list(where[best]):
                s, f = words[wi], freqs[wi]
                for p in zip(s, s[1:]):                    # 1. 减掉旧的相邻对
                    pair_counts[p] -= f
                    where[p].discard(wi)
                out, i = [], 0                             # 2. 合并
                while i < len(s):
                    if i + 1 < len(s) and s[i] == best[0] and s[i + 1] == best[1]:
                        out.append(new_id)
                        i += 2
                    else:
                        out.append(s[i])
                        i += 1
                words[wi] = out
                for p in zip(out, out[1:]):                # 3. 加上新的相邻对
                    pair_counts[p] += f
                    where[p].add(wi)
        self._cache = {}
        return self

    def _encode_chunk(self, chunk):
        if chunk in self._cache:
            return self._cache[chunk]
        seq = list(chunk.encode("utf-8"))
        while len(seq) >= 2:
            rank, pos = min((self.ranks.get((a, b), float("inf")), i) for i, (a, b) in enumerate(zip(seq, seq[1:])))
            if rank == float("inf"):
                break
            a, b = self.merges[rank]
            out, i = [], 0
            while i < len(seq):
                if i + 1 < len(seq) and seq[i] == a and seq[i + 1] == b:
                    out.append(256 + rank)
                    i += 2
                else:
                    out.append(seq[i])
                    i += 1
            seq = out
        self._cache[chunk] = seq
        return seq

    def encode(self, text):
        ids = []
        for chunk in pretokenize(text):
            ids.extend(self._encode_chunk(chunk))
        return ids

    def decode(self, ids):
        return b"".join(self.token_bytes[i] for i in ids).decode("utf-8", errors="replace")

    def token_text(self, tid):
        raw = self.token_bytes[tid]
        try:
            return raw.decode("utf-8")
        except UnicodeDecodeError:
            return "".join(f"<{b:02x}>" for b in raw)

    def tokens(self, text):
        return [self.token_text(t) for t in self.encode(text)]

    def vocab_size(self):
        return 256 + len(self.merges)


__all__ = ["char_bpe_merges", "FastBPE", "pretokenize", "show"]