"""
Lecture 19 -- Byte Pair Encoding (BPE) merge-training from scratch.
Toy implementation on the instructor's example corpus. Uses classic
BPE (pair FREQUENCY, summed across word counts) -- the algorithm
GPT-style tokenizers use. (WordPiece, used by BERT, scores pairs
differently -- see the lecture notes.)
"""
from collections import defaultdict, Counter

corpus = ["huggingface", "hugging", "face", "hug",
          "hugger", "learning", "learner", "learners", "learn"]
word_freq = Counter(corpus)   # each word appears once here

def split_word(w):
    """Represent a word as a list of symbols; '##' marks a
    non-initial character, following the lecture's convention."""
    return [w[0]] + ["##" + c for c in w[1:]]

splits = {w: split_word(w) for w in word_freq}

def get_vocab(splits):
    v = set()
    for symbols in splits.values():
        v.update(symbols)
    return v

def pair_freqs(splits, word_freq):
    freqs = defaultdict(int)
    for w, symbols in splits.items():
        for a, b in zip(symbols, symbols[1:]):
            freqs[(a, b)] += word_freq[w]
    return freqs

def merge_pair(pair, splits):
    a, b = pair
    merged = a + (b[2:] if b.startswith("##") else b)
    new_splits = {}
    for w, symbols in splits.items():
        out, i = [], 0
        while i < len(symbols):
            if i < len(symbols) - 1 and symbols[i] == a and symbols[i+1] == b:
                out.append(merged); i += 2
            else:
                out.append(symbols[i]); i += 1
        new_splits[w] = out
    return new_splits, merged

print("Initial vocab:", sorted(get_vocab(splits)))

NUM_MERGES = 6
for step in range(1, NUM_MERGES + 1):
    freqs = pair_freqs(splits, word_freq)
    if not freqs:
        break
    best_pair = max(freqs, key=freqs.get)
    splits, merged_token = merge_pair(best_pair, splits)
    print(f"Merge {step}: {best_pair} (freq={freqs[best_pair]}) -> '{merged_token}'")

print("\nFinal splits:")
for w in corpus:
    print(f"  {w:12s} -> {splits[w]}")