Hands-on lab — Day 101: Matrix Multiplication

"""Three from-scratch implementations and NumPy's @, asserted equal.

Run from inside this directory:

    ../.venv/bin/python3 01_matmul_from_scratch.py

Every claim printed below is also asserted. If an assertion fails the script
stops with a traceback instead of printing a reassuring final line.
"""

import numpy as np

from dataset import HIGHLIGHT_COLUMN, HIGHLIGHT_ROW, HIGHLIGHT_TERMS, HIGHLIGHT_VALUE, W, X, XW
from matmul import (
    ShapeMismatch,
    dot,
    identity,
    matmul_columns,
    matmul_dots,
    matmul_loops,
    matvec,
    multiplication_count,
    shape,
    transpose,
)


def show(M):
    return "[" + ", ".join("[" + ", ".join(f"{v:g}" for v in row) + "]" for row in M) + "]"


print("=" * 74)
print("1. The dot product: multiply pairwise, then add")
print("=" * 74)

u = [3, 4]
v = [4, 3]
w = [-4, 3]

print(f"  u = {u}, v = {v}, w = {w}")
print(f"  u . v = 3*4 + 4*3 = 12 + 12 = {dot(u, v)}")
print(f"  u . u = 3*3 + 4*4 =  9 + 16 = {dot(u, u)}   <- the length of u, squared")
print(f"  u . w = 3*-4 + 4*3 = -12 + 12 = {dot(u, w)}   <- zero means perpendicular")
assert dot(u, v) == 24
assert dot(u, u) == 25
assert dot(u, w) == 0
assert dot(u, u) == round(float(np.linalg.norm(u)) ** 2)
print("  A vector dotted with itself is its squared length, every time.")
print(f"  NumPy agrees: np.dot(u, v) = {np.dot(u, v)}, u @ v = {np.array(u) @ np.array(v)}")

print()
print("=" * 74)
print("2. Matrix times vector, as a weighted sum of the matrix's COLUMNS")
print("=" * 74)

A = [[2, 0], [-1, 1], [0, 4]]  # shape (3, 2)
c = [3, 5]

print(f"  A = {show(A)}, shape {shape(A)}")
print(f"  c = {c}")
print("  A @ c takes 3 copies of A's first column plus 5 copies of its second:")
print("      3 * [2, -1, 0]  =  [6, -3, 0]")
print("      5 * [0,  1, 4]  =  [0,  5, 20]")
print("      sum             =  [6,  2, 20]")
result = matvec(A, c)
print(f"  matvec(A, c) = {result}")
assert result == [6, 2, 20]
assert result == (np.array(A) @ np.array(c)).tolist()
print("  NumPy returns the same three numbers.")
print("  Note what this guarantees: the answer is ALWAYS a combination of A's")
print("  columns, so it can never land anywhere those columns cannot reach.")

print()
print("=" * 74)
print("3. Matrix times matrix: the same thing, once per column")
print("=" * 74)

print(f"  X = {show(X)}   shape {shape(X)}")
print(f"  W = {show(W)}   shape {shape(W)}")
print()
print("  The highlighted cell, entry (1, 1) of X @ W, in full:")
print(f"      row 1 of X    = {HIGHLIGHT_ROW}")
print(f"      column 1 of W = {HIGHLIGHT_COLUMN}")
print(
    f"      0*0 + 1*1 + 3*4 = {HIGHLIGHT_TERMS[0]} + {HIGHLIGHT_TERMS[1]}"
    f" + {HIGHLIGHT_TERMS[2]} = {HIGHLIGHT_VALUE}"
)
assert sum(HIGHLIGHT_TERMS) == HIGHLIGHT_VALUE

by_loops = matmul_loops(X, W)
by_dots = matmul_dots(X, W)
by_columns = matmul_columns(X, W)
by_numpy = (np.array(X) @ np.array(W)).tolist()

print()
print(f"  three nested loops   {show(by_loops)}")
print(f"  list of dot products {show(by_dots)}")
print(f"  A applied per column {show(by_columns)}")
print(f"  NumPy's @            {show(by_numpy)}")
print(f"  worked out by hand   {show(XW)}")
assert by_loops == by_dots == by_columns == by_numpy == XW
assert by_loops[1][1] == HIGHLIGHT_VALUE
print("  All five agree, including the one a human derived with a pen.")

print()
print("  And the cost, counted rather than guessed:")
m, n = shape(X)
_, p = shape(W)
print(f"      ({m}, {n}) @ ({n}, {p}) costs m*n*p = {m}*{n}*{p} = {multiplication_count(m, n, p)}")
print("      multiplications, one per (row, column, inner step).")
assert multiplication_count(m, n, p) == 12

print()
print("=" * 74)
print("4. Several shapes, all four implementations, all equal")
print("=" * 74)

rng = np.random.default_rng(101)  # seeded, so this script is reproducible
for m, n, p in [(1, 1, 1), (2, 3, 2), (3, 2, 4), (4, 4, 4), (5, 1, 3), (1, 6, 2)]:
    L = rng.integers(-9, 10, size=(m, n)).tolist()
    R = rng.integers(-9, 10, size=(n, p)).tolist()
    expected = (np.array(L) @ np.array(R)).tolist()
    got_loops = matmul_loops(L, R)
    got_dots = matmul_dots(L, R)
    got_cols = matmul_columns(L, R)
    assert got_loops == expected, (m, n, p)
    assert got_dots == expected, (m, n, p)
    assert got_cols == expected, (m, n, p)
    count = multiplication_count(m, n, p)
    print(
        f"  ({m}, {n}) @ ({n}, {p}) -> ({m}, {p})   "
        f"loops = dots = columns = NumPy   "
        f"({count} multiplication{'' if count == 1 else 's'})"
    )

print()
print("=" * 74)
print("5. The identity matrix: the transformation that does nothing")
print("=" * 74)

I3 = identity(3)
I2 = identity(2)
print(f"  identity(3) = {show(I3)}")
print(f"  I2 @ X      = {show(matmul_loops(I2, X))}    (X unchanged)")
print(f"  X  @ I3     = {show(matmul_loops(X, I3))}    (X unchanged)")
assert matmul_loops(I2, X) == X
assert matmul_loops(X, I3) == X
assert I3 == np.eye(3, dtype=int).tolist()
print("  Note the two different sizes. X is (2, 3), so the identity that fits")
print("  on the left is 2 by 2 and the one that fits on the right is 3 by 3.")
print("  'The' identity matrix is really one per size.")

print()
print("=" * 74)
print("6. The shape rule, and what happens when it is broken")
print("=" * 74)

try:
    matmul_loops(X, X)
except ShapeMismatch as exc:
    print("  matmul_loops(X, X) raises ShapeMismatch:")
    print(f"    {exc}")
else:  # pragma: no cover - only reached if the guard is broken
    raise AssertionError("(2, 3) @ (2, 3) should not have been allowed")

try:
    np.array(X) @ np.array(X)
except ValueError as exc:
    print("  NumPy raises ValueError for the same thing:")
    print(f"    {type(exc).__name__}: {str(exc).splitlines()[0]}")
else:  # pragma: no cover
    raise AssertionError("NumPy should have refused (2, 3) @ (2, 3)")

fixed = matmul_loops(X, transpose(X))
print(f"  X @ X.T is (2, 3) @ (3, 2) -> {shape(fixed)}: {show(fixed)}")
assert shape(fixed) == (2, 2)
assert fixed == (np.array(X) @ np.array(X).T).tolist()

print()
print("01_matmul_from_scratch.py: every assertion held.")

"""Matrix multiplication IS composition — and that is why the rule is the rule.

Run from inside this directory:

    ../.venv/bin/python3 02_composition.py

Two transformations of the plane, applied in both orders, with real coordinates
at every step. Then the three algebraic facts that follow: multiplication is
not commutative, it is associative, and it distributes over addition.
"""

import numpy as np

from dataset import (
    FLIP_AFTER_ROT,
    FLIP_AFTER_ROT_V,
    FLIP_V,
    FLIP_X,
    P,
    P_AT_Q,
    Q,
    Q_AT_P,
    ROT90,
    ROT_AFTER_FLIP,
    ROT_AFTER_FLIP_V,
    ROT_V,
    V,
)
from matmul import chain_costs, matmul_loops, matvec, multiplication_count


def show(M):
    return "[" + ", ".join("[" + ", ".join(f"{v:g}" for v in row) + "]" for row in M) + "]"


A = ROT90  # a quarter turn anticlockwise
B = FLIP_X  # a reflection in the horizontal axis

print("=" * 74)
print("1. Two transformations, read off their columns")
print("=" * 74)
print(f"  A = {show(A)}   a quarter turn anticlockwise")
print(f"  B = {show(B)}   a reflection in the horizontal axis")
print()
print("  How to read a transformation matrix without doing any arithmetic:")
print("  its columns are where the basis vectors land.")
print(f"      A sends (1, 0) to {matvec(A, [1, 0])}  <- A's first column")
print(f"      A sends (0, 1) to {matvec(A, [0, 1])}  <- A's second column")
print(f"      B sends (1, 0) to {matvec(B, [1, 0])}  <- B's first column")
print(f"      B sends (0, 1) to {matvec(B, [0, 1])}  <- B's second column")
assert matvec(A, [1, 0]) == [0, 1]
assert matvec(A, [0, 1]) == [-1, 0]
assert matvec(B, [1, 0]) == [1, 0]
assert matvec(B, [0, 1]) == [0, -1]
print("  That is not a coincidence and it is not a mnemonic. A @ v is a weighted")
print("  sum of A's columns, so A @ (1, 0) takes one copy of column 0 and none")
print("  of column 1. The columns ARE the images of the basis vectors.")

print()
print("=" * 74)
print("2. Doing B and then A, one step at a time")
print("=" * 74)
print(f"  Start at v = {V}.")
step1 = matvec(B, V)
step2 = matvec(A, step1)
print(f"      B @ v          = {step1}      (reflected: y flipped sign)")
print(f"      A @ (B @ v)    = {step2}      (then turned a quarter anticlockwise)")
assert step1 == FLIP_V
assert step2 == ROT_AFTER_FLIP_V

print()
print("  Now the single matrix that does both in one step:")
AB = matmul_loops(A, B)
print(f"      A @ B          = {show(AB)}")
print(f"      (A @ B) @ v    = {matvec(AB, V)}")
assert AB == ROT_AFTER_FLIP
assert matvec(AB, V) == step2
print("  Same destination, one multiplication instead of two.")
print("  A @ B is the reflection in the line y = x: it swaps the coordinates,")
print(f"  and {V} becoming {matvec(AB, V)} is exactly that.")
print()
print("  Read the order carefully, because it is the thing people get wrong:")
print("  in A @ B, B runs FIRST. It is the one standing next to the vector in")
print("  A @ (B @ v). Matrices compose right to left, like nested function")
print("  calls: A(B(v)). English says 'A times B'; the arithmetic says 'B, then A'.")

print()
print("=" * 74)
print("3. The same two operations in the other order")
print("=" * 74)
other1 = matvec(A, V)
other2 = matvec(B, other1)
print(f"      A @ v          = {other1}     (turned first)")
print(f"      B @ (A @ v)    = {other2}    (then reflected)")
BA = matmul_loops(B, A)
print(f"      B @ A          = {show(BA)}")
print(f"      (B @ A) @ v    = {matvec(BA, V)}")
assert other1 == ROT_V
assert other2 == FLIP_AFTER_ROT_V
assert BA == FLIP_AFTER_ROT
assert matvec(BA, V) == other2

print()
print(f"  A @ B = {show(AB)}   reflection in the line y = x")
print(f"  B @ A = {show(BA)}   reflection in the line y = -x")
print(f"  The same starting point {V} ends at {matvec(AB, V)} one way and "
      f"{matvec(BA, V)} the other.")
assert AB != BA
assert matvec(AB, V) != matvec(BA, V)
print()
print("  Matrix multiplication is NOT commutative. That is not a defect and it")
print("  is not a subtlety: it is the honest consequence of what it means.")
print("  Putting your socks on and then your shoes is not the same as putting")
print("  your shoes on and then your socks, and no amount of algebra will make")
print("  it so.")
print()
print("  NumPy, on the same numbers, to be sure the from-scratch code is right:")
npA, npB = np.array(A), np.array(B)
print(f"      npA @ npB = {(npA @ npB).tolist()}")
print(f"      npB @ npA = {(npB @ npA).tolist()}")
assert (npA @ npB).tolist() == AB
assert (npB @ npA).tolist() == BA

print()
print("=" * 74)
print("4. A second pair, where neither answer is a tidy reflection")
print("=" * 74)
PQ = matmul_loops(P, Q)
QP = matmul_loops(Q, P)
elementwise = [[P[i][j] * Q[i][j] for j in range(2)] for i in range(2)]
print(f"  P = {show(P)}    Q = {show(Q)}")
print(f"      P @ Q = {show(PQ)}     1*5 + 2*7 = 19, and so on")
print(f"      Q @ P = {show(QP)}     5*1 + 6*3 = 23, and so on")
print(f"      P * Q = {show(elementwise)}      no summing anywhere; section 5 of")
print("                                       03_star_versus_at.py is about this")
assert PQ == P_AT_Q
assert QP == Q_AT_P
assert PQ != QP != elementwise
print("  Two matrices, three different answers, and only one of them is what")
print("  'multiply these matrices' means in linear algebra.")

print()
print("=" * 74)
print("5. What IS true: associativity")
print("=" * 74)
C = [[1, 1], [0, 2]]
left_first = matmul_loops(matmul_loops(A, B), C)
right_first = matmul_loops(A, matmul_loops(B, C))
print(f"  C = {show(C)}")
print(f"      (A @ B) @ C = {show(left_first)}")
print(f"      A @ (B @ C) = {show(right_first)}")
assert left_first == right_first
print("  Identical. The BRACKETS may move freely; the ORDER may not. Those are")
print("  two different statements and confusing them is the usual mistake.")
print()
print("  Associativity is not a curiosity. It is the only reason you are allowed")
print("  to choose the cheaper way to evaluate a chain, and the choice is worth")
print("  a great deal:")
for label, (m, n, p, q) in [
    ("(10, 100) @ (100, 5) @ (5, 50)", (10, 100, 5, 50)),
    ("(1024, 4096) @ (4096, 8) @ (8, 4096)", (1024, 4096, 8, 4096)),
]:
    left_first, right_first = chain_costs(m, n, p, q)
    ratio = right_first / left_first
    print(f"      {label}")
    print(f"        (AB)C : {left_first:>15,} multiplications")
    print(f"        A(BC) : {right_first:>15,} multiplications   ({ratio:.0f}x more)")
assert chain_costs(10, 100, 5, 50) == (7_500, 75_000)
assert chain_costs(1024, 4096, 8, 4096) == (67_108_864, 17_314_086_912)
assert 17_314_086_912 // 67_108_864 == 258
print("  Same answer to the last digit, 258 times the arithmetic. The second")
print("  set of shapes is a low-rank adapter on a 4096-wide layer, which is a")
print("  real thing people train, and this cost gap is most of why they can.")

print()
print("=" * 74)
print("6. What IS true: distributivity")
print("=" * 74)
D = [[2, 0], [1, 1]]
sum_then_mult = matmul_loops(A, [[B[i][j] + D[i][j] for j in range(2)] for i in range(2)])
mult_then_sum_parts = (matmul_loops(A, B), matmul_loops(A, D))
mult_then_sum = [
    [mult_then_sum_parts[0][i][j] + mult_then_sum_parts[1][i][j] for j in range(2)]
    for i in range(2)
]
print(f"  D = {show(D)}")
print(f"      A @ (B + D)     = {show(sum_then_mult)}")
print(f"      A @ B  +  A @ D = {show(mult_then_sum)}")
assert sum_then_mult == mult_then_sum
print("  Equal. This is what lets you split a layer's weights into a base part")
print("  and a small correction and add the two results — the other half of why")
print("  adapters work.")

print()
print("=" * 74)
print("7. The cost of the chain, counted rather than asserted")
print("=" * 74)
print("  (m, n) @ (n, p) runs its innermost line m*n*p times, once per")
print("  (row, column, inner step). Nothing about that count is an estimate:")
for m, n, p in [(2, 3, 2), (10, 100, 5), (200, 200, 200)]:
    label = f"({m}, {n}) @ ({n}, {p})"
    print(f"      {label:<26} {multiplication_count(m, n, p):>11,} multiplications")
assert multiplication_count(200, 200, 200) == 8_000_000
print("  Eight million for a pair of 200 by 200 matrices, which is small by any")
print("  modern standard. 05_cost_and_speed.py times that one for real.")

print()
print("02_composition.py: every assertion held.")

"""`*` and `@` are different operations, and NumPy will not warn you.

Run from inside this directory:

    ../.venv/bin/python3 03_star_versus_at.py

The trap is not that one of them is wrong. It is that on the right operands
BOTH are legal, both return an array, and only the shape gives you away — and
sometimes not even that. This script meets it deliberately, then covers the
three spellings of matrix multiplication and how to read a shape error.
"""

import numpy as np

from dataset import P, P_AT_Q, P_TIMES_Q, Q, U, W, X, X_AT_U, X_AT_XT, X_TIMES_U, XT_AT_X

print("=" * 74)
print("1. Same two operands, both operations legal, different answers")
print("=" * 74)

npP, npQ = np.array(P), np.array(Q)
print(f"  P = {npP.tolist()}   Q = {npQ.tolist()}   both shape {npP.shape}")
print()
print(f"  P * Q  -> shape {(npP * npQ).shape}  {(npP * npQ).tolist()}")
print("             entry by entry: 1*5=5, 2*6=12, 3*7=21, 4*8=32. Nothing summed.")
print(f"  P @ Q  -> shape {(npP @ npQ).shape}  {(npP @ npQ).tolist()}")
print("             row dotted with column: 1*5 + 2*7 = 19, and so on.")
assert (npP * npQ).tolist() == P_TIMES_Q
assert (npP @ npQ).tolist() == P_AT_Q
assert (npP * npQ).shape == (npP @ npQ).shape == (2, 2)
print()
print("  Note the shapes are IDENTICAL. Two square matrices of the same size")
print("  give a same-sized answer either way, so a shape check catches nothing")
print("  here. Only the numbers differ, and nothing will tell you which you got.")

print()
print("=" * 74)
print("2. The version where the shapes do give you away")
print("=" * 74)

