"""Follow one short text through the tokenizer, then price a request. Run it from this folder: python3 trace.py

It uses the lesson's own functions and adds no rules of its own. For each piece of the text it calls
encode() with one more merge allowed each time and prints what changed. The bill and the passes come
from cost.py.
"""
from bpe import encode, load_corpus, pieces, train, vocab_bytes
from cost import decode_positions_read, request_cost, sequential_passes

TEXT = "The ferry leaves at nine"
USE_MERGES = 512  # how many of the learned merges the tokenizer may use

merges = train(load_corpus())[:USE_MERGES]
table = vocab_bytes(merges)


def show(token_id):
    """The text a token stands for, or its raw bytes when it is only part of a character."""
    try:
        return table[token_id].decode("utf-8")
    except UnicodeDecodeError:
        return table[token_id]


n_pieces = len(pieces(TEXT))
print(f"text {TEXT!r}: {len(TEXT.encode('utf-8'))} bytes, {n_pieces} piece{'' if n_pieces == 1 else 's'}, {len(merges)} merges allowed")
ids = []
for piece in pieces(TEXT):
    state = encode(piece, [])
    print(f"\npiece {piece!r}: bytes {state}")
    for k in range(len(merges)):
        now = encode(piece, merges[: k + 1])
        if now != state:
            left, right = merges[k]
            print(f"  merge {k}: {left} + {right} -> {256 + k} {table[256 + k]!r}, now {now}")
            state = now
    ids.extend(state)
print(f"\ntokens: {len(ids)}, from {len(TEXT.encode('utf-8'))} bytes")
print("blocks: " + " ".join(repr(show(i)) for i in ids))

print("\nprice per million tokens (examples): new input 2, cached input 0.2, output 10 dollars")
new, cached, out = 500, 1500, 400
parts = (request_cost(new, 0, 0, 2, 0, 0), request_cost(0, cached, 0, 0, "0.2", 0), request_cost(0, 0, out, 0, 0, 10))
total = request_cost(new, cached, out, 2, "0.2", 10)
print(f"  {new:>5,} new input     at  2   = {parts[0]} dollars")
print(f"  {cached:>5,} cached input  at  0.2 = {parts[1]} dollars")
print(f"  {out:>5,} output        at 10   = {parts[2]} dollars")
print(f"  total {total} dollars; with nothing cached {request_cost(new + cached, 0, out, 2, '0.2', 10)} dollars")

print("\nsequential passes (the same 2,400 tokens, split two ways)")
for n_in, n_out in ((2000, 400), (400, 2000)):
    print(f"  {n_in:,} in + {n_out:,} out: {sequential_passes(n_in, n_out):,} passes in a row, "
          f"{decode_positions_read(n_in, n_out):,} positions read")