"""Byte-level byte pair encoding (BPE), small enough to read in one sitting.

Ids 0 to 255 are the 256 possible bytes. Training adds one new id per merge: merge number k
(counting from 0) makes id 256 + k out of a pair of existing ids. Encoding then applies the
same merges to new text, so any string can be written as ids and read back exactly.

This file is lesson 1's bpe.py, copied so that this lesson stands alone. The widget in ../widget.html
carries the merges learned here and runs encode() in JavaScript. Two things must stay identical in
both: the splitting pattern PIECE and the merge order that encode() follows.
"""
import re
from collections import Counter
from pathlib import Path

# Split text into pieces first, so a merge never joins across a word and the punctuation after it.
# Letters, digits, punctuation and whitespace each form their own runs; a single leading space
# stays with the word that follows, so " the" and "the" are different pieces. Any character
# outside ASCII counts as a letter. The six alternatives cover every character exactly once.
PIECE = re.compile(
    r" ?[A-Za-z\u0080-\U0010ffff]+"
    r"| ?[0-9]+"
    r"| ?[!-/:-@\[-`{-~]+"
    r"|[ \t\r\n]+(?![^ \t\r\n])"
    r"|[ \t\r\n]+"
    r"|[\x00-\x1f\x7f]"
)

MIN_COUNT = 2  # a pair seen only once is not worth a vocabulary slot
NUM_MERGES = 512  # the lesson's vocabulary: 256 bytes + 512 merges = 768 ids
CORPUS_PATH = Path(__file__).with_name("corpus.txt")


def load_corpus():
    """The training text: written for this lesson, about 15 KB of prose, code, numbers and Korean."""
    return CORPUS_PATH.read_text(encoding="utf-8")


def pieces(text):
    """Split text into the pieces that BPE works inside. Joined back together they equal the text."""
    return PIECE.findall(text)


def train_counted(text, num_merges=NUM_MERGES, min_count=MIN_COUNT):
    """Learn up to num_merges merges. Returns a list of ((left_id, right_id), count) in order.

    Each round counts every adjacent pair of ids inside every piece and merges the most frequent
    one. A tie goes to the smaller (left_id, right_id), so training is deterministic. Training
    stops early when the best pair appears fewer than min_count times.
    """
    words = Counter(tuple(p.encode("utf-8")) for p in pieces(text))
    learned = []
    for _ in range(num_merges):
        pairs = Counter()
        for word, freq in words.items():
            for pair in zip(word, word[1:]):
                pairs[pair] += freq
        if not pairs:
            break
        best = min(pairs, key=lambda pair: (-pairs[pair], pair))
        if pairs[best] < min_count:
            break
        learned.append((best, pairs[best]))
        words = _merge_words(words, best, 255 + len(learned))
    return learned


def train(text, num_merges=NUM_MERGES, min_count=MIN_COUNT):
    """Like train_counted, keeping only the pairs."""
    return [pair for pair, _ in train_counted(text, num_merges, min_count)]


def _apply(word, pair, new_id):
    """Replace every left-to-right occurrence of pair in a tuple of ids with new_id."""
    out = []
    i = 0
    while i < len(word):
        if i + 1 < len(word) and (word[i], word[i + 1]) == pair:
            out.append(new_id)
            i += 2
        else:
            out.append(word[i])
            i += 1
    return tuple(out)


def _merge_words(words, pair, new_id):
    merged = Counter()
    for word, freq in words.items():
        merged[_apply(word, pair, new_id)] += freq
    return merged


def vocab_bytes(merges):
    """The bytes each id stands for: 256 single bytes, then one entry per merge."""
    table = [bytes([b]) for b in range(256)]
    for left, right in merges:
        table.append(table[left] + table[right])
    return table


def encode(text, merges):
    """Turn text into ids. Within a piece, always apply the earliest-learned merge that is present."""
    rank = {pair: k for k, pair in enumerate(merges)}
    cache = {}
    ids = []
    for piece in pieces(text):
        if piece not in cache:
            word = list(piece.encode("utf-8"))
            while len(word) > 1:
                best = min(zip(word, word[1:]), key=lambda pair: rank.get(pair, len(rank)))
                if best not in rank:
                    break
                word = list(_apply(word, best, 256 + rank[best]))
            cache[piece] = word
        ids.extend(cache[piece])
    return ids


def decode(ids, merges):
    """Turn ids back into text. Bytes that are not valid UTF-8 become the replacement character."""
    table = vocab_bytes(merges)
    return b"".join(table[i] for i in ids).decode("utf-8", errors="replace")