"""Tokenise several strings with a model's own tokeniser and show what it did.

Purpose: turn the tokens lesson into measurements: how many tokens the same meaning costs
         in English and Portuguese, what happens to source code and emoji, how the
         byte-level vocabulary displays non-ASCII text, and what the special and chat
         control tokens of this model are.
Platform: all (CPU only; the tokeniser is data, not a model)
Minimum memory: 8 GB (the tokeniser itself needs a few hundred megabytes)
Assumes: a model directory downloaded with `hf download ... --local-dir <dir>` containing
         tokenizer.json and tokenizer_config.json; transformers and tokenizers installed
         in the course environment from Part 1.

Usage: python tokenise-samples.py --model ~/llm-course/models/qwen3-1.7b
           [--guess 4,6,8,2] [--text "your own string"] [--labbook labbook.md]
"""
import argparse
import json
from pathlib import Path

from transformers import AutoTokenizer

# The four strings the tokens lesson asked you to predict, in that order.
PREDICTIONS = [
    ("hello", "Hello, world!"),
    ("hello-pt", "Olá, mundo!"),
    ("signature", "def get_user_by_id(user_id: int):"),
    ("flag-rocket", "🇧🇷🚀"),
]

SAMPLES = [
    ("english", "The quick brown fox jumps over the lazy dog near the river bank."),
    ("portuguese", "A rápida raposa castanha salta sobre o cão preguiçoso perto da margem do rio."),
    ("code", "def get_user_by_id(user_id: int) -> User | None:\n    return session.get(User, user_id)\n"),
    ("emoji", "🇧🇷 🚀 ✨ 🧮"),
]

# Chat control tokens named in the Qwen3 chat template; printed if the vocabulary has them.
CONTROL_TOKENS = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<think>", "</think>"]


def describe(tokeniser, label: str, text: str, guess: int | None = None) -> dict:
    """Tokenise one string and print ids, token strings and the three counts."""
    ids = tokeniser(text, add_special_tokens=False)["input_ids"]
    pieces = tokeniser.convert_ids_to_tokens(ids)
    n_bytes = len(text.encode("utf-8"))
    n_words = len([w for w in text.split() if w])
    print(f"\n[{label}]  {text!r}")
    guessed = f"   your guess {guess:4d}" if guess is not None else ""
    print(f"  tokens {len(ids):4d}   words {n_words:4d}   characters {len(text):4d}   utf-8 bytes {n_bytes:4d}{guessed}")
    print(f"  ids    {ids}")
    print(f"  pieces {pieces}")
    if not text.isascii():
        # Byte-level BPE shows each UTF-8 byte as one printable character, so 'á' (bytes
        # C3 A1) appears as 'á'. Decoding each piece on its own shows the text it holds.
        decoded = [tokeniser.decode([i]) for i in ids]
        print(f"  pieces decoded one by one {decoded}")
    if n_words:
        print(f"  tokens per word {len(ids) / n_words:.2f}   bytes per token {n_bytes / max(len(ids), 1):.2f}")
    # Round-tripping proves the split is lossless for this tokeniser.
    restored = tokeniser.decode(ids)
    print(f"  decode round-trip identical: {restored == text}")
    return {
        "label": label, "tokens": len(ids), "words": n_words, "characters": len(text),
        "utf8_bytes": n_bytes, "round_trip": restored == text, "guess": guess,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model", required=True, help="model directory or Hub id")
    parser.add_argument("--guess", default=None,
                        help="your four predicted token counts for the lesson's strings, e.g. 4,6,8,2")
    parser.add_argument("--text", action="append", default=[], help="an extra string to tokenise; repeatable")
    parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file")
    args = parser.parse_args()

    model_dir = Path(args.model).expanduser()
    if model_dir.is_dir() and not (model_dir / "tokenizer.json").exists():
        raise SystemExit(f"{model_dir} has no tokenizer.json; the download did not finish, or the path is wrong")

    guesses: list[int | None] = [None] * len(PREDICTIONS)
    if args.guess:
        parts = [p.strip() for p in args.guess.split(",")]
        if len(parts) != len(PREDICTIONS) or not all(p.isdigit() for p in parts):
            raise SystemExit(f"--guess needs {len(PREDICTIONS)} integers separated by commas, e.g. 4,6,8,2")
        guesses = [int(p) for p in parts]

    tokeniser = AutoTokenizer.from_pretrained(args.model)
    print(f"tokeniser class:            {type(tokeniser).__name__}")
    print(f"base vocabulary (BPE):      {tokeniser.vocab_size:,}   (tokeniser.vocab_size)")
    print(f"entries incl. added tokens: {len(tokeniser):,}   (len(tokeniser))")
    config_path = model_dir / "config.json"
    if config_path.exists():
        config_vocab = json.loads(config_path.read_text(encoding="utf-8")).get("vocab_size")
        if config_vocab:
            print(f"embedding rows (config):    {config_vocab:,}   (vocab_size in config.json; "
                  f"{config_vocab - len(tokeniser):,} rows no token can select)")
    print(f"model max length:           {tokeniser.model_max_length:,}")

    specials = []
    for name in ("bos_token", "eos_token", "pad_token", "unk_token"):
        token = getattr(tokeniser, name, None)
        token_id = getattr(tokeniser, f"{name}_id", None)
        print(f"  {name:10s} {token!r:18s} id {token_id}")
        specials.append({"name": name, "token": token, "id": token_id})
    print("  chat control tokens in this vocabulary:")
    controls = []
    for token in CONTROL_TOKENS:
        token_id = tokeniser.convert_tokens_to_ids(token)
        if token_id is not None and token_id != tokeniser.unk_token_id:
            print(f"    {token!r:16s} id {token_id}")
            controls.append({"token": token, "id": token_id})

    print("\n=== the four strings you predicted in the tokens lesson ===")
    results = [describe(tokeniser, label, text, guess)
               for (label, text), guess in zip(PREDICTIONS, guesses)]
    print("\n=== four kinds of text ===")
    results += [describe(tokeniser, label, text) for label, text in SAMPLES]
    results += [describe(tokeniser, f"custom-{i + 1}", text) for i, text in enumerate(args.text)]

    print("\nsummary")
    print(f"  {'sample':12s} {'tokens':>7s} {'guess':>6s} {'words':>7s} {'bytes':>7s} {'bytes/token':>12s} {'round-trip':>11s}")
    for r in results:
        per_token = r["utf8_bytes"] / max(r["tokens"], 1)
        guess = str(r["guess"]) if r["guess"] is not None else "-"
        print(f"  {r['label']:12s} {r['tokens']:7d} {guess:>6s} {r['words']:7d} {r['utf8_bytes']:7d} "
              f"{per_token:12.2f} {str(r['round_trip']):>11s}")

    if args.labbook:
        record = {
            "lab": "part-02/tokenise-samples",
            "model": args.model,
            "tokeniser": type(tokeniser).__name__,
            "base_vocab_size": tokeniser.vocab_size,
            "entries_with_added_tokens": len(tokeniser),
            "special_tokens": specials,
            "control_tokens": controls,
            "samples": results,
        }
        with Path(args.labbook).expanduser().open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(record, ensure_ascii=False) + "\n")
        print(f"\nrecorded in {args.labbook}")


if __name__ == "__main__":
    main()