"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
from prefill_sim import (
    A100_BANDWIDTH, A100_PEAK_FLOPS, AXIS_MS, DRAWN_STEPS, EMBED_PARAMS, FP16_BYTES, HEAD_DIM, HIDDEN, KV_BYTES_PER_TOKEN,
    LAYERS, PARAMS, WEIGHT_FILES_PARAMS, attention_flops_prefill, grid_intensity, attention_share_of_linear, decode_step, first_arithmetic_limited_prompt, fmt_ms,
    prefill, prompt_lengths, ridge, view,
)

GIB = 2**30


def main() -> None:
    step = decode_step(1)
    print(f"parameters from the config {PARAMS:,}; listed for the weight files {WEIGHT_FILES_PARAMS:,} (Meta's card rounds it to 8B)")
    print(f"weights in 16 bit: {FP16_BYTES * PARAMS / 1e9:.2f} GB; chip: {A100_PEAK_FLOPS / 1e12:.0f} TFLOP/s, {A100_BANDWIDTH / 1e9:.0f} GB/s")
    print(f"ridge {ridge():.2f} FLOP per byte; first prompt limited by arithmetic: {first_arithmetic_limited_prompt()} tokens")
    print(f"decode step {step.seconds * 1000:.3f} ms = {1 / step.seconds:.1f} tokens/s at most; arithmetic busy {step.busy * 100:.2f}% (1/ridge = {100 / ridge():.2f}%)")
    print(f"FLOPs per token 2P = {2 * PARAMS:,} ({2 * PARAMS / 1e9:.2f} GFLOP)")
    print()
    print("tokens | math floor ms | memory floor ms | time ms | limited by | busy | steps-long")
    for n in (1, 10, 20, 100, 150, 153, 154, 160, 300, 1000, 3000, 8000):
        p = prefill(n)
        print(f"{n:6,} | {p.math_s * 1000:9.3f} | {p.memory_s * 1000:9.3f} | {p.seconds * 1000:8.3f} | {p.bound:10} | {p.busy * 100:5.1f}% | {p.seconds / step.seconds:6.2f}")
    print()
    print("widget lines at the opening state (8,000 tokens) and at 20 tokens:")
    for n in (8000, 20):
        print(view(n))
    print()
    print(f"timeline: axis {AXIS_MS} ms; longest case = prefill {fmt_ms(prefill(8000, True).seconds)} ms + {DRAWN_STEPS} steps of {fmt_ms(decode_step(8000, True).seconds)} ms"
          f" = {(prefill(8000, True).seconds + DRAWN_STEPS * decode_step(8000, True).seconds) * 1000:.1f} ms")
    print(f"slider stops {len(prompt_lengths())}: {prompt_lengths()[:6]} ... {prompt_lengths()[-3:]}")
    print()
    print("left out of the short formula:")
    for n in (1000, 8000):
        a, b = prefill(n), prefill(n, True)
        print(f"  n={n:,}: attention FLOPs (each token against itself and earlier ones) {attention_flops_prefill(n):.3e}"
              f" = {100 * attention_share_of_linear(n):.1f}% of the weight FLOPs; prefill {a.seconds * 1000:.1f} -> {b.seconds * 1000:.1f} ms;"
              f" decode step {step.seconds * 1000:.3f} -> {decode_step(n, True).seconds * 1000:.3f} ms")
    print(f"  all pairs counted instead of the causal half: attention at 8,000 would be {4 * LAYERS * HIDDEN * 8000**2:.3e} FLOPs")
    n5 = next(n for n in range(1, 100000) if attention_share_of_linear(n) >= 0.05)
    n10 = next(n for n in range(1, 100000) if attention_share_of_linear(n) >= 0.10)
    print(f"  attention reaches 5% of the weight FLOPs at {n5:,} tokens and 10% at {n10:,}")
    print(f"  KV cache: {KV_BYTES_PER_TOKEN:,} bytes per token = 2 x {LAYERS} layers x 8 KV heads x {HEAD_DIM} x {FP16_BYTES} bytes")
    for n in (1000, 8000, 8192):
        print(f"  cache at {n:,} tokens {n * KV_BYTES_PER_TOKEN / 1e9:.3f} GB = {n * KV_BYTES_PER_TOKEN / GIB:.3f} GiB = {100 * n * KV_BYTES_PER_TOKEN / (FP16_BYTES * PARAMS):.1f}% of the weight bytes")
    first_grid = next(m for m in range(1, 1000) if grid_intensity(m) >= ridge())
    print(f"  lesson 19's single grid (4096 x 4096, activations counted): first prompt limited by arithmetic {first_grid}; activation bytes {100 * 2 * first_grid * 4096 * 2 / (2 * 4096 ** 2):.1f}% of the weight bytes there")
    print(f"  input table {EMBED_PARAMS:,} weights = {100 * EMBED_PARAMS / PARAMS:.1f}% of all weights (a lookup reads one row, the short formula counts all)")


if __name__ == "__main__":
    main()