"""
Byte-level BPE from scratch (lecture: Tokenization).

Core idea:
  1. Start with 256 tokens = every byte. There is no "unknown token"; anything can be encoded.
  2. Pre-tokenize text with a regex (GPT-2 style) so merges never cross spaces/punctuation.
  3. Repeat: find the most frequent adjacent token pair -> merge it into a new token.
  4. Stop when the vocab reaches its target size.

Training is incremental (heap + pair->word index) so thousands of merges finish in seconds
instead of recounting every pair on every iteration.
"""
from __future__ import annotations

import heapq
import json
from collections import Counter, defaultdict

import regex as re

# GPT-2 pre-tokenization pattern: contractions, words, numbers, punctuation, whitespace
GPT2_SPLIT = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")

EOS = "<|eos|>"
USER = "<|user|>"
ASSISTANT = "<|assistant|>"
SPECIAL_TOKENS = [EOS, USER, ASSISTANT]


class BPETokenizer:
    def __init__(self, merges: list[tuple[int, int]] | None = None, special_tokens: list[str] | None = None):
        self.merges: list[tuple[int, int]] = merges or []
        self.special_tokens = special_tokens or list(SPECIAL_TOKENS)
        self._build()

    # ------------------------------------------------------------------ setup
    def _build(self) -> None:
        self.ranks = {pair: i for i, pair in enumerate(self.merges)}
        self.vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
        for i, (a, b) in enumerate(self.merges):
            self.vocab[256 + i] = self.vocab[a] + self.vocab[b]
        base = 256 + len(self.merges)
        self.special_to_id = {tok: base + i for i, tok in enumerate(self.special_tokens)}
        self.id_to_special = {v: k for k, v in self.special_to_id.items()}
        self._special_re = re.compile("(" + "|".join(re.escape(t) for t in self.special_tokens) + ")")
        self._cache: dict[str, list[int]] = {}

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

    @property
    def eos_id(self) -> int:
        return self.special_to_id[EOS]

    # --------------------------------------------------------------- training
    @classmethod
    def train(cls, text: str, vocab_size: int, verbose: bool = True) -> "BPETokenizer":
        n_merges = vocab_size - 256 - len(SPECIAL_TOKENS)
        assert n_merges > 0, "vocab_size is too small"

        word_freq = Counter(GPT2_SPLIT.findall(text))
        words = [list(w.encode("utf-8")) for w in word_freq]
        freqs = list(word_freq.values())

        pair_counts: dict[tuple[int, int], int] = defaultdict(int)
        where: dict[tuple[int, int], set[int]] = defaultdict(set)
        for idx, w in enumerate(words):
            for p in zip(w, w[1:]):
                pair_counts[p] += freqs[idx]
                where[p].add(idx)

        # max-heap with lazy invalidation: stale entries are skipped on pop
        heap = [(-c, p) for p, c in pair_counts.items()]
        heapq.heapify(heap)

        merges: list[tuple[int, int]] = []
        while len(merges) < n_merges and heap:
            neg_c, pair = heapq.heappop(heap)
            if pair_counts.get(pair, 0) != -neg_c or -neg_c <= 0:
                continue  # stale
            new_id = 256 + len(merges)
            merges.append(pair)

            for idx in list(where[pair]):
                w, f = words[idx], freqs[idx]
                # remove the old pairs' contribution
                for p in zip(w, w[1:]):
                    pair_counts[p] -= f
                    where[p].discard(idx)
                # merge the pair
                out, i = [], 0
                while i < len(w):
                    if i < len(w) - 1 and w[i] == pair[0] and w[i + 1] == pair[1]:
                        out.append(new_id)
                        i += 2
                    else:
                        out.append(w[i])
                        i += 1
                words[idx] = out
                # add the new pairs' contribution
                touched = set()
                for p in zip(out, out[1:]):
                    pair_counts[p] += f
                    where[p].add(idx)
                    touched.add(p)
                for p in touched:
                    heapq.heappush(heap, (-pair_counts[p], p))
            pair_counts.pop(pair, None)
            where.pop(pair, None)

            if verbose and len(merges) % 500 == 0:
                print(f"  merge {len(merges)}/{n_merges}")

        tok = cls(merges)
        if verbose:
            print(f"Tokenizer done: vocab {tok.vocab_size} ({len(merges)} merges)")
        return tok

    # ------------------------------------------------------------ encode/decode
    def _encode_chunk(self, chunk: str) -> list[int]:
        cached = self._cache.get(chunk)
        if cached is not None:
            return cached
        ids = list(chunk.encode("utf-8"))
        while len(ids) >= 2:
            # lowest-rank pair = the merge learned earliest
            best = min(zip(ids, ids[1:]), key=lambda p: self.ranks.get(p, 1 << 30))
            rank = self.ranks.get(best)
            if rank is None:
                break
            new_id = 256 + rank
            out, i = [], 0
            while i < len(ids):
                if i < len(ids) - 1 and ids[i] == best[0] and ids[i + 1] == best[1]:
                    out.append(new_id)
                    i += 2
                else:
                    out.append(ids[i])
                    i += 1
            ids = out
        if len(self._cache) < 500_000:
            self._cache[chunk] = ids
        return ids

    def encode(self, text: str, allow_special: bool = True) -> list[int]:
        ids: list[int] = []
        parts = self._special_re.split(text) if allow_special else [text]
        for part in parts:
            if not part:
                continue
            if allow_special and part in self.special_to_id:
                ids.append(self.special_to_id[part])
                continue
            for chunk in GPT2_SPLIT.findall(part):
                ids.extend(self._encode_chunk(chunk))
        return ids

    def decode(self, ids: list[int]) -> str:
        out = bytearray()
        for i in ids:
            if i in self.id_to_special:
                out.extend(self.id_to_special[i].encode("utf-8"))
            else:
                out.extend(self.vocab[i])
        return out.decode("utf-8", errors="replace")

    # ---------------------------------------------------------------- persist
    def save(self, path: str) -> None:
        with open(path, "w") as f:
            json.dump({"merges": self.merges, "special_tokens": self.special_tokens}, f)

    @classmethod
    def load(cls, path: str) -> "BPETokenizer":
        with open(path) as f:
            d = json.load(f)
        return cls([tuple(m) for m in d["merges"]], d["special_tokens"])