"""A toy simulation of two GPUs that serve chats, laid out two ways: both GPUs do everything ("shared"), or one reads prompts and the other writes replies ("split").

Every cost is a toy value. Two of them echo lesson 23 (why-ai-waits-before-the-first-word): the time to read one prompt token and the time of one writing step are its
best-case numbers for Llama 3.1 8B with 16-bit weights on an NVIDIA A100 80GB SXM (the chip of lesson 19). The cache bytes per token echo lesson 22 (paged-kv-cache). The link
speeds, the slot count and the two targets are invented. Nothing here is a measurement, and test_pd_sim.py checks the copied values against lesson 23's code.

Time moves in whole steps. A GPU that has a waiting prompt and a free slot reads that whole prompt first (nothing else runs; there is no chunking) and the prompt's first word
comes out at the end. Otherwise it runs one writing step: every reply in progress gets one more word. A new reply joins at the start of a step. Python standard library only.
The widget (../widget.html) shows a table that show.py builds from this file, and test_widget.py checks the page against it.
"""
from __future__ import annotations

import random
from itertools import takewhile
from typing import NamedTuple

# --- toy values: the chip, copied from lesson 23 (prefill_sim.py) and lesson 19 (tiled-matmul-roofline)
PARAMS = 8_030_261_248                                  # Llama 3.1 8B, lesson 23
A100_PEAK_FLOPS, A100_BANDWIDTH = 312e12, 2039e9        # A100 80GB SXM data sheet, lesson 19: 312 TFLOPS, 2,039 GB/s (GB is 10**9 bytes)
READ_S_PER_TOKEN = 2 * PARAMS / A100_PEAK_FLOPS         # reading the prompt: 2 FLOPs per weight per token, limited by arithmetic (lesson 23)
STEP_S = 2 * PARAMS / A100_BANDWIDTH                    # one writing step: every weight is read once, limited by memory, the same for any number of replies (lesson 23)
CACHE_BYTES_PER_TOKEN = 131_072                         # 2 x 32 layers x 8 key-value heads x 128 x 2 bytes, lesson 22
# --- toy values: invented
SLOTS = 32                                              # replies one GPU writes at once
LINKS_GBPS = (100, 25, 12.5, 1.25, 0.125)               # link speeds offered, in 10**9 bytes per second (800, 200, 100, 10 and 1 gigabits per second)
FIRST_WORD_TARGET_S, WORD_GAP_TARGET_S = 0.5, 0.1       # a request is good if its first word comes within 0.5 s and no gap between words is over 0.1 s
MIX = (0.5, 1.0, 1.5)                                   # each prompt and each reply is its stated length times one of these, chosen at random
N_REQUESTS, SEED = 600, 7


class Request(NamedTuple):
    arrival: float      # seconds since the first request
    prompt: int         # tokens in the prompt
    reply: int          # words in the reply, at least 1


def make_requests(rps: float, prompt: int, reply: int, n: int = N_REQUESTS, seed: int = SEED) -> list[Request]:
    """Step 1: seeded Poisson arrivals (gaps between requests are exponential, mean 1 / rps)."""
    rng, t, out = random.Random(seed), 0.0, []
    for _ in range(n):
        t += rng.expovariate(rps)
        out.append(Request(t, round(prompt * rng.choice(MIX)), max(1, round(reply * rng.choice(MIX)))))
    return out


