"""Byte-level BPE: turn text into a list of integers and back.

Training starts from the 256 possible bytes and repeatedly glues together the
pair of neighbouring tokens that occurs most often. Every glued pair becomes a
new token. Encoding replays the same merges in the same order.
"""
import json
from collections import Counter, defaultdict

import regex as re

# Split text into chunks first, so a merge never crosses a word boundary:
# contractions, words with their leading space, single digits, punctuation, spaces.
SPLIT = re.compile(r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")

SPECIALS = ['<|endoftext|>', '<|user|>', '<|assistant|>', '<|end|>']


class Tokenizer:
    def __init__(self, merges=(), specials=SPECIALS):
        self.merges = {tuple(pair): 256 + i for i, pair in enumerate(merges)}
        self.vocab = {i: bytes([i]) for i in range(256)}
        for (a, b), idx in self.merges.items():
            self.vocab[idx] = self.vocab[a] + self.vocab[b]
        first_special = 256 + len(self.merges)
        self.specials = {s: first_special + i for i, s in enumerate(specials)}
        for s, idx in self.specials.items():
            self.vocab[idx] = s.encode()
        self.special_split = re.compile('(' + '|'.join(re.escape(s) for s in specials) + ')')
        self.cache = {}

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

    def __getitem__(self, special):
        return self.specials[special]

    # ---- training -------------------------------------------------------
    @classmethod
    def train(cls, texts, vocab_size, specials=SPECIALS, verbose=False):
        n_merges = vocab_size - 256 - len(specials)
        words = Counter()
        for text in texts:
            for chunk in SPLIT.findall(text):
                words[chunk] += 1
        # every distinct chunk is a list of token ids, counted once with its frequency
        seqs = [list(w.encode()) for w in words]
        freqs = list(words.values())
        pairs = Counter()
        where = defaultdict(set)          # pair -> indices of chunks that contain it
        for i, seq in enumerate(seqs):
            for pair in zip(seq, seq[1:]):
                pairs[pair] += freqs[i]
                where[pair].add(i)
        merges = []
        for step in range(n_merges):
            if not pairs:
                break
            best = max(pairs, key=pairs.get)
            new_id = 256 + step
            merges.append(best)
            for i in list(where[best]):
                seq, f = seqs[i], freqs[i]
                for pair in zip(seq, seq[1:]):  # forget this chunk's old pairs
                    pairs[pair] -= f
                    if pairs[pair] <= 0:
                        del pairs[pair]
                    where[pair].discard(i)
                seq = merge(seq, best, new_id)
                seqs[i] = seq
                for pair in zip(seq, seq[1:]):  # and count the new ones
                    pairs[pair] += f
                    where[pair].add(i)
            if verbose and step % 500 == 0:
                tok = cls(merges).vocab[new_id]
                print(f'merge {step:5d}: {tok!r}')
        return cls(merges, specials)

    # ---- encoding -------------------------------------------------------
    def encode_chunk(self, chunk):
        ids = self.cache.get(chunk)
        if ids is not None:
            return ids
        ids = list(chunk.encode())
        while len(ids) > 1:
            # the pair that was learned earliest gets merged first
            pair = min(zip(ids, ids[1:]), key=lambda p: self.merges.get(p, float('inf')))
            if pair not in self.merges:
                break
            ids = merge(ids, pair, self.merges[pair])
        if len(self.cache) < 500_000:
            self.cache[chunk] = ids
        return ids

    def encode(self, text, allow_special=True):
        out = []
        parts = self.special_split.split(text) if allow_special else [text]
        for part in parts:
            if part in self.specials and allow_special:
                out.append(self.specials[part])
                continue
            for chunk in SPLIT.findall(part):
                out.extend(self.encode_chunk(chunk))
        return out

    def decode(self, ids):
        return b''.join(self.vocab[i] for i in ids).decode('utf-8', errors='replace')

    # ---- saving ---------------------------------------------------------
    def save(self, path):
        with open(path, 'w') as f:
            json.dump({'merges': [list(p) for p in self.merges], 'specials': list(self.specials)}, f)

    @classmethod
    def load(cls, path):
        with open(path) as f:
            data = json.load(f)
        return cls(data['merges'], data['specials'])


def merge(ids, pair, new_id):
    """Replace every occurrence of `pair` in `ids` with `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