"""Byte pair encoding, traced step by step as Chapter 5 traces it.

Run this and you get the same corpus tables, the same pair counts, the same
ties and the same five merges the book prints.

    python3 bpe.py                    # the low/lower/newest/widest corpus
    python3 bpe.py --merges 8         # keep going and watch it continue
    python3 bpe.py --encode lowest    # apply the learned rules to a new word
    python3 bpe.py --no-tiebreak      # see why the tie-break rule exists

The two functions that matter are `train` and `encode`. Everything else is
printing.

Install:  nothing, the Python standard library is enough
"""

import argparse
from collections import Counter

END = "</w>"


def pair_counts(words):
    """Every adjacent pair, weighted by the frequency of the word it sits in.

    The weighting is the step most hand simulations get wrong. A pair inside a
    word seen 6 times contributes 6, not 1.
    """
    counts = Counter()
    for symbols, freq in words.items():
        for a, b in zip(symbols, symbols[1:]):
            counts[(a, b)] += freq
    return counts


def merge_pair(words, pair):
    """Rewrite the corpus with `pair` joined into one symbol."""
    a, b = pair
    out = {}
    for symbols, freq in words.items():
        joined, i = [], 0
        while i < len(symbols):
            if i + 1 < len(symbols) and symbols[i] == a and symbols[i + 1] == b:
                joined.append(a + b)
                i += 2
            else:
                joined.append(symbols[i])
                i += 1
        out[tuple(joined)] = out.get(tuple(joined), 0) + freq
    return out


def train(corpus, num_merges, tiebreak=True, trace=None):
    """Learn an ordered merge list from {word: frequency}.

    Returns the merges as "a b" strings, which is the form `encode` expects
    and the form a GPT-2 merges.txt file uses.
    """
    words = {tuple(list(w) + [END]): f for w, f in corpus.items()}
    merges = []
    for step in range(1, num_merges + 1):
        counts = pair_counts(words)
        if not counts:
            break
        top = max(counts.values())
        tied = sorted(p for p, n in counts.items() if n == top)
        # Ties are common in small corpora. Break them lexicographically or
        # the run is not reproducible.
        best = tied[0] if tiebreak else tied[-1]
        if trace is not None:
            trace(step, words, counts, top, tied, best)
        merges.append(" ".join(best))
        words = merge_pair(words, best)
    return merges, words


def encode(word, merges):
    """Apply the merge rules IN LEARNED ORDER to one new word."""
    symbols = list(word) + [END]
    history = [("split", list(symbols), True)]
    for rule in merges:
        a, b = rule.split(" ")
        before = list(symbols)
        symbols = list(merge_pair({tuple(symbols): 1}, (a, b)).keys())[0]
        history.append((rule, list(symbols), list(symbols) != before))
    return list(symbols), history


# ----------------------------------------------------------------- printing

def show_corpus(words, title):
    print(f"\n  {title}")
    print(f"    {'symbols':<28}{'frequency'}")
    for symbols, freq in words.items():
        print(f"    {' '.join(symbols):<28}{freq}")


def show_counts(counts, top, limit=7):
    print(f"\n    {'pair':<16}{'count'}")
    for pair, n in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))[:limit]:
        flag = "   <-- tied for first" if n == top else ""
        print(f"    {' '.join(pair):<16}{n}{flag}")
    if len(counts) > limit:
        print(f"    ... and {len(counts) - limit} more, all lower")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--merges", type=int, default=5)
    ap.add_argument("--encode", default="lowest",
                    help="word to tokenise with the learned rules")
    ap.add_argument("--no-tiebreak", action="store_true",
                    help="take the lexicographically LAST tied pair instead")
    a = ap.parse_args()

    corpus = {"low": 5, "lower": 2, "newest": 6, "widest": 3}
    print("\nCORPUS")
    print(f"  {'word':<10}{'frequency'}")
    for w, f in corpus.items():
        print(f"  {w:<10}{f}")

    def trace(step, words, counts, top, tied, best):
        print(f"\n{'=' * 62}\nSTEP {step}")
        show_corpus(words, "corpus going in:")
        show_counts(counts, top)
        if len(tied) > 1:
            names = ", ".join(f"'{x} {y}'" for x, y in tied)
            print(f"\n    {len(tied)} pairs tie at {top}: {names}")
            print(f"    tie-break takes '{best[0]} {best[1]}'")
        else:
            print(f"\n    clear winner at {top}, no tie")
        print(f"\n    MERGE: '{best[0]} {best[1]}'  ->  {best[0] + best[1]}")

    merges, final = train(corpus, a.merges,
                          tiebreak=not a.no_tiebreak, trace=trace)

    print(f"\n{'=' * 62}")
    show_corpus(final, "corpus after the last merge:")

    print(f"\nTHE MERGE LIST (this ordered list IS the tokeniser)")
    for i, m in enumerate(merges, 1):
        a_, b_ = m.split(" ")
        print(f"  {i}. '{m}'  ->  {a_ + b_}")

    print(f"\nTRAINING WORDS, TOKENISED")
    for w in corpus:
        toks, _ = encode(w, merges)
        print(f"  {w:<10}{' '.join(toks)}")

    print(f"\nAN UNSEEN WORD: {a.encode!r}")
    toks, history = encode(a.encode, merges)
    for rule, state, changed in history:
        label = "split" if rule == "split" else f"rule '{rule}'"
        note = "" if changed or rule == "split" else "   (no match)"
        print(f"  {label:<20}{' '.join(state)}{note}")
    print(f"\n  result: {' '.join(toks)}")


if __name__ == "__main__":
    main()