"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
import numpy as np

from contiguous import ContiguousPool, requests_that_fit_contiguous, requests_that_fit_paged
from kv_sizes import (LLAMA3_8B, OPT_13B, decode_both_ways, kv_bytes, kv_bytes_per_token,
                      pairs_with_cache, pairs_without_cache)
from paged_kv import prefix_sharing_example, worked_example

GIB = 1024**3


def main():
    print("== steps with and without a cache")
    for n in (10, 100, 1000):
        print(f"n={n}: without={pairs_without_cache(n)}  with={pairs_with_cache(n)}  "
              f"ratio={pairs_without_cache(n) / pairs_with_cache(n):.1f}")
    out_a, out_b, pa, pb = decode_both_ways(1000)
    print(f"demo n=1000: counted without={pa} with={pb}  outputs equal={np.allclose(out_a, out_b)}")

    print("== sizes")
    for cfg in (LLAMA3_8B, OPT_13B):
        per = kv_bytes_per_token(cfg)
        print(f"{cfg.name}: {per} bytes/token = {per / 1024:.0f} KiB; "
              f"full context ({cfg.max_tokens}) = {kv_bytes(cfg, cfg.max_tokens)} bytes = "
              f"{kv_bytes(cfg, cfg.max_tokens) / GIB:.4f} GiB")
    print("Llama 3 8B, 1000 tokens:", kv_bytes(LLAMA3_8B, 1000) / 1024**2, "MiB")
    print("Llama 3 8B, 131072 tokens:", kv_bytes(LLAMA3_8B, 131072) / GIB, "GiB")

    print("== contiguous slabs (figure 1): 32 slots")
    pool = ContiguousPool(32)
    pool.allocate("R1", 12); pool.allocate("R2", 8); pool.allocate("R3", 8)
    pool.used.update(R1=5, R2=3, R3=6)
    print("before: internal waste", pool.internal_waste(), "runs", pool.free_runs())
    pool.release("R2")
    print("after R2: runs", pool.free_runs(), "free_total", pool.free_total(),
          "largest", pool.largest_free_run(), "internal waste", pool.internal_waste(),
          "D(12) ->", pool.allocate("D", 12))
    eq = ContiguousPool(32)
    print("12 reserved each:", [eq.allocate(r, 12) for r in "ABC"])

    print("== capacity: 16 GiB of cache, Llama 3 8B, max_len 8192, block size 16")
    pool_tokens = 16 * GIB // kv_bytes_per_token(LLAMA3_8B)
    lengths = np.random.default_rng(0).integers(100, 2001, size=500)
    fit_c = requests_that_fit_contiguous(pool_tokens, 8192)
    paged = requests_that_fit_paged(lengths, pool_tokens, 16)
    print("pool tokens", pool_tokens, "mean length", lengths.mean())
    print("contiguous admits", fit_c, "-> utilisation of reserved slots",
          f"{lengths[:fit_c].sum() / (fit_c * 8192):.3f}")
    print("paged", paged, f"waste fraction of allocated {paged['waste'] / paged['allocated']:.4f}",
          f"gain {paged['admitted'] / fit_c:.1f}x")

    print("== worked example")
    for event, tables, free, lengths_ in worked_example():
        used = sum(lengths_.values())
        allocated = 4 * sum(len(t) for t in tables.values())
        print(event, tables, free, lengths_, f"used={used} allocated={allocated} wasted={allocated - used}")
    print("== prefix sharing")
    for row in prefix_sharing_example():
        print(*row)


if __name__ == "__main__":
    main()