"""A toy AI server, a pile of requests and the numbers a test reports about them: lesson 27 (why-average-speed-is-misleading).

The server is the continuous-batching server of lesson 24 (how-an-ai-server-answers-many-people-at-once): 8 slots on one A100, step times from its
`step_seconds`. Everything runs inside this one Python process. There is no network code: the "load generator" is a list of send times and the
"server" is a loop over steps, so nothing here can be pointed at a real service.

Per request the run records two times: the wait before the first word (from send to the end of the step that reads its prompt) and the total
time (from send to the last word). `summarize` turns a list of such times into an average and the percentiles p50, p95 and p99.

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

import heapq
import importlib.util
import math
import random
import sys
from collections import deque
from pathlib import Path
from typing import NamedTuple

# --- lesson 24's step time, loaded by path (a sibling lesson folder is not a package) --------------------------------

def _load_batching():
    path = Path(__file__).resolve().parent.parent.parent / "how-an-ai-server-answers-many-people-at-once" / "code" / "batching_sim.py"
    spec = importlib.util.spec_from_file_location("batching_sim_lesson24", path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module                                   # a dataclass looks its own module up by name
    keep, sys.dont_write_bytecode = sys.dont_write_bytecode, True
    try:
        spec.loader.exec_module(module)
    finally:
        sys.dont_write_bytecode = keep
    return module


batching = _load_batching()
step_seconds = batching.step_seconds

# --- the invented workload -----------------------------------------------------------------------------------------------
# A seeded draw, not a measurement. Prompts run from 30 to 600 tokens (equal steps on a log scale, as in lesson 24). Replies are a skewed mix:
# most are short (8 to 60 tokens) and 12 in 100 are long (100 to 400 tokens).

SEED = 2027
N_REQUESTS = 2000
SLOTS = 8
LONG_SHARE = 0.12
SHORT_REPLY, LONG_REPLY = (8, 60), (100, 400)
RATES = [4, 8, 12, 15, 17, 19]            # intended requests per second, the load slider's stops
OPENING_RATE = 17
CONNECTIONS = 8                           # the waiting tester keeps this many requests going, one per connection
PAUSE_AT, PAUSE = 40.0, 5.0               # the toy server freezes for PAUSE seconds in the first step that starts at or after PAUSE_AT seconds
PERCENTILES = (50, 95, 99)


def draw_workload(seed: int = SEED, n: int = N_REQUESTS) -> dict:
    """Gaps (seconds between requests at one request a second, rounded to 4 places so JavaScript reads the same doubles), prompts and replies."""
    rng = random.Random(seed)
    gaps = [round(rng.expovariate(1.0), 4) for _ in range(n)]
    prompts = [round(30 * (600 / 30) ** rng.random()) for _ in range(n)]
    long = [rng.random() < LONG_SHARE for _ in range(n)]
    where = [rng.random() for _ in range(n)]
    replies = []
    for is_long, x in zip(long, where):
        lo, hi = LONG_REPLY if is_long else SHORT_REPLY
        replies.append(round(lo * (hi / lo) ** x))
    return {"gaps": gaps, "prompts": prompts, "replies": replies}


WORKLOAD = draw_workload()


def draw_rerolls(seed: int = SEED + 1, count: int = 8, n: int = N_REQUESTS) -> list[tuple[int, int]]:
    """Index maps i -> (a*i + b) % n for the widget's "Shuffle" button: same requests, another order. The first is the identity.
    `a` is odd and not a multiple of 5, so it is coprime with 2000 and the map is a permutation."""
    assert n % 10 == 0
    rng = random.Random(seed)
    out = [(1, 0)]
    while len(out) < count:
        a = rng.randrange(3, n)
        if a % 2 and a % 5:
            out.append((a, rng.randrange(n)))
    return out


REROLLS = draw_rerolls()


def shuffled(workload: dict, a: int, b: int) -> dict:
    n = len(workload["gaps"])
    order = [(a * i + b) % n for i in range(n)]
    return {key: [values[j] for j in order] for key, values in workload.items()}


def intended_times(gaps: list[float], rate: float) -> list[float]:
    """When the tester means to send each request: random gaps that average 1/rate seconds. The first goes at time 0."""
    out, t = [], 0.0
    for gap in gaps:
        out.append(t)
        t += gap / rate
    return out


# --- the toy server ---------------------------------------------------------------------------------------------------------

class Served(NamedTuple):
    intended: list      # when the tester meant to send each request
    sent: list          # when it actually did
    first: list         # when the first word came out
    done: list          # when the last word came out
    steps: int


def serve(workload: dict, rate: float, connections: int = 0, pause: float = 0.0, slots: int = SLOTS) -> Served:
    """Send the requests to the toy server and run it one step at a time, with lesson 24's continuous batching: at the start of every step each empty
    slot takes the next waiting request; a joining prompt is read in the step it joins and its first word comes out at the end of that step.

    connections == 0: "on schedule". Every request is sent at its intended time, whatever the server is doing.
    connections == K: "waits for each reply". K connections; connection c owns requests c, c+K, c+2K, ... and sends the next one when the previous
    reply is complete or at its intended time, whichever is later.
    pause > 0: the server freezes for `pause` seconds in the first step that starts at or after PAUSE_AT."""
    prompts, replies, n = workload["prompts"], workload["replies"], len(workload["gaps"])
    intended = intended_times(workload["gaps"], rate)
    sent, first, done = [0.0] * n, [0.0] * n, [0.0] * n
    ahead = [(intended[i], i) for i in range(n if connections == 0 else min(connections, n))]    # (send time, id), a heap
    heapq.heapify(ahead)
    held, made, waiting = [-1] * slots, [0] * n, deque()
    now, finished, steps, frozen = 0.0, 0, 0, False
    while finished < n:
        if not waiting and all(h < 0 for h in held) and ahead[0][0] > now:
            now = ahead[0][0]                                         # nothing is running: wait for the next request
        while ahead and ahead[0][0] <= now:
            when, i = heapq.heappop(ahead)
            sent[i] = when
            waiting.append(i)
        joined = []
        for s in range(slots):
            if held[s] < 0 and waiting:
                held[s] = waiting.popleft()
                joined.append(held[s])
        writing = [h for h in held if h >= 0 and h not in joined]
        dt = step_seconds([prompts[i] + made[i] for i in writing], sum(prompts[i] for i in joined))
        if pause > 0 and not frozen and now >= PAUSE_AT:
            dt += pause
            frozen = True
        now += dt
        steps += 1
        for s, h in enumerate(held):
            if h < 0:
                continue
            made[h] += 1
            if h in joined:
                first[h] = now
            if made[h] == replies[h]:
                done[h] = now
                held[s] = -1
                finished += 1
                follow = h + connections
                if connections and follow < n:
                    heapq.heappush(ahead, (max(intended[follow], now), follow))
    return Served(intended, sent, first, done, steps)


def capacity(workload: dict = WORKLOAD) -> float:
    """Requests a second the toy server can finish: all requests wait at time 0 and the server never idles."""
    n = len(workload["gaps"])
    run = serve({**workload, "gaps": [0.0] * n}, 1.0)
    return n / max(run.done)


# --- the numbers a test reports ---------------------------------------------------------------------------------------------

def percentile(values: list[float], p: int) -> float:
    """Nearest rank: the smallest value that at least p in 100 of the values do not exceed, i.e. the value at rank ceil(p * n / 100) when sorted."""
    ordered = sorted(values)
    return ordered[-(-p * len(ordered) // 100) - 1]


def mean(values: list[float]) -> float:
    total = 0.0
    for v in values:
        total += v
    return total / len(values)


def summarize(values: list[float]) -> dict:
    return {"average": mean(values), **{f"p{p}": percentile(values, p) for p in PERCENTILES}, "max": max(values)}


def totals(run: Served) -> list[float]:
    """Total time per request: send to last word."""
    return [d - s for d, s in zip(run.done, run.sent)]


def waits(run: Served) -> list[float]:
    """Wait before the first word per request: send to first word."""
    return [f - s for f, s in zip(run.first, run.sent)]


MEASURES = {"total": totals, "wait": waits}


def sweep(workload: dict = WORKLOAD, rates: list = RATES) -> dict:
    """For each rate, the summary of both measures on an on-schedule run."""
    out = {}
    for rate in rates:
        run = serve(workload, rate)
        out[rate] = {name: summarize(measure(run)) for name, measure in MEASURES.items()}
    return out


def pooled_example(workload: dict = WORKLOAD) -> dict:
    """Two servers that each handled the same 2,000 requests, one at a quiet 4 a second and one at a busy 19 a second, measured by the wait before
    the first word. The average of their two p99s is not the p99 of all 4,000 requests together."""
    quiet = waits(serve(workload, 4))
    busy = waits(serve(workload, 19))
    p99s = [percentile(quiet, 99), percentile(busy, 99)]
    return {"requests": len(quiet), "quiet_p99": p99s[0], "busy_p99": p99s[1], "averaged": mean(p99s), "pooled": percentile(quiet + busy, 99)}


def shuffle_spread(rate: int = OPENING_RATE, workload: dict = WORKLOAD) -> dict:
    """Lowest and highest average and p99 total time over the widget's shuffles (REROLLS, the opening order included) at one load."""
    runs = [summarize(totals(serve(shuffled(workload, a, b), rate))) for a, b in REROLLS]
    return {key: (min(r[key] for r in runs), max(r[key] for r in runs)) for key in ("average", "p99")}


