"""Guards for batching_sim.py: the chip and model it copies from lesson 23, the one-step time, the two schedulers, the invented workload, and the numbers the lesson prints."""
import importlib.util
import sys
from pathlib import Path

import pytest

import batching_sim as sim

sys.dont_write_bytecode = True                           # loading lesson 23's file below must not leave a __pycache__ in its folder
LESSON_23 = Path(__file__).resolve().parents[2] / "why-ai-waits-before-the-first-word" / "code" / "prefill_sim.py"
GIB = 2**30
GROUPS = [(rate, mix, slots) for mix in sim.MIXES for slots in sim.SLOT_CHOICES for rate in sim.RATES]


def lesson_23():
    spec = importlib.util.spec_from_file_location("prefill_sim_lesson_23", LESSON_23)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module                      # dataclasses looks the module up by name
    spec.loader.exec_module(module)
    return module


def toy(policy):
    """Four slots, six requests that all arrive at once with replies of 2, 6, 3, 1, 4 and 2 tokens, every step taking one second."""
    requests = [sim.Request(0.0, 10, reply) for reply in (2, 6, 3, 1, 4, 2)]
    return requests, sim.simulate(requests, 4, policy, cost=lambda contexts, joining: 1.0)


def test_model_and_chip_equal_lesson_23s_and_a_lone_user_costs_what_it_did_there():
    """Guards a copy that drifted: this lesson quotes lesson 23's 7.88 ms step and 128 KiB per token, so one user alone must still cost exactly that."""
    ours, theirs = sim, lesson_23()
    for name in ("PARAMS", "KV_BYTES_PER_TOKEN", "A100_PEAK_FLOPS", "A100_BANDWIDTH", "FP16_BYTES"):
        assert getattr(ours, name) == getattr(theirs, name), name
    assert ours.batch_step([0]).seconds == theirs.decode_step(1).seconds == ours.FP16_BYTES * ours.PARAMS / ours.A100_BANDWIDTH
    for n in (1, 100, 154, 600):                          # a prompt that joins an empty server is lesson 23's prefill
        assert ours.batch_step([], n).seconds == theirs.prefill(n).seconds
    for context in (1, 1000, 8000):                       # a user alone reads its cache and nothing else (lesson 23, with the cache counted)
        assert ours.batch_step([context]).memory_s == theirs.decode_step(context, True).memory_s


def test_numbers_the_lesson_prints_for_one_step():
    """Pins the step times, rates and cache sizes in lesson.md and the figures, so a change to the model cannot leave stale prose."""
    step = lambda users, context: sim.batch_step([context] * users)
    ms = lambda users, context: round(step(users, context).seconds * 1000, 2)
    assert [ms(1, 1000), ms(64, 1000), ms(256, 1000)] == [7.94, 11.99, 24.33]
    assert [round(users / step(users, 1000).seconds) for users in (1, 64, 256)] == [126, 5337, 10521]
    assert [round(1 / step(users, 1000).seconds, 1) for users in (1, 64, 256)] == [125.9, 83.4, 41.1]
    assert [round(users * 1000 * sim.KV_BYTES_PER_TOKEN / 1e9, 2) for users in (1, 64, 256)] == [0.13, 8.39, 33.55]
    assert round(sim.FP16_BYTES * sim.PARAMS / sim.A100_BANDWIDTH * 1000, 2) == 7.88 and ms(64, 100) == 8.29
    assert all(step(users, 1000).bound == "memory" for users in (1, 64, 256, 487))
    assert round(sim.batch_step([1000] * 64).busy * 100, 1) == 27.5 and sim.batch_step([1000]).busy < 0.01     # the share of the step the arithmetic is busy


def test_each_user_adds_more_reading_than_arithmetic_at_a_thousand_tokens_but_not_at_a_hundred():
    """Guards the lesson's claim about why the chip stays memory-limited: reading one user's cache costs 64 us at 1,000 tokens and
    51 us of arithmetic is all a user adds, so the step never turns arithmetic-limited. At 100 tokens it does, from 175 users on."""
    us = lambda seconds: round(seconds * 1e6, 1)
    math_per_user = sim.batch_step([0]).math_s
    assert us(math_per_user) == 51.5 and us(sim.KV_BYTES_PER_TOKEN * 1000 / sim.A100_BANDWIDTH) == 64.3
    assert not any(sim.batch_step([1000] * u).bound == "arithmetic" for u in range(1, sim.users_that_fit(1000) + 1))
    assert sim.batch_step([100] * 174).bound == "memory" and sim.batch_step([100] * 175).bound == "arithmetic"


def test_users_that_fit_in_memory_beside_the_weights():
    """Guards the memory cap the lesson states (best case, all the rest of the 80 GB is cache): 487, 243 and 59 users."""
    assert [sim.users_that_fit(c) for c in (1000, 2000, 8192)] == [487, 243, 59]
    assert (sim.GPU_MEMORY_BYTES - sim.FP16_BYTES * sim.PARAMS) / 1e9 == pytest.approx(63.94, abs=5e-3) and 8192 * sim.KV_BYTES_PER_TOKEN == GIB


