#!/usr/bin/env python3
"""Bucket the 8192-token vocab into coarse types.

Used to log the mixing scalar's distribution over token types: knowing that
alpha averages 0.5 overall says nothing, knowing that it sits at 0.8 on digits
and 0.4 on word continuations says a lot about what each tower is doing.

The Mini tokenizer is byte-level BPE (with digits split out by a pre-tokenizer),
so token strings are byte-level encoded and have to be mapped back through the
GPT-2 byte<->unicode table before they can be classified.
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np


TYPE_NAMES = (
    "special",
    "whitespace",
    "digit",
    "word_start",
    "word_cont",
    "punct",
    "other",
)
TYPE_INDEX = {name: i for i, name in enumerate(TYPE_NAMES)}


def _byte_decoder() -> dict[str, int]:
    bs = (
        list(range(ord("!"), ord("~") + 1))
        + list(range(ord("¡"), ord("¬") + 1))
        + list(range(ord("®"), ord("ÿ") + 1))
    )
    cs = bs[:]
    n = 0
    for b in range(256):
        if b not in bs:
            bs.append(b)
            cs.append(256 + n)
            n += 1
    return {chr(c): b for b, c in zip(bs, cs)}


def _decode(token: str, byte_decoder: dict[str, int]) -> str:
    try:
        raw = bytes(byte_decoder[ch] for ch in token)
    except KeyError:
        return token
    return raw.decode("utf-8", errors="replace")


def classify(text: str) -> int:
    if text == "":
        return TYPE_INDEX["other"]
    core = text.lstrip(" \t")
    leading_space = core != text
    if core.strip() == "":
        return TYPE_INDEX["whitespace"]
    if any(ch.isdigit() for ch in core):
        return TYPE_INDEX["digit"]
    letters = core.replace("'", "").replace("’", "")
    if letters and all(ch.isalpha() for ch in letters):
        return TYPE_INDEX["word_start"] if leading_space else TYPE_INDEX["word_cont"]
    if all(not ch.isalnum() for ch in core):
        return TYPE_INDEX["punct"]
    return TYPE_INDEX["other"]


def build_token_type_table(tokenizer_path: Path, vocab_size: int) -> np.ndarray:
    """Returns an int64 array of length vocab_size mapping token id -> type index."""
    path = Path(tokenizer_path)
    if path.is_dir():
        path = path / "tokenizer.json"
    with path.open("r", encoding="utf-8") as f:
        data = json.load(f)

    vocab: dict[str, int] = data["model"]["vocab"]
    specials = {entry["content"] for entry in data.get("added_tokens", [])}
    byte_decoder = _byte_decoder()

    table = np.full(vocab_size, TYPE_INDEX["other"], dtype=np.int64)
    for token, idx in vocab.items():
        if idx >= vocab_size:
            continue
        if token in specials:
            table[idx] = TYPE_INDEX["special"]
            continue
        table[idx] = classify(_decode(token, byte_decoder))
    return table


if __name__ == "__main__":
    import sys
    from collections import Counter

    tok = Path(
        sys.argv[1]
        if len(sys.argv) > 1
        else "/home/banaxi/Desktop/BananaMind/BananaMind-2/tokenizers/fineweb_edu_first_50gib_8k_digits/tokenizer.json"
    )
    table = build_token_type_table(tok, 8192)
    counts = Counter(table.tolist())
    for i, name in enumerate(TYPE_NAMES):
        print(f"{name:12s} {counts.get(i, 0):5d}  {counts.get(i, 0) / len(table):6.2%}")