"""Guards for attention.py: the toy's weights, the mask, weights that add up to 1, an independent NumPy check, the scaling numbers."""
import math
import random
import re
from pathlib import Path

import numpy as np

import attention as att

TOL = 1e-12                                        # agreement with NumPy, in absolute weight or output units
IT = att.TOKENS.index("it")


def test_the_toys_weights_for_it_match_the_hand_worked_numbers():
    """Guards a wrong weight: head 1 gives "it" a score of 3 / sqrt(2) on "cat" and 0 on the others, so cat gets e^2.121 / (e^2.121 + 4)."""
    refers = att.attend(att.VECTORS, att.HEADS["refers"])
    scores, weights = refers["scores"][IT], refers["weights"][IT]
    cat = att.TOKENS.index("cat")
    assert scores[:IT + 1] == [0.0, 3 / math.sqrt(2), 0.0, 0.0, 0.0] and scores[IT + 1:] == [-math.inf] * 2
    assert weights[cat] == math.exp(3 / math.sqrt(2)) / (math.exp(3 / math.sqrt(2)) + 4) and round(weights[cat], 3) == 0.676
    assert [round(w, 3) for w in weights] == [0.081, 0.676, 0.081, 0.081, 0.081, 0.0, 0.0]
    for word, size, weight in (("sat", 3.57, 0.641), ("tired", 5.46, 0.476)):                     # the other words the lesson works by hand
        row = att.TOKENS.index(word)
        assert round(math.exp(refers["scores"][row][cat]), 2) == size and round(refers["weights"][row][cat], 3) == weight
    assert refers["weights"][att.TOKENS.index("because")][:4] == [0.25] * 4                      # a query of all zeros: every score is 0
    before = att.attend(att.VECTORS, att.HEADS["before"])["weights"][IT]
    assert max(range(IT + 1), key=lambda j: before[j]) == att.TOKENS.index("because") and round(max(before), 2) == 0.71
    even = att.attend(att.VECTORS, att.HEADS["even"])["weights"][IT]
    assert even[:IT + 1] == [0.2] * 5
    assert att.view("refers", IT)["sentence"] == "“it” puts 68% of its weight on “cat”."
    assert att.view("refers", 0)["sentence"] == "“The” puts 100% of its weight on “The”."
    assert att.view("even", IT)["sentence"] == "“it” spreads its weight evenly: 20% on each of 5 words."
    assert att.view("refers", IT)["blend"] == "Blended note, “an animal”: 0.68. “it” alone: 0.00."
    assert len(att.VECTORS) == len(att.TOKENS) == 7 and {len(v) for v in att.VECTORS} == {att.D_MODEL}


def test_weights_are_positive_add_to_one_and_a_word_never_looks_ahead():
    """Guards the two failures a reader would see: a weight that is negative or the row not adding to 1, and a word whose answer changes
    when a LATER word changes. Runs every head with the division on and off, then rewrites every word after word i and checks rows 0..i."""
    n = len(att.TOKENS)
    for name, head in att.HEADS.items():
        for scale in (True, False):
            run = att.attend(att.VECTORS, head, scale)
            for i in range(n):
                row = run["weights"][i]
                assert abs(sum(row) - 1) < TOL and all(w > 0 for w in row[:i + 1]) and row[i + 1:] == [0.0] * (n - 1 - i), (name, scale, i)
                changed = [v if j <= i else [x + 1.7 for x in v] for j, v in enumerate(att.VECTORS)]
                other = att.attend(changed, head, scale)
                assert other["outputs"][i] == run["outputs"][i] and other["weights"][i] == row, (name, scale, i)


def reference(x, wq, wk, wv, scale=True):
    """Attention as NumPy matrix operations: softmax(Q K^T / sqrt(d) + mask) V, with the mask as -inf above the diagonal."""
    q, k, v = x @ wq, x @ wk, x @ wv
    s = q @ k.T / (math.sqrt(k.shape[1]) if scale else 1.0)
    s = s + np.triu(np.full(s.shape, -np.inf), k=1)
    e = np.exp(s - s.max(axis=1, keepdims=True))
    w = e / e.sum(axis=1, keepdims=True)
    return w, w @ v