npX, npU = np.array(X), np.array(U)
print(f"  X = {npX.tolist()}   shape {npX.shape}")
print(f"  u = {npU.tolist()}                 shape {npU.shape}")
print()
star = npX * npU
at = npX @ npU
print(f"  X * u  -> shape {star.shape}  {star.tolist()}")
print("             u is broadcast across both rows (Day 100), then multiplied")
print("             entry by entry. Three numbers per row go in, three come out.")
print(f"  X @ u  -> shape {at.shape}     {at.tolist()}")
print("             each row is multiplied by u AND THEN SUMMED:")
print("               row 0: 1*10 + 2*2 + 0*5 = 10 + 4 + 0  = 14")
print("               row 1: 0*10 + 1*2 + 3*5 =  0 + 2 + 15 = 17")
assert star.tolist() == X_TIMES_U
assert at.tolist() == X_AT_U
assert star.shape == (2, 3)
assert at.shape == (2,)
print()
print("  The summing is the whole difference. `*` keeps every product; `@` adds")
print("  them up and loses a dimension doing it. That collapse is what makes it")
print("  a transformation rather than a rescaling.")

print()
print("=" * 74)
print("3. Three spellings of the same operation")
print("=" * 74)

by_at = npX @ np.array(W)
by_matmul = np.matmul(npX, np.array(W))
by_dot = np.dot(npX, np.array(W))
print(f"  X @ W            = {by_at.tolist()}")
print(f"  np.matmul(X, W)  = {by_matmul.tolist()}")
print(f"  np.dot(X, W)     = {by_dot.tolist()}")
assert by_at.tolist() == by_matmul.tolist() == by_dot.tolist()
print()
print("  All three agree on two-dimensional arrays, and `@` is the one to use:")
print("  it says at a glance which operation you meant. The other two exist for")
print("  reasons that matter only when the arrays have more than two dimensions,")
print("  where matmul and dot genuinely differ — matmul treats leading axes as a")
print("  stack of matrices, dot does not. Here is that difference, on shapes")
print("  small enough to read:")

stack = np.arange(8).reshape(2, 2, 2)
plain = np.arange(4).reshape(2, 2)
other = np.arange(8).reshape(2, 2, 2)

print(f"      a stack of two 2x2 matrices, shape {stack.shape}, times one 2x2:")
print(f"        np.matmul -> {np.matmul(stack, plain).shape}")
print(f"        np.dot    -> {np.dot(stack, plain).shape}")
print("        Here they agree, in shape AND in every value. This case was")
print("        checked rather than assumed, and the assumption would have been")
print("        that they differ. They do not.")
assert np.matmul(stack, plain).shape == (2, 2, 2)
assert np.dot(stack, plain).shape == (2, 2, 2)
assert np.array_equal(np.matmul(stack, plain), np.dot(stack, plain))

print(f"      the same stack times ANOTHER stack, {stack.shape} against {other.shape}:")
print(f"        np.matmul -> {np.matmul(stack, other).shape}   two matrices, paired up")
print(f"        np.dot    -> {np.dot(stack, other).shape}   every pairing, all four")
assert np.matmul(stack, other).shape == (2, 2, 2)
assert np.dot(stack, other).shape == (2, 2, 2, 2)
print("        THAT is where they part company, and the gap is a whole extra")
print("        axis. matmul pairs the stacks off matrix by matrix; dot sums over")
print("        the last axis of the left and the second-to-last of the right and")
print("        keeps everything else, so it produces every combination.")
print()
print("      The rule to carry is simpler than the exceptions: use `@` for")
print("      matrix multiplication and `np.dot` only for two plain vectors.")
print(f"      For two vectors, np.dot is exact and readable: np.dot([3,4],[4,3]) = "
      f"{np.dot([3, 4], [4, 3])}")

print()
print("=" * 74)
print("4. A shape error on purpose, and how to read it")
print("=" * 74)

print(f"  X is {npX.shape}. X @ X asks for ({npX.shape[0]}, {npX.shape[1]}) @ "
      f"({npX.shape[0]}, {npX.shape[1]}).")
print(f"  The inner dimensions are {npX.shape[1]} and {npX.shape[0]}, and they disagree.")
try:
    npX @ npX
except ValueError as exc:
    message = str(exc)
    print()
    print("  NumPy raises:")
    print(f"    {type(exc).__name__}: {message}")
    caught = type(exc)
else:  # pragma: no cover - only reached if NumPy changes its rules
    raise AssertionError("(2, 3) @ (2, 3) should have raised")

assert caught is ValueError
assert "size 2 is different from 3" in message
print()
print("  The first thing to check, always, is the two shapes and nothing else.")
print("  Print them before you read another word of the traceback:")
print(f"      left  {npX.shape}     right {npX.shape}")
print("      inner dimensions: 3 and 2 — that is the bug, and it is the whole bug.")

print()
print("=" * 74)
print("5. The transpose fixes it — but there are TWO fixes and they differ")
print("=" * 74)

left_fix = npX @ npX.T
right_fix = npX.T @ npX
print(f"  X @ X.T  is (2, 3) @ (3, 2) -> {left_fix.shape}")
print(f"      {left_fix.tolist()}")
print("      entry (i, j) is example i dotted with example j — a table of how")
print("      alike the two EXAMPLES are. Symmetric, as any such table must be.")
print(f"  X.T @ X  is (3, 2) @ (2, 3) -> {right_fix.shape}")
print(f"      {right_fix.tolist()}")
print("      entry (i, j) is feature i dotted with feature j — a table of how")
print("      alike the three FEATURES are. Also symmetric, and a different object.")
assert left_fix.tolist() == X_AT_XT
assert right_fix.tolist() == XT_AT_X
assert left_fix.shape == (2, 2)
assert right_fix.shape == (3, 3)
assert (left_fix == left_fix.T).all()
assert (right_fix == right_fix.T).all()
print()
print("  Both make the exception go away. Only one of them answers the question")
print("  you had. This is the reason 'just transpose it until it runs' is a bad")
print("  habit: the error was telling you something, and silencing it at random")
print("  swaps a loud failure for a quiet wrong answer.")

print()
print("=" * 74)
print("6. The full comparison table, checked")
print("=" * 74)
rows = [
    ("P * Q", (npP * npQ).shape, "entry by entry, nothing summed"),
    ("P @ Q", (npP @ npQ).shape, "rows dotted with columns"),
    ("X * u", star.shape, "u broadcast across rows, nothing summed"),
    ("X @ u", at.shape, "rows dotted with u, then summed"),
    ("X @ X", "ValueError", "inner dimensions 3 and 2 disagree"),
    ("X @ X.T", left_fix.shape, "example against example"),
    ("X.T @ X", right_fix.shape, "feature against feature"),
]
print(f"  {'expression':<10} {'result shape':<14} what it means")
for name, shp, meaning in rows:
    print(f"  {name:<10} {str(shp):<14} {meaning}")

print()
print("03_star_versus_at.py: every assertion held.")

"""One layer of a neural network is a matrix multiply plus a vector add.

Run from inside this directory:

    ../.venv/bin/python3 04_network_layer.py

Every number below is small enough to work out with a pen, and the hand-worked
answers in dataset.py were derived that way before anything was run. This is
the operation that consumes essentially all the compute in training any model
you will ever use. It is `X @ W + b`, and that is the whole of it.
"""

import numpy as np

from dataset import BIAS, LAYER_OUT, W, X, XW
from matmul import add_bias, matmul_loops, multiplication_count, shape

npX = np.array(X)
npW = np.array(W)
npB = np.array(BIAS)

print("=" * 74)
print("1. The three things a layer is made of")
print("=" * 74)
print(f"  X, the batch      shape {npX.shape}  {npX.tolist()}")
print("      two examples, three features each. One example per ROW — that is")
print("      the convention almost every framework uses, and the reason the")
print("      weights end up on the right-hand side of the multiply.")
print(f"  W, the weights    shape {npW.shape}  {npW.tolist()}")
print("      three inputs in, two outputs out. Column j holds the weights that")
print("      produce output j. Read it that way and the shape rule is obvious")
print("      rather than memorised.")
print(f"  b, the bias       shape {npB.shape}     {npB.tolist()}")
print("      one number per output unit. Not per example — per OUTPUT.")

print()
print("=" * 74)
print("2. The multiply, worked out by hand")
print("=" * 74)
print("  Entry (i, j) of X @ W is example i dotted with the weights of output j.")
print()
print("  example 0 = [1, 2, 0]")
print("    output 0, weights [2, -1, 0]:  1*2 + 2*(-1) + 0*0  =  2 - 2 + 0  =  0")
print("    output 1, weights [0,  1, 4]:  1*0 + 2*1    + 0*4  =  0 + 2 + 0  =  2")
print("  example 1 = [0, 1, 3]")
print("    output 0, weights [2, -1, 0]:  0*2 + 1*(-1) + 3*0  =  0 - 1 + 0  = -1")
print("    output 1, weights [0,  1, 4]:  0*0 + 1*1    + 3*4  =  0 + 1 + 12 = 13")
print()
product = npX @ npW
print(f"  X @ W = {product.tolist()}   shape {product.shape}")
assert product.tolist() == XW
assert product.shape == (2, 2)
print("  Two examples in, two outputs each. The 3 was consumed — it had to")
print("  match, and it is gone from the answer. The 2 on the left survived")
print("  because it is the batch, and the 2 on the right survived because it is")
print("  the width of the layer. Neither of those numbers is the same kind of")
print("  thing, and they are only both 2 here by accident.")

print()
print("=" * 74)
print("3. The bias, broadcast across the rows")
print("=" * 74)
print(f"  b has shape {npB.shape} and X @ W has shape {product.shape}.")
print("  Broadcasting (Day 100) lines the shapes up from the right: (2, 2)")
print("  against (2,) pads to (1, 2), the trailing 2s match, and the 1 stretches")
print("  down the rows. So EVERY example gets the same bias:")
print()
print("    example 0: [ 0,  2] + [5, -2] = [5,  0]")
print("    example 1: [-1, 13] + [5, -2] = [4, 11]")
print()
out = product + npB
print(f"  X @ W + b = {out.tolist()}   shape {out.shape}")
assert out.tolist() == LAYER_OUT

print()
print("  Written out as an explicit loop, with no broadcasting at all:")
by_hand = add_bias(matmul_loops(X, W), BIAS)
print(f"  add_bias(matmul_loops(X, W), b) = {by_hand}")
assert by_hand == LAYER_OUT
print("  Same answer. The broadcast is shorthand for that loop, and knowing it")
print("  is shorthand is what stops it being magic.")

print()
print("  And the mistake worth meeting once: a bias of the wrong length.")
try:
    product + np.array([5, -2, 7])
except ValueError as exc:
    print(f"    (2, 2) + (3,) raises {type(exc).__name__}: {exc}")
    assert "could not be broadcast" in str(exc)
else:  # pragma: no cover
    raise AssertionError("a length-3 bias on a 2-wide layer should have raised")
print("    One bias per OUTPUT. A layer two units wide takes two numbers, and")
print("    the exception is the shape rule from Day 100 doing its job.")

print()
print("=" * 74)
print("3b. A convention clash worth meeting head on")
print("=" * 74)
print("  02_composition.py said matrices compose RIGHT TO LEFT: in A @ B, B runs")
print("  first, because B is the one standing next to the vector in A @ (B @ v).")
print("  Then this script writes a layer as X @ W, with the data on the LEFT and")
print("  the transformation on the right. Those two look like they contradict")
print("  each other. They do not, and the reason is worth knowing.")
print()
print("  It comes down to whether a vector is a column or a row.")
print()
print("    Column convention (textbooks, 02_composition.py):")
print("      v is (n, 1), the matrix goes on the LEFT:      y = A @ v")
print("      chaining reads right to left:                  y = B @ (A @ v)")
print()
print("    Row convention (this script, and every framework you will use):")
print("      each example is a ROW, the matrix goes on the RIGHT: y = x @ A")
print("      chaining reads left to right:                        y = x @ A @ B")
print()
print("  Both compute the same thing. One is the transpose of the other:")
v_row = np.array([1, 2, 0])
A_col = npW.T  # (2, 3): the same layer written for column vectors
print(f"      row form:    x @ W        = {(v_row @ npW).tolist()}")
print(f"      column form: W.T @ x      = {(A_col @ v_row).tolist()}")
assert np.array_equal(v_row @ npW, A_col @ v_row)
print("  Identical, because (x @ W) and (W.T @ x) are the same numbers written")
print("  the two different ways round.")
print()
print("  Why frameworks chose rows: a batch is a stack of examples, and stacking")
print("  them as rows means example i lives at X[i], which is how every dataset,")
print("  CSV file and database table you have met since Day 65 is already laid")
print("  out. The cost is this one moment of confusion, and you have now had it.")

print()
print("=" * 74)
print("4. What changes when the batch grows, and what does not")
print("=" * 74)
bigger = np.array([[1, 2, 0], [0, 1, 3], [2, 0, 1], [1, 1, 1]])
bigger_out = bigger @ npW + npB
print(f"  a batch of {bigger.shape[0]} instead of {npX.shape[0]}: {bigger.shape} @ {npW.shape} "
      f"-> {(bigger @ npW).shape}, plus b -> {bigger_out.shape}")
print(f"  {bigger_out.tolist()}")
assert bigger_out.shape == (4, 2)
assert bigger_out[:2].tolist() == LAYER_OUT
print("  The first two rows are unchanged, because each row is computed")
print("  independently of the others. W did not change shape and b did not")
print("  change shape — only the batch dimension moved. That independence is")
print("  exactly what makes a batch worth having: the same weights, reused")
print("  across every example, in one multiply.")

print()
print("=" * 74)
print("5. Stacking two layers is multiplying three matrices")
print("=" * 74)
W2 = np.array([[1, 0, 2], [3, 1, 0]])  # (2, 3): the 2-wide layer feeds a 3-wide one
b2 = np.array([0, 1, -1])
hidden = npX @ npW + npB
final = hidden @ W2 + b2
print(f"  layer 1: X {npX.shape} @ W {npW.shape} + b {npB.shape} -> {hidden.shape}")
print(f"           {hidden.tolist()}")
print(f"  layer 2: h {hidden.shape} @ W2 {W2.shape} + b2 {b2.shape} -> {final.shape}")
print(f"           {final.tolist()}")
assert hidden.tolist() == LAYER_OUT
assert final.shape == (2, 3)
print()
print("  The shapes chain: 3 features in, 2 hidden, 3 out. Each layer's output")
print("  width must equal the next layer's input width, and that is the shape")
print("  rule again, wearing a different hat.")
print()
print("  Now the honest caveat, and it matters. Without a non-linear function")
print("  between the layers, those two layers COLLAPSE into one:")
collapsed_W = npW @ W2
collapsed_b = npB @ W2 + b2
collapsed = npX @ collapsed_W + collapsed_b
print(f"      W @ W2 has shape {collapsed_W.shape}, and X @ (W @ W2) + (b @ W2 + b2) =")
print(f"      {collapsed.tolist()}")
assert collapsed.tolist() == final.tolist()
print("      — identical to running the two layers separately. That is")
print("      associativity, and it is the reason activation functions exist:")
print("      a stack of pure matrix multiplies is just one matrix multiply, no")
print("      matter how many layers deep you make it.")

print()
print("=" * 74)
print("6. The cost, at a size people actually train")
print("=" * 74)
print("  This layer:")
m, n = shape(X)
_, p = shape(W)
print(f"      ({m}, {n}) @ ({n}, {p}) = {multiplication_count(m, n, p)} multiplications. You could do it on paper.")
for batch, d_in, d_out, label in [
    (32, 768, 768, "one modest layer, batch 32"),
    (1024, 4096, 4096, "one wide layer, batch 1024"),
]:
    count = multiplication_count(batch, d_in, d_out)
    print(f"  {label}:")
    print(f"      ({batch}, {d_in}) @ ({d_in}, {d_out}) = {count:,} multiplications")
assert multiplication_count(32, 768, 768) == 18_874_368
assert multiplication_count(1024, 4096, 4096) == 17_179_869_184
print()
print("  Seventeen billion multiplications, for ONE layer, on ONE batch, in ONE")
print("  forward pass. A model has many layers, training does a backward pass")
print("  too, and you repeat the whole thing for every batch in the dataset,")
print("  many times over. This is where the compute goes. Not somewhere else.")

print()
print("04_network_layer.py: every assertion held.")

"""Association order, and the loop against NumPy at a size where it shows.

Run from inside this directory:

    ../.venv/bin/python3 05_cost_and_speed.py

READ THIS BEFORE READING THE NUMBERS. The durations printed below are from one
machine on one day, and yours will differ — possibly by a lot. They are printed
because a measurement you can see beats an assertion you have to trust, not
because the milliseconds mean anything. What matters, and what the tests
actually assert, is the SHAPE of the gap: a ratio in the hundreds or thousands,
growing with the size of the problem. No test in this lab asserts a duration.
"""

import time

import numpy as np

from dataset import BIG_CHAIN, BIG_LEFT_FIRST, BIG_RATIO, BIG_RIGHT_FIRST
from dataset import SMALL_CHAIN, SMALL_LEFT_FIRST, SMALL_RIGHT_FIRST
from matmul import chain_costs, matmul_loops, multiplication_count

print("=" * 74)
print("1. Where the brackets go changes the arithmetic, not the answer")
print("=" * 74)

for chain, expected in [(SMALL_CHAIN, (SMALL_LEFT_FIRST, SMALL_RIGHT_FIRST)),
                        (BIG_CHAIN, (BIG_LEFT_FIRST, BIG_RIGHT_FIRST))]:
    m, n, p, q = chain
    left_first, right_first = chain_costs(m, n, p, q)
    assert (left_first, right_first) == expected
    print(f"  ({m}, {n}) @ ({n}, {p}) @ ({p}, {q})")
    print(f"      (AB)C  = {m}*{n}*{p} + {m}*{p}*{q}")
    print(f"             = {multiplication_count(m, n, p):,} + {multiplication_count(m, p, q):,}"
          f" = {left_first:,}")
    print(f"      A(BC)  = {n}*{p}*{q} + {m}*{n}*{q}")
    print(f"             = {multiplication_count(n, p, q):,} + {multiplication_count(m, n, q):,}"
          f" = {right_first:,}")
    print(f"      ratio  = {right_first / left_first:,.0f}x")
    print()

assert BIG_RIGHT_FIRST // BIG_LEFT_FIRST == BIG_RATIO
print("  The second chain is a low-rank adapter: a 4096-wide layer with an 8-wide")
print("  detour through A and back out through B, on a batch of 1024. Multiplying")
print("  A and B together first builds a full (4096, 4096) matrix and then hits")
print(f"  the whole batch with it — {BIG_RATIO} times the work for the identical answer.")
print("  Associativity is what makes both spellings legal. Counting is what tells")
print("  you which one to write.")

