"""Print every number the lesson computes. Run it from this folder: python3 show_numbers.py

compare_tiktoken.py prints the numbers that come from a published tokenizer.
"""
from bpe import NUM_MERGES, encode, decode, load_corpus, pieces, train_counted, vocab_bytes
from cost import common_prefix_len, decode_positions_read, request_cost, sequential_passes
from samples import SAMPLES

corpus = load_corpus()
learned = train_counted(corpus)
merges = [pair for pair, _ in learned]
table = vocab_bytes(merges)


def words(text):
    return len(text.split())


print("== Bytes per character in UTF-8")
for ch in ("a", "\u00e9", "\ud55c", "\U0001f600"):
    print(f"  U+{ord(ch):04X} {ch!r}: {len(ch.encode('utf-8'))} bytes {list(ch.encode('utf-8'))}")
print(f"  '한국어' = {len('한국어')} characters, {len('한국어'.encode())} bytes")

print("\n== Worked example: train on the string 'aaabdaaabac' (one piece)")
tiny = train_counted("aaabdaaabac")
for k, (pair, count) in enumerate(tiny):
    tiny_table = vocab_bytes([p for p, _ in tiny[: k + 1]])
    print(f"  merge {k}: id {256 + k} = {pair[0]} + {pair[1]} = {tiny_table[256 + k]!r}, seen {count} times")
tiny_merges = [p for p, _ in tiny]
print(f"  encode -> {encode('aaabdaaabac', tiny_merges)}  (11 bytes become {len(encode('aaabdaaabac', tiny_merges))} tokens)")

print(f"\n== The corpus (corpus.txt, written for this lesson)")
chars, size = len(corpus), len(corpus.encode())
print(f"  {chars:,} characters, {size:,} bytes ({size / 1024:.1f} KiB), {len(pieces(corpus)):,} pieces, {len(set(pieces(corpus))):,} distinct")
hangul = sum(1 for c in corpus if "가" <= c <= "힣")
print(f"  Hangul syllables {hangul:,} = {hangul * 3:,} bytes ({hangul * 3 / size:.1%} of the bytes)")

print(f"\n== Training ({NUM_MERGES} merges asked for, {len(learned)} learned, vocabulary {256 + len(learned)} ids)")
for k in (0, 1, 2, 3, 4, 5, 6, 7, 8, 9):
    pair, count = learned[k]
    print(f"  merge {k}: id {256 + k} = {table[256 + k]!r}, seen {count} times")
for k in (99, 255, 511):
    pair, count = learned[k]
    print(f"  merge {k}: id {256 + k} = {table[256 + k]!r}, seen {count} times")
print(f"  smallest count among the merges: {min(c for _, c in learned)}")
longest = max((t for t in table if t.strip()), key=len)
print(f"  longest entry that is not only whitespace: {len(longest)} bytes {longest!r}")

print("\n== The four samples (none is in the corpus)")
print(f"  {'sample':8} {'chars':>5} {'bytes':>5} {'words':>5}  {'tokens@0':>8} {'tokens@512':>10} {'bytes/token':>11} {'tokens/word':>11}")
for name, text in SAMPLES.items():
    b, n0, n1 = len(text.encode()), len(encode(text, [])), len(encode(text, merges))
    assert decode(encode(text, merges), merges) == text
    print(f"  {name:8} {len(text):>5} {b:>5} {words(text):>5}  {n0:>8} {n1:>10} {b / n1:>11.2f} {n1 / words(text):>11.2f}")

print("\n== Vocabulary size: tokens needed at 0, 64, 128, 256 and 512 merges")
ks = (0, 64, 128, 256, 512)
print(f"  {'text':16} " + " ".join(f"{k:>6}" for k in ks))
rows = {**{f'{n} sample': t for n, t in SAMPLES.items()}, "the corpus": corpus}
for name, text in rows.items():
    print(f"  {name:16} " + " ".join(f"{len(encode(text, merges[:k])):>6}" for k in ks))
print("  tokens saved per extra merge, between steps (corpus | all four samples together):")
joined = "".join(SAMPLES.values())
for a, b in zip(ks, ks[1:]):
    s_corpus = len(encode(corpus, merges[:a])) - len(encode(corpus, merges[:b]))
    s_held = len(encode(joined, merges[:a])) - len(encode(joined, merges[:b]))
    print(f"    {a:>3} to {b:>3}: {s_corpus / (b - a):6.2f} | {s_held / (b - a):5.2f}")
