"""Byte-level Byte-Pair Encoding (BPE) tokenizer, implemented from scratch.

Design choices (optimised for readability over speed):

- **Byte-level.** The base vocabulary is the 256 possible byte values, so any
  UTF-8 text is representable and there are no "unknown token" holes.
- **Word pre-tokenization.** Text is first split into word-ish chunks with a
  regex (GPT-2 style, keeping a word's leading space attached: " the" is one
  chunk). BPE merges only ever happen *within* a chunk, never across word or
  whitespace boundaries — this is what keeps the learned tokens sensible.
- **Special tokens.** Strings like ``<|endoftext|>`` are reserved: they get a
  fixed id, are never split by BPE, and decode back to their literal text.

ID layout:  0..255 raw bytes | 256..256+M-1 merges | then one id per special.

The trainer recounts pair frequencies from scratch on every merge. That is the
simple, legible algorithm; it operates over *unique* pre-tokenized words (not
the raw byte stream), which keeps it fast enough to train a 4k vocab on a
~50 MB sample in a few minutes.
"""

from __future__ import annotations

import json
import re
from collections import Counter
from pathlib import Path

# Tiles the input completely: a word (optionally led by one space), or a run of
# punctuation/symbols (optionally led by one space), or a run of whitespace, or
# — as a catch-all so no character is ever dropped — any single character.
# re.DOTALL lets the final "." also match newlines.
_SPLIT_PATTERN = re.compile(r" ?\w+| ?[^\w\s]+|\s+|.", re.DOTALL)


def pre_tokenize(text: str) -> list[str]:
    """Split text into word-ish chunks. Concatenating the result rebuilds text."""
    return _SPLIT_PATTERN.findall(text)


def _merge(ids: list[int], pair: tuple[int, int], new_id: int) -> list[int]:
    """Replace every non-overlapping occurrence of `pair` in `ids` with `new_id`."""
    out: list[int] = []
    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


class BPETokenizer:
    def __init__(
        self,
        merges: list[tuple[int, int]],
        special_tokens: dict[str, int],
    ) -> None:
        self.merges = merges
        self.special_tokens = special_tokens
        # rank = merge priority (lower merges first); id of a merged pair = 256 + rank
        self.ranks = {pair: rank for rank, pair in enumerate(merges)}
        self.vocab = self._build_vocab(merges, special_tokens)

    @staticmethod
    def _build_vocab(
        merges: list[tuple[int, int]],
        special_tokens: dict[str, int],
    ) -> dict[int, bytes]:
        """Map every id to the byte string it expands to (used for decoding)."""
        vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
        for rank, (a, b) in enumerate(merges):
            vocab[256 + rank] = vocab[a] + vocab[b]
        for text, idx in special_tokens.items():
            vocab[idx] = text.encode("utf-8")
        return vocab

    @property
    def vocab_size(self) -> int:
        return 256 + len(self.merges) + len(self.special_tokens)

    # ------------------------------------------------------------------ training
    @classmethod
    def train(
        cls,
        text: str,
        vocab_size: int,
        special_tokens: list[str] | None = None,
    ) -> "BPETokenizer":
        special_tokens = special_tokens or []
        num_merges = vocab_size - 256 - len(special_tokens)
        if num_merges < 0:
            raise ValueError("vocab_size too small for 256 bytes + special tokens")

        # Count unique pre-tokenized words. Special-token strings are stripped
        # first so BPE never learns to merge pieces of "<|endoftext|>".
        word_freqs: Counter[bytes] = Counter()
        for segment, is_special in cls._split_on_specials(text, special_tokens):
            if is_special:
                continue
            for chunk in pre_tokenize(segment):
                word_freqs[chunk.encode("utf-8")] += 1

        # Each unique word is a list of byte ids that we progressively merge.
        words: dict[bytes, list[int]] = {w: list(w) for w in word_freqs}

        merges: list[tuple[int, int]] = []
        for _ in range(num_merges):
            pair_counts: Counter[tuple[int, int]] = Counter()
            for w, freq in word_freqs.items():
                ids = words[w]
                for pair in zip(ids, ids[1:]):
                    pair_counts[pair] += freq
            if not pair_counts:
                break  # nothing left to merge
            # Highest count wins; ties broken by pair value for reproducibility.
            best = max(pair_counts.items(), key=lambda kv: (kv[1], kv[0]))[0]
            new_id = 256 + len(merges)
            merges.append(best)
            for w in words:
                words[w] = _merge(words[w], best, new_id)

        specials = {tok: 256 + len(merges) + i for i, tok in enumerate(special_tokens)}
        return cls(merges, specials)

    # ------------------------------------------------------------------ encoding
    @staticmethod
    def _split_on_specials(text: str, specials: list[str]):
        """Yield (segment, is_special), splitting on any special-token string."""
        if not specials:
            if text:
                yield text, False
            return
        pattern = "(" + "|".join(re.escape(s) for s in specials) + ")"
        for part in re.split(pattern, text):
            if part == "":
                continue
            yield part, part in specials

    def _encode_chunk(self, chunk: str) -> list[int]:
        """Apply learned merges to one pre-tokenized word, best (lowest) rank first."""
        ids = list(chunk.encode("utf-8"))
        while len(ids) >= 2:
            # Find the present pair with the lowest merge rank.
            best_rank: int | None = None
            best_i = -1
            for i in range(len(ids) - 1):
                rank = self.ranks.get((ids[i], ids[i + 1]))
                if rank is not None and (best_rank is None or rank < best_rank):
                    best_rank, best_i = rank, i
            if best_rank is None:
                break  # no learned merge applies
            ids = _merge(ids, self.merges[best_rank], 256 + best_rank)
        return ids

    def encode(self, text: str) -> list[int]:
        ids: list[int] = []
        for segment, is_special in self._split_on_specials(
            text, list(self.special_tokens)
        ):
            if is_special:
                ids.append(self.special_tokens[segment])
            else:
                for chunk in pre_tokenize(segment):
                    ids.extend(self._encode_chunk(chunk))
        return ids

    def decode(self, ids: list[int]) -> str:
        data = b"".join(self.vocab[i] for i in ids)
        return data.decode("utf-8", errors="replace")

    # ------------------------------------------------------------------ persistence
    def save(self, path: str | Path) -> None:
        """Save as human-inspectable JSON (merges + specials; vocab is derived)."""
        path = Path(path)
        blob = {
            "special_tokens": self.special_tokens,
            "merges": [[a, b] for a, b in self.merges],
        }
        path.write_text(json.dumps(blob))

    @classmethod
    def load(cls, path: str | Path) -> "BPETokenizer":
        blob = json.loads(Path(path).read_text())
        merges = [(a, b) for a, b in blob["merges"]]
        return cls(merges, blob["special_tokens"])