print()
print("  And a proof, on small enough shapes to check, that the answers really")
print("  are identical rather than merely close:")
rng = np.random.default_rng(101)
A = rng.integers(-5, 6, size=(4, 7))
B = rng.integers(-5, 6, size=(7, 2))
C = rng.integers(-5, 6, size=(2, 6))
left = (A @ B) @ C
right = A @ (B @ C)
print(f"      (A @ B) @ C and A @ (B @ C), shapes {left.shape} and {right.shape}")
print(f"      identical in every entry: {np.array_equal(left, right)}")
assert np.array_equal(left, right)
print("      These are integers, so 'identical' is exact. With floating point")
print("      (Day 70) the two orders can differ in the last bits, because")
print("      addition is not associative in floating point even though matrix")
print("      multiplication is associative in mathematics. Worth knowing before")
print("      it surprises you in a test.")
fA = A.astype(np.float64) / 3
fleft = (fA @ B) @ C
fright = fA @ (B @ C)
print(f"      in float64, exactly equal: {np.array_equal(fleft, fright)}; "
      f"close within 1e-9: {np.allclose(fleft, fright, atol=1e-9)}")
assert np.allclose(fleft, fright, atol=1e-9)

print()
print("=" * 74)
print("2. The loop against NumPy")
print("=" * 74)

SIZE = 200
print(f"  Two {SIZE} by {SIZE} matrices. The nested loop's innermost line will run")
print(f"  {multiplication_count(SIZE, SIZE, SIZE):,} times.")
print()

rng = np.random.default_rng(2026)
ints_left = rng.integers(0, 10, size=(SIZE, SIZE))
ints_right = rng.integers(0, 10, size=(SIZE, SIZE))
floats_left = ints_left.astype(np.float64)
floats_right = ints_right.astype(np.float64)
left_lists = ints_left.tolist()
right_lists = ints_right.tolist()


def best_of(fn, repeats=5):
    """Fastest of several runs — the run least disturbed by everything else."""
    best = float("inf")
    result = None
    for _ in range(repeats):
        start = time.perf_counter()
        result = fn()
        best = min(best, time.perf_counter() - start)
    return best, result


start = time.perf_counter()
loop_answer = matmul_loops(left_lists, right_lists)
loop_seconds = time.perf_counter() - start

int_seconds, int_answer = best_of(lambda: ints_left @ ints_right)
float_seconds, float_answer = best_of(lambda: floats_left @ floats_right)

assert loop_answer == int_answer.tolist(), "the two answers must be identical"
assert np.array_equal(float_answer, int_answer.astype(np.float64))
print("  All three answers are identical, entry for entry. Only the time differs.")
print()
print(f"      three nested loops in Python  : {loop_seconds:9.4f} s")
print(f"      NumPy @ on int64  (best of 5) : {int_seconds:9.6f} s"
      f"   {loop_seconds / int_seconds:>9,.0f}x faster than the loop")
print(f"      NumPy @ on float64 (best of 5): {float_seconds:9.6f} s"
      f"   {loop_seconds / float_seconds:>9,.0f}x faster than the loop")
print()
print("  Those durations are from one machine on one day and yours will differ.")
print("  The ratios are the part that travels, and even they vary with hardware")
print("  and with how your NumPy was built. No test in this lab asserts a time.")
assert loop_seconds / int_seconds > 10, "int64 should still beat the loop comfortably"
assert loop_seconds / float_seconds > 200, "float64 should beat the loop enormously"

print()
print("=" * 74)
print("3. The surprise in that table, and what it tells you")
print("=" * 74)
print("  The two NumPy rows are not the same speed, and the difference is not")
print("  small. On this machine, on this run:")
print()
print(f"      int64 divided by float64: {int_seconds / float_seconds:,.0f}x")
print()
print("  Same shapes, same values, same operator, same library. The only thing")
print("  that changed was the dtype, and the float version was dramatically")
print("  faster. That is not a quirk to file away. It is the single best piece")
print("  of evidence for what NumPy is actually doing:")
print()
print("  **BLAS only handles floating point.** BLAS — Basic Linear Algebra")
print("  Subprograms — is a decades-old interface with several competing")
print("  implementations, all compiled, all tuned to the exact processor they run")
print("  on, using vector instructions and cache-aware blocking and often several")
print("  cores. Its matrix-multiply routines are defined for float and complex")
print("  types and nothing else. So a float64 `@` is handed straight to BLAS,")
print("  while an int64 `@` falls back to NumPy's own compiled C loop — still far")
print("  better than interpreted Python, and still nowhere near BLAS.")
print()
print("  This is why the answer to 'why is NumPy fast?' is not 'because it is C'.")
print("  The int64 row IS C, and it is the slow NumPy row. NumPy is fast because")
print("  for the types that matter it stops being NumPy too, and calls out to a")
print("  library that people have been optimising since the 1970s.")
print()
config = np.__config__.show(mode="dicts") if hasattr(np.__config__, "show") else None
blas = {}
if isinstance(config, dict):
    blas = config.get("Build Dependencies", {}).get("blas", {}) or {}
if blas.get("name"):
    print("  This installation reports its own BLAS, read from the build config")
    print("  rather than assumed:")
    for key in ("name", "found", "detection method"):
        if key in blas:
            print(f"      {key:<17} {blas[key]}")
else:  # pragma: no cover - depends entirely on how NumPy was built
    print("  This installation did not report a BLAS name in its build config,")
    print("  so none is claimed here. That is a gap in what can be shown, not")
    print("  evidence that no BLAS is present.")
print()
print("  The practical consequence, and it is worth carrying: if a matrix")
print("  multiply is slower than you expected, check the dtype before you check")
print("  anything else. This is also why every framework you will meet stores")
print("  weights as float32 or a smaller float and never as integers.")

print()
print("=" * 74)
print("4. Does the gap hold as the problem grows?")
print("=" * 74)
print("  Both implementations do work proportional to n cubed, so the RATIO")
print("  should stay in the same broad range as n grows — the loop's overhead is")
print("  per operation, not per call. Whether it actually does is a measurement,")
print("  not a deduction, so here it is:")
print()
header = f"  {'n':>5}  {'operations':>14}  {'loop (s)':>10}  {'float64 (s)':>12}  {'ratio':>10}"
print(header)
ratios = []
for n in (40, 80, 160):
    a = rng.integers(0, 10, size=(n, n)).astype(np.float64)
    b = rng.integers(0, 10, size=(n, n)).astype(np.float64)
    al, bl = a.tolist(), b.tolist()
    start = time.perf_counter()
    matmul_loops(al, bl)
    t_loop = time.perf_counter() - start
    t_np, _ = best_of(lambda: a @ b)
    ratios.append(t_loop / t_np)
    print(f"  {n:>5}  {multiplication_count(n, n, n):>14,}  {t_loop:>10.4f}  "
          f"{t_np:>12.6f}  {t_loop / t_np:>9,.0f}x")

print()
print("  Read the ratio column, not the two before it. At the smallest size the")
print("  ratio is held down by NumPy's own fixed per-call overhead, which is a")
print("  real cost that simply stops mattering once the matrices are big enough.")
print("  If your machine shows something else, believe your machine — and then")
print("  work out why, which is a better exercise than the one this script set.")
assert all(r > 20 for r in ratios), "even the smallest size should show a wide gap"

print()
print("05_cost_and_speed.py: every assertion held.")

"""Make this directory's own matmul.py the one its tests import.

Both `examples/` and `starter/` contain a module called `matmul`, and pytest
imports test files by putting their directory on `sys.path`. Without this file,
running `pytest` across both directories at once would import whichever
`matmul` was seen first and then reuse it for the other suite — so the starter
tests would silently pass against the reference solution instead of skipping.
That is a wrong answer with a green tick on it, which is the worst kind.

This is not a hypothetical. It happened while building the Day 100 lab, where
eleven unwritten exercises reported as passing, and it was caught only because
the number of skips changed between two runs that should have agreed. Section 4
of `tests/run_tests.sh` now asserts that the skip count is identical whether the
suites are collected separately or together, so the same mistake cannot come
back quietly.

So: put this directory first on the import path, and drop any already-imported
`matmul`, `dataset` or `answers` that came from somewhere else.
"""

import sys
from pathlib import Path

HERE = str(Path(__file__).parent.resolve())

if HERE in sys.path:
    sys.path.remove(HERE)
sys.path.insert(0, HERE)

for name in ("matmul", "dataset", "answers"):
    module = sys.modules.get(name)
    origin = getattr(module, "__file__", "") or ""
    if module is not None and not origin.startswith(HERE):
        del sys.modules[name]

"""The numbers every script in this lab shares, and the answers worked by hand.

Everything here is invented and everything here is small. That is deliberate:
each answer below was worked out with a pen before it was ever run, and you can
do the same. A lab about an operation you cannot check by hand is a lab that
teaches you to trust output, which is the opposite of the point.

The hand-worked answers are stored beside the inputs so that the reference
tests can assert against a number a human derived, not against whatever NumPy
happened to return. If NumPy and the pen disagree, that is a finding, and the
tests are arranged so you would see it rather than absorb it.
"""

# ---------------------------------------------------------------------------
# One layer of a neural network: X @ W + b
# ---------------------------------------------------------------------------

# A batch of two examples, three features each. Shape (2, 3).
# Row 0 is the first example, row 1 is the second. Rows are examples and
# columns are features — the table reading from Day 100.
X = [
    [1, 2, 0],
    [0, 1, 3],
]

# The weights of one layer: three inputs in, two outputs out. Shape (3, 2).
# Column j holds the weights that produce output j. Read it that way and the
# shape rule stops needing to be memorised.
W = [
    [2, 0],
    [-1, 1],
    [0, 4],
]

# One bias per output unit. Shape (2,). It is added to every row of X @ W,
# which is broadcasting from Day 100 doing exactly the job it was built for.
BIAS = [5, -2]

# X @ W, worked by hand:
#   row 0 = [1, 2, 0]
#     column 0 of W is [2, -1, 0]:  1*2 + 2*(-1) + 0*0  =  2 - 2 + 0  =  0
#     column 1 of W is [0,  1, 4]:  1*0 + 2*1    + 0*4  =  0 + 2 + 0  =  2
#   row 1 = [0, 1, 3]
#     column 0 of W is [2, -1, 0]:  0*2 + 1*(-1) + 3*0  =  0 - 1 + 0  = -1
#     column 1 of W is [0,  1, 4]:  0*0 + 1*1    + 3*4  =  0 + 1 + 12 = 13
XW = [
    [0, 2],
    [-1, 13],
]

# Then the bias is added to every row: [5, -2] on top of each.
#   row 0: [0 + 5,  2 - 2] = [5,  0]
#   row 1: [-1 + 5, 13 - 2] = [4, 11]
LAYER_OUT = [
    [5, 0],
    [4, 11],
]

# The single output cell the architecture diagram highlights: row 1, column 1
# of X @ W. Row 1 of X is [0, 1, 3]; column 1 of W is [0, 1, 4].
HIGHLIGHT_CELL = (1, 1)
HIGHLIGHT_ROW = [0, 1, 3]
HIGHLIGHT_COLUMN = [0, 1, 4]
HIGHLIGHT_TERMS = [0, 1, 12]  # 0*0, 1*1, 3*4
HIGHLIGHT_VALUE = 13

# ---------------------------------------------------------------------------
# Composition: two transformations of the plane, in both orders
# ---------------------------------------------------------------------------

# A quarter turn anticlockwise. It sends (1, 0) to (0, 1) and (0, 1) to (-1, 0),
# and those two images ARE its columns — which is the whole trick for reading a
# transformation matrix off a picture.
ROT90 = [
    [0, -1],
    [1, 0],
]

# A reflection in the horizontal axis: x stays, y flips sign.
FLIP_X = [
    [1, 0],
    [0, -1],
]

# ROT90 @ FLIP_X means "flip first, then rotate" — the rightmost matrix meets
# the vector first, because that is the one standing next to it in A @ (B @ v).
# The product turns out to be the reflection in the line y = x.
ROT_AFTER_FLIP = [
    [0, 1],
    [1, 0],
]

# FLIP_X @ ROT90 means "rotate first, then flip". Same two operations, opposite
# order, and the result is the reflection in the line y = -x. A different
# transformation entirely, which is what "not commutative" means in practice.
FLIP_AFTER_ROT = [
    [0, -1],
    [-1, 0],
]

# The test vector the flow diagram follows, and every waypoint on its journey.
V = [3, 1]
FLIP_V = [3, -1]  # FLIP_X @ V
ROT_AFTER_FLIP_V = [1, 3]  # ROT90 @ (FLIP_X @ V), and also ROT_AFTER_FLIP @ V
ROT_V = [-1, 3]  # ROT90 @ V
FLIP_AFTER_ROT_V = [-1, -3]  # FLIP_X @ (ROT90 @ V), and also FLIP_AFTER_ROT @ V

# The elementwise product of those same two matrices is all zeros, because
# every entry of one lines up with a zero of the other. Same operands, same
# shape out, and not a single number in common with the matrix product.
ROT_TIMES_FLIP_ELEMENTWISE = [
    [0, 0],
    [0, 0],
]

# ---------------------------------------------------------------------------
# A second, less tidy pair: three different answers from two matrices
# ---------------------------------------------------------------------------

P = [
    [1, 2],
    [3, 4],
]
Q = [
    [5, 6],
    [7, 8],
]

# P @ Q, worked by hand:
#   [0][0] = 1*5 + 2*7 = 5 + 14 = 19      [0][1] = 1*6 + 2*8 = 6 + 16 = 22
#   [1][0] = 3*5 + 4*7 = 15 + 28 = 43     [1][1] = 3*6 + 4*8 = 18 + 32 = 50
P_AT_Q = [
    [19, 22],
    [43, 50],
]

# Q @ P, worked by hand:
#   [0][0] = 5*1 + 6*3 = 5 + 18 = 23      [0][1] = 5*2 + 6*4 = 10 + 24 = 34
#   [1][0] = 7*1 + 8*3 = 7 + 24 = 31      [1][1] = 7*2 + 8*4 = 14 + 32 = 46
Q_AT_P = [
    [23, 34],
    [31, 46],
]

# P * Q, entry by entry, no summing anywhere:
P_TIMES_Q = [
    [5, 12],
    [21, 32],
]

# ---------------------------------------------------------------------------
# The `*` versus `@` trap, where both are legal and neither warns you
# ---------------------------------------------------------------------------

# A vector of length 3 meeting the (2, 3) batch X.
U = [10, 2, 5]

# X * U broadcasts U across both rows and multiplies entry by entry.
# Shape (2, 3) out — the same shape it went in with, nothing summed.
X_TIMES_U = [
    [10, 4, 0],
    [0, 2, 15],
]

# X @ U multiplies and then SUMS along each row. Shape (2,) out.
#   row 0: 1*10 + 2*2 + 0*5 = 10 + 4 + 0  = 14
#   row 1: 0*10 + 1*2 + 3*5 = 0 + 2 + 15  = 17
X_AT_U = [14, 17]

# ---------------------------------------------------------------------------
# The deliberate shape error, and the two different repairs
# ---------------------------------------------------------------------------

# X @ X is (2, 3) @ (2, 3). The inner dimensions are 3 and 2 and they disagree,
# so NumPy raises. Transposing one side fixes it — but WHICH side you transpose
# changes the answer, its shape and its meaning, and nothing will tell you if
# you pick the one you did not want.

# X @ X.T is (2, 3) @ (3, 2) -> (2, 2). Entry (i, j) is row i dotted with row j,
# so this is a table of how alike the two EXAMPLES are.
#   [0][0] = 1*1 + 2*2 + 0*0 = 5      [0][1] = 1*0 + 2*1 + 0*3 = 2
#   [1][0] = 0*1 + 1*2 + 3*0 = 2      [1][1] = 0*0 + 1*1 + 3*3 = 10
X_AT_XT = [
    [5, 2],
    [2, 10],
]

# X.T @ X is (3, 2) @ (2, 3) -> (3, 3). Entry (i, j) is column i dotted with
# column j, so this is a table of how alike the three FEATURES are.
#   columns of X are [1, 0], [2, 1] and [0, 3]
#   [0][0] = 1*1 + 0*0 = 1     [0][1] = 1*2 + 0*1 = 2     [0][2] = 1*0 + 0*3 = 0
#   [1][1] = 2*2 + 1*1 = 5     [1][2] = 2*0 + 1*3 = 3
#   [2][2] = 0*0 + 3*3 = 9
XT_AT_X = [
    [1, 2, 0],
    [2, 5, 3],
    [0, 3, 9],
]

# ---------------------------------------------------------------------------
# The dot product, arithmetically and geometrically
# ---------------------------------------------------------------------------

DOT_U = [3, 4]  # length 5, from Day 99
DOT_V = [4, 3]  # length 5 as well
DOT_W = [-4, 3]  # length 5, and at right angles to DOT_U

DOT_U_U = 25  # 3*3 + 4*4 — a vector dotted with itself is its length squared
DOT_U_V = 24  # 3*4 + 4*3
DOT_U_W = 0  # 3*(-4) + 4*3 = -12 + 12 — perpendicular, and the zero says so

# A pair whose angle has a name you can state exactly.
ANGLE_A = [2, 0]
ANGLE_B = [1, 1]
ANGLE_A_B = 2  # 2*1 + 0*1
# |ANGLE_A| = 2, |ANGLE_B| = sqrt(2), so cos(theta) = 2 / (2 * sqrt(2))
# = 1 / sqrt(2), and that angle is 45 degrees exactly.
ANGLE_DEGREES = 45.0

# ---------------------------------------------------------------------------
# Association order: same answer, wildly different arithmetic
# ---------------------------------------------------------------------------

# A small chain you can count on paper: (10, 100) @ (100, 5) @ (5, 50).
SMALL_CHAIN = (10, 100, 5, 50)
SMALL_LEFT_FIRST = 7_500  # (AB)C  = 10*100*5 + 10*5*50
SMALL_RIGHT_FIRST = 75_000  # A(BC)  = 100*5*50 + 10*100*50

# A chain with the shapes of a real low-rank adapter sitting on a 4096-wide
# layer, applied to a batch of 1024: (1024, 4096) @ (4096, 8) @ (8, 4096).
BIG_CHAIN = (1024, 4096, 8, 4096)
BIG_LEFT_FIRST = 67_108_864  # (XA)B
BIG_RIGHT_FIRST = 17_314_086_912  # X(AB)
BIG_RATIO = 258  # exactly, on these shapes

# ---------------------------------------------------------------------------
# Shapes used to exercise the shape rule
# ---------------------------------------------------------------------------

# (left shape, right shape, expected result shape or the string "error")
SHAPE_CASES = [
    ((2, 3), (3, 2), (2, 2)),
    ((3, 2), (2, 3), (3, 3)),
    ((2, 3), (2, 3), "error"),
    ((1, 4), (4, 1), (1, 1)),
    ((4, 1), (1, 4), (4, 4)),
    ((5, 5), (5, 5), (5, 5)),
    ((2, 3), (4, 2), "error"),
]

