"""Tests on the serve-batch recording: Lab 12's load test (09_batching.py) rerun on Qwen3-4B 4-bit (Apple M4 Max),
40 simulated requests arriving at about 4 a second (the same arrival times for every server), up to 128 tokens each,
with every request's timing kept; and the lab's batch-size runs. Run: pytest test_batch_practice.py"""
from batch_practice import one_at_a_time, p50, static_batch, summarise, throughput

# Per request, in seconds from the start of the test: (arrival, start, first word, done).
ONE = [
    (0.16, 0.17, 0.3, 1.09), (0.43, 1.09, 1.21, 2.0), (0.7, 2.0, 2.1, 2.89), (0.74, 2.89, 2.99, 3.79),
    (0.74, 3.79, 3.88, 4.68), (0.86, 4.68, 4.81, 5.62), (0.94, 5.62, 5.73, 6.54), (1.36, 6.54, 6.65, 7.45),
    (1.65, 7.45, 7.54, 8.34), (1.88, 8.34, 8.48, 9.29), (2.08, 9.29, 9.4, 10.2), (2.36, 10.2, 10.34, 11.15),
    (2.39, 11.15, 11.26, 12.06), (2.54, 12.06, 12.17, 12.98), (2.58, 12.98, 13.1, 13.92), (3.18, 13.92, 14.01, 14.82),
    (3.19, 14.82, 14.96, 15.77), (3.62, 15.77, 15.9, 16.72), (3.64, 16.72, 16.84, 17.65), (3.93, 17.65, 17.75, 18.55),
    (4.03, 18.55, 18.65, 19.46), (4.16, 19.46, 19.58, 20.4), (4.62, 20.4, 20.57, 21.39), (4.63, 21.39, 21.49, 22.24),
    (4.64, 22.24, 22.36, 23.19), (5.26, 23.19, 23.32, 24.15), (5.44, 24.15, 24.24, 25.05), (5.46, 25.05, 25.15, 25.73),
    (6.55, 25.73, 25.85, 26.67), (7.28, 26.67, 26.82, 27.64), (7.31, 27.64, 27.76, 28.58), (7.45, 28.58, 28.69, 29.51),
    (7.48, 29.51, 29.66, 30.48), (7.58, 30.48, 30.6, 31.35), (7.82, 31.35, 31.47, 32.29), (7.87, 32.29, 32.46, 33.33),
    (8.16, 33.33, 33.47, 34.31), (8.18, 34.31, 34.41, 35.22), (8.22, 35.22, 35.34, 36.16), (8.65, 36.16, 36.28, 37.1),
]
STATIC = [
    (0.16, 0.67, 0.88, 1.74), (0.43, 0.67, 0.88, 1.74), (0.7, 1.74, 2.47, 4.85), (0.74, 1.74, 2.47, 4.85),
    (0.74, 1.74, 2.47, 4.85), (0.86, 1.74, 2.47, 4.85), (0.94, 1.74, 2.47, 4.85), (1.36, 1.74, 2.47, 4.85),
    (1.65, 1.74, 2.47, 4.85), (1.88, 4.85, 5.76, 8.23), (2.08, 4.85, 5.76, 8.23), (2.36, 4.85, 5.76, 8.23),
    (2.39, 4.85, 5.76, 8.21), (2.54, 4.85, 5.76, 8.23), (2.58, 4.85, 5.76, 8.23), (3.18, 4.85, 5.76, 8.23),
    (3.19, 4.85, 5.76, 8.23), (3.62, 8.23, 9.08, 11.33), (3.64, 8.23, 9.08, 11.33), (3.93, 8.23, 9.08, 11.33),
    (4.03, 8.23, 9.08, 11.33), (4.16, 8.23, 9.08, 11.33), (4.62, 8.23, 9.08, 11.33), (4.63, 8.23, 9.08, 11.13),
    (4.64, 8.23, 9.08, 11.33), (5.26, 11.33, 12.16, 14.45), (5.44, 11.33, 12.16, 14.45), (5.46, 11.33, 12.16, 13.8),
    (6.55, 11.33, 12.16, 14.45), (7.28, 11.33, 12.16, 14.45), (7.31, 11.33, 12.16, 14.45), (7.45, 11.33, 12.16, 14.45),
    (7.48, 11.33, 12.16, 14.45), (7.58, 14.45, 15.34, 17.57), (7.82, 14.45, 15.34, 17.72), (7.87, 14.45, 15.34, 17.72),
    (8.16, 14.45, 15.34, 17.72), (8.18, 14.45, 15.34, 17.72), (8.22, 14.45, 15.34, 17.72), (8.65, 14.45, 15.34, 17.72),
]
CONTINUOUS = [
    (0.16, 0.16, 0.28, 5.08), (0.43, 0.43, 0.57, 6.24), (0.7, 0.7, 1.19, 7.12), (0.74, 0.8, 1.21, 7.17),
    (0.74, 0.8, 1.21, 7.17), (0.86, 0.97, 1.24, 7.21), (0.94, 0.97, 1.24, 7.21), (1.36, 1.36, 1.53, 8.88),
    (1.65, 1.66, 1.82, 9.26), (1.88, 1.89, 2.1, 9.49), (2.08, 2.1, 2.28, 9.62), (2.36, 2.36, 2.88, 9.85),
    (2.39, 2.52, 2.93, 9.92), (2.54, 2.66, 2.97, 9.96), (2.58, 2.66, 2.97, 9.96), (3.18, 3.18, 3.51, 10.33),
    (3.19, 3.3, 3.55, 10.38), (3.62, 3.65, 4.07, 10.61), (3.64, 3.65, 4.07, 10.61), (3.93, 3.96, 4.33, 10.74),
    (4.03, 4.07, 4.38, 10.8), (4.16, 4.19, 4.42, 10.94), (4.62, 4.64, 5.13, 11.45), (4.63, 4.64, 5.13, 11.1),
    (4.64, 4.64, 5.13, 11.45), (5.26, 5.27, 5.56, 11.74), (5.44, 5.45, 5.73, 11.85), (5.46, 5.56, 5.77, 10.21),
    (6.55, 6.58, 6.8, 12.74), (7.28, 7.3, 7.95, 13.33), (7.31, 7.46, 8.27, 13.46), (7.45, 7.46, 8.27, 13.46),
    (7.48, 7.68, 8.61, 13.55), (7.58, 7.68, 8.61, 13.06), (7.82, 7.95, 8.67, 13.63), (7.87, 7.95, 8.67, 13.63),
    (8.16, 8.27, 8.82, 13.67), (8.18, 8.27, 8.82, 13.67), (8.22, 8.27, 8.82, 13.67), (8.65, 8.67, 8.93, 13.69),
]


