"""A step-by-step scheduler for one GPU that serves many users, on top of the time model of lesson 23 (why-ai-waits-before-the-first-word).

The model is Llama 3.1 8B with 16-bit weights on an NVIDIA A100 80GB SXM. One step of the server is one pass of the model over every token in it:
the next token of each user who is already writing, plus the whole prompt of each user who joins in that step. The pass reads every weight once,
whatever the number of users, and reads the attention cache of each user who is writing. Its time is the larger of two floors, as in lessons 19 and 23:

    memory floor = (2 P + cache bytes per token * sum of the writing users' contexts) / bandwidth
    math floor   = 2 P * (writing users + joining prompt tokens) / peak          (2 FLOPs per weight per token)

Both floors assume the chip runs at its data-sheet peak, so these are best cases, not measurements. The scheduler runs two policies over the same
requests, one step at a time: fixed groups ("fixed") and continuous batching ("continuous"). The workload is seeded and invented.

Python 3 standard library only. The widget (../widget.html) repeats this scheduler in JavaScript; test_widget.py runs both and compares them.
"""
from __future__ import annotations

import math
import random
from dataclasses import dataclass
from typing import Callable, NamedTuple

# --- the chip and the model: copied from lesson 23 (prefill_sim.py), which copied the chip from lesson 19.
# test_batching_sim.py checks every value against ../../why-ai-waits-before-the-first-word/code/prefill_sim.py.
A100_PEAK_FLOPS = 312e12          # NVIDIA A100 data sheet, revision "Jun21", 80GB SXM column: 312 TFLOPS dense FP16
A100_BANDWIDTH = 2039e9           # the same column: 2,039 GB/s. GB is 10**9 bytes.
GPU_MEMORY_BYTES = 80e9           # the same column: "80GB HBM2e"; the sheet does not say which GB, read here as 80 billion bytes
FP16_BYTES = 2
HIDDEN, FFN, LAYERS, HEADS, KV_HEADS, VOCAB = 4096, 14336, 32, 32, 8, 128256     # Llama 3.1 8B, config.json
HEAD_DIM = HIDDEN // HEADS


def parameter_count() -> int:
    attention = 2 * HIDDEN * HIDDEN + 2 * HIDDEN * KV_HEADS * HEAD_DIM
    layer = attention + 3 * HIDDEN * FFN + 2 * HIDDEN
    return LAYERS * layer + 2 * VOCAB * HIDDEN + HIDDEN


PARAMS = parameter_count()
KV_BYTES_PER_TOKEN = 2 * LAYERS * KV_HEADS * HEAD_DIM * FP16_BYTES        # lesson 22: 128 KiB for each token of context


# --- one step ---------------------------------------------------------------------------------

@dataclass(frozen=True)
class Pass:
    flops: int
    nbytes: int

    @property
    def math_s(self) -> float:
        return self.flops / A100_PEAK_FLOPS

    @property
    def memory_s(self) -> float:
        return self.nbytes / A100_BANDWIDTH

    @property
    def seconds(self) -> float:
        return max(self.math_s, self.memory_s)

    @property
    def bound(self) -> str:
        """A tie counts as arithmetic-limited, as in lessons 19 and 23."""
        return "arithmetic" if self.flops / self.nbytes >= A100_PEAK_FLOPS / A100_BANDWIDTH else "memory"

    @property
    def busy(self) -> float:
        """Share of the chip's peak arithmetic this step reaches."""
        return self.math_s / self.seconds


def batch_step(contexts: list[int], joining_tokens: int = 0) -> Pass:
    """One step for users who are writing (`contexts` holds how many tokens each one has in its cache) and for prompts that join (`joining_tokens` in all).
    Left out: the attention arithmetic, the cache a joining prompt writes and the attention among prompt tokens (lesson 23's "extras"), all small next to the rest."""
    tokens = len(contexts) + joining_tokens
    return Pass(2 * PARAMS * tokens, FP16_BYTES * PARAMS + KV_BYTES_PER_TOKEN * sum(contexts))


def step_seconds(contexts: list[int], joining_tokens: int = 0) -> float:
    return batch_step(contexts, joining_tokens).seconds