"""Matrix multiplication from first principles, written three different ways.

The three functions below compute exactly the same thing and are asserted equal
to each other and to NumPy's `@` in `01_matmul_from_scratch.py`. They exist
separately because each one makes a different fact obvious:

* `matmul_loops`   — three nested loops. The definition, transcribed. It makes
                     the COST obvious: the body runs m * n * p times.
* `matmul_dots`    — a list of dot products. It makes the DEFINITION obvious:
                     entry (i, j) is row i of A dotted with column j of B, and
                     nothing else.
* `matmul_columns` — column j of the answer is A applied to column j of B, and
                     A applied to a vector is a weighted sum of A's columns. It
                     makes the MEANING obvious, and it is the reading that
                     makes every later idea in linear algebra easy.

Nothing here imports NumPy. That is the point of a from-scratch build: you
should be able to read every line and see where each number came from.

Matrices are plain nested lists of numbers, rows outermost, exactly as they are
written on paper. Vectors are flat lists.
"""

from __future__ import annotations


class ShapeMismatch(ValueError):
    """Raised when two matrices cannot be multiplied.

    Subclasses ValueError deliberately, so that code which only knows to catch
    ValueError still catches this, while code that wants to be specific can be.
    NumPy raises its own ValueError for the same situation.
    """


def shape(M: list[list[float]]) -> tuple[int, int]:
    """Return (rows, columns), checking that the grid is rectangular."""
    if not M or not isinstance(M, list) or not isinstance(M[0], list):
        raise ValueError("a matrix is a non-empty list of rows, each a list")
    n_cols = len(M[0])
    for i, row in enumerate(M):
        if len(row) != n_cols:
            raise ValueError(
                f"row {i} has {len(row)} entries but row 0 has {n_cols}; "
                "a matrix is rectangular"
            )
    return (len(M), n_cols)


def check_multipliable(A: list[list[float]], B: list[list[float]]) -> tuple[int, int, int]:
    """Check the shape rule and return (m, n, p) for an (m, n) @ (n, p).

    The rule is not a convention to memorise. A @ B means "do B, then do A".
    B has shape (n, p), so it turns a vector of length p into one of length n.
    A has shape (m, n), so it accepts a vector of length n and returns one of
    length m. A's column count is the length of what it ACCEPTS; B's row count
    is the length of what it PRODUCES. The inner dimensions must match because
    the second thing has to accept what the first one hands it. The outer
    dimensions survive because they are the two ends of the pipeline: what goes
    in at one end, and what comes out at the other.
    """
    m, n = shape(A)
    n2, p = shape(B)
    if n != n2:
        raise ShapeMismatch(
            f"cannot multiply ({m}, {n}) by ({n2}, {p}): "
            f"the inner dimensions {n} and {n2} disagree. "
            f"In A @ B the right-hand matrix runs first and returns vectors of "
            f"length {n2}, and the left-hand matrix only accepts vectors of "
            f"length {n}."
        )
    return m, n, p


def dot(u: list[float], v: list[float]) -> float:
    """Multiply pairwise, then add. The whole of the dot product."""
    if len(u) != len(v):
        raise ShapeMismatch(
            f"cannot dot a vector of length {len(u)} with one of length {len(v)}; "
            "there is no sensible partner for the leftover entries"
        )
    total = 0
    for a, b in zip(u, v):
        total += a * b
    return total