def serve(jobs: list[tuple]) -> tuple[dict, dict]:
    """One GPU. A job is (ready_time, read_seconds, words_left_to_write, id). Returns {id: time of first word} and {id: [gaps between its words]}.
    A job with read_seconds 0 arrives with its cache and first word already made (the split's writing GPU); the others read their prompt here."""
    waiting, active, first, gaps, last, t = sorted(jobs), [], {}, {}, {}, 0.0
    while waiting or active:
        if not active and waiting[0][0] > t:
            t = waiting[0][0]                                         # an idle GPU waits for the next job
        ready = list(takewhile(lambda j: j[0] <= t, waiting))         # `waiting` is sorted, so the jobs that are ready are its first few
        if ready and len(active) < SLOTS and ready[0][1] > 0:         # step 2: a prompt waits and a slot is free, so read the whole prompt now
            ready_t, read, left, i = ready[0]
            waiting.remove(ready[0])
            t += read
            first[i], gaps[i], last[i] = t, [], t
            if left > 0:
                active.append([i, left])
            continue
        for j in ready[: SLOTS - len(active)]:                        # a job that comes with its cache joins at the start of a step
            if j[1] == 0:
                waiting.remove(j)
                first[j[3]], gaps[j[3]], last[j[3]] = j[0], [], j[0]
                if j[2] > 0:
                    active.append([j[3], j[2]])
        if not active:
            continue
        t += STEP_S                                                   # one writing step: every reply in progress gets one word
        for a in active:
            gaps[a[0]].append(t - last[a[0]])
            last[a[0]], a[1] = t, a[1] - 1
        active = [a for a in active if a[1] > 0]
    return first, gaps


def handover_seconds(prompt_tokens: int, link_gbps: float) -> float:
    """Step 5: the prompt's cache is prompt tokens x bytes per token, and it crosses the link at the link's speed."""
    return prompt_tokens * CACHE_BYTES_PER_TOKEN / (link_gbps * 1e9)


def shared(requests: list[Request]) -> tuple[dict, dict]:
    """Step 2: two GPUs that each read prompts and write replies. Requests take turns: 0 goes to the first GPU, 1 to the second, 2 to the first."""
    first, gaps = {}, {}
    for g in (0, 1):
        f, w = serve([(r.arrival, r.prompt * READ_S_PER_TOKEN, r.reply - 1, i) for i, r in enumerate(requests) if i % 2 == g])
        first.update(f)
        gaps.update(w)
    return first, gaps


def split(requests: list[Request], link_gbps: float) -> tuple[dict, dict]:
    """Step 4: one GPU only reads prompts, one at a time in order; each prompt's cache crosses the link (step 5); the other GPU only writes. The first word reaches the
    reader of the reply when the cache lands on the writing GPU. The link carries any number of caches at once (no contention)."""
    free, jobs = 0.0, []
    for i, r in enumerate(requests):
        free = max(free, r.arrival) + r.prompt * READ_S_PER_TOKEN
        jobs.append((free + handover_seconds(r.prompt, link_gbps), 0.0, r.reply - 1, i))
    return serve(jobs)


def percentile(values: list[float], p: int) -> float:
    """Nearest rank, as in lesson 27 (why-average-speed-is-misleading)."""
    ordered = sorted(values)
    return ordered[-(-p * len(ordered) // 100) - 1]


def goodput(requests: list[Request], run: tuple[dict, dict], first_target: float = FIRST_WORD_TARGET_S, gap_target: float = WORD_GAP_TARGET_S) -> float:
    """Step 6: the share of requests whose first word came within `first_target` seconds and whose every gap between words was at most `gap_target`."""
    first, gaps = run
    good = sum(1 for i, r in enumerate(requests) if first[i] - r.arrival <= first_target and max(gaps[i], default=0.0) <= gap_target)
    return good / len(requests)


def measure(requests: list[Request], run: tuple[dict, dict]) -> dict:
    """Steps 3 and 6: time to first word per request, the gaps between words pooled over all requests, and goodput at the default targets."""
    first, gaps = run
    waits = [first[i] - r.arrival for i, r in enumerate(requests)]
    pooled = [g for i in range(len(requests)) for g in gaps[i]]
    return {"first_p50": percentile(waits, 50), "first_p99": percentile(waits, 99), "gap_p50": percentile(pooled, 50), "gap_p99": percentile(pooled, 99), "goodput": goodput(requests, run)}


def compare(rps: float, prompt: int, reply: int, link_gbps: float) -> dict:
    """The same requests through both layouts."""
    requests = make_requests(rps, prompt, reply)
    return {"shared": measure(requests, shared(requests)), "split": measure(requests, split(requests, link_gbps))}