def test_throughput_of_a_batch():
    assert throughput(128, 0.926, 1) == (138, 138)  # one reply alone
    assert throughput(4058, 8.416, 32) == (482, 15)  # 32 at once: 3.5x the total, each far slower
    assert throughput(8124, 17.572, 64) == (462, 7)  # 64: the total falls


def test_p50_is_nearest_rank():
    assert p50([3, 1, 2]) == 2
    assert p50([4, 1, 3, 2]) == 2  # the lower middle value, not 2.5
    assert p50([5.0]) == 5.0


def test_summarise_one_at_a_time():
    s = summarise(ONE)
    assert round(s["first"], 2) == 13.82  # the typical customer waited ~14 s for the first word
    assert round(s["maxwait"], 1) == 27.6
    assert round(s["write"], 2) == 0.81  # ...then got the reply in under a second
    assert round(s["whole"], 1) == 14.6


def test_summarise_batches_and_continuous():
    s, c = summarise(STATIC), summarise(CONTINUOUS)
    assert round(s["first"], 2) == 4.46 and round(c["first"], 2) == 0.43
    assert round(s["first"] / c["first"]) == 10 and round(summarise(ONE)["first"] / c["first"]) == 32
    # the whole reply: only about 2x sooner than one at a time, and the waiting bus was nearly as quick
    assert round(c["whole"], 1) == 6.4 and round(s["whole"], 1) == 6.7
    assert round(c["write"], 2) == 5.97  # writing is slower: everyone on board shares the chip


def test_one_at_a_time_queue_matches_the_recording():
    got = one_at_a_time([r[0] for r in ONE], [r[3] - r[1] for r in ONE])
    assert len(got) == 40
    for (start, done), r in zip(got, ONE):
        assert abs(start - r[1]) < 0.02 and abs(done - r[3]) < 0.02
    assert one_at_a_time([0, 0, 5], [1, 1, 1]) == [(0, 1), (1, 2), (5, 6)]


def test_static_batch_rules():
    assert static_batch([0.0, 0.1, 2.0], 0) == (0.5, 2)  # timer: the oldest waited 0.5 s
    assert static_batch([0.1 * i for i in range(10)], 0) == (0.5, 6)  # timer fires before the 8th arrives at 0.7
    assert static_batch([0.01 * i for i in range(10)], 0) == (0.07, 8)  # full first
    assert static_batch([1.0, 1.2], 3.0) == (3.0, 2)  # both waited long enough already
    assert static_batch([1.0], 1.0) == (1.0, 1)  # nobody else is coming: leave now


def test_static_batches_match_the_recording():
    waiting, now, i, starts = [r[0] for r in STATIC], 0.0, 0, []
    while i < len(STATIC):
        t, n = static_batch(waiting[i:], now)
        group = STATIC[i:i + n]
        assert len({r[1] for r in group}) == 1 and abs(t - group[0][1]) < 0.02  # they all started together, at t
        assert i + n == len(STATIC) or STATIC[i + n][1] > group[0][1]  # and the next one didn't
        starts.append(n)
        now, i = max(r[3] for r in group), i + n  # the bus is back when its last reply is done
    assert starts == [2, 7, 8, 8, 8, 7]