@pytest.mark.parametrize("policy", ["fixed", "continuous"])
def test_toy_example_slot_by_slot(policy):
    """Guards the hand-worked example in the lesson (figure 1 and the Go-deeper table): which request sits in which slot at each step."""
    requests, run = toy(policy)
    slots = [step.slots for step in run.steps]
    if policy == "fixed":
        assert slots == [[0, 1, 2, 3], [0, 1, 2, -1], [-1, 1, 2, -1], [-1, 1, -1, -1], [-1, 1, -1, -1], [-1, 1, -1, -1],
                         [4, 5, -1, -1], [4, 5, -1, -1], [4, -1, -1, -1], [4, -1, -1, -1]]
        assert [run.first_token[i] for i in (4, 5)] == [7.0, 7.0] and run.steps[-1].end == 10.0
    else:
        assert slots == [[0, 1, 2, 3], [0, 1, 2, 4], [5, 1, 2, 4], [5, 1, -1, 4], [-1, 1, -1, 4], [-1, 1, -1, -1]]
        assert [run.first_token[i] for i in (4, 5)] == [2.0, 3.0] and run.steps[-1].end == 6.0
    assert run.finished[1] == 6.0 and [s.finished for s in run.steps][:3] == [[3], [0], [2]]


def test_snapshot_reads_a_lane_in_the_middle_of_a_step():
    """Guards the widget's picture: half way through step 0 each slot shows the share of its reply written (a request of 2, 6, 3 and 1 tokens),
    and two requests are counted as waiting until the step they join starts."""
    requests, run = toy("continuous")
    assert sim.snapshot(run, requests, 0.5) == {"slots": [0.5, 1 / 6, 1 / 3, 1.0], "in_use": 4, "waiting": 2, "done": 0}
    assert sim.snapshot(run, requests, 1.5)["slots"] == [1.0, 2 / 6, 2 / 3, 0.25] and sim.snapshot(run, requests, 1.5)["done"] == 1
    assert sim.snapshot(run, requests, 5.9) == {"slots": [-1.0, 6 / 6, -1.0, -1.0], "in_use": 1, "waiting": 0, "done": 5}
    assert sim.snapshot(run, requests, 6.0) == {"slots": [-1.0] * 4, "in_use": 0, "waiting": 0, "done": 6}


def test_an_idle_server_waits_for_the_next_request_and_a_lone_prompt_costs_lesson_23s_prefill():
    """Guards the clock jump: a request that arrives at 10 s finds nothing running, so its first token comes one joining step later, not at 0 s plus anything."""
    prefill = lesson_23().prefill
    requests = [sim.Request(0.0, 100, 3), sim.Request(10.0, 600, 2)]
    for policy in ("fixed", "continuous"):
        run = sim.simulate(requests, 2, policy)
        assert run.steps[-2].start == 10.0 and run.first_token[1] - 10.0 == pytest.approx(prefill(600).seconds, rel=1e-9)
        assert run.first_token[0] == prefill(100).seconds and run.steps[-1].end == run.finished[1]


@pytest.mark.parametrize("policy", ["fixed", "continuous"])
def test_every_setting_obeys_the_rules_of_its_policy(policy):
    """Guards scheduler slips a reader would see as a bug in the picture, in every traffic, reply and slot setting: a request in two places, a reply with
    the wrong number of steps, a start before arrival, a first token out of order, a fixed group that takes a newcomer early, or a continuous lane
    with an empty slot while someone waits."""
    for rate, mix, slots in GROUPS:
        requests = sim.make_requests(rate, mix)
        run = sim.simulate(requests, slots, policy)
        where = {i: [] for i in range(len(requests))}
        for k, step in enumerate(run.steps):
            assert len(step.slots) == slots and len(set(h for h in step.slots if h >= 0)) == sum(1 for h in step.slots if h >= 0), (rate, mix, slots, k)
            assert step.end > step.start and (k == 0 or step.start >= run.steps[k - 1].end)
            for s, h in enumerate(step.slots):
                if h >= 0:
                    where[h].append((k, s))
            writing = [h for h in step.slots if h >= 0 and h not in step.joined]
            if policy == "fixed" and step.joined:
                assert not writing, (rate, mix, slots, k)                       # a new group starts only into an empty lane
            if policy == "continuous" and -1 in step.slots:
                assert all(run.join_step[i] <= k for i, r in enumerate(requests) if r.arrival <= step.start), (rate, mix, slots, k)
        for i, r in enumerate(requests):
            assert [k for k, _ in where[i]] == list(range(run.join_step[i], run.join_step[i] + r.reply)), (rate, mix, slots, i)
            assert len({s for _, s in where[i]}) == 1                          # a reply stays in its slot
            assert run.steps[run.join_step[i]].start >= r.arrival and run.first_token[i] == run.steps[run.join_step[i]].end
            assert run.first_token[i] <= run.finished[i] == run.steps[run.join_step[i] + r.reply - 1].end


