"""Baselines for comparison against SPIRA.

Compression: gzip -9, HuggingFace byte-level BPE (matched vocabulary sizes).
Prediction: unigram, patch-bigram with unigram backoff.
"""
from __future__ import annotations

import gzip
import time

import numpy as np

from tokenizers import Tokenizer, models, trainers, pre_tokenizers


# ------------------------------------------------------------- compression
def gzip_bits_per_byte(data: bytes) -> tuple[float, float]:
    """(bits/byte, seconds) for gzip level 9."""
    t0 = time.perf_counter()
    comp = gzip.compress(data, 9)
    dt = time.perf_counter() - t0
    return len(comp) * 8 / len(data), dt


def bpe_encode(train: bytes, test: bytes, vocab_size: int) -> dict:
    """Train a byte-level BPE on `train` and encode `test` with it.

    Returns dict(ids, vocab_size, vocab_bytes, train_s, encode_s,
                 MB_per_s_train).
    """
    t0 = time.perf_counter()
    tok = Tokenizer(models.BPE())
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        special_tokens=[],
    )
    tok.train_from_iterator([train.decode("utf-8", "replace")], trainer=trainer)
    train_s = time.perf_counter() - t0
    t1 = time.perf_counter()
    enc = tok.encode(test.decode("utf-8", "replace"))
    encode_s = time.perf_counter() - t1
    vocab = tok.get_vocab()
    vocab_bytes = sum(len(s.encode("utf-8")) for s in vocab)
    return {
        "ids": np.asarray(enc.ids, dtype=np.int64),
        "vocab_size": len(vocab),
        "vocab_bytes": vocab_bytes,
        "train_s": train_s,
        "encode_s": encode_s,
    }


def stream_entropy_bits(ids: np.ndarray, order: int = 0) -> float:
    """Empirical entropy of a token stream in bits/token (order 0 or 1)."""
    # remap arbitrary ids (e.g. uint64 patch ids) to dense indices
    uniq, inv = np.unique(ids, return_inverse=True)
    ids = inv.astype(np.int64)
    if order == 0:
        _, cnt = np.unique(ids, return_counts=True)
        p = cnt / cnt.sum()
        return float(-(p * np.log2(p)).sum())
    # order 1: H(next | prev) estimated from the same stream
    prev = ids[:-1]
    cur = ids[1:]
    span = len(uniq)
    pair = prev * span + cur
    uniq_pair, cnt_pair = np.unique(pair, return_counts=True)
    pair_prev = uniq_pair // span
    prev_ids, cnt_prev = np.unique(prev, return_counts=True)
    lookup = np.zeros(span, dtype=np.int64)
    lookup[prev_ids] = cnt_prev
    denom = lookup[pair_prev].astype(np.float64)
    p_pair = cnt_pair / denom
    bits = -(cnt_pair * np.log2(p_pair)).sum()
    return float(bits / len(cur))


def dict_bits(vocab_size: int, vocab_bytes: int, overhead_bits: float = 32.0) -> float:
    """Decode-dictionary storage: raw token bytes + id/length overhead."""
    return vocab_bytes * 8 + vocab_size * overhead_bits


# ------------------------------------------------------------- prediction
class UnigramBaseline:
    def __init__(self, train_cols: np.ndarray, V: int | None = None) -> None:
        if V is not None:
            train_cols = train_cols[train_cols != V]   # exclude the UNK column
        self.best = int(np.bincount(train_cols).argmax())

    def accuracy(self, cur: np.ndarray) -> float:
        return float((cur == self.best).mean())


class BigramBaseline:
    """Patch-bigram predictor with unigram backoff (sparse pair counts).

    Dense (V, V) storage is infeasible at V=16k (1+ GB); pairs are counted
    via packed int64 keys instead, and the argmax continuation per preceding
    column is precomputed by a (prev, -count) lexsort.
    """

    def __init__(self, train_cols: np.ndarray, V: int) -> None:
        prev = train_cols[:-1]
        cur = train_cols[1:]
        ok = (prev != V) & (cur != V)           # in-vocab pairs only
        p = prev[ok].astype(np.int64)
        c = cur[ok].astype(np.int64)
        if len(p) == 0:
            self.map_prev = np.array([-1], dtype=np.int64)
            self.map_pred = np.array([0], dtype=np.int64)
        else:
            key = p * np.int64(V + 1) + c
            uk, cnt = np.unique(key, return_counts=True)   # sorted by key
            prev_u = (uk // np.int64(V + 1)).astype(np.int64)
            pred_c = (uk % np.int64(V + 1)).astype(np.int64)
            order = np.lexsort((-cnt, prev_u))             # prev asc, count desc
            pu, pc = prev_u[order], pred_c[order]
            first = np.ones(len(pu), dtype=bool)
            first[1:] = pu[1:] != pu[:-1]
            self.map_prev = pu[first]
            self.map_pred = pc[first]
        self.uni_best = int(np.bincount(train_cols[train_cols != V]).argmax())

    def accuracy(self, prev: np.ndarray, cur: np.ndarray) -> float:
        pos = np.searchsorted(self.map_prev, prev)
        pos_c = np.minimum(pos, len(self.map_prev) - 1)
        hit = (pos < len(self.map_prev)) & (self.map_prev[pos_c] == prev)
        pred = np.where(hit, self.map_pred[pos_c], self.uni_best)
        return float((pred == cur).mean())