"""Byte-level BPE tokenizer, written from scratch (minbpe-style).

How BPE works, in one paragraph:
Start from raw bytes (256 possible values — covers ANY text, so there is no
such thing as an "unknown" character). Count which PAIR of tokens appears most
often in the data, merge that pair into one new token, repeat. After a few
thousand merges the vocabulary contains the most useful chunks of THIS dataset
("the", " said", "once upon"), and any text encodes into far fewer tokens.

Contract notes: A1 (lossless round-trip) holds by construction — decode is the
exact inverse of encode. A2 (no unknowns) holds because the base is all 256 bytes.
"""

import json


class BPETokenizer:
    def __init__(self):
        # merges: (token_a, token_b) -> new_token_id, in the ORDER they were learned
        self.merges = {}
        # vocab: token_id -> bytes it represents
        self.vocab = {i: bytes([i]) for i in range(256)}

    # ---------------- training ----------------

    def train(self, text, vocab_size, progress_every=200):
        """Learn merges until the vocabulary reaches vocab_size."""
        assert vocab_size > 256, "vocab must exceed the 256 base bytes"
        ids = list(text.encode("utf-8"))
        num_merges = vocab_size - 256

        for i in range(num_merges):
            pairs = self._count_pairs(ids)
            if not pairs:
                break  # nothing left to merge (tiny input)
            best = max(pairs, key=lambda p: pairs[p])  # most frequent pair
            new_id = 256 + i
            ids = self._merge(ids, best, new_id)      # replace pair everywhere
            self.merges[best] = new_id
            self.vocab[new_id] = self.vocab[best[0]] + self.vocab[best[1]]
            if (i + 1) % progress_every == 0:
                print(f"  merge {i + 1}/{num_merges}: {best} -> {new_id} "
                      f"({self.vocab[new_id]!r}), data length {len(ids):,}")

    @staticmethod
    def _count_pairs(ids):
        counts = {}
        for a, b in zip(ids, ids[1:]):
            counts[(a, b)] = counts.get((a, b), 0) + 1
        return counts

    @staticmethod
    def _merge(ids, pair, new_id):
        out = []
        i = 0
        while i < len(ids):
            if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
                out.append(new_id)
                i += 2
            else:
                out.append(ids[i])
                i += 1
        return out

    # ---------------- encode / decode ----------------

    def encode(self, text):
        """Text -> token ids, applying merges in the order they were learned."""
        ids = list(text.encode("utf-8"))
        while len(ids) >= 2:
            pairs = self._count_pairs(ids)
            # among pairs present, pick the one learned EARLIEST (lowest new id)
            candidate = min(pairs, key=lambda p: self.merges.get(p, float("inf")))
            if candidate not in self.merges:
                break  # no learned merge applies anymore
            ids = self._merge(ids, candidate, self.merges[candidate])
        return ids

    def decode(self, ids):
        """Token ids -> text. Exact inverse of encode."""
        data = b"".join(self.vocab[i] for i in ids)
        return data.decode("utf-8", errors="replace")

    # ---------------- save / load ----------------

    def save(self, path):
        payload = {
            "merges": [[a, b, new] for (a, b), new in self.merges.items()],
        }
        with open(path, "w", encoding="utf-8") as f:
            json.dump(payload, f)

    @classmethod
    def load(cls, path):
        tok = cls()
        with open(path, "r", encoding="utf-8") as f:
            payload = json.load(f)
        for a, b, new in payload["merges"]:
            tok.merges[(a, b)] = new
            tok.vocab[new] = tok.vocab[a] + tok.vocab[b]
        return tok

    @property
    def vocab_size(self):
        return len(self.vocab)


if __name__ == "__main__":
    # self-check: train on a tiny sample, round-trip must be exact
    sample = "Once upon a time there was a little dog. The dog liked to play."
    tok = BPETokenizer()
    tok.train(sample * 20, vocab_size=300, progress_every=50)
    ids = tok.encode(sample)
    assert tok.decode(ids) == sample, "round-trip failed"
    print(f"self-check OK: {len(sample)} chars -> {len(ids)} tokens, "
          f"vocab {tok.vocab_size}")