"""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.

The widget in ../widget.html runs the same rules in JavaScript. Two things must stay identical in
both: the splitting pattern PIECE and the tie-break rule in train_counted().
"""
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")