"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
import statistics

from batching_sim import (
    A100_BANDWIDTH, A100_PEAK_FLOPS, FP16_BYTES, GPU_MEMORY_BYTES, KV_BYTES_PER_TOKEN, MIXES, N_REQUESTS, OPENING, PARAMS, RATES, SEED, SLOT_CHOICES,
    WORKLOAD, Request, batch_step, make_requests, mix_names, now_text, report, simulate, snapshot, state_line, summary_line, time_at, users_that_fit,
)


def toy() -> None:
    requests = [Request(0.0, 1, r) for r in (2, 6, 3, 1, 4, 2)]
    for policy in ("fixed", "continuous"):
        run = simulate(requests, 4, policy, cost=lambda contexts, joining: 1.0)
        print(f"  {policy}: {len(run.steps)} steps; first tokens at {[run.first_token[i] for i in range(6)]}; last tokens at {run.finished}")
        for k, st in enumerate(run.steps):
            print(f"    step {k + 1}: slots {st.slots}")


def main() -> None:
    one = batch_step([1000])
    print(f"parameters {PARAMS:,}; weights {FP16_BYTES * PARAMS / 1e9:.2f} GB; chip {A100_PEAK_FLOPS / 1e12:.0f} TFLOP/s, {A100_BANDWIDTH / 1e9:.0f} GB/s, {GPU_MEMORY_BYTES / 1e9:.0f} GB")
    print(f"cache per token {KV_BYTES_PER_TOKEN:,} bytes ({KV_BYTES_PER_TOKEN / 1024:.0f} KiB); weights alone take {FP16_BYTES * PARAMS / A100_BANDWIDTH * 1000:.3f} ms to read")
    print(f"arithmetic per user {2 * PARAMS / A100_PEAK_FLOPS * 1e6:.1f} us; cache reading per user at 1,000 tokens {1000 * KV_BYTES_PER_TOKEN / A100_BANDWIDTH * 1e6:.1f} us, at 100 tokens {100 * KV_BYTES_PER_TOKEN / A100_BANDWIDTH * 1e6:.1f} us")
    print()
    print("one step, every user with the same context")
    print("context | users | step ms | limited by | server tokens/s | per user tokens/s | cache GB read | busy")
    for context in (1000, 100):
        for users in (1, 8, 64, 175, 256, 487):
            if users > users_that_fit(context):
                continue
            p = batch_step([context] * users)
            print(f"{context:7,} | {users:5} | {p.seconds * 1000:7.2f} | {p.bound:10} | {users / p.seconds:15,.0f} | {1 / p.seconds:17.1f} | {users * context * KV_BYTES_PER_TOKEN / 1e9:6.2f} | {p.busy * 100:4.1f}%")
    print(f"  one user at 1,000 tokens: {one.seconds * 1000:.2f} ms (weights alone {FP16_BYTES * PARAMS / A100_BANDWIDTH * 1000:.2f} ms, so the cache adds {(one.seconds - FP16_BYTES * PARAMS / A100_BANDWIDTH) * 1000:.2f} ms)")
    first = next(n for n in range(1, 10000) if batch_step([100] * n).bound == "arithmetic")
    print(f"  at 100 tokens of context the step turns arithmetic-limited from {first} users; at 1,000 tokens never up to {users_that_fit(1000)}")
    print(f"  64 users against 1 user at 1,000 tokens: step time x{batch_step([1000] * 64).seconds / one.seconds:.2f}, server tokens/s x{64 * one.seconds / batch_step([1000] * 64).seconds:.1f}")
    print()
    print("users that fit beside the weights (best case: all the rest is cache)")
    for context in (100, 1000, 2000, 8192):
        print(f"  context {context:6,}: cache {context * KV_BYTES_PER_TOKEN / 1e9:6.3f} GB each -> {users_that_fit(context):,} users; room for cache {(GPU_MEMORY_BYTES - FP16_BYTES * PARAMS) / 1e9:.2f} GB")
    print(f"  4,000 tokens each: {users_that_fit(4000)} users")
    print()
    print("the worked table, and the steps that fill the memory")
    for contexts in ([1000], [1000] * 8, [1000] * 64):
        p = batch_step(contexts)
        print(f"  {len(contexts):3} users at 1,000 tokens: math {p.math_s * 1000:.4f} ms, memory {p.memory_s * 1000:.4f} ms, step {p.seconds * 1000:.2f} ms,"
              f" {len(contexts) / p.seconds:,.1f} tokens/s in all, {1 / p.seconds:.2f} per user, {p.nbytes / 1e9:.2f} GB read")
    for users, context in ((users_that_fit(1000), 1000), (users_that_fit(2000), 2000)):
        p = batch_step([context] * users)
        print(f"  memory full, {users} users at {context:,} tokens: step {p.seconds * 1000:.2f} ms, {p.nbytes / 1e9:.2f} GB read")
    joiner = batch_step([300] * 7, 577)
    print(f"  a 577-token prompt joins 7 users at 300 tokens: step {joiner.seconds * 1000:.2f} ms ({joiner.bound}) against {batch_step([300] * 7).seconds * 1000:.2f} ms without it")
    print()
    print("toy example: six requests at time zero, four slots, unit steps, replies 2 6 3 1 4 2")
    toy()
    print()
    prompts, replies = [r.prompt for r in make_requests(1, "mixed")], {m: [r.reply for r in make_requests(1, m)] for m in MIXES}
    print(f"workload: seed {SEED}, {N_REQUESTS} requests; prompts {min(prompts)} to {max(prompts)} (mean {statistics.fmean(prompts):.1f})")
    for m in MIXES:
        print(f"  {mix_names()[m]}: replies {min(replies[m])} to {max(replies[m])}, mean {statistics.fmean(replies[m]):.1f}, total {sum(replies[m]):,}")
    print(f"  first arrivals at traffic 1: {[round(r.arrival, 2) for r in make_requests(1, 'mixed')[:5]]}; gaps sum {sum(WORKLOAD['gaps']):.1f} s at one a second")
    slow = make_requests(0.5, "mixed")
    print(f"  at 0.5 a second the last request arrives at {slow[-1].arrival:.1f} s and its reply is {slow[-1].reply} tokens")
    print()
    print("results (8 slots unless noted): seconds | tokens/s | average wait | slowest wait | slots in use")
    for mix in MIXES:
        for slots in (8,):
            for rate in RATES:
                requests = make_requests(rate, mix)
                a, b = (report(simulate(requests, slots, p), requests, slots) for p in ("fixed", "continuous"))
                print(f"  {mix:10} S={slots:2} rate {rate:4}: fixed {a['seconds']:6.1f} s {a['tokens_per_second']:5.0f}/s wait {a['average_wait']:5.2f}/{a['slowest_wait']:5.2f} use {a['slots_in_use'] * 100:3.0f}%"
                      f" | continuous {b['seconds']:6.1f} s {b['tokens_per_second']:5.0f}/s wait {b['average_wait']:5.2f}/{b['slowest_wait']:5.2f} use {b['slots_in_use'] * 100:3.0f}%")
    never_worse = all(
        report(simulate(make_requests(r, m), s, "continuous"), make_requests(r, m), s)[k] <= report(simulate(make_requests(r, m), s, "fixed"), make_requests(r, m), s)[k]
        for m in MIXES for s in SLOT_CHOICES for r in RATES for k in ("seconds", "average_wait", "slowest_wait"))
    print(f"  continuous never slower or longer-waiting than fixed in all {len(MIXES) * len(SLOT_CHOICES) * len(RATES)} settings: {never_worse}")
    print()
    print("opening state of the widget")
    requests = make_requests(OPENING["rate"], OPENING["mix"])
    runs = [simulate(requests, OPENING["slots"], p) for p in ("fixed", "continuous")]
    reps = [report(run, requests, OPENING["slots"]) for run in runs]
    end = max(rep["seconds"] for rep in reps)
    t = time_at(round(OPENING["share"] * 1000), end)
    print(f"  {now_text(t, end)}")
    for p, run, rep in zip(("fixed", "continuous"), runs, reps):
        print(f"  {p}: {state_line(snapshot(run, requests, t), OPENING['slots'], N_REQUESTS)}")
        print(f"    {summary_line(rep)}")


if __name__ == "__main__":
    main()