def matmul_loops(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
    """Version 1: three nested loops. The definition with nothing hidden.

    Count the work: the innermost line runs once for every (i, k, j), which is
    m * n * p times. For two 200 by 200 matrices that is eight million
    multiply-and-adds, and Python will do every one of them as an interpreted
    step. That number is the reason the timing script exists.
    """
    m, n, p = check_multipliable(A, B)
    C = [[0] * p for _ in range(m)]
    for i in range(m):  # each row of the answer
        for j in range(p):  # each column of the answer
            total = 0
            for k in range(n):  # walk A's row and B's column together
                total += A[i][k] * B[k][j]
            C[i][j] = total
    return C


def matmul_dots(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
    """Version 2: entry (i, j) is row i of A dotted with column j of B."""
    check_multipliable(A, B)
    B_columns = transpose(B)  # so a column is a list we can hand to dot()
    return [[dot(row, column) for column in B_columns] for row in A]


def matvec(A: list[list[float]], v: list[float]) -> list[float]:
    """A applied to one vector, computed as a weighted sum of A's COLUMNS.

    This is the picture worth keeping. A @ v does not really "dot v with the
    rows". It takes v[0] copies of A's first column, v[1] copies of the second,
    and so on, and adds them up. The answer is therefore always a combination
    of A's columns and can never leave the space they span — which is why the
    output has as many entries as A has rows, and why the length of v must
    equal the number of columns: one weight per column, no spares.
    """
    m, n = shape(A)
    if len(v) != n:
        raise ShapeMismatch(
            f"cannot apply an ({m}, {n}) matrix to a vector of length {len(v)}: "
            f"there is one weight per column and the matrix has {n} columns"
        )
    out = [0] * m
    for j in range(n):  # for each column of A
        weight = v[j]
        for i in range(m):  # add that many copies of it into the running total
            out[i] += weight * A[i][j]
    return out


def matmul_columns(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
    """Version 3: column j of the answer is A applied to column j of B.

    Matrix-matrix multiplication is matrix-vector multiplication done once per
    column, and nothing more. Once you believe `matvec`, this needs no separate
    justification — which is the reason to write it this way.
    """
    check_multipliable(A, B)
    columns = [matvec(A, column) for column in transpose(B)]
    return transpose(columns)


def transpose(M: list[list[float]]) -> list[list[float]]:
    """Swap rows and columns: an (r, c) matrix becomes (c, r). From Day 100."""
    rows, cols = shape(M)
    return [[M[i][j] for i in range(rows)] for j in range(cols)]


def identity(n: int) -> list[list[float]]:
    """The n by n matrix that does nothing: 1 on the diagonal, 0 elsewhere.

    "Does nothing" is the definition worth carrying, not the picture. Applied
    to a vector it returns that same vector, because the weighted sum of its
    columns picks out exactly one column per coordinate.
    """
    if n < 1:
        raise ValueError("the identity matrix needs at least one row")
    return [[1 if i == j else 0 for j in range(n)] for i in range(n)]


def add_bias(M: list[list[float]], bias: list[float]) -> list[list[float]]:
    """Add one bias vector to EVERY row. NumPy would do this by broadcasting.

    Written out, so that the two-character version in the NumPy script is
    visibly shorthand for a loop rather than magic.
    """
    rows, cols = shape(M)
    if len(bias) != cols:
        raise ShapeMismatch(
            f"the bias has {len(bias)} entries but each row has {cols}; "
            "there must be one bias per output column"
        )
    return [[M[i][j] + bias[j] for j in range(cols)] for i in range(rows)]


def multiplication_count(m: int, n: int, p: int) -> int:
    """How many multiplications an (m, n) @ (n, p) costs: one per (i, k, j)."""
    return m * n * p


def chain_costs(m: int, n: int, p: int, q: int) -> tuple[int, int]:
    """Cost of ((A B) C) against (A (B C)) for (m,n) @ (n,p) @ (p,q).

    Both associations give the identical answer — that is associativity, and it
    is a theorem, not a coincidence. They do not give the identical amount of
    work, and on realistic shapes the gap is not a rounding difference.
    """
    left_first = multiplication_count(m, n, p) + multiplication_count(m, p, q)
    right_first = multiplication_count(n, p, q) + multiplication_count(m, n, q)
    return left_first, right_first

"""The reference suite: real values, real shapes, real exception types.

Run from the LAB DIRECTORY:

    .venv/bin/pytest examples -q -p no:cacheprovider

Nothing here asserts a duration. Timings live in 05_cost_and_speed.py, where
they are printed as observations from one machine; a test that asserts a
millisecond figure is a test that fails on somebody else's laptop for no good
reason. The one performance claim made here is a wide margin, and it is in
test_the_gap_is_wide_not_marginal at the foot of the file.
"""

import numpy as np
import pytest

from dataset import (
    ANGLE_A,
    ANGLE_A_B,
    ANGLE_B,
    ANGLE_DEGREES,
    BIAS,
    BIG_CHAIN,
    BIG_LEFT_FIRST,
    BIG_RATIO,
    BIG_RIGHT_FIRST,
    DOT_U,
    DOT_U_U,
    DOT_U_V,
    DOT_U_W,
    DOT_V,
    DOT_W,
    FLIP_AFTER_ROT,
    FLIP_AFTER_ROT_V,
    FLIP_V,
    FLIP_X,
    HIGHLIGHT_CELL,
    HIGHLIGHT_COLUMN,
    HIGHLIGHT_ROW,
    HIGHLIGHT_TERMS,
    HIGHLIGHT_VALUE,
    LAYER_OUT,
    P,
    P_AT_Q,
    P_TIMES_Q,
    Q,
    Q_AT_P,
    ROT90,
    ROT_AFTER_FLIP,
    ROT_AFTER_FLIP_V,
    ROT_TIMES_FLIP_ELEMENTWISE,
    ROT_V,
    SHAPE_CASES,
    SMALL_CHAIN,
    SMALL_LEFT_FIRST,
    SMALL_RIGHT_FIRST,
    U,
    V,
    W,
    X,
    X_AT_U,
    X_AT_XT,
    X_TIMES_U,
    XT_AT_X,
    XW,
)
from matmul import (
    ShapeMismatch,
    add_bias,
    chain_costs,
    dot,
    identity,
    matmul_columns,
    matmul_dots,
    matmul_loops,
    matvec,
    multiplication_count,
    shape,
    transpose,
)

TOL = 1e-12

IMPLEMENTATIONS = [matmul_loops, matmul_dots, matmul_columns]


# -- The dot product ---------------------------------------------------------


def test_dot_multiplies_pairwise_and_adds():
    assert dot(DOT_U, DOT_V) == DOT_U_V == 24


def test_dot_of_a_vector_with_itself_is_its_squared_length():
    assert dot(DOT_U, DOT_U) == DOT_U_U == 25
    assert dot(DOT_U, DOT_U) == pytest.approx(float(np.linalg.norm(DOT_U)) ** 2, abs=1e-9)


def test_dot_is_zero_exactly_when_the_vectors_are_perpendicular():
    assert dot(DOT_U, DOT_W) == DOT_U_W == 0
    cosine = dot(DOT_U, DOT_W) / (np.linalg.norm(DOT_U) * np.linalg.norm(DOT_W))
    assert np.degrees(np.arccos(cosine)) == pytest.approx(90.0, abs=1e-9)


def test_dot_agrees_with_the_geometric_formula_on_a_named_angle():
    """|a| |b| cos(theta) and the pairwise sum are the same number."""
    assert dot(ANGLE_A, ANGLE_B) == ANGLE_A_B == 2
    cosine = dot(ANGLE_A, ANGLE_B) / (np.linalg.norm(ANGLE_A) * np.linalg.norm(ANGLE_B))
    assert np.degrees(np.arccos(cosine)) == pytest.approx(ANGLE_DEGREES, abs=1e-9)
    geometric = np.linalg.norm(ANGLE_A) * np.linalg.norm(ANGLE_B) * cosine
    assert geometric == pytest.approx(dot(ANGLE_A, ANGLE_B), abs=1e-9)


def test_dot_is_commutative_even_though_matrix_multiplication_is_not():
    assert dot(DOT_U, DOT_V) == dot(DOT_V, DOT_U)


def test_dot_refuses_mismatched_lengths():
    with pytest.raises(ShapeMismatch):
        dot([1, 2, 3], [1, 2])


def test_dot_agrees_with_numpy():
    assert dot(DOT_U, DOT_V) == np.dot(DOT_U, DOT_V)
    assert dot(DOT_U, DOT_V) == np.array(DOT_U) @ np.array(DOT_V)


# -- Matrix times vector, as a combination of columns ------------------------


def test_matvec_is_a_weighted_sum_of_the_columns():
    A = [[2, 0], [-1, 1], [0, 4]]
    assert matvec(A, [3, 5]) == [6, 2, 20]


def test_matvec_agrees_with_numpy():
    A = [[2, 0], [-1, 1], [0, 4]]
    assert matvec(A, [3, 5]) == (np.array(A) @ np.array([3, 5])).tolist()


def test_matvec_of_a_basis_vector_returns_that_column():
    """The columns of a matrix ARE the images of the basis vectors."""
    assert matvec(ROT90, [1, 0]) == [0, 1]
    assert matvec(ROT90, [0, 1]) == [-1, 0]
    assert matvec(ROT90, [1, 0]) == [row[0] for row in ROT90]
    assert matvec(ROT90, [0, 1]) == [row[1] for row in ROT90]


def test_matvec_output_length_is_the_row_count():
    A = [[2, 0], [-1, 1], [0, 4]]
    assert shape(A) == (3, 2)
    assert len(matvec(A, [3, 5])) == 3


def test_matvec_refuses_a_vector_of_the_wrong_length():
    with pytest.raises(ShapeMismatch):
        matvec([[2, 0], [-1, 1], [0, 4]], [3, 5, 7])


# -- The three implementations agree with each other and with NumPy ----------


@pytest.mark.parametrize("implementation", IMPLEMENTATIONS)
def test_every_implementation_reproduces_the_hand_worked_product(implementation):
    assert implementation(X, W) == XW


@pytest.mark.parametrize("implementation", IMPLEMENTATIONS)
def test_every_implementation_agrees_with_numpy(implementation):
    assert implementation(X, W) == (np.array(X) @ np.array(W)).tolist()


@pytest.mark.parametrize("m, n, p", [(1, 1, 1), (2, 3, 2), (3, 2, 4), (4, 4, 4), (5, 1, 3), (1, 6, 2)])
def test_all_three_implementations_agree_on_many_shapes(m, n, p):
    rng = np.random.default_rng(m * 100 + n * 10 + p)
    left = rng.integers(-9, 10, size=(m, n)).tolist()
    right = rng.integers(-9, 10, size=(n, p)).tolist()
    expected = (np.array(left) @ np.array(right)).tolist()
    assert matmul_loops(left, right) == expected
    assert matmul_dots(left, right) == expected
    assert matmul_columns(left, right) == expected


def test_the_highlighted_cell_is_row_dotted_with_column():
    i, j = HIGHLIGHT_CELL
    assert [X[i][k] for k in range(3)] == HIGHLIGHT_ROW
    assert [W[k][j] for k in range(3)] == HIGHLIGHT_COLUMN
    assert HIGHLIGHT_TERMS == [X[i][k] * W[k][j] for k in range(3)]
    assert sum(HIGHLIGHT_TERMS) == HIGHLIGHT_VALUE == 13
    assert matmul_loops(X, W)[i][j] == HIGHLIGHT_VALUE


# -- The shape rule ----------------------------------------------------------


@pytest.mark.parametrize("left_shape, right_shape, expected", SHAPE_CASES)
def test_the_shape_rule_predicts_the_result_or_the_failure(left_shape, right_shape, expected):
    rng = np.random.default_rng(sum(left_shape) * 31 + sum(right_shape))
    left = rng.integers(0, 5, size=left_shape)
    right = rng.integers(0, 5, size=right_shape)
    if expected == "error":
        with pytest.raises(ValueError):
            left @ right
        with pytest.raises(ShapeMismatch):
            matmul_loops(left.tolist(), right.tolist())
    else:
        assert (left @ right).shape == expected
        assert shape(matmul_loops(left.tolist(), right.tolist())) == expected


def test_shape_mismatch_is_a_valueerror_so_broad_handlers_still_catch_it():
    assert issubclass(ShapeMismatch, ValueError)
    with pytest.raises(ValueError):
        matmul_loops(X, X)


def test_the_error_message_names_both_shapes_and_the_inner_dimensions():
    with pytest.raises(ShapeMismatch) as caught:
        matmul_loops(X, X)
    message = str(caught.value)
    assert "(2, 3)" in message
    assert "inner dimensions 3 and 2" in message


def test_numpy_raises_valueerror_for_the_same_mismatch():
    with pytest.raises(ValueError) as caught:
        np.array(X) @ np.array(X)
    assert "size 2 is different from 3" in str(caught.value)


def test_the_two_transpose_repairs_give_different_shapes_and_different_answers():
    npX = np.array(X)
    assert (npX @ npX.T).tolist() == X_AT_XT
    assert (npX.T @ npX).tolist() == XT_AT_X
    assert (npX @ npX.T).shape == (2, 2)
    assert (npX.T @ npX).shape == (3, 3)


def test_both_transpose_repairs_are_symmetric():
    npX = np.array(X)
    for repaired in (npX @ npX.T, npX.T @ npX):
        assert np.array_equal(repaired, repaired.T)


# -- Composition and non-commutativity ---------------------------------------


def test_the_product_is_the_composition_for_a_real_vector():
    step = matvec(FLIP_X, V)
    assert step == FLIP_V
    assert matvec(ROT90, step) == ROT_AFTER_FLIP_V
    assert matmul_loops(ROT90, FLIP_X) == ROT_AFTER_FLIP
    assert matvec(ROT_AFTER_FLIP, V) == ROT_AFTER_FLIP_V


def test_the_rightmost_matrix_acts_on_the_vector_first():
    """A @ (B @ v) equals (A @ B) @ v, which is what fixes the order."""
    two_steps = matvec(ROT90, matvec(FLIP_X, V))
    one_step = matvec(matmul_loops(ROT90, FLIP_X), V)
    assert two_steps == one_step == ROT_AFTER_FLIP_V


def test_the_other_order_gives_a_genuinely_different_matrix():
    assert matmul_loops(FLIP_X, ROT90) == FLIP_AFTER_ROT
    assert matmul_loops(ROT90, FLIP_X) != matmul_loops(FLIP_X, ROT90)


def test_the_other_order_sends_the_same_vector_somewhere_else():
    assert matvec(ROT90, V) == ROT_V
    assert matvec(FLIP_X, matvec(ROT90, V)) == FLIP_AFTER_ROT_V
    assert ROT_AFTER_FLIP_V != FLIP_AFTER_ROT_V


def test_non_commutativity_on_a_second_untidy_pair():
    assert matmul_loops(P, Q) == P_AT_Q == [[19, 22], [43, 50]]
    assert matmul_loops(Q, P) == Q_AT_P == [[23, 34], [31, 46]]
    assert P_AT_Q != Q_AT_P


def test_numpy_agrees_that_the_order_matters():
    npP, npQ = np.array(P), np.array(Q)
    assert (npP @ npQ).tolist() == P_AT_Q
    assert (npQ @ npP).tolist() == Q_AT_P
    assert not np.array_equal(npP @ npQ, npQ @ npP)


def test_associativity_holds_on_integers_exactly():
    rng = np.random.default_rng(7)
    A = rng.integers(-5, 6, size=(4, 7))
    B = rng.integers(-5, 6, size=(7, 2))
    C = rng.integers(-5, 6, size=(2, 6))
    assert np.array_equal((A @ B) @ C, A @ (B @ C))


def test_associativity_holds_for_the_from_scratch_implementation_too():
    C = [[1, 1], [0, 2]]
    assert matmul_loops(matmul_loops(ROT90, FLIP_X), C) == matmul_loops(
        ROT90, matmul_loops(FLIP_X, C)
    )


def test_distributivity_over_addition():
    D = [[2, 0], [1, 1]]
    combined = [[FLIP_X[i][j] + D[i][j] for j in range(2)] for i in range(2)]
    left = matmul_loops(ROT90, combined)
    part_one = matmul_loops(ROT90, FLIP_X)
    part_two = matmul_loops(ROT90, D)
    right = [[part_one[i][j] + part_two[i][j] for j in range(2)] for i in range(2)]
    assert left == right


# -- Elementwise versus matrix multiplication --------------------------------


def test_star_and_at_give_different_values_at_the_same_shape():
    npP, npQ = np.array(P), np.array(Q)
    assert (npP * npQ).tolist() == P_TIMES_Q
    assert (npP @ npQ).tolist() == P_AT_Q
    assert (npP * npQ).shape == (npP @ npQ).shape == (2, 2)
    assert not np.array_equal(npP * npQ, npP @ npQ)


def test_star_and_at_give_different_shapes_on_a_matrix_and_a_vector():
    npX, npU = np.array(X), np.array(U)
    assert (npX * npU).shape == (2, 3)
    assert (npX @ npU).shape == (2,)
    assert (npX * npU).tolist() == X_TIMES_U
    assert (npX @ npU).tolist() == X_AT_U


def test_at_is_star_followed_by_a_sum_along_the_last_axis():
    """The one sentence that separates them: `@` sums, `*` does not."""
    npX, npU = np.array(X), np.array(U)
    assert np.array_equal((npX * npU).sum(axis=1), npX @ npU)


def test_the_elementwise_product_of_the_two_transformations_is_all_zeros():
    npA, npB = np.array(ROT90), np.array(FLIP_X)
    assert (npA * npB).tolist() == ROT_TIMES_FLIP_ELEMENTWISE
    assert (npA @ npB).tolist() == ROT_AFTER_FLIP
    assert not np.array_equal(npA * npB, npA @ npB)


def test_dot_matmul_and_the_operator_agree_on_two_dimensional_arrays():
    npX, npW = np.array(X), np.array(W)
    assert np.array_equal(npX @ npW, np.matmul(npX, npW))
    assert np.array_equal(npX @ npW, np.dot(npX, npW))


def test_matmul_and_dot_part_company_on_two_stacks():
    """Checked rather than assumed — they agree on 3-D against 2-D."""
    stack = np.arange(8).reshape(2, 2, 2)
    plain = np.arange(4).reshape(2, 2)
    assert np.matmul(stack, plain).shape == (2, 2, 2)
    assert np.dot(stack, plain).shape == (2, 2, 2)
    assert np.array_equal(np.matmul(stack, plain), np.dot(stack, plain))
    other = np.arange(8).reshape(2, 2, 2)
    assert np.matmul(stack, other).shape == (2, 2, 2)
    assert np.dot(stack, other).shape == (2, 2, 2, 2)


# -- The identity matrix -----------------------------------------------------


def test_identity_is_ones_on_the_diagonal():
    assert identity(3) == [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    assert identity(3) == np.eye(3, dtype=int).tolist()


def test_identity_leaves_a_matrix_alone_from_either_side():
    assert matmul_loops(identity(2), X) == X
    assert matmul_loops(X, identity(3)) == X


def test_the_identity_that_fits_depends_on_the_side():
    """X is (2, 3), so it takes a 2x2 on the left and a 3x3 on the right."""
    with pytest.raises(ShapeMismatch):
        matmul_loops(identity(3), X)
    with pytest.raises(ShapeMismatch):
        matmul_loops(X, identity(2))


def test_identity_leaves_a_vector_alone():
    v = [1.5, -2.0, 0.25]
    assert np.allclose(matvec(identity(3), v), v, atol=TOL)


def test_identity_rejects_a_size_below_one():
    with pytest.raises(ValueError):
        identity(0)


# -- One layer of a neural network -------------------------------------------


def test_the_layer_reproduces_the_hand_worked_output():
    assert (np.array(X) @ np.array(W) + np.array(BIAS)).tolist() == LAYER_OUT


def test_the_layer_from_scratch_matches_the_numpy_version():
    assert add_bias(matmul_loops(X, W), BIAS) == LAYER_OUT


def test_the_bias_is_broadcast_across_rows_not_columns():
    product = np.array(X) @ np.array(W)
    assert product.shape == (2, 2)
    with_bias = product + np.array(BIAS)
    for row in with_bias.tolist():
        assert [row[j] - BIAS[j] for j in range(2)] in product.tolist()


def test_a_bias_of_the_wrong_length_raises():
    product = np.array(X) @ np.array(W)
    with pytest.raises(ValueError) as caught:
        product + np.array([5, -2, 7])
    assert "could not be broadcast" in str(caught.value)
    with pytest.raises(ShapeMismatch):
        add_bias(XW, [5, -2, 7])


def test_growing_the_batch_leaves_the_earlier_rows_unchanged():
    bigger = np.array([[1, 2, 0], [0, 1, 3], [2, 0, 1], [1, 1, 1]])
    out = bigger @ np.array(W) + np.array(BIAS)
    assert out.shape == (4, 2)
    assert out[:2].tolist() == LAYER_OUT


def test_two_layers_without_a_nonlinearity_collapse_into_one():
    npX, npW, npB = np.array(X), np.array(W), np.array(BIAS)
    W2 = np.array([[1, 0, 2], [3, 1, 0]])
    b2 = np.array([0, 1, -1])
    two_layers = (npX @ npW + npB) @ W2 + b2
    collapsed = npX @ (npW @ W2) + (npB @ W2 + b2)
    assert np.array_equal(two_layers, collapsed)


# -- Cost -------------------------------------------------------------------


def test_the_multiplication_count_is_m_times_n_times_p():
    assert multiplication_count(2, 3, 2) == 12
    assert multiplication_count(200, 200, 200) == 8_000_000
    assert multiplication_count(1024, 4096, 4096) == 17_179_869_184


def test_the_small_chain_costs_what_the_hand_count_said():
    assert chain_costs(*SMALL_CHAIN) == (SMALL_LEFT_FIRST, SMALL_RIGHT_FIRST)
    assert SMALL_RIGHT_FIRST // SMALL_LEFT_FIRST == 10


def test_the_adapter_chain_costs_what_the_hand_count_said():
    assert chain_costs(*BIG_CHAIN) == (BIG_LEFT_FIRST, BIG_RIGHT_FIRST)
    assert BIG_RIGHT_FIRST // BIG_LEFT_FIRST == BIG_RATIO == 258


def test_both_associations_of_the_adapter_chain_give_the_same_answer():
    """Cheap and expensive are the same computation, on shapes small enough to run."""
    rng = np.random.default_rng(11)
    batch = rng.integers(0, 4, size=(6, 12))
    down = rng.integers(0, 4, size=(12, 2))
    up = rng.integers(0, 4, size=(2, 12))
    assert np.array_equal((batch @ down) @ up, batch @ (down @ up))


# -- Utilities the rest of the lab leans on ----------------------------------


def test_numpy_integer_products_overflow_silently():
    """NumPy int64 wraps; Python's own integers do not. NumPy is the wrong one.

    Asserted rather than described, including the absence of a warning, because
    "it fails silently" is exactly the kind of claim that rots when a library
    changes. If NumPy ever starts warning here, this test says so.
    """
    import warnings

    big = [[3037000500, 0], [0, 1]]
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        wrapped = (np.array(big) @ np.array(big)).tolist()
    exact = matmul_loops(big, big)

    assert exact[0][0] == 3037000500**2 == 9223372037000250000
    assert wrapped[0][0] == -9223372036709301616
    assert wrapped[0][0] < 0 < exact[0][0]
    assert caught == [], "NumPy raised no warning — that is the dangerous part"


def test_shape_rejects_a_ragged_grid():
    with pytest.raises(ValueError):
        shape([[1, 2, 3], [4, 5]])


def test_transpose_swaps_the_shape_and_undoes_itself():
    assert shape(transpose(X)) == (3, 2)
    assert transpose(transpose(X)) == X
    assert transpose(X) == np.array(X).T.tolist()


# -- The one performance claim, stated as a wide margin ----------------------


def test_the_gap_is_wide_not_marginal():
    """No duration is asserted. Only that the loop loses by a large factor.

    The threshold is deliberately far below what was measured while writing
    this lab, so that a slow or busy machine still passes. If this ever fails,
    the interesting question is not the number but why NumPy is not reaching
    BLAS on your installation.
    """
    import time

    rng = np.random.default_rng(2026)
    size = 120
    left = rng.integers(0, 10, size=(size, size)).astype(np.float64)
    right = rng.integers(0, 10, size=(size, size)).astype(np.float64)
    left_lists, right_lists = left.tolist(), right.tolist()

    start = time.perf_counter()
    loop_answer = matmul_loops(left_lists, right_lists)
    loop_seconds = time.perf_counter() - start

    best = float("inf")
    for _ in range(5):
        start = time.perf_counter()
        numpy_answer = left @ right
        best = min(best, time.perf_counter() - start)

    assert np.allclose(np.array(loop_answer), numpy_answer, atol=1e-9)
    assert loop_seconds / best > 50

lesson_id: D101
day: 101
kind: guided-build
languages: [python, bash]
setup_commands:
  - cd labs/sections/math-statistics-and-data/day-101-matrix-multiplication
  - python3 -m venv .venv
  - .venv/bin/pip install -r requirements/requirements.txt
  - .venv/bin/python3 -c "import numpy; print(numpy.__version__)"
run_commands:
  - 'cd examples && ../.venv/bin/python3 01_matmul_from_scratch.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 02_composition.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 03_star_versus_at.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 04_network_layer.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 05_cost_and_speed.py && cd ..'
  - .venv/bin/pytest examples -q -p no:cacheprovider
  - .venv/bin/pytest starter -q -p no:cacheprovider
test_commands:
  - bash tests/run_tests.sh
cleanup_commands:
  - "find . -type d -name '__pycache__' -prune -exec rm -rf -- {} +"
  - rm -rf .pytest_cache
  - 'rm -rf .venv  # optional: removes the lab virtual environment'
  - 'git checkout -- starter/  # optional: reset your work'
requires_network: true
requires_api_key: false
estimated_minutes: 30
last_executed: '2026-08-16'
executed_on: 'macOS 26.5.2 (Apple Silicon, arm64), Python 3.14.0, numpy 2.5.2, pytest 9.1.1, bash 3.2.57 — bash tests/run_tests.sh -> 58 checks, 0 failure(s), exit 0; pytest examples -> 71 passed; pytest starter -> 1 passed, 56 skipped on an untouched checkout, and 57 passed against a fully solved copy of starter/ kept outside the lab. All five reference scripts exit 0 with every internal assertion holding. Network is needed once to install numpy and pytest; nothing else in the lab opens a socket, and section 7 of the harness greps the sources to prove it. Section 6 re-runs the harness with the network-layer output deliberately swapped for a wrong value and asserts that the re-run exits non-zero and reports exactly one failure, so the suite is demonstrated to be capable of failing rather than merely claimed to be. Timings observed on this machine and asserted only as wide ratios, never as durations: the three-nested-loop implementation took 0.1957 s on two 200x200 matrices, NumPy @ on int64 took 0.002479 s (79x faster) and NumPy @ on float64 took 0.000037 s (5223x faster). The 66x gap between the two NumPy rows was not anticipated when the lab was planned and is the most interesting measurement in it: BLAS implements no integer matrix-multiply routine, so an int64 @ falls back to NumPy own compiled loop while a float64 @ is handed to BLAS. This NumPy reports its BLAS as accelerate, read from numpy.__config__ rather than assumed.'

# Dependencies for the Day 101 lab

Two packages, both free and open source, both installed from the Python Package
Index with `pip`, both running entirely on your own machine.

| Package | Pinned version | Why this lab needs it |
| --- | --- | --- |
| `numpy` | `2.5.2` | The array library the lesson compares against. It supplies `@`, `np.matmul`, `np.dot`, the broadcasting that adds a bias across rows, and — the point of the timing script — the route out to BLAS. |
| `pytest` | `9.1.1` | The test runner from Days 071–074. Nothing new here except what it is pointed at. |

Nothing else is required. Both suites and all five reference scripts run on
those two packages plus the standard library.

## Why the version is pinned, and what is actually checked

The pin is there because the version is *checked* rather than assumed. Section 1
of `tests/run_tests.sh` reads the installed version and compares it against this
file, so a mismatch is reported at the top of the run instead of surfacing later
as a confusing diff.

The version was read from the installed package rather than guessed:

```bash
.venv/bin/python3 -c "from importlib.metadata import version; print(version('numpy'))"
```

On the authoring machine, on 16 August 2026, that printed `2.5.2`.

Every array this lab prints goes through `.tolist()` first, so none of the
captured output in `../expected-output/` depends on how a particular NumPy
version formats an array. A newer NumPy should reproduce every number in this
lab exactly. Only the two timing sections of `05_cost_and_speed.py` will differ,
and they are labelled as machine-specific in `../expected-output/FIELDS.md`.

## The one dependency this lab has that is not in this file

`05_cost_and_speed.py` reports which **BLAS** implementation your NumPy was
built against, read from `numpy.__config__`. BLAS is not a Python package and
you do not install it separately — it arrives inside the NumPy wheel, or NumPy
finds one already on your system.

On the authoring machine that reported `accelerate`, which is Apple's own
implementation. A Linux wheel from the Package Index will usually report
OpenBLAS instead. **Neither is better for the purposes of this lab and neither
is claimed to be**; what matters is that some BLAS is present, because that is
what makes the float64 timing so much faster than the int64 one. If your
installation reports no BLAS at all, the script says so rather than inventing a
name.

## Licences

NumPy is distributed under the BSD 3-Clause licence and pytest under the MIT
licence, each stated on that project's own documentation site. Both are
maintained in the open, cost nothing, and need no account, no key and no
signup — personally or commercially.

## One-time install

From the lab directory:

```bash
python3 -m venv .venv
.venv/bin/pip install -r requirements/requirements.txt
.venv/bin/python3 -c "import numpy; print(numpy.__version__)"
```

Expect `2.5.2`. Day 43 covered `python3 -m venv` in full; this is the same
pattern. The environment lives in `.venv/` inside the lab, is already excluded
from version control, and can be deleted at any time with `rm -rf .venv`.

## Network

Installing needs the network, once. **Nothing else in this lab does.** No
script opens a socket, reads a URL or contacts a service, and section 7 of
`tests/run_tests.sh` greps every file under `examples/` and `starter/` for the
patterns that would indicate otherwise.

If you are offline and already have NumPy available somewhere, you do not need
the install at all — see the next section.

## Running without a lab-local environment

If NumPy and pytest are already available in an environment you have activated,
the harness will find `pytest` on your `PATH`. You can also point it at a
specific binary:

```bash
PYTEST=/path/to/pytest bash tests/run_tests.sh
```

The harness uses the `python3` that sits beside that `pytest`, because that is
the interpreter NumPy is installed into. If NumPy is not importable from it, the
harness says so and stops rather than skipping checks quietly.

## What you would give up without NumPy

Exercise 1 — the whole from-scratch build, all nine functions — needs nothing
but the standard library, and you can complete it on a bare `python3`. That is
most of the lab's work.

Everything after it compares your implementation against NumPy, or demonstrates
something (`*` against `@`, the shape error's exact type, the BLAS timing) that
exists only because NumPy exists. Those parts cannot be faked, and the lab does
not pretend otherwise.

numpy==2.5.2
pytest==9.1.1

# Multiply It Yourself — the six exercises, in order

Work top to bottom. Check yourself at any point from the **lab directory**:

```bash
.venv/bin/pytest starter -q
```

On an untouched checkout that prints `1 passed, 56 skipped`. A skip means
"not attempted". A failure means "attempted and wrong", and it prints both your
answer and the real one. When it prints `57 passed`, you are finished.

Two rules that make this worth doing:

1. **Write the answer down before you run anything.** Every prediction in
   `answers.py` can be done with a pen in under a minute. If you run the code
   first and then fill them in, all of them will be right and you will have
   learned nothing.
2. **Do not import NumPy in `matmul.py`.** Exercise 1 is a from-scratch build.
   Its whole value is that you can read every line and see where each number
   came from.

The matrices used throughout, so you never have to hunt for them:

```
X    = [[1, 2, 0],        (2, 3)   a batch: two examples, three features
        [0, 1, 3]]

W    = [[ 2, 0],          (3, 2)   a layer's weights: three in, two out
        [-1, 1],
        [ 0, 4]]

bias = [5, -2]            (2,)     one per output unit

u    = [10, 2, 5]         (3,)

P    = [[1, 2],           (2, 2)
        [3, 4]]

Q    = [[5, 6],           (2, 2)
        [7, 8]]

A = ROT90  = [[0, -1],    (2, 2)   a quarter turn anticlockwise
              [1,  0]]

B = FLIP_X = [[1,  0],    (2, 2)   a reflection in the horizontal axis
              [0, -1]]

v    = [3, 1]
```

---

## Exercise 1 — build it (`matmul.py`)

Nine functions, each marked `EXERCISE` in the file. Written for you already:
`shape`, `transpose`, and the `ShapeMismatch` exception.

| # | Function | The idea |
| --- | --- | --- |
| 1.1 | `dot` | Multiply pairwise, then add. Six lines. |
| 1.2 | `check_multipliable` | The shape rule, with an error message that names both shapes. |
| 1.3 | `matmul_loops` | Three nested loops. The definition, transcribed. |
| 1.4 | `matvec` | A vector, as a weighted sum of the matrix's **columns**. |
| 1.5 | `matmul_dots` | The same product, as a list of dot products. |
| 1.6 | `identity` | The matrix that does nothing. |
| 1.7 | `add_bias` | Broadcasting, written out as a loop. |
| 1.8 | `multiplication_count` | Count the work: one expression. |
| 1.9 | `chain_costs` | Both associations of a three-matrix chain. |

Two traps worth naming before you meet them.

**In 1.3, build the result grid with `[[0] * p for _ in range(m)]`.** The
shorter-looking `[[0] * p] * m` makes `m` references to **one** row, so writing
to `C[0][0]` changes every row at once. That is the view-versus-copy lesson from
Day 100 turning up in plain Python, and there is a test that catches it by name.

**In 1.4, resist writing "dot v with each row".** It gives the right numbers and
the wrong picture. Write it as columns: take `v[0]` copies of column 0, plus
`v[1]` copies of column 1, and add them up. Then try `matvec(A, [1, 0])` and see
which column comes back. That is the fact everything in Week 15 rests on.

---

## Exercise 2 — the shape rule (`answers.py`)

Eight predictions. For each expression, give the shape of the result, or the
string `"error"`. Derive them from the rule rather than recalling them:

> An `(m, n) @ (n, p)` is legal only when the inner dimensions agree, and the
> result is `(m, p)`.

And one question with a different flavour: which exception **class** does NumPy
raise when they do not agree? Give the class, not its name as a string.

---

## Exercise 3 — composition and order (`answers.py`)

Seven predictions, all doable with a pen — each is four multiplications and two
additions.

Apply `B` to `v`. Then apply `A` to that. Then work out the single matrix
`A @ B` and check it takes `v` to the same place in one step. Then do `B @ A`
and see that it does not.

The question that decides whether you have understood it: **in `A @ B @ v`,
which matrix meets the vector first?** Think about `A @ (B @ v)` before you
answer.

---

## Exercise 4 — `*` against `@` (`answers.py`)

Seven predictions. `P * Q` and `P @ Q` are both legal, both return a `(2, 2)`
array, and are not the same numbers — so no shape check will save you.

Then `X * u` and `X @ u`, where the shapes finally do differ, and one question
that states the whole distinction as code: **`@` is `*` followed by a sum along
which axis?**

---

## Exercise 5 — one layer of a neural network (`answers.py`)

Compute `X @ W + bias` entirely by hand and record both the product and the
final output. Four dot products of length three, then two additions per row.

This is the operation that consumes essentially all the compute in training any
model you will ever use. Doing it once with a pen is the point of the day.

Then three questions about what the shapes mean: is the bias one per example or
one per output; which shapes change when the batch grows; and whether two
layers with no activation function between them collapse into one.

---

## Exercise 6 — cost (`answers.py`)

Seven predictions. Count the multiplications for a small layer and for a pair of
200 by 200 matrices. Then count both associations of the chain
`(10, 100) @ (100, 5) @ (5, 50)` and say which is cheaper — and confirm that
they nonetheless give the same answer.

The last question is the one the timing script is really about. NumPy's `@` on a
`float64` array is dramatically faster than on an `int64` array of the same
shape and the same values. Four explanations are offered; one is true. Run
`05_cost_and_speed.py` if you want the evidence before committing.

---

## When you are done

Read the reference, which is written to be read rather than merely to be
correct. Each script prints its working and asserts every claim it makes:

```bash
cd examples
../.venv/bin/python3 01_matmul_from_scratch.py
../.venv/bin/python3 02_composition.py
../.venv/bin/python3 03_star_versus_at.py
../.venv/bin/python3 04_network_layer.py
../.venv/bin/python3 05_cost_and_speed.py
cd ..
```

Then the full harness:

```bash
bash tests/run_tests.sh
echo "exit=$?"
```

"""Exercises 2 to 6 — predictions. Replace each None with your answer.

The rule that makes this worth doing: WRITE THE ANSWER DOWN BEFORE YOU RUN
ANYTHING. A prediction you check is worth ten outputs you read. If you run the
code first and then fill these in, every one will be right and you will have
learned nothing.

Anything still `None` is reported as SKIPPED, not failed. A failure means you
committed to an answer and it was wrong, and the failure prints both numbers.

All the matrices referred to here are defined at the top of test_starter.py and
repeated in 00_brief.md, so you never have to guess what X or W is.
"""

# ===========================================================================
# EXERCISE 2 — the shape rule
#
#   X is (2, 3)      W is (3, 2)      P and Q are both (2, 2)
#
# For each expression, give the SHAPE of the result as a tuple, or the string
# "error" if it raises. Work them out from the rule, not from memory:
# an (m, n) @ (n, p) is legal only when the inner dimensions agree, and the
# result is (m, p).
# ===========================================================================

SHAPE_OF_X_AT_W = None  # e.g. (2, 2) or "error"
SHAPE_OF_W_AT_X = None
SHAPE_OF_X_AT_X = None
SHAPE_OF_X_AT_X_T = None  # X @ X.T
SHAPE_OF_X_T_AT_X = None  # X.T @ X
SHAPE_OF_P_AT_Q = None

# X is (2, 3) and u is the length-3 vector [10, 2, 5]. What shape is X @ u?
# Careful: u is one-dimensional, so the answer is not a pair.
SHAPE_OF_X_AT_U = None

# Which exception class does NumPy raise for X @ X? Give the CLASS itself,
# not its name as a string — e.g. TypeError, not "TypeError".
SHAPE_ERROR_EXCEPTION = None


# ===========================================================================
# EXERCISE 3 — composition and order
#
#   A = ROT90  = [[0, -1], [1, 0]]    a quarter turn anticlockwise
#   B = FLIP_X = [[1,  0], [0, -1]]   a reflection in the horizontal axis
#   v = [3, 1]
#
# Do these with a pen. Each one is four multiplications and two additions.
# ===========================================================================

# Apply B to v. (Reflection in the horizontal axis flips the sign of y.)
B_TIMES_V = None  # a list of two numbers

# Now apply A to THAT result — so this is A @ (B @ v).
A_TIMES_B_TIMES_V = None

# The single matrix A @ B, worked out entry by entry.
A_AT_B = None  # a list of two lists

# And B @ A, the other order.
B_AT_A = None

# Does (A @ B) @ v land in the same place as A @ (B @ v)? True or False.
COMPOSITION_MATCHES = None

# Does A @ B equal B @ A? True or False.
ORDER_DOES_NOT_MATTER = None

# In the expression A @ B @ v, which matrix meets the vector first?
# Answer with the string "A" or the string "B".
WHICH_ACTS_FIRST = None


# ===========================================================================
# EXERCISE 4 — `*` against `@`
#
#   P = [[1, 2], [3, 4]]     Q = [[5, 6], [7, 8]]
#   X = [[1, 2, 0], [0, 1, 3]]      u = [10, 2, 5]
# ===========================================================================

# P * Q — entry by entry, nothing summed.
P_STAR_Q = None  # a list of two lists

# P @ Q — rows dotted with columns.
P_AT_Q = None

# Are those two the same SHAPE? True or False. (Think before you answer: this
# is the reason the `*` versus `@` mistake survives so long in real code.)
P_STAR_AND_AT_SAME_SHAPE = None

# X * u — u is broadcast across both rows, then multiplied entry by entry.
X_STAR_U_SHAPE = None  # a tuple

# X @ u — each row is multiplied by u and then SUMMED.
X_AT_U_SHAPE = None  # a tuple
X_AT_U_VALUES = None  # a list of numbers

# One sentence, expressed as code: `@` is `*` followed by a sum along which
# axis? Give the integer axis number that turns (X * u) into (X @ u).
AXIS_THAT_TURNS_STAR_INTO_AT = None


# ===========================================================================
# EXERCISE 5 — one layer of a neural network
#
#   X    = [[1, 2, 0],      the batch: two examples, three features each
#            [0, 1, 3]]
#   W    = [[ 2, 0],        the weights: three inputs, two outputs
#           [-1, 1],
#           [ 0, 4]]
#   bias = [5, -2]
#
# Compute X @ W + bias entirely by hand. This is the operation that consumes
# essentially all the compute in training any model, and it is worth having
# done once with a pen.
# ===========================================================================

X_AT_W = None  # a list of two lists
LAYER_OUTPUT = None  # X @ W + bias, a list of two lists

# The bias has 2 entries. Is that one per EXAMPLE or one per OUTPUT?
# Answer with the string "example" or the string "output".
BIAS_IS_ONE_PER = None

# If the batch grew from 2 examples to 64, which of the three shapes changes?
# Answer with a list of the names that change, drawn from "X", "W" and "bias".
# For example ["X", "W"] — but that is not the answer.
SHAPES_THAT_CHANGE_WITH_THE_BATCH = None

# Two layers with no activation function between them collapse into a single
# layer. True or False?
TWO_LINEAR_LAYERS_COLLAPSE = None


# ===========================================================================
# EXERCISE 6 — cost
# ===========================================================================

# How many multiplications does a (2, 3) @ (3, 2) cost?
COST_OF_THE_SMALL_LAYER = None

# How many does a (200, 200) @ (200, 200) cost?
COST_OF_200_SQUARED = None

# For the chain (10, 100) @ (100, 5) @ (5, 50), count both associations.
#   (AB)C  = 10*100*5 + 10*5*50
#   A(BC)  = 100*5*50 + 10*100*50
CHAIN_LEFT_FIRST = None
CHAIN_RIGHT_FIRST = None

# Which association is cheaper here? Answer "(AB)C" or "A(BC)".
CHEAPER_ASSOCIATION = None

# Do the two associations give the same ANSWER, ignoring cost? True or False.
ASSOCIATIONS_AGREE = None

# NumPy's `@` on a float64 array is far faster than on an int64 array of the
# same shape and values. What is the reason? Answer with one of these strings:
#   "floats are smaller"
#   "BLAS only implements floating point"
#   "integers overflow so numpy checks every entry"
#   "numpy converts integers to Python objects"
WHY_FLOAT_BEATS_INT = None

"""Make this directory's own matmul.py the one its tests import.

Both `examples/` and `starter/` contain a module called `matmul`, and pytest
imports test files by putting their directory on `sys.path`. Without this file,
running `pytest` across both directories at once would import whichever
`matmul` was seen first and then reuse it for the other suite — so the starter
tests would silently pass against the reference solution instead of skipping.
That is a wrong answer with a green tick on it, which is the worst kind.

This is not a hypothetical. It happened while building the Day 100 lab, where
eleven unwritten exercises reported as passing, and it was caught only because
the number of skips changed between two runs that should have agreed. Section 4
of `tests/run_tests.sh` now asserts that the skip count is identical whether the
suites are collected separately or together, so the same mistake cannot come
back quietly.

So: put this directory first on the import path, and drop any already-imported
`matmul`, `dataset` or `answers` that came from somewhere else.
"""

import sys
from pathlib import Path

HERE = str(Path(__file__).parent.resolve())

if HERE in sys.path:
    sys.path.remove(HERE)
sys.path.insert(0, HERE)

for name in ("matmul", "dataset", "answers"):
    module = sys.modules.get(name)
    origin = getattr(module, "__file__", "") or ""
    if module is not None and not origin.startswith(HERE):
        del sys.modules[name]

"""Exercise 1 — matrix multiplication from first principles. Your work.

Seven functions to write, each marked EXERCISE. Everything else is written for
you, including `shape`, `transpose` and the `ShapeMismatch` exception, so that
you spend your time on the operation itself rather than on plumbing.

Check yourself from the LAB DIRECTORY at any point:

    .venv/bin/pytest starter -q

A function you have not written raises NotImplementedError, and its tests SKIP
rather than fail. A skip means "not attempted". A failure means "attempted and
wrong", and it prints both your answer and the real one.

Do not import numpy in this file. The whole point of exercise 1 is that you can
read every line and see where each number came from.
"""

from __future__ import annotations


class ShapeMismatch(ValueError):
    """Raised when two matrices cannot be multiplied. Written for you.

    It subclasses ValueError on purpose, so that code which only knows to catch
    ValueError still catches this. NumPy raises a plain ValueError for the same
    situation.
    """


def shape(M: list[list[float]]) -> tuple[int, int]:
    """Return (rows, columns), checking the grid is rectangular. Written for you."""
    if not M or not isinstance(M, list) or not isinstance(M[0], list):
        raise ValueError("a matrix is a non-empty list of rows, each a list")
    n_cols = len(M[0])
    for i, row in enumerate(M):
        if len(row) != n_cols:
            raise ValueError(
                f"row {i} has {len(row)} entries but row 0 has {n_cols}; "
                "a matrix is rectangular"
            )
    return (len(M), n_cols)


def transpose(M: list[list[float]]) -> list[list[float]]:
    """Swap rows and columns. Written for you; you built this on Day 100."""
    rows, cols = shape(M)
    return [[M[i][j] for i in range(rows)] for j in range(cols)]


# ---------------------------------------------------------------------------
# EXERCISE 1.1 — the dot product
# ---------------------------------------------------------------------------


def dot(u: list[float], v: list[float]) -> float:
    """Multiply the two vectors pairwise, then add up the products.

    EXERCISE 1.1

        dot([3, 4], [4, 3])  ->  3*4 + 4*3  ->  24
        dot([3, 4], [3, 4])  ->  3*3 + 4*4  ->  25   (the squared length)
        dot([3, 4], [-4, 3]) ->  -12 + 12   ->  0    (perpendicular)

    Raise ShapeMismatch if the two lengths differ — there is no sensible
    partner for the leftover entries, and returning something anyway would be
    worse than stopping.

    Hint: `zip(u, v)` walks both lists together.
    """
    raise NotImplementedError("exercise 1.1: dot")


# ---------------------------------------------------------------------------
# EXERCISE 1.2 — the shape rule
# ---------------------------------------------------------------------------


def check_multipliable(A: list[list[float]], B: list[list[float]]) -> tuple[int, int, int]:
    """Check that A @ B is legal and return (m, n, p) for an (m, n) @ (n, p).

    EXERCISE 1.2

    Use `shape` on each. If A is (m, n) and B is (n2, p), then n and n2 must be
    equal; raise ShapeMismatch if they are not, and put BOTH shapes and both
    inner dimensions in the message. An error message that does not tell you
    the two numbers that disagreed has wasted the exception.

    The test looks for the substrings "(2, 3)" and "inner dimensions 3 and 2"
    when called with two (2, 3) matrices, so include those exact phrasings.

    Return (m, n, p) when it is legal.
    """
    raise NotImplementedError("exercise 1.2: check_multipliable")


# ---------------------------------------------------------------------------
# EXERCISE 1.3 — three nested loops
# ---------------------------------------------------------------------------


def matmul_loops(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
    """Multiply two matrices with three nested loops. The definition, verbatim.

    EXERCISE 1.3

    Call `check_multipliable` first, then build an m by p grid of zeros and
    fill it. Entry (i, j) is the sum over k of A[i][k] * B[k][j].

        C = [[0] * p for _ in range(m)]

    Write that line exactly. `[[0] * p] * m` looks equivalent and is a bug: it
    makes m references to ONE row, so writing to C[0][0] changes every row at
    once. That is the Day 100 view-versus-copy lesson turning up in plain
    Python, and it catches people every year.

    Count as you write: the innermost line runs m*n*p times. That number is why
    exercise 6 exists.
    """
    raise NotImplementedError("exercise 1.3: matmul_loops")


# ---------------------------------------------------------------------------
# EXERCISE 1.4 — a matrix applied to a vector, as a sum of COLUMNS
# ---------------------------------------------------------------------------


def matvec(A: list[list[float]], v: list[float]) -> list[float]:
    """Apply A to v as a weighted sum of A's columns.

    EXERCISE 1.4

    You could write this as "dot v with each row", and you would get the right
    numbers. Write it the other way instead, because the other way is the
    picture that makes everything later obvious:

        take v[0] copies of A's first column,
        plus v[1] copies of A's second column,
        plus ... and add them all up.

    So: start with a list of m zeros, loop over the columns j, and add
    v[j] * A[i][j] into out[i] for every row i.

    Raise ShapeMismatch if len(v) is not the number of COLUMNS of A. There is
    one weight per column and no spares.

    When you have it, try `matvec(A, [1, 0])` and see which column comes back.
    That is not a trick; it is the reason the columns of a transformation
    matrix are where the basis vectors land.
    """
    raise NotImplementedError("exercise 1.4: matvec")


# ---------------------------------------------------------------------------
# EXERCISE 1.5 — the same product, as a list of dot products
# ---------------------------------------------------------------------------


def matmul_dots(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
    """Entry (i, j) is row i of A dotted with column j of B. Nothing else.

    EXERCISE 1.5

    Call `check_multipliable`, get B's columns with `transpose(B)`, then build
    the answer with your own `dot`. This can be a single comprehension, and it
    should produce exactly the same numbers as `matmul_loops` on every input —
    the tests check that on six different shapes.
    """
    raise NotImplementedError("exercise 1.5: matmul_dots")


# ---------------------------------------------------------------------------
# EXERCISE 1.6 — the identity matrix, and adding a bias
# ---------------------------------------------------------------------------


def identity(n: int) -> list[list[float]]:
    """The n by n matrix with 1 on the main diagonal and 0 everywhere else.

    EXERCISE 1.6

    Raise ValueError for n < 1. Then check the property that actually defines
    it: multiplying by it changes nothing. The tests check both sides, and note
    that X being (2, 3) means the identity that fits on its left is 2 by 2 and
    the one that fits on its right is 3 by 3.
    """
    raise NotImplementedError("exercise 1.6: identity")


def add_bias(M: list[list[float]], bias: list[float]) -> list[list[float]]:
    """Add one bias vector to EVERY row of M.

    EXERCISE 1.7

    NumPy does this with `M + b` and calls it broadcasting. Write it out as a
    loop once, so that the two-character version is visibly shorthand rather
    than magic.

    Raise ShapeMismatch if len(bias) is not the number of COLUMNS of M: there
    is one bias per output, not one per example.
    """
    raise NotImplementedError("exercise 1.7: add_bias")


# ---------------------------------------------------------------------------
# EXERCISE 1.8 — counting the work
# ---------------------------------------------------------------------------


def multiplication_count(m: int, n: int, p: int) -> int:
    """How many multiplications does an (m, n) @ (n, p) cost?

    EXERCISE 1.8

    Look at your own `matmul_loops` and count how many times the innermost
    line runs. One expression, no loops needed.
    """
    raise NotImplementedError("exercise 1.8: multiplication_count")


def chain_costs(m: int, n: int, p: int, q: int) -> tuple[int, int]:
    """Cost of ((A B) C) against (A (B C)) for (m,n) @ (n,p) @ (p,q).

    EXERCISE 1.9

    Return the two totals as a tuple, left-first then right-first.

        (A B) C : the (m, n) @ (n, p) costs one lot, and multiplying that
                  (m, p) result by C costs another.
        A (B C) : the (n, p) @ (p, q) costs one lot, and multiplying A by
                  that (n, q) result costs another.

    Both give the identical answer. They do not cost the identical amount, and
    on the shapes in exercise 6 the gap is not a rounding difference.
    """
    raise NotImplementedError("exercise 1.9: chain_costs")

"""Your running score. Run from the LAB DIRECTORY:

    .venv/bin/pytest starter -q

Anything you have not written yet is SKIPPED, not failed. A skip means "not
attempted"; a failure means "attempted and wrong", and the failure prints both
your answer and the real one.

Float comparisons use numpy.allclose with atol=TOL, stated below.
"""

import numpy as np
import pytest

import answers
from matmul import ShapeMismatch, shape, transpose

TOL = 1e-12

X = [[1, 2, 0], [0, 1, 3]]
W = [[2, 0], [-1, 1], [0, 4]]
BIAS = [5, -2]
U = [10, 2, 5]
P = [[1, 2], [3, 4]]
Q = [[5, 6], [7, 8]]
ROT90 = [[0, -1], [1, 0]]
FLIP_X = [[1, 0], [0, -1]]
V = [3, 1]


def written(fn, *args, **kwargs):
    """Run part of your work, or skip the test if it is not written yet."""
    try:
        return fn(*args, **kwargs)
    except NotImplementedError as exc:
        pytest.skip(f"not written yet: {exc}")


def predicted(name):
    """Read one prediction from answers.py, or skip if it is still None."""
    value = getattr(answers, name)
    if value is None:
        pytest.skip(f"answers.{name} is still unanswered")
    return value


# -- Exercise 0: the environment ---------------------------------------------


def test_0_the_environment_is_ready():
    """Always passes once the install worked. Everything below is your work."""
    assert np.__version__, "numpy is importable"
    # Written for you, and already refusing a ragged grid: a matrix is
    # rectangular, and that is not a detail you get to skip.
    with pytest.raises(ValueError):
        shape([[1, 2, 3], [4, 5]])
    assert transpose(X) == np.array(X).T.tolist()
    assert issubclass(ShapeMismatch, ValueError)


# -- Exercise 1.1: the dot product -------------------------------------------


def test_1_1_dot_multiplies_pairwise_and_adds():
    from matmul import dot

    assert written(dot, [3, 4], [4, 3]) == 24
    assert dot([3, 4], [3, 4]) == 25
    assert dot([3, 4], [-4, 3]) == 0


def test_1_1_dot_agrees_with_numpy_on_longer_vectors():
    from matmul import dot

    written(dot, [1], [1])
    rng = np.random.default_rng(3)
    for _ in range(5):
        a = rng.integers(-9, 10, size=7)
        b = rng.integers(-9, 10, size=7)
        assert dot(a.tolist(), b.tolist()) == int(np.dot(a, b))


def test_1_1_dot_refuses_mismatched_lengths():
    from matmul import dot

    written(dot, [1], [1])
    with pytest.raises(ShapeMismatch):
        dot([1, 2, 3], [1, 2])


# -- Exercise 1.2: the shape rule --------------------------------------------


def test_1_2_check_multipliable_returns_m_n_p():
    from matmul import check_multipliable

    assert written(check_multipliable, X, W) == (2, 3, 2)
    assert check_multipliable(W, X) == (3, 2, 3)


def test_1_2_check_multipliable_rejects_a_mismatch():
    from matmul import check_multipliable

    written(check_multipliable, X, W)
    with pytest.raises(ShapeMismatch):
        check_multipliable(X, X)


def test_1_2_the_message_names_both_shapes_and_the_inner_dimensions():
    from matmul import check_multipliable

    written(check_multipliable, X, W)
    with pytest.raises(ShapeMismatch) as caught:
        check_multipliable(X, X)
    message = str(caught.value)
    assert "(2, 3)" in message, "the message should name the shapes"
    assert "inner dimensions 3 and 2" in message, "and the two numbers that disagreed"


# -- Exercise 1.3: three nested loops ----------------------------------------


def test_1_3_matmul_loops_reproduces_the_hand_worked_product():
    from matmul import matmul_loops

    assert written(matmul_loops, X, W) == [[0, 2], [-1, 13]]


def test_1_3_matmul_loops_agrees_with_numpy_on_six_shapes():
    from matmul import matmul_loops

    written(matmul_loops, X, W)
    for m, n, p in [(1, 1, 1), (2, 3, 2), (3, 2, 4), (4, 4, 4), (5, 1, 3), (1, 6, 2)]:
        rng = np.random.default_rng(m * 100 + n * 10 + p)
        left = rng.integers(-9, 10, size=(m, n)).tolist()
        right = rng.integers(-9, 10, size=(n, p)).tolist()
        assert matmul_loops(left, right) == (np.array(left) @ np.array(right)).tolist(), (
            f"disagreed on ({m}, {n}) @ ({n}, {p})"
        )


def test_1_3_the_rows_of_the_result_are_independent_objects():
    """Catches `[[0] * p] * m`, which makes m references to ONE row."""
    from matmul import matmul_loops

    result = written(matmul_loops, X, W)
    assert result[0] is not result[1], (
        "the two rows are the same object — you built the grid with "
        "[[0] * p] * m, which repeats one row rather than making m of them"
    )


def test_1_3_matmul_loops_refuses_a_shape_mismatch():
    from matmul import matmul_loops

    written(matmul_loops, X, W)
    with pytest.raises(ShapeMismatch):
        matmul_loops(X, X)


# -- Exercise 1.4: a matrix applied to a vector ------------------------------


def test_1_4_matvec_is_a_weighted_sum_of_the_columns():
    from matmul import matvec

    A = [[2, 0], [-1, 1], [0, 4]]
    assert written(matvec, A, [3, 5]) == [6, 2, 20]


def test_1_4_matvec_of_a_basis_vector_returns_that_column():
    from matmul import matvec

    written(matvec, ROT90, [1, 0])
    assert matvec(ROT90, [1, 0]) == [0, 1], "column 0 of ROT90"
    assert matvec(ROT90, [0, 1]) == [-1, 0], "column 1 of ROT90"


def test_1_4_matvec_agrees_with_numpy():
    from matmul import matvec

    written(matvec, X, U)
    assert matvec(X, U) == (np.array(X) @ np.array(U)).tolist()


def test_1_4_matvec_refuses_a_vector_of_the_wrong_length():
    from matmul import matvec

    written(matvec, X, U)
    with pytest.raises(ShapeMismatch):
        matvec(X, [1, 2])


# -- Exercise 1.5: the same product as a list of dot products ----------------


def test_1_5_matmul_dots_agrees_with_matmul_loops():
    from matmul import matmul_dots, matmul_loops

    assert written(matmul_dots, X, W) == [[0, 2], [-1, 13]]
    for m, n, p in [(1, 1, 1), (2, 3, 2), (3, 2, 4), (4, 4, 4), (5, 1, 3), (1, 6, 2)]:
        rng = np.random.default_rng(m * 100 + n * 10 + p)
        left = rng.integers(-9, 10, size=(m, n)).tolist()
        right = rng.integers(-9, 10, size=(n, p)).tolist()
        expected = (np.array(left) @ np.array(right)).tolist()
        assert matmul_dots(left, right) == expected, f"({m}, {n}) @ ({n}, {p})"
        assert matmul_loops(left, right) == expected


# -- Exercise 1.6 and 1.7: identity and bias ---------------------------------


def test_1_6_identity_is_ones_on_the_diagonal():
    from matmul import identity

    assert written(identity, 3) == [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    assert identity(1) == [[1]]


def test_1_6_identity_does_nothing_from_either_side():
    from matmul import identity, matmul_loops

    written(identity, 2)
    written(matmul_loops, X, W)
    assert matmul_loops(identity(2), X) == X
    assert matmul_loops(X, identity(3)) == X


def test_1_6_the_identity_that_fits_depends_on_the_side():
    from matmul import identity, matmul_loops

    written(identity, 2)
    written(matmul_loops, X, W)
    with pytest.raises(ShapeMismatch):
        matmul_loops(identity(3), X)
    with pytest.raises(ShapeMismatch):
        matmul_loops(X, identity(2))


def test_1_6_identity_rejects_a_size_below_one():
    from matmul import identity

    written(identity, 1)
    with pytest.raises(ValueError):
        identity(0)


def test_1_7_add_bias_adds_the_same_vector_to_every_row():
    from matmul import add_bias

    assert written(add_bias, [[0, 2], [-1, 13]], BIAS) == [[5, 0], [4, 11]]


def test_1_7_add_bias_matches_numpy_broadcasting():
    from matmul import add_bias

    written(add_bias, [[0, 2]], BIAS)
    product = np.array(X) @ np.array(W)
    assert add_bias(product.tolist(), BIAS) == (product + np.array(BIAS)).tolist()


def test_1_7_add_bias_refuses_a_bias_of_the_wrong_length():
    from matmul import add_bias

    written(add_bias, [[0, 2]], BIAS)
    with pytest.raises(ShapeMismatch):
        add_bias([[0, 2], [-1, 13]], [5, -2, 7])


# -- Exercise 1.8 and 1.9: counting the work ---------------------------------


def test_1_8_multiplication_count():
    from matmul import multiplication_count

    assert written(multiplication_count, 2, 3, 2) == 12
    assert multiplication_count(200, 200, 200) == 8_000_000
    assert multiplication_count(1024, 4096, 4096) == 17_179_869_184


def test_1_9_chain_costs():
    from matmul import chain_costs

    assert written(chain_costs, 10, 100, 5, 50) == (7_500, 75_000)
    assert chain_costs(1024, 4096, 8, 4096) == (67_108_864, 17_314_086_912)


# -- Exercise 2: the shape rule ----------------------------------------------


def test_2_1_shape_of_X_at_W():
    assert predicted("SHAPE_OF_X_AT_W") == (np.array(X) @ np.array(W)).shape


def test_2_2_shape_of_W_at_X():
    assert predicted("SHAPE_OF_W_AT_X") == (np.array(W) @ np.array(X)).shape


def test_2_3_X_at_X_is_an_error():
    guess = predicted("SHAPE_OF_X_AT_X")
    assert guess == "error", "the inner dimensions are 3 and 2"


def test_2_4_shape_of_X_at_X_transposed():
    assert predicted("SHAPE_OF_X_AT_X_T") == (np.array(X) @ np.array(X).T).shape


def test_2_5_shape_of_X_transposed_at_X():
    assert predicted("SHAPE_OF_X_T_AT_X") == (np.array(X).T @ np.array(X)).shape


def test_2_6_shape_of_P_at_Q():
    assert predicted("SHAPE_OF_P_AT_Q") == (np.array(P) @ np.array(Q)).shape


def test_2_7_shape_of_matrix_times_vector():
    assert predicted("SHAPE_OF_X_AT_U") == (np.array(X) @ np.array(U)).shape


def test_2_8_the_exception_class():
    guess = predicted("SHAPE_ERROR_EXCEPTION")
    assert isinstance(guess, type), "give the class itself, not its name as a string"
    with pytest.raises(guess):
        np.array(X) @ np.array(X)


# -- Exercise 3: composition and order ---------------------------------------


def test_3_1_B_times_v():
    assert predicted("B_TIMES_V") == (np.array(FLIP_X) @ np.array(V)).tolist()


def test_3_2_A_times_B_times_v():
    expected = (np.array(ROT90) @ (np.array(FLIP_X) @ np.array(V))).tolist()
    assert predicted("A_TIMES_B_TIMES_V") == expected


def test_3_3_A_at_B():
    assert predicted("A_AT_B") == (np.array(ROT90) @ np.array(FLIP_X)).tolist()


def test_3_4_B_at_A():
    assert predicted("B_AT_A") == (np.array(FLIP_X) @ np.array(ROT90)).tolist()


def test_3_5_composition_matches():
    two_steps = np.array(ROT90) @ (np.array(FLIP_X) @ np.array(V))
    one_step = (np.array(ROT90) @ np.array(FLIP_X)) @ np.array(V)
    assert predicted("COMPOSITION_MATCHES") == bool(np.array_equal(two_steps, one_step))


def test_3_6_order_does_not_matter_is_false():
    same = np.array_equal(np.array(ROT90) @ np.array(FLIP_X), np.array(FLIP_X) @ np.array(ROT90))
    assert predicted("ORDER_DOES_NOT_MATTER") == bool(same)


def test_3_7_which_matrix_acts_first():
    guess = predicted("WHICH_ACTS_FIRST")
    assert guess in ("A", "B"), 'answer with the string "A" or "B"'
    assert guess == "B", "the rightmost matrix is the one standing next to the vector"


# -- Exercise 4: `*` against `@` ---------------------------------------------


def test_4_1_P_star_Q():
    assert predicted("P_STAR_Q") == (np.array(P) * np.array(Q)).tolist()


def test_4_2_P_at_Q():
    assert predicted("P_AT_Q") == (np.array(P) @ np.array(Q)).tolist()


def test_4_3_star_and_at_have_the_same_shape_here():
    same = (np.array(P) * np.array(Q)).shape == (np.array(P) @ np.array(Q)).shape
    assert predicted("P_STAR_AND_AT_SAME_SHAPE") == bool(same)


def test_4_4_X_star_u_shape():
    assert predicted("X_STAR_U_SHAPE") == (np.array(X) * np.array(U)).shape


def test_4_5_X_at_u_shape_and_values():
    assert predicted("X_AT_U_SHAPE") == (np.array(X) @ np.array(U)).shape
    assert predicted("X_AT_U_VALUES") == (np.array(X) @ np.array(U)).tolist()


def test_4_6_the_axis_that_turns_star_into_at():
    axis = predicted("AXIS_THAT_TURNS_STAR_INTO_AT")
    summed = (np.array(X) * np.array(U)).sum(axis=axis)
    assert np.array_equal(summed, np.array(X) @ np.array(U))


# -- Exercise 5: one layer of a neural network -------------------------------


def test_5_1_X_at_W():
    assert predicted("X_AT_W") == (np.array(X) @ np.array(W)).tolist()


def test_5_2_layer_output():
    assert predicted("LAYER_OUTPUT") == (np.array(X) @ np.array(W) + np.array(BIAS)).tolist()


def test_5_3_the_bias_is_one_per_output():
    guess = predicted("BIAS_IS_ONE_PER")
    assert guess in ("example", "output"), 'answer "example" or "output"'
    assert guess == "output", (
        "the bias has as many entries as the layer has output units; grow the "
        "batch and it does not change"
    )


def test_5_4_only_the_batch_shape_changes():
    guess = predicted("SHAPES_THAT_CHANGE_WITH_THE_BATCH")
    assert isinstance(guess, list), "give a list of names"
    assert guess == ["X"], "W and the bias belong to the layer, not to the batch"


def test_5_5_two_linear_layers_collapse():
    npX, npW, npB = np.array(X), np.array(W), np.array(BIAS)
    W2 = np.array([[1, 0, 2], [3, 1, 0]])
    b2 = np.array([0, 1, -1])
    two_layers = (npX @ npW + npB) @ W2 + b2
    collapsed = npX @ (npW @ W2) + (npB @ W2 + b2)
    assert predicted("TWO_LINEAR_LAYERS_COLLAPSE") == bool(np.array_equal(two_layers, collapsed))


# -- Exercise 6: cost --------------------------------------------------------


def test_6_1_cost_of_the_small_layer():
    assert predicted("COST_OF_THE_SMALL_LAYER") == 2 * 3 * 2


def test_6_2_cost_of_200_squared():
    assert predicted("COST_OF_200_SQUARED") == 200 * 200 * 200


def test_6_3_the_two_associations():
    assert predicted("CHAIN_LEFT_FIRST") == 10 * 100 * 5 + 10 * 5 * 50
    assert predicted("CHAIN_RIGHT_FIRST") == 100 * 5 * 50 + 10 * 100 * 50


def test_6_4_which_association_is_cheaper():
    guess = predicted("CHEAPER_ASSOCIATION")
    assert guess in ("(AB)C", "A(BC)"), 'answer "(AB)C" or "A(BC)"'
    left = 10 * 100 * 5 + 10 * 5 * 50
    right = 100 * 5 * 50 + 10 * 100 * 50
    assert guess == ("(AB)C" if left < right else "A(BC)")


def test_6_5_the_associations_agree_on_the_answer():
    rng = np.random.default_rng(11)
    A = rng.integers(0, 4, size=(6, 12))
    B = rng.integers(0, 4, size=(12, 2))
    C = rng.integers(0, 4, size=(2, 12))
    assert predicted("ASSOCIATIONS_AGREE") == bool(np.array_equal((A @ B) @ C, A @ (B @ C)))


def test_6_6_why_float_beats_int():
    guess = predicted("WHY_FLOAT_BEATS_INT")
    assert guess == "BLAS only implements floating point", (
        "BLAS has no integer matrix-multiply routine, so an int64 `@` falls "
        "back to NumPy's own compiled loop. That loop is still far faster than "
        "interpreted Python and still far slower than BLAS."
    )

#!/usr/bin/env bash
# Tests for the Day 101 lab. Run from the lab directory:
#   bash tests/run_tests.sh
#
# The harness proves the lesson's claims by running code and reading real
# values, never by reading source:
#
#   * three from-scratch implementations of matrix multiplication — nested
#     loops, a list of dot products, and a weighted sum of columns — agree with
#     each other and with NumPy's @ on six different shapes;
#   * one output cell equals the row-times-column arithmetic a human did by
#     hand, digit for digit;
#   * A @ B and B @ A are both defined, both computed in full, and different;
#   * `*` and `@` on the same operands give different values at the same shape,
#     and different SHAPES on a matrix and a vector;
#   * a deliberate shape error raises ValueError, and each of the two transpose
#     repairs gives a different shape and a different answer;
#   * one network layer, X @ W + b, matches the hand-computed output, and a
#     wrong-length bias raises;
#   * the loop loses to NumPy by a wide margin — a margin, never a duration;
#   * nothing is left behind on disk.
#
# Everything except the one-time install runs offline. Nothing binds a port,
# nothing writes outside the lab, nothing needs a key. Deterministic,
# non-interactive, exits 0 only if every check passes.
set -u

export PYTHONDONTWRITEBYTECODE=1

lab_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

# Bytecode left by an EARLIER command is not this run's litter. The README
# documents `pytest starter -q`, and running it writes .pyc files that would
# then fail the cleanliness check at the end of this script -- failing the
# reader for following the instructions. Clearing them here makes that final
# check measure what it claims to: what THIS run left behind. `.venv` is
# untouched, because the packages' own bytecode is theirs, not ours.
find "${lab_dir}" -name '.venv' -prune -o -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null || true
find "${lab_dir}" -name '.venv' -prune -o -type d -name '.pytest_cache' -exec rm -rf {} + 2>/dev/null || true

failures=0
checks=0

check() {
  local label="$1" ok="$2"
  checks=$((checks + 1))
  if [ "${ok}" = "yes" ]; then
    echo "  ok: ${label}"
  else
    echo "  FAIL: ${label}"
    failures=$((failures + 1))
  fi
}

check_eq() {
  # check_eq <label> <expected> <actual>
  if [ "$2" = "$3" ]; then
    check "$1" "yes"
  else
    check "$1 (expected [$2], got [$3])" "no"
  fi
}

# Resolve pytest: an explicit override, then this lab's .venv, then PATH.
# Fails loudly with instructions rather than silently skipping checks.
resolve_tool() {
  local tool="$1" override="$2"
  if [ -n "${override}" ] && [ -x "${override}" ]; then echo "${override}"; return 0; fi
  if [ -x "${lab_dir}/.venv/bin/${tool}" ]; then echo "${lab_dir}/.venv/bin/${tool}"; return 0; fi
  if command -v "${tool}" >/dev/null 2>&1; then command -v "${tool}"; return 0; fi
  return 1
}

pytest_bin="$(resolve_tool pytest "${PYTEST:-}")" || {
  echo "FAIL: pytest not found." >&2
  echo "  Install the lab's dependencies with:" >&2
  echo "    python3 -m venv .venv" >&2
  echo "    .venv/bin/pip install -r requirements/requirements.txt" >&2
  echo "  Or point this suite at an existing pytest:" >&2
  echo "    PYTEST=/path/to/pytest bash tests/run_tests.sh" >&2
  exit 1
}

# The Python that owns that pytest is the one with numpy installed.
python_bin="$(dirname "${pytest_bin}")/python3"
if [ ! -x "${python_bin}" ]; then
  python_bin="$(command -v python3 || true)"
fi
if [ -z "${python_bin}" ]; then
  echo "FAIL: python3 not found on PATH." >&2
  exit 1
fi

if ! "${python_bin}" -c "import numpy" >/dev/null 2>&1; then
  echo "FAIL: numpy is not importable from ${python_bin}." >&2
  echo "  Install the lab's dependencies with:" >&2
  echo "    python3 -m venv .venv" >&2
  echo "    .venv/bin/pip install -r requirements/requirements.txt" >&2
  exit 1
fi

echo "Day 101 — Multiply It Yourself"
echo

# --------------------------------------------------------------------------
echo "1. The tools and the versions this lab was written against"
# --------------------------------------------------------------------------

versions="$("${python_bin}" - <<'PY'
import platform
import sys
from importlib.metadata import version

print(f"python   {platform.python_version()}")
for name in ("numpy", "pytest"):
    print(f"{name:<8} {version(name)}")
print(f"platform {platform.platform()}")
print(f"exe      {sys.executable.rsplit('/', 3)[-1]}")
PY
)"
echo "${versions}" | sed 's/^/  /'

pinned_numpy="$(grep -E '^numpy==' "${lab_dir}/requirements/requirements.txt" | cut -d= -f3)"
installed_numpy="$("${python_bin}" -c "from importlib.metadata import version; print(version('numpy'))")"
check_eq "installed numpy matches requirements.txt" "${pinned_numpy}" "${installed_numpy}"

major="$("${python_bin}" -c "import numpy; print(numpy.__version__.split('.')[0])")"
check_eq "numpy is version 2 or later" "2" "${major}"

# --------------------------------------------------------------------------
echo
echo "2. Every reference script runs and every assertion inside it holds"
# --------------------------------------------------------------------------

for script in 01_matmul_from_scratch 02_composition 03_star_versus_at \
              04_network_layer 05_cost_and_speed; do
  out="$(cd "${lab_dir}/examples" && "${python_bin}" "${script}.py" 2>&1)"
  status=$?
  if [ "${status}" -ne 0 ]; then
    check "${script}.py exits 0" "no"
    echo "${out}" | tail -5 | sed 's/^/      /'
  else
    check "${script}.py exits 0" "yes"
  fi
  case "${out}" in
    *"${script}.py: every assertion held."*)
      check "${script}.py reports every assertion held" "yes" ;;
    *) check "${script}.py reports every assertion held" "no" ;;
  esac
done

# --------------------------------------------------------------------------
echo
echo "3. The reference pytest suite: real values, real shapes"
# --------------------------------------------------------------------------

ref_out="$(cd "${lab_dir}" && "${pytest_bin}" examples -q -p no:cacheprovider 2>&1)"
ref_status=$?
echo "${ref_out}" | tail -3 | sed 's/^/  /'
if [ "${ref_status}" -eq 0 ]; then
  check "pytest examples exits 0" "yes"
else
  check "pytest examples exits 0" "no"
fi
case "${ref_out}" in
  *" failed"*) check "no test in the reference suite failed" "no" ;;
  *)           check "no test in the reference suite failed" "yes" ;;
esac
ref_passed="$(printf '%s\n' "${ref_out}" | grep -o '[0-9][0-9]* passed' | head -1 | cut -d' ' -f1)"
if [ "${ref_passed:-0}" -ge 60 ]; then
  check "the reference suite ran at least 60 tests (ran ${ref_passed})" "yes"
else
  check "the reference suite ran at least 60 tests (ran ${ref_passed:-0})" "no"
fi

# --------------------------------------------------------------------------
echo
echo "4. The starter suite skips unattempted work instead of failing it"
# --------------------------------------------------------------------------

start_out="$(cd "${lab_dir}" && "${pytest_bin}" starter -q -p no:cacheprovider 2>&1)"
start_status=$?
echo "${start_out}" | tail -3 | sed 's/^/  /'
if [ "${start_status}" -eq 0 ]; then
  check "pytest starter exits 0 on an untouched checkout" "yes"
else
  check "pytest starter exits 0 on an untouched checkout" "no"
fi
case "${start_out}" in
  *" failed"*) check "the starter suite reports no failures" "no" ;;
  *)           check "the starter suite reports no failures" "yes" ;;
esac
case "${start_out}" in
  *skipped*) check "unwritten exercises are reported as skipped, not passed" "yes" ;;
  *) check "unwritten exercises are reported as skipped, not passed" "no" ;;
