"""Run: python3 -B -m pytest -q -p no:cacheprovider code/ -- each test guards a number the lesson or the widget shows a reader."""
import math
import re
import statistics
from functools import lru_cache
from pathlib import Path

import settings_calc as sc
from settings_calc import BLOCK_BYTES, BLOCK_TOKENS, GPU_BYTES, RESERVE_BYTES, WEIGHT_BYTES, Request, blocks_for, room, serve


LESSON = Path(__file__).resolve().parent.parent


@lru_cache(maxsize=None)
def swept(share=sc.OPENING["share"], longest=sc.OPENING["longest"]):
    return sc.sweep(share, longest)


@lru_cache(maxsize=None)
def ran(share, longest, cap, reserved=False):
    return serve(sc.backlog(), cap, room(share, longest)["blocks"], longest, reserved)


def test_the_parts_add_up_to_the_card_and_the_pool_is_whole_blocks():
    """A bar that does not add up to 80 GB, or a pool counted in half blocks, would put wrong numbers on every line of the widget."""
    for share in sc.SHARES + [1, 20, 23]:
        r = room(share, 4000)
        assert r["weights"] + r["reserve"] + r["pool"] + r["unused"] == GPU_BYTES
        assert r["pool"] == r["blocks"] * BLOCK_BYTES and r["pool"] <= max(0, r["allowed"] - WEIGHT_BYTES - RESERVE_BYTES)
        assert r["allowed"] - WEIGHT_BYTES - RESERVE_BYTES - r["pool"] < BLOCK_BYTES or r["blocks"] == 0
    assert (room(90, 8000)["blocks"], room(90, 8000)["pool_tokens"]) == (25_720, 411_520)
    assert room(20, 8000)["blocks"] == 0 and room(23, 8000)["blocks"] == 161       # the card holds 18.06 GB of weights and reserve: 22.6% is the lowest share with any cache


def test_chats_that_fit_and_the_server_that_will_not_start():
    r = room(90, 8000)
    assert (r["typical"], r["at_longest"], r["starts"]) == (205, 51, True)          # 25,720 // 125 and 25,720 // 500
    assert room(90, 8000, reserved=True)["typical"] == 51
    for share in sc.SHARES:                                                          # starts exactly when the pool can hold one chat of the longest length
        for longest in sc.LONGEST:
            r = room(share, longest)
            assert r["starts"] == (r["blocks"] >= blocks_for(longest))
    assert not room(40, 128_000)["starts"] and room(44, 128_000)["starts"]
    assert not room(23, 128_000)["starts"] and "will not start" in sc.room_line(room(40, 128_000), 128_000, False)


def test_blocks_hold_more_chats_than_reserving_the_longest_length():
    """With blocks the longest chat is a ceiling; if each chat held a full-length slice, long ceilings would shrink the number of chats."""
    for longest in sc.LONGEST[1:]:
        r = room(90, longest)
        assert r["typical"] == room(90, 8000)["typical"] > room(90, longest, True)["typical"] == r["at_longest"]
    blocks = room(90, 32_000)["blocks"]
    grown = serve(sc.backlog(), 512, blocks, 32_000)
    held = serve(sc.backlog(), 512, blocks, 32_000, reserved=True)
    assert held["paused"] == 0 and held["lowest_free"] >= 0 and held["running"] < 13 < grown["running"]
    assert grown["tokens_per_second"] > held["tokens_per_second"]


def test_the_server_matches_lesson_24_when_the_pool_is_unlimited():
    """The new scheduler must be lesson 24's continuous batching plus a memory limit, or its numbers would not carry on from that lesson."""
    for mix in ("same", "very mixed"):
        requests = sc.L24.make_requests(8, mix)
        theirs = sc.L24.simulate(requests, 8, "continuous")
        mine = serve(requests, 8, 10**9, 10**5)
        assert mine["first_token"] == theirs.first_token and mine["finished"] == theirs.finished
        assert mine["steps"] == len(theirs.steps) and mine["paused"] == 0


def test_a_full_pool_pauses_the_newest_chat_and_rebuilds_it_later():
    """Two chats of 30 tokens each need a third block at token 32 and the pool has 4: the newer pauses, the older finishes, the newer reads its 32 tokens again. A chat of exactly the longest length is served; one token more is turned away."""
    requests = [Request(0.0, 30, 10), Request(0.0, 30, 10)]
    seen = []
    run = serve(requests, 2, 4, 64, cost=lambda contexts, joining: seen.append((sorted(contexts), joining)) or 1.0)
    assert run["paused"] == 1 and run["rebuilt_tokens"] == 32 and run["lowest_free"] == 0
    assert run["finished"][0] < run["finished"][1] and run["first_token"][1] == 1.0   # the paused chat keeps its first-token time
    assert seen[0] == ([], 60) and seen[1] == ([31, 31], 0) and seen[2] == ([32], 0)          # step 3: only the older chat writes
    assert seen[10] == ([], 32) and run["steps"] == 18                                       # step 11: the newer chat is read in again, then writes 8 more tokens
    edge = serve([Request(0.0, 54, 10), Request(0.0, 55, 10)], 2, 4, 64)                         # "up to 64" means a chat of 64 tokens is served and one of 65 is turned away
    assert (edge["served"], edge["refused"]) == (1, 1)