DEMO_RATE = 12            # the load of the two-testers comparison: well below capacity, so a pause is the only trouble
SLOW_WAIT = 1.0           # seconds; a wait before the first word above this counts as slow in the comparison


def generator_demo(workload: dict = WORKLOAD) -> dict:
    """The same server and the same 5 s freeze, measured by two testers. For each: the summary of the wait before the first word and of the total
    time as that tester records them, and how many requests it saw waiting more than SLOW_WAIT seconds."""
    out = {}
    for name, connections in (("on schedule", 0), ("waits for each reply", CONNECTIONS)):
        run = serve(workload, DEMO_RATE, connections, PAUSE)
        wait = waits(run)
        out[name] = {"wait": summarize(wait), "total": summarize(totals(run)), "slow": sum(1 for w in wait if w > SLOW_WAIT)}
    return out


def session_chance(slow_share: float, requests: int) -> float:
    """Chance that a session of `requests` independent requests meets at least one that is slower than the cutoff, when `slow_share` of requests are."""
    return 1 - (1 - slow_share) ** requests


# --- the words the widget shows (widget.html repeats these in JavaScript; test_widget.py compares them) --------------------------

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


def headline(stats: dict, measure: str) -> str:
    """The one line above the widget's chart: the average, then what 1 request in 20 and 1 in 100 meet."""
    avg, p95, p99 = (seconds_text(stats[k]) for k in ("average", "p95", "p99"))
    if measure == "total":
        return f"Average {avg} s. But 1 request in 20 takes more than {p95} s, and 1 in 100 more than {p99} s."
    return f"Average {avg} s. But 1 request in 20 waits more than {p95} s for its first word, and 1 in 100 more than {p99} s."


KEY_FIELDS = (("Average", "average"), ("p50", "p50"), ("p95", "p95"), ("p99", "p99"))


def key_rows(stats: dict) -> list[tuple[str, str]]:
    """The widget's key under the chart: name and value in seconds."""
    return [(name, f"{seconds_text(stats[field])} s") for name, field in KEY_FIELDS]