def users_that_fit(context: int) -> int:
    """How many users, each with `context` tokens in the cache, fit in the chip's memory beside the weights (best case: all the rest is cache)."""
    return int((GPU_MEMORY_BYTES - FP16_BYTES * PARAMS) // (KV_BYTES_PER_TOKEN * context))


# --- the scheduler ----------------------------------------------------------------------------

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


class Step(NamedTuple):
    start: float
    end: float
    slots: list          # for each slot, the id of the request running in it during this step, or -1
    joined: list         # ids that joined at the start of this step (prefill happens in it)
    finished: list       # ids whose last token came out at the end of this step


class Run(NamedTuple):
    steps: list
    first_token: list    # for each request, the time its first token came out
    finished: list       # for each request, the time its last token came out
    join_step: list      # for each request, the index of the step it joined in


def simulate(requests: list[Request], slots: int, policy: str, cost: Callable[[list[int], int], float] = step_seconds) -> Run:
    """Run `requests` (sorted by arrival) through `slots` slots, one step at a time. A request's id is its index.

    fixed:      new requests are taken only when every slot is empty, so a group runs until its longest reply ends and finished users' slots sit empty.
    continuous: at the start of every step, each empty slot takes the next waiting request.
    Either way, an idle server waits for the next arrival. A request that joins has its prompt read in the step it joins, and its first token comes
    out at the end of that step; each later step writes one more token. `cost(contexts, joining_tokens)` is the step time."""
    assert policy in ("fixed", "continuous")
    n = len(requests)
    held, made = [-1] * slots, [0] * n
    first, last, steps, join_step = [0.0] * n, [0.0] * n, [], [0] * n
    now, nxt, done = 0.0, 0, 0
    while done < n:
        joined = []
        if policy == "continuous" or all(h < 0 for h in held):
            if all(h < 0 for h in held) and requests[nxt].arrival > now:
                now = requests[nxt].arrival                         # nothing is running: wait for the next request
            for s in range(slots):
                if held[s] < 0 and nxt < n and requests[nxt].arrival <= now:
                    held[s] = nxt
                    joined.append(nxt)
                    join_step[nxt] = len(steps)
                    nxt += 1
        writing = [h for h in held if h >= 0 and h not in joined]
        dt = cost([requests[i].prompt + made[i] for i in writing], sum(requests[i].prompt for i in joined))
        start, now = now, now + dt
        snapshot, finished = list(held), []
        for s, h in enumerate(held):
            if h < 0:
                continue
            made[h] += 1
            if h in joined:
                first[h] = now
            if made[h] == requests[h].reply:
                last[h] = now
                held[s] = -1
                finished.append(h)
        steps.append(Step(start, now, snapshot, joined, finished))
        done += len(finished)
    return Run(steps, first, last, join_step)


def report(run: Run, requests: list[Request], slots: int) -> dict:
    """What the lesson and the widget say about a run. Sums are plain left-to-right loops so that JavaScript gives the same bits.
    `slots_in_use` is the share of slot-seconds in use while the server was running steps (waiting for the next request is not counted)."""
    waits = [run.first_token[i] - requests[i].arrival for i in range(len(requests))]
    total_wait = tokens = 0.0
    for w in waits:
        total_wait += w
    for r in requests:
        tokens += r.reply
    used = busy = 0.0
    for st in run.steps:
        used += sum(1 for h in st.slots if h >= 0) * (st.end - st.start)
        busy += st.end - st.start
    seconds = run.steps[-1].end
    return {"replies": len(requests), "seconds": seconds, "tokens_per_second": tokens / seconds, "average_wait": total_wait / len(requests),
            "slowest_wait": max(waits), "slots_in_use": used / (slots * busy), "steps": len(run.steps)}


def snapshot(run: Run, requests: list[Request], t: float) -> dict:
    """What one lane shows at simulated time t (seconds): for each slot how much of its reply is written (0 to 1, or -1 when empty),
    how many requests have arrived and wait for a slot, and how many replies are finished."""
    k = len(run.steps) - 1
    while k >= 0 and run.steps[k].start > t:
        k -= 1
    shares = [-1.0] * len(run.steps[0].slots)
    if k >= 0 and t < run.steps[k].end:
        for s, h in enumerate(run.steps[k].slots):
            if h >= 0:
                shares[s] = (k - run.join_step[h] + 1) / requests[h].reply
    waiting = sum(1 for i, r in enumerate(requests) if r.arrival <= t < run.steps[run.join_step[i]].start)
    return {"slots": shares, "in_use": sum(1 for x in shares if x >= 0), "waiting": waiting, "done": sum(1 for f in run.finished if f <= t)}


# --- the invented workload ----------------------------------------------------------------------
# A seeded draw, not a measurement. Each request has a prompt of 30 to 600 tokens (equal steps on a log scale, mean about 190) and a reply whose length
# depends on the "mix" setting. Requests arrive at random, with the gaps between them drawn once; a traffic setting only scales the gaps.

SEED = 2026
N_REQUESTS = 48
RATES = [0.5, 1, 2, 4, 8, 16]                      # new requests per second, the traffic slider's stops
SLOT_CHOICES = [4, 8, 16, 32]
MIXES = {"same": (60, 60), "mixed": (10, 300), "very mixed": (5, 600)}   # reply length from low to high, equal steps on a log scale


def draw_workload(seed: int = SEED, n: int = N_REQUESTS) -> dict:
    rng = random.Random(seed)
    gaps = [rng.expovariate(1.0) for _ in range(n)]                   # seconds between requests at one request a second
    prompts = [round(30 * (600 / 30) ** rng.random()) for _ in range(n)]
    where = [rng.random() for _ in range(n)]
    replies = {name: [round(lo * (hi / lo) ** x) for x in where] for name, (lo, hi) in MIXES.items()}
    return {"gaps": gaps, "prompts": prompts, "replies": replies}


WORKLOAD = draw_workload()


def make_requests(rate: float, mix: str, workload: dict = WORKLOAD) -> list[Request]:
    """The same requests at any traffic: only the gaps scale. The first request arrives at time 0."""
    out, t = [], 0.0
    for gap, prompt, reply in zip(workload["gaps"], workload["prompts"], workload["replies"][mix]):
        out.append(Request(t, prompt, reply))
        t += gap / rate
    return out


# --- what the widget starts with, and the words it shows (widget.html repeats these in JavaScript; test_widget.py compares them) --------

OPENING = {"rate": 8, "slots": 8, "mix": "mixed", "share": 0.4}     # the opening state: heavy-ish traffic, 40% of the way through the longer run
QUEUE_CAP = 16                                                       # little squares drawn for waiting requests; more are counted as "+N"
TIME_TICKS = 1000                                                    # the time slider has this many positions
PLAY_SECONDS = 10                                                    # Play takes about this long for the longer run


def mix_names(workload: dict = WORKLOAD) -> dict:
    names = {"same": "Same length", "mixed": "Mixed", "very mixed": "Very mixed"}
    out = {}
    for key, replies in workload["replies"].items():
        lo, hi = min(replies), max(replies)
        out[key] = f"{names[key]} ({lo} tokens)" if lo == hi else f"{names[key]} ({lo} to {hi} tokens)"
    return out


def commas(n: int) -> str:
    return f"{n:,}"


def fmt_seconds(s: float) -> str:
    return f"{s:.2f}" if s < 10 else f"{s:.1f}"


def rate_text(rate: float) -> str:
    return f"{rate:g}"


def time_at(tick: int, longest_end: float) -> float:
    """Simulated seconds at a position of the time slider."""
    return tick / TIME_TICKS * longest_end


def now_text(t: float, longest_end: float) -> str:
    return f"{fmt_seconds(t)} s of {fmt_seconds(longest_end)} s"


def state_line(snap: dict, slots: int, total: int) -> str:
    return f"In use {snap['in_use']} of {slots}. {snap['waiting']} waiting. {snap['done']} of {total} done."


def summary_line(rep: dict) -> str:
    tokens = commas(int(f"{rep['tokens_per_second']:.0f}"))
    return (f"Whole run: done after {fmt_seconds(rep['seconds'])} s, {tokens} tokens a second. "
            f"First word after {fmt_seconds(rep['average_wait'])} s on average, {fmt_seconds(rep['slowest_wait'])} s at worst. "
            f"Slots in use: {math.floor(rep['slots_in_use'] * 100 + 0.5)}% on average.")