def test_equal_replies_leave_no_slot_empty_inside_a_fixed_group_so_fixed_only_waits():
    """Guards the lesson's caveat that with same-length replies fixed groups waste no slots (a group finishes all at once) and lose only by making newcomers wait."""
    for rate in sim.RATES:
        requests = sim.make_requests(rate, "same")
        run = sim.simulate(requests, 8, "fixed")
        assert all(len(step.finished) in (0, sum(1 for h in step.slots if h >= 0)) for step in run.steps), rate


def test_continuous_is_never_slower_and_never_makes_anyone_wait_longer_in_any_setting_here():
    """Guards the lesson's claim that continuous batching is at least as good on the same requests: the same replies finish no later and the first word comes no later,
    on average and at worst, in every one of the 72 settings (it is not a law: a different workload could break it)."""
    for rate, mix, slots in GROUPS:
        requests = sim.make_requests(rate, mix)
        fixed, continuous = (sim.report(sim.simulate(requests, slots, p), requests, slots) for p in ("fixed", "continuous"))
        for key in ("seconds", "average_wait", "slowest_wait"):
            assert continuous[key] <= fixed[key] * (1 + 1e-12), (rate, mix, slots, key)


def test_numbers_the_lesson_prints_for_the_two_schedulers():
    """Pins the results in lesson.md and figure 4 for 8 slots and mixed replies: light traffic (0.5 a second), the opening state (8) and heavy traffic (16)."""
    def results(rate, mix="mixed", slots=8):
        requests = sim.make_requests(rate, mix)
        return [sim.report(sim.simulate(requests, slots, p), requests, slots) for p in ("fixed", "continuous")]
    line = lambda rep: (round(rep["seconds"], 1), round(rep["tokens_per_second"]), round(rep["average_wait"], 2), round(rep["slowest_wait"], 2), round(rep["slots_in_use"] * 100))
    assert [line(r) for r in results(0.5)] == [(103.9, 47, 0.19, 1.91, 13), (103.9, 47, 0.01, 0.03, 14)]
    assert [line(r) for r in results(8)] == [(14.1, 349, 3.43, 6.18, 38), (8.4, 582, 0.03, 0.33, 61)]
    assert [line(r) for r in results(16)] == [(14.1, 349, 5.02, 9.1, 38), (6.3, 777, 0.53, 1.27, 80)]
    assert [line(r)[:4] for r in results(8, "same")] == [(7.2, 400, 0.31, 0.54), (6.9, 420, 0.02, 0.04)]


def test_workload_is_the_seeded_draw_the_lesson_describes():
    """Guards a changed draw: the lesson says 48 requests, prompts of 30 to 600 tokens averaging about 190, and three reply mixes of 60, 10 to 268 and 5 to 513 tokens."""
    w = sim.WORKLOAD
    assert sim.draw_workload() == w and len(w["gaps"]) == len(w["prompts"]) == 48 and all(len(r) == 48 for r in w["replies"].values())
    assert min(w["prompts"]) >= 30 and max(w["prompts"]) <= 600 and round(sum(w["prompts"]) / 48) == 191
    assert [(min(r), max(r)) for r in w["replies"].values()] == [(60, 60), (10, 268), (5, 513)]
    assert [round(sum(r) / 48, 1) for r in w["replies"].values()] == [60.0, 102.4, 157.4]
    assert sim.mix_names() == {"same": "Same length (60 tokens)", "mixed": "Mixed (10 to 268 tokens)", "very mixed": "Very mixed (5 to 513 tokens)"}


def test_traffic_only_scales_the_gaps_between_arrivals():
    """Guards the widget's promise that the slider changes only the traffic: the same prompts and replies arrive at every rate, 8 times as fast at 8 times the rate."""
    slow, fast = sim.make_requests(1, "mixed"), sim.make_requests(8, "mixed")
    assert slow[0].arrival == fast[0].arrival == 0.0
    assert [(r.prompt, r.reply) for r in slow] == [(r.prompt, r.reply) for r in fast]
    assert all(s.arrival / 8 == pytest.approx(f.arrival, rel=1e-12) for s, f in zip(slow, fast))
    assert all(a.arrival <= b.arrival for a, b in zip(slow, slow[1:]))


def test_the_opening_state_shows_the_point_at_once():
    """Guards the first screen: at the opening setting, 40% in, the fixed lane has most slots empty while 22 wait, and the continuous lane has most slots full and nobody waiting."""
    o = sim.OPENING
    requests = sim.make_requests(o["rate"], o["mix"])
    runs = [sim.simulate(requests, o["slots"], p) for p in ("fixed", "continuous")]
    t = sim.time_at(round(o["share"] * sim.TIME_TICKS), max(r.steps[-1].end for r in runs))
    fixed, continuous = (sim.snapshot(r, requests, t) for r in runs)
    assert (fixed["in_use"], fixed["waiting"], fixed["done"]) == (2, 22, 15) and (continuous["in_use"], continuous["waiting"], continuous["done"]) == (7, 0, 32)
    assert sim.state_line(fixed, 8, 48) == "In use 2 of 8. 22 waiting. 15 of 48 done."
    assert sim.fmt_seconds(t) == "5.63"