esac

# The import guard, and the reason it exists. Both directories contain a module
# called `matmul`, and pytest imports test files by putting their directory on
# sys.path — so collecting both suites at once would otherwise let the starter
# tests import the REFERENCE solution and report unwritten exercises as
# passing. That is exactly what happened while building the Day 100 lab: eleven
# unwritten exercises passed against the reference, and it was caught only
# because the skip count changed. Each directory's conftest.py prevents it.
# This check proves it still does: across both suites the skip count must be
# unchanged.
both_out="$(cd "${lab_dir}" && "${pytest_bin}" -q -p no:cacheprovider 2>&1)"
start_skipped="$(printf '%s\n' "${start_out}" | grep -o '[0-9][0-9]* skipped' | head -1 | cut -d' ' -f1)"
both_skipped="$(printf '%s\n' "${both_out}" | grep -o '[0-9][0-9]* skipped' | head -1 | cut -d' ' -f1)"
check_eq "collecting both suites at once does not turn skips into passes" \
  "${start_skipped:-none}" "${both_skipped:-none}"

# --------------------------------------------------------------------------
echo
echo "5. The lesson's claims, checked one value at a time"
# --------------------------------------------------------------------------

facts="$(cd "${lab_dir}/examples" && "${python_bin}" - <<'PY'
import time

