"""Guards for gpu_split.py: memory per GPU, read time, the all-reduce, the best GPU count and the numbers the lesson quotes."""
import ast
import re
from pathlib import Path

import pytest

import gpu_split as g
from show_numbers import gib, half_up, halved_link_times, ms

LESSON = Path(__file__).resolve().parent.parent
PUBLISHED = {"8b": 8_030_261_248, "70b": 70_553_706_496, "405b": 405_853_388_800}   # same totals lesson 18 pins


def test_weight_counts_match_the_published_totals():
    """Guards a wrong shape file or counting formula: every size on the page starts from these totals."""
    assert {k: g.param_count(g.load_config(k)) for k in g.MODELS} == PUBLISHED


def test_memory_per_gpu_by_hand():
    """Guards the memory sum: weights shared, cache shared by kv head, working memory not shared (70B on 4 GPUs, added up by hand)."""
    m = g.memory_per_gpu("70b", 4)
    assert m["weights"] == 2 * PUBLISHED["70b"] // 4 == 35_276_853_248
    assert m["cache"] == 2 * 80 * 8 * 128 * 2 * 8192 // 4 == 671_088_640
    assert m["working"] == 2048 * (3 * 28672 + 4 * 8192) * 2 + g.GIB == 486_539_264 + g.GIB
    assert m["total"] == m["weights"] + m["cache"] + m["working"]
    # the cache is shared by key/value head, so 16 GPUs hold as much of it as 8 do
    assert g.memory_per_gpu("70b", 16)["cache"] == g.memory_per_gpu("70b", 8)["cache"]


def test_which_settings_fit():
    """Guards a wrong 'fits' verdict, the first thing the page shows: 70B needs 2 GPUs, 405B needs 11 and the page offers 16."""
    assert [g.fits("8b", n) for n in g.COUNTS] == [True] * 5
    assert [g.fits("70b", n) for n in g.COUNTS] == [False, True, True, True, True]
    assert [g.fits("405b", n) for n in g.COUNTS] == [False, False, False, False, True]
    assert {k: g.least_n(k) for k in g.MODELS} == {"8b": 1, "70b": 2, "405b": 11}


def test_read_time_falls_in_step_with_the_gpu_count():
    """Guards a split that does not shorten the read: each GPU reads 1/n of the weights, so n x the read time stays one GPU's read time."""
    assert g.read_time("8b", 1) == 2 * PUBLISHED["8b"] / g.GPU_BANDWIDTH
    for key in g.MODELS:
        for n in g.COUNTS:
            assert g.read_time(key, n) * n == pytest.approx(g.read_time(key, 1), rel=1e-6)
    for key, n, link, split, words in g.settings():   # the page calls the busy part 'reading': it must be the longer of the two
        assert g.compute_time(key, n, words) < g.read_time(key, n)


def ring_all_reduce(vectors: list[list[int]]):
    """Sum the vectors the way a ring does: n-1 steps of reduce-scatter, then n-1 steps of all-gather. Returns (result per GPU, steps, items sent per GPU)."""
    n, size = len(vectors), len(vectors[0])
    chunk = size // n
    data = [list(v) for v in vectors]
    sent = [0] * n
    steps = 0
    for phase in ("reduce", "gather"):
        for step in range(n - 1):
            messages = []
            for gpu in range(n):
                c = (gpu - step) % n if phase == "reduce" else (gpu + 1 - step) % n
                messages.append(((gpu + 1) % n, c, data[gpu][c * chunk:(c + 1) * chunk]))
                sent[gpu] += chunk
            for dest, c, part in messages:
                for i, x in enumerate(part):
                    data[dest][c * chunk + i] = data[dest][c * chunk + i] + x if phase == "reduce" else x
            steps += 1
    return data, steps, sent


