"""Cut the lesson's words into chunks with two published tokenizers, and with the toy for comparison.

Needs the tiktoken package (pip install tiktoken). The first run of each encoding downloads its
vocabulary file, so it needs network access once; later runs use the cache. Run it from this folder:
python3 compare_tiktoken.py
"""
import sys
from importlib.metadata import version

from letters import DEFAULT_WORD, EXAMPLES, SEPARATOR, chunks, toy_merges

try:
    import tiktoken
except ImportError:
    print("tiktoken is not installed: pip install tiktoken, then run this again.")
    sys.exit(0)

ENCODINGS = ("cl100k_base", "o200k_base")
VARIANTS = (DEFAULT_WORD, " " + DEFAULT_WORD, DEFAULT_WORD.capitalize(), DEFAULT_WORD.upper())

published = {}
for name in ENCODINGS:
    try:
        published[name] = tiktoken.get_encoding(name)
    except Exception as err:  # first use downloads the vocabulary; no network means no encoding
        print(f"skipped {name}: {type(err).__name__}")
if not published:
    sys.exit(0)


def pieces(enc, text):
    """The chunks of text as (id, spelling) pairs. Every text here is ASCII, so each chunk decodes alone."""
    return [(i, enc.decode([i])) for i in enc.encode(text)]


def shown(spelling):
    return spelling.replace(" ", "_")


print(f"tiktoken {version('tiktoken')}; " + ", ".join(f"{n} has {e.n_vocab:,} ids" for n, e in published.items()))

print("\n== Chunks of each word (_ is a space; one line per tokenizer, with its ids)")
for text in dict.fromkeys(tuple(EXAMPLES) + VARIANTS):
    print(f"  {text!r}  ({len(text)} characters)")
    for name, enc in published.items():
        got = pieces(enc, text)
        assert "".join(s for _, s in got) == text and enc.decode(enc.encode(text)) == text
        print(f"    {name:12}{len(got):>3} chunk{'' if len(got) == 1 else 's'}  {' | '.join(shown(s) for _, s in got)}   ids {[i for i, _ in got]}")

print("\n== Spelled out: most letters in one chunk, and the number of chunks")
for text in EXAMPLES:
    for sep in (SEPARATOR, " "):
        spelled = sep.join(text)
        row = []
        for name, enc in published.items():
            got = pieces(enc, spelled)
            row.append(f"{name} {len(got)} chunks, at most {max(sum(c.isalpha() for c in s) for _, s in got)} letter per chunk")
        print(f"  {spelled!r:32}" + "   ".join(row))

print("\n== The toy for comparison")
print(f"  {DEFAULT_WORD!r}: " + " | ".join(shown(s) for _, s in chunks(DEFAULT_WORD, toy_merges())))