import numpy as np

from dataset import BIAS, FLIP_X, P, Q, ROT90, U, V, W, X
from matmul import (
    ShapeMismatch,
    add_bias,
    chain_costs,
    dot,
    identity,
    matmul_columns,
    matmul_dots,
    matmul_loops,
    matvec,
    multiplication_count,
)

npX, npW, npB = np.array(X), np.array(W), np.array(BIAS)
npP, npQ = np.array(P), np.array(Q)
npA, npF = np.array(ROT90), np.array(FLIP_X)

print("dot_pairwise", dot([3, 4], [4, 3]))
print("dot_self_is_squared_length", dot([3, 4], [3, 4]), round(float(np.linalg.norm([3, 4])) ** 2))
print("dot_perpendicular", dot([3, 4], [-4, 3]))

print("three_implementations",
      matmul_loops(X, W), matmul_dots(X, W), matmul_columns(X, W), (npX @ npW).tolist())
print("highlighted_cell", matmul_loops(X, W)[1][1], 0 * 0 + 1 * 1 + 3 * 4)
print("matvec_is_column_combination", matvec([[2, 0], [-1, 1], [0, 4]], [3, 5]))
print("columns_are_basis_images", matvec(ROT90, [1, 0]), matvec(ROT90, [0, 1]))