def test_matches_numpy_on_the_toy_and_on_random_inputs():
    """Guards a slip in any one step (the division, the mask, the softmax, the average): both routes must agree within 1e-12."""
    rng = random.Random(0)
    cases = [(att.VECTORS, h) for h in att.HEADS.values()]
    for n, d_model, d_k, d_v in [(1, 3, 2, 2), (5, 4, 3, 4), (9, 8, 8, 5), (12, 16, 4, 6)]:
        def table(rows, cols):
            return [[rng.gauss(0, 1) for _ in range(cols)] for _ in range(rows)]
        cases.append((table(n, d_model), {"wq": table(d_model, d_k), "wk": table(d_model, d_k), "wv": table(d_model, d_v)}))
    for vectors, head in cases:
        for scale in (True, False):
            run = att.attend(vectors, head, scale)
            w, out = reference(np.array(vectors), np.array(head["wq"]), np.array(head["wk"]), np.array(head["wv"]), scale)
            assert np.allclose(run["weights"], w, rtol=0, atol=TOL) and np.allclose(run["outputs"], out, rtol=0, atol=TOL)


def test_two_heads_are_joined_side_by_side_without_mixing():
    """Guards heads sharing numbers: each word's joined output is head 1's output followed by head 2's, each as when run alone."""
    both = att.multi_head(att.VECTORS, [att.HEADS["refers"], att.HEADS["before"]])
    one, two = att.attend(att.VECTORS, att.HEADS["refers"])["outputs"], att.attend(att.VECTORS, att.HEADS["before"])["outputs"]
    assert both == [a + b for a, b in zip(one, two)] and len(both[IT]) == 2 * att.D_V


def test_scaling_experiment_gives_the_numbers_the_lesson_prints():
    """Guards the lesson's claim: without the division the scores spread out in step with sqrt(d) and the biggest word takes almost all
    the weight; with it the spread stays near 1 at every length."""
    rows = {r["d"]: r for r in att.scaling_experiment()}
    assert {d: round(r["spread_raw"], 1) for d, r in rows.items()} == {4: 2.2, 64: 8.1, 1024: 31.6}
    assert {d: round(r["spread_scaled"], 2) for d, r in rows.items()} == {4: 1.08, 64: 1.01, 1024: 0.99}
    assert {d: round(r["top_raw"], 2) for d, r in rows.items()} == {4: 0.53, 64: 0.90, 1024: 0.96}
    assert {d: round(r["top_scaled"], 2) for d, r in rows.items()} == {4: 0.34, 64: 0.36, 1024: 0.36}
    assert all(abs(r["spread_raw"] / math.sqrt(d) - 1) < 0.1 for d, r in rows.items())


def test_cost_counts_match_a_brute_force_count_and_lesson_23():
    """Guards the n-squared-over-2 claim and the copied model shape: the formula equals the number of unmasked cells, the three printed
    figures are right, and the layer and head counts are the ones in lesson 23's prefill_sim.py."""
    for n in range(1, 25):
        assert att.score_count(n) == sum(1 for i in range(n) for j in range(n) if j <= i)
    assert [att.score_count(n) for n in (1000, 8000)] == [500_500, 32_004_000]
    assert att.score_count(8000) * att.LAYERS * att.HEADS_PER_LAYER == 32_772_096_000
    other = (Path(__file__).resolve().parents[2] / "why-ai-waits-before-the-first-word" / "code" / "prefill_sim.py").read_text(encoding="utf-8")
    assert int(re.search(r"^LAYERS = (\d+)", other, re.M).group(1)) == att.LAYERS
    assert int(re.search(r"^HEADS = (\d+)", other, re.M).group(1)) == att.HEADS_PER_LAYER