@pytest.mark.parametrize("n", [2, 4, 8, 16])
def test_all_reduce_sums_and_costs_two_n_minus_one_steps(n):
    """Guards the all-reduce story: every GPU ends with the sum of all the partial results, in 2(n-1) steps and 2(n-1)/n x size sent."""
    size = n * 3
    vectors = [[gpu * 100 + i for i in range(size)] for gpu in range(n)]
    result, steps, sent = ring_all_reduce(vectors)
    assert all(row == [sum(v[i] for v in vectors) for i in range(size)] for row in result)
    assert steps == 2 * (n - 1)
    assert set(sent) == {2 * (n - 1) * size // n}
    # the time formula charges exactly those steps and those bytes
    link = g.LINKS["fast"]
    expected = (g.BASE_US + steps * link[2]) * 1e-6 + sent[0] / link[1]
    assert g.allreduce_time(size, n, "fast") == pytest.approx(expected, rel=1e-12)


def test_one_gpu_waits_on_no_one():
    assert g.allreduce_time(8192, 1, "slow") == 0 and g.word_time("8b", 1, "slow")["transfer"] == 0


def test_a_word_makes_two_swaps_per_layer():
    """Guards the swap count the page states (160 for 70B) and that the wait is that many all-reduces of one hidden vector."""
    cfg = g.load_config("70b")
    t = g.word_time("70b", 4, "fast")
    assert t["transfer"] == pytest.approx(2 * 80 * g.allreduce_time(8192 * 2, 4, "fast"), rel=1e-12)
    assert {k: 2 * g.load_config(k)["num_hidden_layers"] for k in g.MODELS} == {"8b": 64, "70b": 160, "405b": 252}
    assert g.message_bytes(cfg) == 16_384
    assert t["total"] == t["read"] + t["transfer"] and t["share"] == pytest.approx(t["transfer"] / t["total"])


def test_waiting_grows_with_gpus_and_shrinks_with_a_faster_link():
    for key in g.MODELS:
        shares = [g.word_time(key, n, "fast")["share"] for n in g.COUNTS]
        assert shares == sorted(shares)
        for n in g.COUNTS[1:]:
            assert g.word_time(key, n, "slow")["transfer"] > g.word_time(key, n, "fast")["transfer"]


def test_more_gpus_stop_helping():
    """Guards the lesson's turning point: the small model is faster on 8 GPUs than on 16, on both links, and each doubling gains less."""
    for link in g.LINKS:
        times = [g.word_time("8b", n, link)["total"] for n in g.COUNTS]
        gains = [a / b for a, b in zip(times, times[1:])]
        assert gains == sorted(gains, reverse=True) and gains[-1] < 1 < gains[-2]


def test_best_gpu_count_depends_on_the_link():
    """Guards the claim 'a faster link lets more GPUs help': best count per model on each link, any whole number that fits, up to 64."""
    best = {k: {link: g.best_n(k, link) for link in g.LINKS} for k in g.MODELS}
    assert best == {"8b": {"fast": 10, "slow": 8}, "70b": {"fast": 19, "slow": 15}, "405b": {"fast": 36, "slow": 28}}
    assert g.best_n("70b", "slow", words=64) == 14 < g.best_n("70b", "slow")
    # among the counts the page offers, the best is the same on both links, so the lesson does not claim otherwise
    offered = {k: {link: min(g.COUNTS, key=lambda n: g.word_time(k, n, link)["total"] if g.fits(k, n) else float("inf")) for link in g.LINKS} for k in g.MODELS}
    assert offered == {"8b": {"fast": 8, "slow": 8}, "70b": {"fast": 16, "slow": 16}, "405b": {"fast": 16, "slow": 16}}
    # slices are cut by attention head, so a usable count divides the heads: 19 for the 70B model (64 heads) is not one
    usable = {k: {link: g.best_n(k, link, divide_heads=True) for link in g.LINKS} for k in g.MODELS}
    assert usable == {"8b": {"fast": 8, "slow": 8}, "70b": {"fast": 16, "slow": 16}, "405b": {"fast": 32, "slow": 32}}
    assert all(g.load_config(k)["num_attention_heads"] % n == 0 for k in g.MODELS for n in usable[k].values())
    assert g.load_config("70b")["num_attention_heads"] % best["70b"]["fast"] != 0


def test_layer_groups_fit_the_model_but_do_not_speed_up_one_user():
    """Guards the pipeline claim: for one user a word still passes through every weight once, so the time stays near one GPU's full read."""
    for key in g.MODELS:
        whole = g.read_time(key, 1)
        for n in g.COUNTS:
            for link in g.LINKS:
                total = g.word_time(key, n, link, "groups")["total"]
                assert whole <= total <= 1.02 * whole
                assert g.word_time(key, n, link, "groups")["share"] < 0.02
    with pytest.raises(ValueError):   # a split the toy does not know must fail, not be priced as free
        g.word_time("8b", 2, "fast", "diagonal")


def test_many_words_at_once_make_the_link_matter():
    """Guards the 'words at once' setting: the slow link costs 5x more extra time at 64 words than at 1 (70B on 8 GPUs)."""
    gap = {w: g.word_time("70b", 8, "slow", words=w)["total"] - g.word_time("70b", 8, "fast", words=w)["total"] for w in g.WORDS}
    assert gap[64] > 5 * gap[1] and list(gap.values()) == sorted(gap.values())


def test_no_colour_outside_the_site_palette_in_the_lesson_folder():
    """Guards a stray blue or vermillion from the other palette; this file is the only place the codes are written."""
    banned = [bytes(c, "ascii").lower() for c in ("0072B2", "56B4E9", "D55E00")]
    for path in LESSON.rglob("*"):
        if path.is_file() and path.suffix in {".svg", ".html", ".md", ".py", ".json"} and path.name != Path(__file__).name:
            text = path.read_bytes().lower()
            assert not any(b in text for b in banned), path.name


def _pct(x: float, places: str = "0.1") -> str:
    return half_up(x * 100, places) + "%"


def test_every_number_the_page_and_figures_quote_comes_from_the_code():
    """Guards a lesson that drifts from its code: each figure of the prose, the tables and the SVGs is rebuilt here from
    gpu_split and must appear in the file that quotes it. One assert names every fragment that is missing."""
    lesson = (LESSON / "lesson.md").read_text()
    figs = {p.name: re.sub(r"<[^>]+>", " ", p.read_text()) for p in (LESSON / "figs").glob("*.svg")}
    wt = lambda key, n, link="fast", **kw: g.word_time(key, n, link, **kw)
    mem4 = g.memory_per_gpu("70b", 4)
    whole70 = g.weight_bytes(g.load_config("70b"))
    gains = [wt("8b", a)["total"] / wt("8b", b)["total"] for a, b in zip(g.COUNTS, g.COUNTS[1:])]
    wps = lambda n, link: str(round(wt("8b", n, link)["words_per_s"]))
    in_lesson = [
        f"{gib(whole70)} GiB", f"{g.USABLE_BYTES / g.GIB:.1f} GiB", f"{gib(mem4['weights'])} GiB", f"{gib(mem4['total'])} GiB",
        f"{gib(g.weight_bytes(g.load_config('405b')))} GiB", f"{g.least_n('405b')} GPUs", f"{gib(g.memory_per_gpu('405b', 16)['total'])} GiB",
        f"{ms(g.read_time('70b', 1))} ms", f"{ms(g.read_time('70b', 4))} ms", f"{ms(wt('70b', 4)['transfer'])} ms",
        f"{ms(wt('70b', 4)['total'])} ms", f"{g.message_bytes(g.load_config('70b')):,} bytes", f"{g.message_bytes(g.load_config('405b')):,} bytes",
        f"{g.ALL_REDUCES_PER_LAYER * g.load_config('70b')['num_hidden_layers']} swaps", f"{g.ALL_REDUCES_PER_LAYER * g.load_config('405b')['num_hidden_layers']} for each word",
        f"{mem4['total']:,} bytes", f"{mem4['weights']:,}", f"{mem4['cache']:,}", f"{mem4['working']:,}",
        ", ".join(_pct(wt("70b", n)["share"]) for n in (2, 4, 8)) + " and " + _pct(wt("70b", 16)["share"]),
        f"{gains[0]:.2f}, {gains[1]:.2f} and {gains[2]:.2f} times", f"{gains[3]:.2f} times",
        f"{ms(wt('70b', 8)['total'])} ms a word on the fast link and {ms(wt('70b', 8, 'slow')['total'])} ms on the slow one",
        f"{ms(wt('70b', 8, words=64)['total'])} ms on the fast link and {ms(wt('70b', 8, 'slow', words=64)['total'])} ms on the slow one",
        f"{ms(wt('70b', 4, split='groups')['total'])} ms a word on four GPUs", f"{_pct(wt('70b', 4, split='groups')['share'], '0.01')} of it waiting",
        f"{ms(wt('8b', 2)['read'])} ms", f"{ms(wt('8b', 2)['transfer'])} ms", f"{ms(wt('8b', 2)['total'])} ms",
        f"{half_up(g.allreduce_time(g.message_bytes(g.load_config('8b')), 2, 'fast') * 1e6)} in all",
        f"so a swap takes {half_up(g.allreduce_time(g.message_bytes(g.load_config('70b')), 4, 'fast') * 1e6)}.",
        f"{half_up(g.allreduce_time(g.message_bytes(g.load_config('70b')) * 64, 8, 'fast') * 1e6)} microseconds on the fast link and "
        f"{half_up(g.allreduce_time(g.message_bytes(g.load_config('70b')) * 64, 8, 'slow') * 1e6)} on the slow one",
        _pct((g.BASE_US + 6 * 0.6) / (g.allreduce_time(g.message_bytes(g.load_config("70b")), 4, "fast") * 1e6)),   # start-up share of one swap
        f"{ms(g.compute_time('70b', 8, 64))} ms", f"{ms(g.read_time('70b', 8))} ms", f"{ms(g.compute_time('70b', 1))} ms",
        f"{wt('70b', 8)['total'] / wt('70b', 16)['total']:.2f} times faster, from {ms(wt('70b', 8)['total'])} to {ms(wt('70b', 16)['total'])} ms",
        f"from {wps(8, 'fast')} to {wps(16, 'fast')} words a second",
        f"{ms(wt('70b', 8, 'slow', words=64)['total'] - wt('70b', 8, words=64)['total'])} ms more at 64",
        f"{ms(wt('70b', 8, 'slow')['total'] - wt('70b', 8)['total'])} ms more at 1 word",
        "10 on the fast link and 8 on the slow one", "19 and 15", "36 and 28",
        "Each model has {0} key/value heads, so past {0} GPUs".format(*{g.load_config(k)["num_key_value_heads"] for k in g.MODELS}),
    ]
    shares = lambda rows: ", ".join(_pct(t["share"]) for t in rows[:-1]) + " and " + _pct(rows[-1]["share"])
    for n in g.COUNTS:       # the 8B table: one row a count, milliseconds on each link
        in_lesson.append(f"| {n} | {ms(wt('8b', n)['total'])} | {ms(wt('8b', n, 'slow')['total'])} |")
    for words in g.WORDS:    # the words-at-once table
        in_lesson.append(f"| {words} | {ms(wt('70b', 8, words=words)['total'])} | {ms(wt('70b', 8, 'slow', words=words)['total'])} |")
    for link in g.LINKS:     # the waiting shares the tables leave to the sentence below each
        in_lesson.append(shares([wt("8b", n, link) for n in g.COUNTS[1:]]))
        in_lesson.append(shares([wt("70b", 8, link, words=w) for w in g.WORDS]))
    for words in (1, 64):   # the link-direction sentence: as written, then with both link rates halved
        halved = halved_link_times("70b", 8, words)
        slow, slow_half = wt("70b", 8, "slow", words=words)["total"], halved["slow"]
        in_lesson.append(f"{ms(wt('70b', 8, words=words)['total'])} to {ms(halved['fast'])} ms on the fast link and from {ms(slow)} to {ms(slow_half)} ms on the slow one"
                         if words == 1 else f"the slow link goes from {ms(slow)} to {ms(slow_half)} ms, {_pct(slow_half / slow - 1, '1')} more")
        if words == 64:
            assert 0.25 < slow_half / slow - 1 < 0.30
    heads = [g.load_config(k)["num_attention_heads"] for k in g.MODELS]
    usable = [g.best_n(k, "fast", divide_heads=True) for k in g.MODELS]   # the same on both links, as test_best_gpu_count_depends_on_the_link asserts
    in_lesson += [f"the {heads[0]}, {heads[1]} and {heads[2]} the three models have", f"so {g.best_n('70b', 'fast')} GPUs is not a usable count",
                  f"the fastest are {usable[0]} for the 8B model, {usable[1]} for the 70B and {usable[2]} for the 405B, on both links",
                  f"they are {usable[0]}, {usable[1]} and {usable[2]}.", f"({g.USABLE * 100:.0f}% of its {g.GPU_GIB} GiB)"]
    rows70 = [wt("70b", n) for n in (2, 4, 8, 16)]
    in_figs = {
        "weights-split.svg": [gib(whole70), f"{g.USABLE_BYTES / g.GIB:.1f}", gib(mem4["weights"])],
        "time-per-word.svg": [ms(t["total"]) + " ms" for t in rows70] + [", ".join(_pct(t["share"]) for t in rows70)],
        "stops-helping.svg": [f"{wps(8, 'fast')} to {wps(16, 'fast')} and the slow link {wps(8, 'slow')} to {wps(16, 'slow')}"],
        "slices-vs-groups.svg": [ms(wt("70b", 4)["total"]), ms(wt("70b", 4, split="groups")["total"]), ms(g.read_time("70b", 4))],
    }
    missing = [f"lesson.md: {x}" for x in in_lesson if x not in lesson]
    missing += [f"{name}: {x}" for name, parts in in_figs.items() for x in parts if x not in figs[name]]
    assert not missing, "quoted but not found:\n" + "\n".join(missing)


def _without_docstring(source, node):
    """The lines of one function, minus its docstring."""
    lines = source.splitlines()[node.lineno - 1:node.end_lineno]
    first = node.body[0]
    if isinstance(first, ast.Expr) and isinstance(first.value, ast.Constant) and isinstance(first.value.value, str):
        del lines[first.lineno - node.lineno:first.end_lineno - node.lineno + 1]
    return "\n".join(lines)


def test_every_python_block_in_the_lesson_is_a_verbatim_excerpt_of_the_source():
    """The lesson shows `memory_per_gpu` with its docstring left out. If the code changes and the page does not,
    the reader copies a function that was never tested."""
    lesson = (LESSON / "lesson.md").read_text()
    source = Path(g.__file__).read_text()
    in_source = {n.name: _without_docstring(source, n) for n in ast.parse(source).body if isinstance(n, ast.FunctionDef)}
    blocks = re.findall(r"^```python\n(.*?)^```", lesson, flags=re.S | re.M)
    shown = [node for block in blocks for node in ast.parse(block).body]
    assert {node.name for node in shown} == {"memory_per_gpu"}
    for block in blocks:
        for node in ast.parse(block).body:
            assert _without_docstring(block, node) == in_source[node.name], f"{node.name} differs from gpu_split.py"