corpus_counts = [len(encode(corpus, merges[:k])) for k in range(len(merges) + 1)]
saved = [a - b for a, b in zip(corpus_counts, corpus_counts[1:])]
short = [k for k, ((left, right), count) in enumerate(learned) if saved[k] != count]
rises = [k for k in range(1, len(saved)) if saved[k] > saved[k - 1]]
print(f"  merge counts never rise: {all(a[1] >= b[1] for a, b in zip(learned, learned[1:]))}")
print(f"  merges that save fewer tokens than their count: {len(short)}; each joins a pair with itself: "
      f"{all(learned[k][0][0] == learned[k][0][1] for k in short)}")
print(f"  a merge saves more than the one before it in {len(rises)} of {len(saved) - 1} steps; "
      f"each follows a merge that saved less than its count: {all(k - 1 in short for k in rises)}")
print("  example embedding-table size at width 4096 (an example width, not any model's):")
for name, vocab in (("this toy", 256 + len(learned)), ("50,257 ids", 50_257), ("100,277 ids", 100_277), ("200,019 ids", 200_019)):
    print(f"    {name:12}: {vocab:>7,} rows x 4096 = {vocab * 4096:>13,} numbers")

print("\n== One request (example prices, dollars per million tokens: input 2, cached input 0.2, output 10)")
cost = request_cost(500, 1500, 400, 2, "0.2", 10)
plain = request_cost(2000, 0, 400, 2, "0.2", 10)
print(f"  500 new + 1,500 cached input + 400 output tokens: {cost} dollars")
print(f"    parts: input {request_cost(500, 0, 0, 2, 0, 0)}, cached {request_cost(0, 1500, 0, 0, '0.2', 0)}, output {request_cost(0, 0, 400, 0, 0, 10)}")
print(f"    output share of the cost: {request_cost(0, 0, 400, 0, 0, 10) / cost:.1%}")
print(f"    output share of the tokens: {400 / (500 + 1500 + 400):.1%} ({400} of {500 + 1500 + 400})")
print(f"  same request, nothing cached: {plain} dollars (cost {plain / cost:.2f}x as much)")
print(f"  1,000 requests: {cost * 1000:.2f} dollars vs {plain * 1000:.2f} dollars")
print(f"  output tokens cost {10 / 2:g}x input tokens at these prices, so 1 output token = 5 input tokens")
print(f"  2,000 input + 400 output tokens: {sequential_passes(2000, 400)} sequential passes, "
      f"{decode_positions_read(2000, 400):,} positions read by the decode passes")
print(f"  400 input + 2,000 output tokens: {sequential_passes(400, 2000)} sequential passes, "
      f"{decode_positions_read(400, 2000):,} positions read by the decode passes")

print("\n== Cache prefix: how much of a prompt two requests share (toy tokenizer, 512 merges)")
SYSTEM = ("You are a support assistant for a small bakery. Answer in two sentences. "
          "Opening hours are 7 to 18 on weekdays and 8 to 14 on Saturdays. Prices are in dollars.\n")
questions = ("Do you sell rye bread?", "Is the shop open on Sunday?")
stamps = ("2026-10-18 09:14:02", "2026-10-18 09:14:09")
layouts = {
    "stable text first, then time and question": lambda q, t: SYSTEM + f"Time: {t}\nQuestion: {q}\n",
    "time first": lambda q, t: f"Time: {t}\n" + SYSTEM + f"Question: {q}\n",
}
for name, build in layouts.items():
    a, b = (encode(build(q, t), merges) for q, t in zip(questions, stamps))
    shared = common_prefix_len(a, b)
    print(f"  {name}: prompts of {len(a)} and {len(b)} tokens share {shared} ({shared / len(a):.0%} of the first)")

print("\n== Check yourself: the worked answers")
learned = train_counted("ababab")
print(f"  'ababab' merges {learned}, encodes to {encode('ababab', [p for p, _ in learned])}")
print(f"  2,000 new + 6,000 cached input + 500 output at 2 / 0.2 / 10: {request_cost(2000, 6000, 500, 2, '0.2', 10)} dollars; "
      f"nothing cached: {request_cost(8000, 0, 500, 2, '0.2', 10)} dollars")
print(f"  1,000 input + 1,000 output: {sequential_passes(1000, 1000)} passes, {decode_positions_read(1000, 1000):,} positions; "
      f"2,000 input + 1,000 output: {sequential_passes(2000, 1000)} passes, {decode_positions_read(2000, 1000):,} positions")