def test_pool_is_never_overdrawn_and_every_chat_gets_its_whole_reply():
    for share in (30, 35, 50, 90):
        r = room(share, 8000)
        for cap in (16, 128, 512):
            run = serve(sc.backlog(), cap, r["blocks"], 8000)
            assert run["lowest_free"] >= 0 and run["served"] == sc.N_CHATS and all(t > 0 for t in run["finished"])
    cut = serve(sc.backlog(), 64, room(90, 4000)["blocks"], 4000)
    assert cut["refused"] == sum(p + rep > 4000 for p, rep in zip(sc.CHATS["prompts"], sc.CHATS["replies"])) == 89
    assert cut["refused"] + cut["served"] == sc.N_CHATS


def test_raising_the_cap_stops_paying_where_the_lesson_says():
    rows = {row["cap"]: row for row in swept()}
    gains = [rows[2 * c]["tokens_per_second"] / rows[c]["tokens_per_second"] for c in sc.CAPS[:6]]       # 1->2 ... 32->64
    assert all(g > 1.10 for g in gains) and sc.stops_paying(list(rows.values())) == 64
    assert 1.05 < rows[128]["tokens_per_second"] / rows[64]["tokens_per_second"] < 1.10
    assert all(rows[2 * c]["seconds_per_word"] > rows[c]["seconds_per_word"] for c in sc.CAPS[:7])         # waiting keeps climbing up to 128
    assert {**rows[256], "cap": 0} == {**rows[512], "cap": 0} and rows[256]["paused"] > 0 and rows[128]["paused"] == 0                 # past what the pool holds, the cap changes nothing
    assert rows[256]["running"] < 205 and rows[1]["running"] == 1
    small = room(50, 8000)["blocks"]                                                                      # a smaller pool: a higher cap is worse
    low, high = (serve(sc.backlog(), cap, small, 8000) for cap in (64, 128))
    assert high["tokens_per_second"] < low["tokens_per_second"] and high["paused"] > 0 and low["paused"] == 0