print("result_shape", (npX @ npW).shape)
try:
    npX @ npX
except ValueError as exc:
    print("shape_error", type(exc).__name__, "size 2 is different from 3" in str(exc))
else:
    print("shape_error", "NOTHING_RAISED", False)
try:
    matmul_loops(X, X)
except ShapeMismatch as exc:
    print("scratch_shape_error", "ShapeMismatch", "inner dimensions 3 and 2" in str(exc))
else:
    print("scratch_shape_error", "NOTHING_RAISED", False)
print("transpose_repairs", (npX @ npX.T).shape, (npX.T @ npX).shape,
      (npX @ npX.T).tolist(), (npX.T @ npX).tolist())

print("not_commutative", (npA @ npF).tolist(), (npF @ npA).tolist())
print("not_commutative_untidy", (npP @ npQ).tolist(), (npQ @ npP).tolist())
print("composition_two_steps", matvec(ROT90, matvec(FLIP_X, V)))
print("composition_one_step", matvec(matmul_loops(ROT90, FLIP_X), V))
C = [[1, 1], [0, 2]]
print("associative", matmul_loops(matmul_loops(ROT90, FLIP_X), C) ==
      matmul_loops(ROT90, matmul_loops(FLIP_X, C)))

print("star_vs_at_same_shape", (npP * npQ).tolist(), (npP @ npQ).tolist())
print("star_vs_at_elementwise_zeros", (npA * npF).tolist(), (npA @ npF).tolist())
print("star_shape", (npX * np.array(U)).shape, "at_shape", (npX @ np.array(U)).shape)
print("star_values", (npX * np.array(U)).tolist(), "at_values", (npX @ np.array(U)).tolist())
print("at_is_star_then_sum", np.array_equal((npX * np.array(U)).sum(axis=1), npX @ np.array(U)))

print("identity_noop", matmul_loops(identity(2), X) == X, matmul_loops(X, identity(3)) == X)

print("layer_out", (npX @ npW + npB).tolist())
print("layer_out_scratch", add_bias(matmul_loops(X, W), BIAS))
try:
    (npX @ npW) + np.array([5, -2, 7])
except ValueError as exc:
    print("bad_bias", type(exc).__name__, "could not be broadcast" in str(exc))
else:
    print("bad_bias", "NOTHING_RAISED", False)

W2 = np.array([[1, 0, 2], [3, 1, 0]])
b2 = np.array([0, 1, -1])
print("layers_collapse", bool(np.array_equal(
    (npX @ npW + npB) @ W2 + b2, npX @ (npW @ W2) + (npB @ W2 + b2))))

print("count_small", multiplication_count(2, 3, 2))
print("count_200", multiplication_count(200, 200, 200))
print("chain_small", chain_costs(10, 100, 5, 50))
print("chain_big", chain_costs(1024, 4096, 8, 4096))
print("chain_big_ratio", 17314086912 // 67108864)

# The timing. A RATIO is asserted, never a duration. The threshold is far
# below what this machine measured, so a slow or busy machine still passes.
size = 120
rng = np.random.default_rng(2026)
left = rng.integers(0, 10, size=(size, size)).astype(np.float64)
right = rng.integers(0, 10, size=(size, size)).astype(np.float64)
start = time.perf_counter()
loop_answer = matmul_loops(left.tolist(), right.tolist())
loop_seconds = time.perf_counter() - start
best = float("inf")
for _ in range(5):
    start = time.perf_counter()
    numpy_answer = left @ right
    best = min(best, time.perf_counter() - start)
print("timing_answers_agree", bool(np.allclose(np.array(loop_answer), numpy_answer, atol=1e-9)))
print("timing_ratio_over_50", bool(loop_seconds / best > 50))
PY
)"

get() { printf '%s\n' "${facts}" | grep "^$1 " | cut -d' ' -f2-; }

check_eq "the dot product multiplies pairwise and adds" "24" "$(get dot_pairwise)"
check_eq "a vector dotted with itself is its squared length" "25 25" \
  "$(get dot_self_is_squared_length)"
check_eq "perpendicular vectors have a dot product of zero" "0" "$(get dot_perpendicular)"
check_eq "loops, dot products, columns and NumPy all agree" \
  "[[0, 2], [-1, 13]] [[0, 2], [-1, 13]] [[0, 2], [-1, 13]] [[0, 2], [-1, 13]]" \
  "$(get three_implementations)"
check_eq "the highlighted cell equals the hand arithmetic 0*0 + 1*1 + 3*4" "13 13" \
  "$(get highlighted_cell)"
check_eq "matrix times vector is a weighted sum of the columns" "[6, 2, 20]" \
  "$(get matvec_is_column_combination)"
check_eq "the columns of a matrix are where the basis vectors land" "[0, 1] [-1, 0]" \
  "$(get columns_are_basis_images)"
check_eq "(2, 3) @ (3, 2) gives (2, 2)" "(2, 2)" "$(get result_shape)"
check_eq "(2, 3) @ (2, 3) raises ValueError naming the mismatch" "ValueError True" \
  "$(get shape_error)"
check_eq "the from-scratch version raises ShapeMismatch naming both inner dimensions" \
  "ShapeMismatch True" "$(get scratch_shape_error)"
check_eq "the two transpose repairs give different shapes and different answers" \
  "(2, 2) (3, 3) [[5, 2], [2, 10]] [[1, 2, 0], [2, 5, 3], [0, 3, 9]]" \
  "$(get transpose_repairs)"
check_eq "A @ B and B @ A are both defined and genuinely different" \
  "[[0, 1], [1, 0]] [[0, -1], [-1, 0]]" "$(get not_commutative)"
check_eq "the same holds on a second, untidier pair" \
  "[[19, 22], [43, 50]] [[23, 34], [31, 46]]" "$(get not_commutative_untidy)"
check_eq "flip then rotate, one step at a time, lands at [1, 3]" "[1, 3]" \
  "$(get composition_two_steps)"
check_eq "the single product matrix lands at the same point" "[1, 3]" \
  "$(get composition_one_step)"
check_eq "multiplication is associative" "True" "$(get associative)"
check_eq "* and @ give different values at the SAME shape" \
  "[[5, 12], [21, 32]] [[19, 22], [43, 50]]" "$(get star_vs_at_same_shape)"
check_eq "the elementwise product of the two transformations is all zeros" \
  "[[0, 0], [0, 0]] [[0, 1], [1, 0]]" "$(get star_vs_at_elementwise_zeros)"
check_eq "* and @ give different SHAPES on a matrix and a vector" \
  "(2, 3) at_shape (2,)" "$(get star_shape)"
check_eq "* and @ give different values there too, not only different shapes" \
  "[[10, 4, 0], [0, 2, 15]] at_values [14, 17]" "$(get star_values)"
check_eq "@ is * followed by a sum along the last axis" "True" "$(get at_is_star_then_sum)"
check_eq "the identity matrix leaves a matrix alone from either side" "True True" \
  "$(get identity_noop)"
# Section 6 re-runs this script with D101_SELF_TEST=1, which swaps ONE
# expectation below for a deliberately wrong one. That is how the harness
# proves it can fail rather than merely asserting that it could.
expected_layer="[[5, 0], [4, 11]]"
if [ -n "${D101_SELF_TEST:-}" ]; then
  expected_layer="[[5, 0], [4, 10]]"   # off by one in the last cell, deliberately
fi
check_eq "one network layer, X @ W + b, matches the hand-computed output" \
  "${expected_layer}" "$(get layer_out)"
check_eq "the from-scratch layer matches it too" "[[5, 0], [4, 11]]" \
  "$(get layer_out_scratch)"
check_eq "a bias of the wrong length raises ValueError" "ValueError True" "$(get bad_bias)"
check_eq "two layers with no activation between them collapse into one" "True" \
  "$(get layers_collapse)"
check_eq "an (m, n) @ (n, p) costs m*n*p multiplications" "12" "$(get count_small)"
check_eq "two 200 by 200 matrices cost eight million" "8000000" "$(get count_200)"
check_eq "the small chain costs 7500 one way and 75000 the other" "(7500, 75000)" \
  "$(get chain_small)"
check_eq "the adapter chain costs 67108864 one way and 17314086912 the other" \
  "(67108864, 17314086912)" "$(get chain_big)"
check_eq "which is a factor of 258" "258" "$(get chain_big_ratio)"
check_eq "the loop and NumPy return the same answer" "True" "$(get timing_answers_agree)"
check_eq "NumPy beats the loop by a wide margin (a ratio, never a duration)" "True" \
  "$(get timing_ratio_over_50)"

# --------------------------------------------------------------------------
echo
echo "6. The harness can actually fail"
# --------------------------------------------------------------------------

# A green test suite proves nothing until you have watched it go red. This
# section re-runs the whole script with one expectation deliberately swapped
# for a wrong layer output, and asserts that the re-run reports the failure and
# exits non-zero. If this section passes, section 5 is not decorative.
if [ -z "${D101_SELF_TEST:-}" ]; then
  self_out="$(D101_SELF_TEST=1 bash "${BASH_SOURCE[0]}" 2>&1)"
  self_status=$?
  if [ "${self_status}" -ne 0 ]; then
    check "a deliberately wrong expectation makes the harness exit non-zero (${self_status})" "yes"
  else
    check "a deliberately wrong expectation makes the harness exit non-zero" "no"
  fi
  case "${self_out}" in
    *"FAIL: one network layer, X @ W + b"*)
      check "the failing check is named in the output with both values" "yes" ;;
    *) check "the failing check is named in the output with both values" "no" ;;
  esac
  case "${self_out}" in
    *", 1 failure(s)."*)
      check "the summary line counts exactly one failure" "yes" ;;
    *) check "the summary line counts exactly one failure" "no" ;;
  esac
else
  echo "  (self-test run: section 6 does not recurse)"
fi

# --------------------------------------------------------------------------
echo
echo "7. Nothing was left behind"
# --------------------------------------------------------------------------

# `.venv` is pruned from the searches below. A virtual environment ships the
# installed packages' own precompiled bytecode -- hundreds of __pycache__
# directories that came with NumPy or pytest and have nothing to do with
# whether THIS lab tidied up after itself. Without the prune, following the
# README's own setup instructions makes this check fail, which reports a
# problem the reader cannot fix and did not cause.
if find "${lab_dir}" -name '.venv' -prune -o -type d -name '__pycache__' -print -quit 2>/dev/null | grep -q .; then
  check "no __pycache__ directory anywhere under the lab after a full run" "no"
else
  check "no __pycache__ directory anywhere under the lab after a full run" "yes"
fi

if find "${lab_dir}" -name '.venv' -prune -o -type d -name '.pytest_cache' -print -quit 2>/dev/null | grep -q .; then
  check "no .pytest_cache directory left under the lab" "no"
else
  check "no .pytest_cache directory left under the lab" "yes"
fi

if grep -rqE 'urlopen|requests\.|socket\.|http://|https://' \
     "${lab_dir}/examples" "${lab_dir}/starter" 2>/dev/null; then
  check "no lab source opens a network connection" "no"
else
  check "no lab source opens a network connection" "yes"
fi

echo
echo "${checks} checks, ${failures} failure(s)."
[ "${failures}" -eq 0 ]