def test_every_number_in_the_lesson_and_figures_comes_from_the_code():
    """The prose and the four figures quote numbers by hand; a change to the constants or the scheduler must show up here, not in a reader's check."""
    prose = (LESSON / "lesson.md").read_text(encoding="utf-8")
    figures = {path.name: path.read_text(encoding="utf-8") for path in sorted((LESSON / "figs").glob("*.svg"))}
    assert set(figures) == set(re.findall(r"\]\(figs/([\w-]+\.svg)\)", prose)) and len(figures) == 4
    gb = lambda value: f"{value / 1e9:.1f}"
    r, rows = room(90, 8000), {row["cap"]: row for row in swept()}
    tps = lambda row: f"{row['tokens_per_second']:,.0f}"
    ms = lambda row: f"{row['seconds_per_word'] * 1000:.1f}"
    gains = [rows[2 * c]["tokens_per_second"] / rows[c]["tokens_per_second"] for c in sc.CAPS[:6]]
    grow = rows[128]["tokens_per_second"] / rows[64]["tokens_per_second"]
    wait = rows[128]["seconds_per_word"] / rows[64]["seconds_per_word"]
    held = {longest: ran(90, longest, 512, True) for longest in (8000, 32000)}
    grown = ran(90, 8000, 512)
    prompts, replies = sc.CHATS["prompts"], sc.CHATS["replies"]
    seconds = sum(prompts) * 2 * sc.L24.PARAMS / sc.L24.A100_PEAK_FLOPS
    gib = sc.memory_parts(90, sc.GPU_BYTES_GIB)                                                  # the other reading of the card, quoted at the end of the pool sum
    wanted = {"lesson.md": [
        gb(r["pool"]), gb(r["unused"]), f"{r['weights'] / 1e9:.2f} GB", gb(RESERVE_BYTES), f"{r['pool_tokens']:,}", f"{r['blocks']:,}", f"{r['typical']} chats",
        f"{room(90, 8000, True)['typical']}", f"{room(90, 32000, True)['typical']}", f"{room(90, 128000, True)['typical']}", f"{room(90, 4000, True)['typical']}",
        f"{room(50, 8000)['pool_tokens']:,}", gb(room(50, 8000)["pool"]), f"{room(98, 8000)['typical']}", gb(room(98, 8000)["pool"]), gb(room(30, 8000)["pool"]),
        f"{room(30, 8000)['typical']}", f"{room(50, 8000)['typical']}", f"{room(35, 8000)['pool_tokens']:,}",
        f"{room(40, 128000)['blocks']:,} blocks", gb(room(40, 128000)["pool"]), f"{room(44, 128000)['blocks']:,} blocks", f"{blocks_for(128000):,}",
        *(f"{next(s for s in range(1, 100) if room(s, longest)['starts'])}%" for longest in (8000, 16000, 32000, 64000, 128000)),
        *(f"| {c} | {tps(rows[c])} | {ms(rows[c])} ms |" for c in (1, 8, 32, 64, 128, 256)),
        f"{rows[256]['paused']} pauses", f"{rows[256]['rebuilt_tokens']:,}", f"{rows[256]['running']:.1f}", f"{(min(gains) - 1) * 100:.1f}%", f"{(grow - 1) * 100:.1f}%", f"{(grow - 1) * 100:.0f}%",
        f"{wait:.2f} times", f"{wait:.1f} times", f"cap of {sc.stops_paying(list(rows.values()))}",
        f"{sum(prompts):,}", f"{seconds:.1f} s", f"{sum(replies) / seconds:,.0f}", f"{sum(replies):,}", f"{min(prompts)} to {max(prompts):,}", f"mean {statistics.fmean(prompts):,.0f}",
        f"{min(replies)} to {max(replies)}", f"mean {statistics.fmean(replies):.0f}", f"{statistics.fmean(prompts) * 2 * sc.L24.PARAMS / sc.L24.A100_PEAK_FLOPS * 1000:.1f} ms",
        f"writes {tps(grown)} tokens", f"{grown['running']:.1f} chats", f"{tps(held[8000])} at 8,000", f"{held[8000]['running']:.1f} running", f"{held[32000]['running']:.1f} running",
        f"{held[32000]['tokens_per_second']:,.0f} at 32,000", f"{grown['seconds_per_word'] * 1000:.1f} ms", f"{held[8000]['seconds_per_word'] * 1000:.1f} ms",
        f"{held[32000]['seconds_per_word'] * 1000:.1f} ms",
        f"{tps(ran(50, 8000, 128))} with {ran(50, 8000, 128)['paused']}", f"{tps(ran(50, 8000, 64))} tokens a second with no pauses",
        f"{tps(ran(35, 8000, 32))} with {ran(35, 8000, 32)['paused']} pauses", f"{tps(ran(35, 8000, 64))} with {ran(35, 8000, 64)['paused']}",
        f"{sc.BLOCK_BYTES:,}", f"{sc.WEIGHT_BYTES:,}", f"{sc.KV_BYTES:,}", f"{room(90, 8000)['per_typical']} blocks", f"{sc.OPENING['share']}%",
        f"{gb(gib['pool'])} GB, or {gib['blocks']:,} blocks, with room for {gib['blocks'] // r['per_typical']} chats of 2,000 tokens, against {gb(r['pool'])} GB, {r['blocks']:,} and {r['typical']}",
        f"The card is {(sc.GPU_BYTES_GIB / GPU_BYTES - 1) * 100:.1f}% larger and the pool {(gib['pool'] / r['pool'] - 1) * 100:.1f}%",
    ], "memory-parts.svg": [gb(r["pool"]), gb(r["unused"]), f"{r['weights'] / 1e9:.2f}", gb(RESERVE_BYTES), f"{r['pool_tokens']:,}", "72 GB", "90%", "80 GB"],
        "blocks-vs-reserved.svg": [f"{room(90, 8000)['typical']} chats", f"{room(90, 8000, True)['typical']} chats", f"{room(90, 32000, True)['typical']} chats", gb(r["pool"]), f"{r['blocks']:,}"],
        "cap-sweep.svg": [tps(rows[1]), tps(rows[64]), tps(rows[128]), ms(rows[1]), ms(rows[64]), ms(rows[128]), ms(rows[256]), tps(rows[256])],
        "pool-moves-the-knee.svg": [f"{room(share, 8000)['pool_tokens']:,}" for share in (35, 50, 90)] + [tps(ran(share, 8000, cap)) for share in (35, 50, 90) for cap in (32, 64, 128)]}
    missing = [f"{name}: {fragment!r}" for name, fragments in wanted.items() for fragment in fragments if fragment not in (prose if name == "lesson.md" else figures[name])]
    assert not missing, "not found: " + "; ".join(missing)
    for name, text in figures.items():                                                         # the palette is fixed and colour-blind safe: none of the blue or vermillion of the other charts
        assert not re.search(r"0072B2|56B4E9|D55E00", text, re.I) and "<title>" in text and "<desc>" in text and 'fill="#ffffff"' in text, name