Hands-on lab — Day 104: NumPy: Arrays and Vectorized Thinking

"""What an ndarray is, and what it costs, measured rather than asserted.

Run from inside examples/:

    ../.venv/bin/python3 01_list_versus_array.py

The claim under test: a list of a million integers and an array of a million
integers hold the same numbers and do not cost the same memory, because one is
a million separate objects with a million pointers to them and the other is one
typed block.

This script also refuses to make the measurement the easy way, and says why.
"""

import sys

import numpy as np

import dataset
from vectorize import array_bytes, describe, list_bytes


def main() -> None:
    print("01_list_versus_array.py")
    print("=" * 70)

    N = dataset.N_BIG

    # -- 1. The same numbers, two ways ----------------------------------------
    print()
    print("1. A million integers, held two ways")
    print("-" * 70)
    values = list(range(N))
    array = np.arange(N, dtype=np.int64)
    print(f"  list  : {values[:5]} ... {values[-3:]}   len {len(values):,}")
    print(f"  array : {array[:5]} ... {array[-3:]}   size {array.size:,}")
    print(f"  same numbers: {values[:1000] == array[:1000].tolist()} (first thousand)")
    assert values[:1000] == array[:1000].tolist()

    # -- 2. The naive measurement, and why it lies ----------------------------
    print()
    print("2. The measurement almost everyone makes first")
    print("-" * 70)
    naive_list = sys.getsizeof(values)
    print(f"  sys.getsizeof(list)   {naive_list:>12,} bytes")
    print(f"  array.nbytes          {array.nbytes:>12,} bytes")
    print(f"  ratio                 {naive_list / array.nbytes:>12.4f}")
    print()
    print("  Read that ratio again. It says the list costs the SAME as the")
    print("  array, which is the opposite of what every NumPy tutorial")
    print("  promises -- including this one. The measurement is wrong, not the")
    print("  promise.")
    print()
    print("  sys.getsizeof measures the LIST OBJECT: a header plus one 8-byte")
    print("  pointer per element. It does not measure the integers, because")
    print("  the list does not own them. They are a million separate objects")
    print("  sitting elsewhere in memory, and they are where the cost is.")
    assert abs(naive_list / array.nbytes - 1.0) < 0.01, (
        "on this build the pointer array and the int64 block are the same size"
    )

    # -- 3. The honest measurement --------------------------------------------
    print()
    print("3. Counting what the list actually costs")
    print("-" * 70)
    one_int = sys.getsizeof(values[999])
    print(f"  sys.getsizeof(one Python int)     {one_int:>12,} bytes")
    print(f"  one int64 element in the array    {array.itemsize:>12,} bytes")
    print()
    honest_list = list_bytes(values)
    honest_array = array_bytes(array)
    payload = honest_list - naive_list
    print(f"  the list's pointers               {naive_list:>12,} bytes")
    print(f"  the integers they point at        {payload:>12,} bytes")
    print(f"  list total                        {honest_list:>12,} bytes")
    print(f"  array total                       {honest_array:>12,} bytes")
    print(f"  the array is                      {honest_list / honest_array:>12.2f}x smaller")
    assert honest_list == 36_000_056, honest_list
    assert honest_array == 8_000_000, honest_array
    assert honest_list / honest_array > 4.0

    print()
    print("  A Python int is 28 bytes because it is a full object: a reference")
    print("  count, a pointer to its type, a length, and only then the digits.")
    print("  An int64 in an array is 8 bytes because it is 8 bytes.")

    # -- 4. Why the small integers are counted once ---------------------------
    print()
    print("4. A detail that would otherwise overcount")
    print("-" * 70)
    a, b = int("100"), int("100")
    c, d = int("1000"), int("1000")
    print(f'  int("100")  is int("100")   -> {a is b}   (cached: -5 to 256)')
    print(f'  int("1000") is int("1000")  -> {c is d}   (built fresh each time)')
    print()
    print("  The values are built from strings on purpose. Writing the literal")
    print("  1000 twice in one function gives you the same object both times,")
    print("  because the compiler folds equal constants in a code object into")
    print("  one -- which would have made the cache look far bigger than it is.")
    print()
    print("  list_bytes counts each DISTINCT integer object once, so the")
    print("  cached ones are not charged a million times over.")
    assert a is b, "small integers are cached"
    assert c is not d, "1000 is above the cache, so these are two objects"
    distinct = len({id(x) for x in values})
    print(f"  distinct integer objects in the list: {distinct:,} of {N:,}")
    assert distinct == N, "range() builds a new object for every value above 256"

    # -- 5. The three things that make an ndarray different -------------------
    print()
    print("5. dtype, shape, strides -- the three facts a list does not have")
    print("-" * 70)
    grid = np.arange(12).reshape(3, 4)
    print("  a 3 by 4 array of the numbers 0 to 11:")
    for row in grid:
        print(f"    {row}")
    print()
    print(f"  {describe(grid)}")
    print()
    print("  dtype   int64: every element is the same type, decided once,")
    print("          which is what lets the loop live in C.")
    print("  shape   (3, 4): 3 rows of 4, laid out end to end in ONE block.")
    print("  strides (32, 8): to step one row, skip 32 bytes; one column, 8.")
    print("          Four int64 is 32 bytes, so a row IS four elements along.")
    print("  A list of lists has none of these. It has pointers to lists of")
    print("  pointers to integers, scattered wherever the allocator put them.")
    assert grid.dtype == np.int64
    assert grid.shape == (3, 4)
    assert grid.strides == (32, 8)
    assert grid.flags["C_CONTIGUOUS"]

    # -- 6. What contiguity buys ----------------------------------------------
    print()
    print("6. One block, so a transpose costs nothing")
    print("-" * 70)
    transposed = grid.T
    print(f"  grid.T shape      {transposed.shape}")
    print(f"  grid.T strides    {transposed.strides}   <- the two swapped over")
    print(f"  shares memory     {np.shares_memory(grid, transposed)}")
    print(f"  C contiguous      {transposed.flags['C_CONTIGUOUS']}")
    print(f"  F contiguous      {transposed.flags['F_CONTIGUOUS']}")
    print()
    print("  Nothing was copied. NumPy swapped two numbers in the strides and")
    print("  handed back a new way of reading the same bytes. That is what a")
    print("  view is, and section 6 of script 06 is about the bug it causes.")
    assert np.shares_memory(grid, transposed)
    assert transposed.strides == (8, 32)

    print()
    print("=" * 70)
    print("01_list_versus_array.py: every assertion held.")


if __name__ == "__main__":
    main()

"""The dtype is a promise, and a promise you can break by accident.

Run from inside examples/:

    ../.venv/bin/python3 02_dtypes_and_overflow.py

The claim under test: a Python int grows to whatever size it needs and an int8
does not. 127 + 1 in an int8 array is -128. Nothing raises, nothing warns, and
the wrong number goes on down the pipeline.
"""

import warnings

import numpy as np

import dataset
from vectorize import wrap_int8


def main() -> None:
    print("02_dtypes_and_overflow.py")
    print("=" * 70)

    # -- 1. Python integers do not overflow -----------------------------------
    print()
    print("1. What you are used to")
    print("-" * 70)
    big = 2 ** 200
    print(f"  2 ** 200 = {big}")
    print(f"  that is {big.bit_length()} bits, and Python simply allocated them")
    print("  A Python int is a variable-length object. It grows. It has no")
    print("  maximum. That convenience is exactly what an array gives up.")
    assert big.bit_length() == 201

    # -- 2. The dtype is a fixed-width promise --------------------------------
    print()
    print("2. What an array promises instead")
    print("-" * 70)
    info = np.iinfo(np.int8)
    wide = np.iinfo(np.int64)
    print(f"  np.iinfo(np.int8)   min {info.min}   max {info.max}")
    print(f"  np.iinfo(np.int64)  min {wide.min}")
    print(f"                      max {wide.max}")
    print()
    for dtype in (np.int8, np.int16, np.int32, np.int64, np.float32, np.float64):
        a = np.zeros(3, dtype=dtype)
        print(f"  {str(a.dtype):<10} itemsize {a.itemsize} bytes   3 elements = {a.nbytes:>2} bytes")
    assert info.min == dataset.INT8_MIN
    assert info.max == dataset.INT8_MAX

    # -- 3. The wrap ----------------------------------------------------------
    print()
    print("3. Adding 1 to 127")
    print("-" * 70)
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        wrapped = wrap_int8(dataset.INT8_MAX, 1)
        warning_names = [w.category.__name__ for w in caught]
    print(f"  np.array([127], dtype=np.int8) + np.array([1], dtype=np.int8)")
    print(f"  gives                {wrapped}")
    print(f"  warnings raised      {warning_names if warning_names else 'none'}")
    print()
    print("  No exception. No warning. On numpy 2.5.2 the value simply wraps")
    print("  from the top of the range round to the bottom, the way a car")
    print("  odometer rolls from 999999 to 000000, and the next line of your")
    print("  program carries on with -128 as though it were the answer.")
    assert wrapped == dataset.INT8_MIN
    assert warning_names == [], "measured on numpy 2.5.2: no warning is emitted"

    # -- 4. Why -128 and not something arbitrary ------------------------------
    print()
    print("4. Where -128 comes from")
    print("-" * 70)
    print("  An int8 is 8 bits in two's complement. 127 is 0111 1111. Add 1")
    print("  and the carry ripples the whole way:")
    print()
    print("      0111 1111    = 127")
    print("    + 0000 0001    =   1")
    print("      ---------")
    print("      1000 0000    = -128, because the top bit means 'negative'")
    print()
    for value in dataset.INT8_DOUBLING_INPUT:
        doubled = int((np.array([value], dtype=np.int8) * np.int8(2))[0])
        note = "" if doubled == value * 2 else f"   <- wrapped, the true answer is {value * 2}"
        print(f"  int8 {value:>4} doubled -> {doubled:>5}{note}")
    doubled_all = (np.array(dataset.INT8_DOUBLING_INPUT, dtype=np.int8) * np.int8(2)).tolist()
    print(f"  all three at once: {doubled_all}")
    assert doubled_all == [-16, -6, -2]

    # -- 5. A plain Python int does not rescue it -----------------------------
    print()
    print("5. The rule NumPy 2 applies when the types differ")
    print("-" * 70)
    from_python = np.array([127], dtype=np.int8) + 1
    print(f"  np.array([127], dtype=np.int8) + 1  ->  {from_python}  dtype {from_python.dtype}")
    print()
    print("  The plain 1 did not drag the result up to int64. Since NumPy 2 a")
    print("  Python scalar takes the array's dtype rather than the other way")
    print("  round, so the array's promise wins and the result wraps. The")
    print("  array's dtype is the thing to check when a number looks wrong.")
    assert int(from_python[0]) == -128
    assert from_python.dtype == np.int8

    # -- 6. Asking for a wider type on purpose --------------------------------
    print()
    print("6. The fix, which is to say what you meant")
    print("-" * 70)
    widened = np.array([127], dtype=np.int8).astype(np.int16) + 1
    print(f"  .astype(np.int16) + 1  ->  {widened}  dtype {widened.dtype}")
    print("  Two bytes per element instead of one, and 128 fits.")
    print("  This is a decision with a cost. On a million elements it is a")
    print("  megabyte. On a model's weights it can be a gigabyte, which is why")
    print("  half-precision floats exist and why anyone talks about them.")
    assert int(widened[0]) == 128
    assert widened.dtype == np.int16

    # -- 7. Floats lose precision instead of wrapping -------------------------
    print()
    print("7. The float version of the same problem")
    print("-" * 70)
    as32 = np.float32(0.1)
    print(f"  float64 0.1  ->  {0.1!r}")
    print(f"  float32 0.1  ->  {float(as32)!r}")
    print("  Neither is 0.1. One tenth is not representable in binary at all,")
    print("  and float32 has 24 bits of significand where float64 has 53, so")
    print("  it is wrong sooner.")
    print()
    blind = np.float32(dataset.FLOAT32_BLIND_SPOT)
    plus_one = blind + np.float32(1.0)
    print(f"  float32 {dataset.FLOAT32_BLIND_SPOT:.0f} + 1 == {dataset.FLOAT32_BLIND_SPOT:.0f}  ->  {bool(plus_one == blind)}")
    print("  At 2**24 the gap between neighbouring float32 values is exactly")
    print("  1, so adding 1 lands back on the same value. A float does not")
    print("  wrap round like an int8; it stops being able to tell two numbers")
    print("  apart. The failure is quieter and harder to spot.")
    assert bool(plus_one == blind) is True
    assert float(np.float64(dataset.FLOAT32_BLIND_SPOT) + 1.0) == 16777217.0

    print()
    print("=" * 70)
    print("02_dtypes_and_overflow.py: every assertion held.")


if __name__ == "__main__":
    main()

"""Three operations, each written twice. Same answer, different speed.

Run from inside examples/:

    ../.venv/bin/python3 03_same_answer_faster.py

The claim under test: a vectorised expression is the SAME computation as the
loop, not a different one. Every element goes through the same IEEE-754
operation in the same order; only the machinery around it changes. So the two
results are compared with `==`, elementwise, over a million elements -- not
with a tolerance, because a tolerance would hide the very thing being shown.

The timings are then measured, printed with their spread, and deliberately NOT
asserted on. One machine, one day. The test suite asserts the SHAPE of the gap
-- at least twenty times -- which is a claim that survives a slower laptop.
"""

import platform
import sys

import numpy as np

import dataset
from vectorize import (
    clip_loop,
    clip_vec,
    median_seconds,
    roots_loop,
    roots_vec,
    scale_and_offset_loop,
    scale_and_offset_vec,
    speedup,
    time_call,
)

REPEATS = 5


def report(name: str, loop_result, vec_result, loop_times, vec_times) -> float:
    """Print one operation's agreement and its two timings."""
    loop_array = np.array(loop_result)
    identical = bool(np.array_equal(loop_array, vec_result))
    factor = speedup(loop_times, vec_times)
    print(f"  {name}")
    print(f"    elementwise identical over {loop_array.size:,} elements : {identical}")
    print(f"    loop  ms  {[round(t * 1000, 2) for t in loop_times]}")
    print(f"    array ms  {[round(t * 1000, 3) for t in vec_times]}")
    print(f"    median loop {median_seconds(loop_times) * 1000:8.2f} ms")
    print(f"    median array{median_seconds(vec_times) * 1000:8.3f} ms")
    print(f"    speedup     {factor:8.1f}x")
    print()
    assert identical, f"{name}: the two routes must agree exactly"
    return factor


def main() -> None:
    print("03_same_answer_faster.py")
    print("=" * 70)

    # -- 1. The machine this was measured on ----------------------------------
    print()
    print("1. What this was measured on, so the numbers can be read honestly")
    print("-" * 70)
    print(f"  python    {platform.python_version()}")
    print(f"  numpy     {np.__version__}")
    print(f"  platform  {platform.platform()}")
    print(f"  machine   {platform.machine()}")
    print(f"  elements  {dataset.N_BIG:,}")
    print(f"  repeats   {REPEATS} per operation, median reported")
    print()
    print("  Your figures will differ. The ratio is the durable part; the")
    print("  milliseconds are one machine on one day.")

    values = dataset.big_values()
    as_list = values.tolist()
    assert len(as_list) == dataset.N_BIG

    # -- 2. The three operations ----------------------------------------------
    print()
    print("2. Three operations, each computed twice")
    print("-" * 70)
    print()

    factors = []

    factors.append(
        report(
            f"scale and offset:  {dataset.SCALE_M} * x + {dataset.SCALE_C}",
            scale_and_offset_loop(as_list, dataset.SCALE_M, dataset.SCALE_C),
            scale_and_offset_vec(values, dataset.SCALE_M, dataset.SCALE_C),
            time_call(
                lambda: scale_and_offset_loop(as_list, dataset.SCALE_M, dataset.SCALE_C),
                REPEATS,
            ),
            time_call(
                lambda: scale_and_offset_vec(values, dataset.SCALE_M, dataset.SCALE_C),
                REPEATS,
            ),
        )
    )

    factors.append(
        report(
            "square root:       math.sqrt(x)  vs  np.sqrt(a)",
            roots_loop(as_list),
            roots_vec(values),
            time_call(lambda: roots_loop(as_list), REPEATS),
            time_call(lambda: roots_vec(values), REPEATS),
        )
    )

    factors.append(
        report(
            f"clip:              hold x inside [{dataset.CLIP_LO}, {dataset.CLIP_HI}]",
            clip_loop(as_list, dataset.CLIP_LO, dataset.CLIP_HI),
            clip_vec(values, dataset.CLIP_LO, dataset.CLIP_HI),
            time_call(lambda: clip_loop(as_list, dataset.CLIP_LO, dataset.CLIP_HI), REPEATS),
            time_call(lambda: clip_vec(values, dataset.CLIP_LO, dataset.CLIP_HI), REPEATS),
        )
    )

    print(f"  slowest speedup measured here: {min(factors):.1f}x")
    print(f"  fastest speedup measured here: {max(factors):.1f}x")
    assert min(factors) > 20.0, (
        "the tests assert 20x, which is a claim about the shape of the gap"
    )

    # -- 3. Why identical rather than close -----------------------------------
    print()
    print("3. Why those comparisons used == and not a tolerance")
    print("-" * 70)
    print(f"  2.5  is exactly representable in binary: {2.5 == float.fromhex('0x1.4p+1')}")
    print(f"  1.25 is exactly representable in binary: {1.25 == float.fromhex('0x1.4p+0')}")
    print()
    print("  Each element goes through one multiply and one add, in the same")
    print("  order, on the same 64 bits, on the same processor. There is no")
    print("  room for a difference and so none appears. If a vectorised")
    print("  rewrite of yours needs a tolerance, that is worth a second look:")
    print("  it means the two versions are not doing the same arithmetic.")

    # -- 4. Where the loop's time actually goes -------------------------------
    print()
    print("4. What the loop spends its time on")
    print("-" * 70)
    one = as_list[0]
    print(f"  one element as a Python float object : {sys.getsizeof(one)} bytes")
    print(f"  one element inside the array         : {values.itemsize} bytes")
    print()
    print("  Per element, the loop does roughly this: fetch a pointer, follow")
    print("  it, check the object's type, unbox the double, multiply, add, box")
    print("  the result into a NEW float object, store a pointer to it. Seven")
    print("  operations of bookkeeping around one of arithmetic.")
    print()
    print("  The array version fetches eight bytes at a known offset,")
    print("  multiplies, adds, stores eight bytes. No type check, because the")
    print("  dtype already settled that question once for the whole array.")
    print("  THAT is what a dtype buys, and it is why the two facts -- fixed")
    print("  dtype and contiguous block -- are the same fact wearing two hats.")

    # -- 5. It is still a loop ------------------------------------------------
    print()
    print("5. The loop did not disappear")
    print("-" * 70)
    print("  np.sqrt(a) still visits every one of the million elements. The")
    print("  loop moved from CPython's bytecode interpreter into compiled C")
    print("  inside NumPy, where the processor can also work on several")
    print("  elements per instruction. Vectorised does not mean 'no loop', it")
    print("  means 'not YOUR loop'.")
    print()
    print("  Which is also the cost: you can no longer put a print, a")
    print("  breakpoint or an early exit inside it. Script 07 is about when")
    print("  that trade is a bad one.")

    print()
    print("=" * 70)
    print("03_same_answer_faster.py: every assertion held.")


if __name__ == "__main__":
    main()

"""Making arrays without a loop, and operating on them without a loop.

Run from inside examples/:

    ../.venv/bin/python3 04_creating_and_ufuncs.py

The claim under test: you almost never need to build an array by appending to a
list and converting. There is a constructor for the shape you want, and once
you have it, every mathematical function you want already applies to the whole
thing at once.
"""

import math

import numpy as np

import dataset


def main() -> None:
    print("04_creating_and_ufuncs.py")
    print("=" * 70)

    # -- 1. The eight constructors worth knowing by heart ---------------------
    print()
    print("1. Eight ways to make an array, and when each is the right one")
    print("-" * 70)

    from_list = np.array([1.5, 2.5, 3.5])
    zeros = np.zeros(4)
    ones = np.ones((2, 3))
    full = np.full(3, 7)
    arange = np.arange(0, 10, 2)
    linspace = np.linspace(0.0, 1.0, 5)
    eye = np.eye(3)
    rng = np.random.default_rng(dataset.SEED)
    randoms = rng.random(3)

    rows = [
        ("np.array([1.5, 2.5, 3.5])", from_list, "data you already have"),
        ("np.zeros(4)", zeros, "an accumulator to fill in"),
        ("np.ones((2, 3))", ones, "a 2 by 3 block of ones"),
        ("np.full(3, 7)", full, "any constant, dtype taken from it"),
        ("np.arange(0, 10, 2)", arange, "a COUNT: start, stop, step"),
        ("np.linspace(0, 1, 5)", linspace, "a RANGE: start, stop, how many"),
        ("np.eye(3)", eye, "the identity matrix, from Day 102"),
        ("rng.random(3)", randoms, "reproducible pseudo-random values"),
    ]
    for call, value, why in rows:
        flat = np.array2string(value.ravel(), precision=6, separator=", ")
        print(f"  {call:<26} {str(value.dtype):<8} shape {str(value.shape):<7} {flat}")
        print(f"  {'':<26} {why}")
    print()
    print("  The two that get confused: arange counts in steps and EXCLUDES")
    print("  the stop, exactly like Python's range. linspace takes how many")
    print("  points you want and INCLUDES both ends. Ask for a step, use")
    print("  arange; ask for a count, use linspace.")

    assert from_list.dtype == np.float64
    assert zeros.tolist() == [0.0, 0.0, 0.0, 0.0]
    assert ones.shape == (2, 3)
    assert full.tolist() == [7, 7, 7] and full.dtype == np.int64
    assert arange.tolist() == [0, 2, 4, 6, 8]
    assert linspace.tolist() == [0.0, 0.25, 0.5, 0.75, 1.0]
    assert eye.tolist() == [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]

    # -- 2. Seeded randomness -------------------------------------------------
    print()
    print("2. Random, and reproducible, which are not opposites")
    print("-" * 70)
    again = np.random.default_rng(dataset.SEED).random(3)
    print(f"  default_rng({dataset.SEED}).random(3)   {randoms}")
    print(f"  and a second generator, same seed  {again}")
    print(f"  identical: {bool(np.array_equal(randoms, again))}")
    print()
    print("  numpy.random.default_rng is the modern interface. The older")
    print("  numpy.random.seed sets ONE global generator that every library")
    print("  in the process shares, so a call you did not write can move your")
    print("  sequence. A generator object you pass around cannot be moved by")
    print("  anyone else, which is why every number in this lab is stable.")
    assert bool(np.array_equal(randoms, again))

    # -- 3. Universal functions -----------------------------------------------
    print()
    print("3. A ufunc: one call, every element")
    print("-" * 70)
    a = np.array([0.0, 1.0, 4.0, 9.0, 16.0])
    print(f"  a          = {a}")
    print(f"  np.sqrt(a) = {np.sqrt(a)}")
    print(f"  the comprehension [math.sqrt(x) for x in a] gives the same:")
    print(f"               {np.array([math.sqrt(x) for x in a])}")
    print(f"  identical: {bool(np.array_equal(np.sqrt(a), [math.sqrt(x) for x in a]))}")
    assert bool(np.array_equal(np.sqrt(a), [math.sqrt(x) for x in a]))
    print()
    print("  math.sqrt cannot take an array at all -- it wants one number.")
    try:
        math.sqrt(a)
    except TypeError as exc:
        print(f"  math.sqrt(a) raises TypeError: {exc}")
    else:  # pragma: no cover - documents an outcome that would falsify the claim
        raise AssertionError("math.sqrt accepted an array, which it should not")

    # -- 4. The ufuncs you will use --------------------------------------------
    print()
    print("4. The ones that come up daily")
    print("-" * 70)
    small = np.array([-2.0, -0.5, 0.0, 0.5, 2.0])
    for name, result in (
        ("np.abs", np.abs(small)),
        ("np.sqrt(np.abs(x))", np.sqrt(np.abs(small))),
        ("np.exp", np.exp(small)),
        ("np.sign", np.sign(small)),
        ("np.round(x, 2)", np.round(small, 2)),
        ("x ** 2", small ** 2),
        ("np.maximum(x, 0)", np.maximum(small, 0.0)),
    ):
        print(f"  {name:<20} {np.array2string(result, precision=6, separator=', ')}")
    print()
    print("  np.maximum(x, 0) is the ReLU from Day 102, written as a ufunc.")
    print("  It takes TWO arrays and compares them elementwise. np.max takes")
    print("  ONE array and reduces it to a single number. Confusing the two is")
    print("  a rite of passage; the longer name is the elementwise one.")
    assert np.maximum(small, 0.0).tolist() == [0.0, 0.0, 0.0, 0.5, 2.0]
    assert float(np.max(small)) == 2.0
    assert np.abs(small).tolist() == [2.0, 0.5, 0.0, 0.5, 2.0]

    # -- 5. Two arrays at once, elementwise -----------------------------------
    print()
    print("5. Two arrays, elementwise, no loop")
    print("-" * 70)
    left = np.array([1.0, 2.0, 3.0])
    right = np.array([10.0, 20.0, 30.0])
    print(f"  left  = {left}")
    print(f"  right = {right}")
    print(f"  left + right  = {left + right}")
    print(f"  left * right  = {left * right}   <- elementwise, NOT a dot product")
    print(f"  left @ right  = {left @ right}   <- the dot product, from Day 103")
    assert (left * right).tolist() == [10.0, 40.0, 90.0]
    assert float(left @ right) == 140.0
    print()
    print("  `*` is elementwise and `@` is the matrix product. In a language")
    print("  that gives you one symbol for multiplication, this is the single")
    print("  most common source of a silently wrong shape.")

    # -- 6. Shapes must agree, or broadcast -----------------------------------
    print()
    print("6. When the shapes do not match")
    print("-" * 70)
    try:
        np.array([1.0, 2.0, 3.0]) + np.array([1.0, 2.0])
    except ValueError as exc:
        print(f"  (3,) + (2,) raises ValueError: {exc}")
    else:  # pragma: no cover - documents an outcome that would falsify the claim
        raise AssertionError("mismatched shapes must raise")
    print()
    print(f"  but (3,) + a single number works: {left + 100.0}")
    print("  The scalar was BROADCAST: stretched, conceptually, to match. No")
    print("  copy of 100.0 was ever made. Broadcasting is the next section's")
    print("  subject and the reason `2.5 * a + 1.25` in script 03 was legal.")
    assert (left + 100.0).tolist() == [101.0, 102.0, 103.0]

    print()
    print("=" * 70)
    print("04_creating_and_ufuncs.py: every assertion held.")


if __name__ == "__main__":
    main()

"""Boolean masking: the idea that replaces `if` inside a loop.

Run from inside examples/:

    ../.venv/bin/python3 05_masks_and_selection.py

The claim under test: a comparison on an array is an array, one True or False
per element, and that array can be used to count, to select, to choose between
two values and to assign. Almost every filtering loop you would have written
becomes one line.

And the one that trips everyone: combining two masks needs `&` and `|`, not
`and` and `or`, and the reason is a ValueError you can read.
"""

import numpy as np

import dataset
from vectorize import count_above, mask_between, select


def main() -> None:
    print("05_masks_and_selection.py")
    print("=" * 70)

    readings = dataset.small_readings()

    # -- 1. The data, small enough to check by eye ----------------------------
    print()
    print("1. Twenty readings from the seeded generator")
    print("-" * 70)
    print(f"  {readings.tolist()}")
    print(f"  dtype {readings.dtype}   shape {readings.shape}")
    print()
    print("  Written out in dataset.py as SMALL_READINGS_EXPECTED so you can")
    print("  check any answer below without running anything.")
    assert readings.tolist() == dataset.SMALL_READINGS_EXPECTED

    # -- 2. A comparison returns an array -------------------------------------
    print()
    print("2. `readings > 50` is not a yes or a no")
    print("-" * 70)
    mask = readings > 50
    print(f"  readings > 50  ->  {mask}")
    print(f"  dtype {mask.dtype}   shape {mask.shape}   size {mask.size}")
    print()
    print("  Twenty answers, one per element. That single fact is what every")
    print("  other line in this script is built on.")
    assert mask.dtype == np.bool_
    assert mask.shape == (dataset.N_SMALL,)
    assert mask.tolist()[:4] == [True, True, False, True]

    # -- 3. Counting, without a counter ---------------------------------------
    print()
    print("3. Counting")
    print("-" * 70)
    how_many = count_above(readings, 50)
    print(f"  mask.sum()          {int(mask.sum())}")
    print(f"  count_above(a, 50)  {how_many}")
    print("  True is 1 and False is 0 when summed, so a count is a sum. The")
    print("  loop with a `total += 1` inside it does not need writing again.")
    print()
    print(f"  mask.any()   {bool(mask.any())}    is anything above 50?")
    print(f"  mask.all()   {bool(mask.all())}   is EVERYTHING above 50?")
    print(f"  mask.mean()  {float(mask.mean())}    what fraction? (a sum divided by n)")
    assert how_many == 9
    assert bool(mask.any()) is True
    assert bool(mask.all()) is False
    assert float(mask.mean()) == 0.45

    # -- 4. Selecting ---------------------------------------------------------
    print()
    print("4. Selecting the elements the mask marks")
    print("-" * 70)
    chosen = select(readings, mask)
    print(f"  readings[readings > 50]  ->  {chosen.tolist()}")
    print(f"  shape {chosen.shape}, which is the count from section 3")
    print()
    print(f"  np.nonzero(mask)[0]      ->  {np.nonzero(mask)[0].tolist()}")
    print("  ...if it is the POSITIONS you want rather than the values.")
    assert chosen.tolist() == [70, 83, 69, 65, 75, 73, 97, 64, 82]
    assert chosen.shape == (9,)
    assert np.nonzero(mask)[0].tolist() == [0, 1, 3, 8, 11, 13, 16, 17, 19]

    # -- 5. Combining masks, and the error everybody meets --------------------
    print()
    print("5. Two conditions at once")
    print("-" * 70)
    between = mask_between(readings, 30, 70)
    print(f"  (readings > 30) & (readings < 70)")
    print(f"  count {int(between.sum())}   values {readings[between].tolist()}")
    outer = (readings < 10) | (readings > 90)
    print(f"  (readings < 10) | (readings > 90)")
    print(f"  count {int(outer.sum())}   values {readings[outer].tolist()}")
    print(f"  ~between  (the negation)  count {int((~between).sum())}")
    assert int(between.sum()) == 7
    assert readings[between].tolist() == [34, 69, 65, 37, 37, 41, 64]
    assert int(outer.sum()) == 1
    assert readings[outer].tolist() == [97]
    assert int((~between).sum()) == 13

    print()
    print("  Now the same thing with `and`:")
    try:
        (readings > 30) and (readings < 70)
    except ValueError as exc:
        print(f"    ValueError: {exc}")
        message = str(exc)
    else:  # pragma: no cover - documents an outcome that would falsify the claim
        raise AssertionError("`and` on two arrays must raise")
    assert "truth value of an array" in message
    print()
    print("  `and` is not an operator NumPy can define. It is a control-flow")
    print("  keyword: Python asks the left operand 'are you true?' and an")
    print("  array of twenty answers cannot say. `&` IS an operator, so NumPy")
    print("  defines it to mean elementwise-and, which is what you wanted.")
    print()
    print("  The same refusal, more directly:")
    try:
        bool(readings > 30)
    except ValueError as exc:
        print(f"    bool(readings > 30) -> ValueError: {exc}")
    else:  # pragma: no cover - documents an outcome that would falsify the claim
        raise AssertionError("bool() on a multi-element array must raise")
    print()
    print("  It even tells you the two ways out, .any() and .all(), which are")
    print("  the two questions that DO have a single answer.")

    # -- 6. The parentheses are not optional ----------------------------------
    print()
    print("6. Why those brackets are load-bearing")
    print("-" * 70)
    print("  `&` binds TIGHTER than `>` in Python, so")
    print("      readings > 30 & readings < 70")
    print("  parses as")
    print("      readings > (30 & readings) < 70")
    print("  which is a bitwise-and of 30 with every reading, then a chained")
    print("  comparison. Here is what that actually raises -- the expression")
    print("  below is written out literally, brackets and all left off:")
    try:
        readings > 30 & readings < 70  # noqa: B015 - the point is the exception
    except ValueError as exc:
        print(f"    ValueError: {exc}")
    else:  # pragma: no cover - documents an outcome that would falsify the claim
        raise AssertionError("the unbracketed form must raise")
    print()
    print("  Not a syntax error, which would be kinder. It is a chained")
    print("  comparison, and chaining calls bool() on the first half.")
    print(f"  And 30 & readings really is a bitwise-and, element by element:")
    print(f"    {(30 & readings).tolist()[:6]} ...")

    # -- 7. Choosing between two values ---------------------------------------
    print()
    print("7. np.where: the vectorised if-else")
    print("-" * 70)
    labels = np.where(readings > 50, 1, 0)
    print(f"  np.where(readings > 50, 1, 0)")
    print(f"    {labels.tolist()}")
    capped = np.where(readings > 50, 50, readings)
    print(f"  np.where(readings > 50, 50, readings)   <- cap at 50")
    print(f"    {capped.tolist()}")
    print()
    print("  Three arguments: the mask, the value where True, the value where")
    print("  False. Either of the last two may be an array, and then it is")
    print("  read elementwise. This is `x if cond else y` for a whole array.")
    assert labels.sum() == 9
    assert int(capped.max()) == 50
    assert capped.tolist()[:4] == [50, 50, 34, 50]

    # -- 8. Assigning through a mask ------------------------------------------
    print()
    print("8. Writing through a mask")
    print("-" * 70)
    working = readings.copy()
    working[working > 90] = 90
    print(f"  a[a > 90] = 90   ->  max is now {int(working.max())}")
    print(f"  and the original is untouched: max {int(readings.max())}")
    print("  (untouched only because .copy() was called first -- script 06)")
    assert int(working.max()) == 90
    assert int(readings.max()) == 97

    # -- 9. Fancy indexing ----------------------------------------------------
    print()
    print("9. Fancy indexing: an array of positions")
    print("-" * 70)
    wanted = np.array([0, 5, 19, 5])
    print(f"  readings[[0, 5, 19, 5]]  ->  {readings[wanted].tolist()}")
    print()
    print("  Three differences from a boolean mask, all of them useful:")
    print("    * the result has the shape of the INDEX array, not the source")
    print("    * you choose the order")
    print("    * you may ask for the same element twice, as 5 is here")
    print("  This is how a batch of rows is pulled out of a dataset, and it is")
    print("  what the next script's top-k does with the argsort result.")
    assert readings[wanted].tolist() == [70, 21, 82, 21]
    assert readings[wanted].shape == (4,)

    print()
    print("=" * 70)
    print("05_masks_and_selection.py: every assertion held.")


if __name__ == "__main__":
    main()

"""Axes, shapes, views, and the ranking that Day 103 needed.

Run from inside examples/:

    ../.venv/bin/python3 06_axes_views_and_ranking.py

Four claims under test:

  * the axis you name is the one that DISAPPEARS;
  * `np.newaxis` turns a row into a column so broadcasting can pair everything
    with everything;
  * a slice is a VIEW, so writing to it writes to the original, and this is
    where a beginner's hardest bug lives;
  * `argsort` and not `sort` is what a search needs, because the indices are
    the answer and the scores are only how you got there.
"""

import numpy as np

import dataset
from vectorize import cosine_similarities, top_k_indices


def main() -> None:
    print("06_axes_views_and_ranking.py")
    print("=" * 70)

    # -- 1. The axis rule -----------------------------------------------------
    print()
    print("1. Aggregating, and the one rule that makes axis make sense")
    print("-" * 70)
    grid = np.arange(12).reshape(3, 4)
    print("  a 3 by 4 array:")
    for row in grid:
        print(f"    {row}")
    print()
    print(f"  grid.sum()          {grid.sum():<20} shape {np.shape(grid.sum())}")
    print(f"  grid.sum(axis=0)    {str(grid.sum(axis=0)):<20} shape {grid.sum(axis=0).shape}")
    print(f"  grid.sum(axis=1)    {str(grid.sum(axis=1)):<20} shape {grid.sum(axis=1).shape}")
    print()
    print("  THE RULE: the axis you name is the one that disappears.")
    print("    shape (3, 4), axis=0 -> shape (4,)   the 3 went")
    print("    shape (3, 4), axis=1 -> shape (3,)   the 4 went")
    print()
    print("  So axis=0 collapses DOWN the rows and gives one number per")
    print("  column; axis=1 collapses ACROSS the columns and gives one number")
    print("  per row. Reading it as 'which axis do I want to keep' is the")
    print("  mistake, and it is off by exactly one every time.")
    assert int(grid.sum()) == 66
    assert grid.sum(axis=0).tolist() == [12, 15, 18, 21]
    assert grid.sum(axis=1).tolist() == [6, 22, 38]
    assert grid.sum(axis=0).shape == (4,)
    assert grid.sum(axis=1).shape == (3,)

    print()
    print("  The same rule for every aggregation:")
    for name, fn in (("min", np.min), ("max", np.max), ("mean", np.mean)):
        print(
            f"    {name:<5} whole {str(fn(grid)):<8} axis=0 {str(fn(grid, axis=0)):<22}"
            f" axis=1 {fn(grid, axis=1)}"
        )
    assert np.max(grid, axis=1).tolist() == [3, 7, 11]

    print()
    print(f"  keepdims=True holds the shape open: {grid.sum(axis=1, keepdims=True).shape}")
    print("  which is what you want when the result has to broadcast back")
    print("  against the array it came from -- normalising every row, say.")
    assert grid.sum(axis=1, keepdims=True).shape == (3, 1)

    # -- 2. newaxis -----------------------------------------------------------
    print()
    print("2. np.newaxis: making a row into a column")
    print("-" * 70)
    v = np.array([1.0, 2.0, 3.0])
    print(f"  v                shape {v.shape}")
    print(f"  v[:, np.newaxis] shape {v[:, np.newaxis].shape}   a column")
    print(f"  v[np.newaxis, :] shape {v[np.newaxis, :].shape}   a row")
    print()
    print("  A column against a row broadcasts to a full table, every pairing")
    print("  at once, with no loop over pairs:")
    table = v[:, np.newaxis] - v[np.newaxis, :]
    for row in table:
        print(f"    {row}")
    print(f"  shape {table.shape}: every difference of every pair.")
    print("  This is how a whole distance matrix gets built in one line, and")
    print("  section 7 of script 07 is about when that line is a bad idea.")
    assert v[:, np.newaxis].shape == (3, 1)
    assert table.shape == (3, 3)
    assert table.tolist() == [[0.0, -1.0, -2.0], [1.0, 0.0, -1.0], [2.0, 1.0, 0.0]]
    print()
    print(f"  reshape does the same job explicitly: {v.reshape(3, 1).shape}")
    print(f"  and -1 means 'work it out': v.reshape(-1, 1) -> {v.reshape(-1, 1).shape}")
    assert v.reshape(-1, 1).shape == (3, 1)

    # -- 3. A view is not a copy ----------------------------------------------
    print()
    print("3. A slice is a VIEW, and this is where the bug lives")
    print("-" * 70)
    original = np.arange(12).reshape(3, 4)
    print("  original:")
    for row in original:
        print(f"    {row}")
    row_one = original[1]
    print()
    print(f"  row_one = original[1]  ->  {row_one}")
    print(f"  shares memory with the original: {np.shares_memory(original, row_one)}")
    print(f"  row_one.base is None: {row_one.base is None}   (a view knows its owner)")
    row_one[0] = 999
    print()
    print("  row_one[0] = 999")
    print("  the ORIGINAL now reads:")
    for row in original:
        print(f"    {row}")
    print()
    print("  Nothing was copied, so nothing was protected. In a list, `b =")
    print("  a[1:3]` hands you a new list and you can do what you like to it.")
    print("  In NumPy it hands you a different way of reading the same bytes.")
    assert np.shares_memory(original, row_one)
    assert int(original[1, 0]) == 999
    assert row_one.base is not None

    print()
    print("  .copy() breaks the link:")
    detached = original[2].copy()
    detached[0] = -1
    print(f"    detached = original[2].copy(); detached[0] = -1")
    print(f"    detached      {detached}")
    print(f"    original[2]   {original[2]}   <- untouched")
    print(f"    shares memory {np.shares_memory(original, detached)}")
    assert int(original[2, 0]) == 8
    assert not np.shares_memory(original, detached)

    # -- 4. Which operations give a view --------------------------------------
    print()
    print("4. Which of them hand back a view, and which a copy")
    print("-" * 70)
    base = np.arange(12).reshape(3, 4)
    cases = [
        ("base[1]           row slice", base[1]),
        ("base[:, 1]        column slice", base[:, 1]),
        ("base[0:2, 1:3]    block slice", base[0:2, 1:3]),
        ("base.T            transpose", base.T),
        ("base.reshape(4,3) reshape", base.reshape(4, 3)),
        ("base.ravel()      flatten (view when it can)", base.ravel()),
        ("base[base > 5]    boolean mask", base[base > 5]),
        ("base[[0, 2]]      fancy index", base[[0, 2]]),
        ("base.copy()       explicit copy", base.copy()),
        ("base + 0          arithmetic", base + 0),
    ]
    print(f"  {'expression':<44} {'view?':<6}")
    for label, result in cases:
        is_view = np.shares_memory(base, result)
        print(f"  {label:<44} {'VIEW' if is_view else 'copy'}")
    print()
    print("  The pattern: if the elements you asked for are evenly spaced,")
    print("  NumPy can describe them with a stride and gives you a view. If")
    print("  they are not -- a mask, a list of positions -- it has no choice")
    print("  but to copy. So the cheap operations are the dangerous ones.")
    assert np.shares_memory(base, base[:, 1])
    assert np.shares_memory(base, base.T)
    assert not np.shares_memory(base, base[base > 5])
    assert not np.shares_memory(base, base[[0, 2]])
    assert not np.shares_memory(base, base + 0)

    # -- 5. sort, and why it is the wrong tool --------------------------------
    print()
    print("5. sort loses the thing you were looking for")
    print("-" * 70)
    scores = np.array([5.0, 1.0, 9.0, 3.0])
    print(f"  scores          {scores}")
    print(f"  np.sort(scores) {np.sort(scores)}   <- the values, in order")
    print(f"  scores after    {scores}   <- np.sort returns a NEW array")
    print(f"  np.argsort      {np.argsort(scores)}   <- the POSITIONS, in order")
    print()
    print("  argsort answers 'which element would come first, then which'.")
    print("  scores[np.argsort(scores)] rebuilds the sorted values:")
    print(f"    {scores[np.argsort(scores)]}")
    print()
    print("  When the rows mean something -- an article, a customer, a token --")
    print("  the sorted values are useless on their own and the indices are")
    print("  the entire answer. That is why argsort is the one to reach for.")
    assert np.sort(scores).tolist() == [1.0, 3.0, 5.0, 9.0]
    assert np.argsort(scores).tolist() == [1, 3, 0, 2]
    assert scores.tolist() == [5.0, 1.0, 9.0, 3.0]
    assert scores[np.argsort(scores)].tolist() == [1.0, 3.0, 5.0, 9.0]
    print()
    print("  a.sort() -- the method, no np. -- sorts IN PLACE and returns None:")
    in_place = scores.copy()
    returned = in_place.sort()
    print(f"    returned {returned}, array now {in_place}")
    assert returned is None
    assert in_place.tolist() == [1.0, 3.0, 5.0, 9.0]

    # -- 6. Day 103's search, done with argsort -------------------------------
    print()
    print("6. Day 103's search, ranked with argsort")
    print("-" * 70)
    sims = cosine_similarities(dataset.CATALOGUE, dataset.QUERY)
    print(f"  catalogue shape {dataset.CATALOGUE.shape}   query shape {dataset.QUERY.shape}")
    print(f"  query: 'training for a race and what to eat' = {dataset.QUERY}")
    print()
    print("  all six similarities, from one matrix-vector product and one")
    print("  norm along axis=1 -- no loop over articles:")
    for name, score in zip(dataset.ARTICLE_NAMES, sims):
        print(f"    {name:<20} {score:.6f}")
    assert sims.shape == (6,)

    print()
    ordering = np.argsort(sims)
    top = top_k_indices(sims, dataset.TOP_K)
    print(f"  np.argsort(sims)          {ordering.tolist()}   <- worst first")
    print(f"  reversed, first {dataset.TOP_K}         {top.tolist()}   <- best first")
    print()
    print(f"  top {dataset.TOP_K}:")
    for rank, index in enumerate(top, start=1):
        print(f"    {rank}. {dataset.ARTICLE_NAMES[index]:<20} {sims[index]:.6f}")
    assert ordering.tolist() == [4, 5, 1, 0, 2, 3]
    assert top.tolist() == [3, 2, 0]
    assert [dataset.ARTICLE_NAMES[i] for i in top] == [
        "race-day-nutrition",
        "marathon-plan",
        "roast-chicken",
    ]

    print()
    margin = float(sims[top[0]] - sims[top[1]])
    print(f"  margin between first and second: {margin:.6f}")
    print("  Day 103 called this a close call and it still is. The ranking is")
    print("  reported with its margin rather than as a verdict, because a gap")
    print("  of two thousandths is not evidence of much.")
    assert 0.0 < margin < 0.01

    print()
    print("  One more way to say the same thing, and the one you will meet in")
    print("  model code, where only the top few matter out of a hundred")
    print("  thousand:")
    partitioned = np.argpartition(-sims, dataset.TOP_K)[: dataset.TOP_K]
    ordered = partitioned[np.argsort(-sims[partitioned])]
    print(f"    np.argpartition(-sims, {dataset.TOP_K})[:{dataset.TOP_K}] then sorted -> {ordered.tolist()}")
    print("  argpartition does not sort everything; it only guarantees that")
    print("  the k best are in the first k places, in no particular order. On")
    print("  six articles that saves nothing. On a hundred thousand it is the")
    print("  difference between sorting them all and not.")
    assert ordered.tolist() == top.tolist()

    print()
    print("=" * 70)
    print("06_axes_views_and_ranking.py: every assertion held.")


if __name__ == "__main__":
    main()

"""Missing values, and the honest limits of the habit this lab is teaching.

Run from inside examples/:

    ../.venv/bin/python3 07_nan_and_when_not_to_vectorise.py

Two claims under test.

First: `nan` is not equal to itself, which sounds like a bug and is a rule, and
it is the reason `x == np.nan` never finds anything and `np.isnan(x)` does.

Second, and more important: vectorising is a trade, not an upgrade. There are
three situations where the loop is the better code, and this script measures
all three rather than asserting them.
"""

import math
import time

import numpy as np

import dataset
from vectorize import nan_aware_mean, roots_vec


def main() -> None:
    print("07_nan_and_when_not_to_vectorise.py")
    print("=" * 70)

    # -- 1. nan is not equal to itself ----------------------------------------
    print()
    print("1. The one comparison that surprises everybody")
    print("-" * 70)
    print(f"  np.nan == np.nan   ->  {np.nan == np.nan}")
    print(f"  np.nan != np.nan   ->  {np.nan != np.nan}")
    print(f"  np.nan  is np.nan  ->  {np.nan is np.nan}   (it is one object)")
    print()
    print("  Not a NumPy decision. IEEE-754 says nan compares unequal to")
    print("  everything including itself, because nan means 'not a number' --")
    print("  the result of 0/0, of sqrt of a negative, of a reading that was")
    print("  never taken. Two unknowns are not known to be the same unknown.")
    assert (np.nan == np.nan) is False
    assert (np.nan != np.nan) is True

    # -- 2. What that costs you -----------------------------------------------
    print()
    print("2. So this does not work")
    print("-" * 70)
    holed = dataset.WITH_A_HOLE
    print(f"  a = {holed}")
    a_eq = holed == np.nan
    print(f"  a == np.nan   ->  {a_eq}")
    print("  Every answer False, including for the element that IS nan. A")
    print("  filter written this way finds nothing and reports success.")
    print()
    print(f"  np.isnan(a)   ->  {np.isnan(holed)}")
    print(f"  np.isnan(a).sum()  ->  {int(np.isnan(holed).sum())}")
    print("  np.isnan asks about the bit pattern rather than about equality,")
    print("  which is the only question with a useful answer here.")
    assert not a_eq.any()
    assert np.isnan(holed).tolist() == [False, False, True, False]
    assert int(np.isnan(holed).sum()) == 1

    # -- 3. nan is contagious, and that is a feature --------------------------
    print()
    print("3. One hole poisons the aggregate, deliberately")
    print("-" * 70)
    print(f"  a.sum()        {holed.sum()}")
    print(f"  a.mean()       {holed.mean()}")
    print(f"  a.max()        {holed.max()}")
    print()
    print(f"  np.nansum(a)   {np.nansum(holed)}")
    print(f"  np.nanmean(a)  {np.nanmean(holed)}")
    print(f"  np.nanmax(a)   {np.nanmax(holed)}")
    print()
    print(f"  nan_aware_mean(a) = {nan_aware_mean(holed)}")
    print("  which is 7 / 3, the mean of the three readings that exist.")
    print()
    print("  The plain versions are not broken. They are telling you that a")
    print("  value is missing, loudly, at the point where it starts to matter.")
    print("  Reaching for the nan- version is a decision to ignore that, and")
    print("  it should be a decision rather than a reflex: the mean of the")
    print("  three you have is not the mean of the four you wanted.")
    assert math.isnan(float(holed.mean()))
    assert float(np.nansum(holed)) == 7.0
    assert nan_aware_mean(holed) == 7.0 / 3.0
    assert float(np.nanmax(holed)) == 4.0

    print()
    print("  Where a nan comes from in the first place:")
    with np.errstate(invalid="ignore", divide="ignore"):
        zero_over_zero = np.float64(0.0) / np.float64(0.0)
        root_of_negative = np.sqrt(np.array([-1.0]))[0]
        inf_minus_inf = np.float64(np.inf) - np.float64(np.inf)
    print(f"    0.0 / 0.0        {zero_over_zero}")
    print(f"    np.sqrt(-1.0)    {root_of_negative}")
    print(f"    inf - inf        {inf_minus_inf}")
    print("  Each of those emits a RuntimeWarning by default, which np.errstate")
    print("  is silencing here only because the point is the VALUE. In your")
    print("  own code, leave the warning on.")
    assert math.isnan(float(zero_over_zero))
    assert math.isnan(float(root_of_negative))
    assert math.isnan(float(inf_minus_inf))

    # -- 4. When NOT to vectorise: the array is small -------------------------
    print()
    print("4. When not to vectorise, case one: the array is small")
    print("-" * 70)
    small = [1.0, 2.0, 3.0, 4.0]
    reps = 20000

    start = time.perf_counter()
    for _ in range(reps):
        [math.sqrt(x) for x in small]
    loop_us = (time.perf_counter() - start) / reps * 1e6

    start = time.perf_counter()
    for _ in range(reps):
        np.sqrt(np.array(small))
    numpy_us = (time.perf_counter() - start) / reps * 1e6

    print(f"  four elements, {reps:,} repetitions, microseconds per call")
    print(f"    [math.sqrt(x) for x in xs]     {loop_us:7.3f} us")
    print(f"    np.sqrt(np.array(xs))          {numpy_us:7.3f} us")
    print(f"    the comprehension is           {numpy_us / loop_us:7.2f}x faster here")
    print()
    print("  Every NumPy call has a fixed cost -- work out the dtypes, work")
    print("  out the output shape, allocate it -- before any arithmetic")
    print("  happens. On four elements that setup is the whole bill. The")
    print("  crossover on this machine is in the low hundreds of elements.")
    print("  One machine, one day; measure yours rather than trusting this.")
    assert numpy_us > loop_us, "on four elements the NumPy call costs more here"

    # -- 5. When NOT to vectorise: the loop is clearer -------------------------
    print()
    print("5. When not to vectorise, case two: the loop is clearer")
    print("-" * 70)
    print("  A running balance where each step depends on the last one:")
    print()
    balances = [100.0]
    for change in (-30.0, 50.0, -200.0, 20.0):
        nxt = balances[-1] + change
        balances.append(max(nxt, 0.0))
    print(f"    start 100, changes -30, +50, -200, +20, floored at zero")
    print(f"    balances {balances}")
    print()
    print("  There is no one-line NumPy for that, because step four depends on")
    print("  the floor applied at step three. np.cumsum would give you the")
    print("  running total, and the floor would be wrong:")
    naive = 100.0 + np.cumsum([-30.0, 50.0, -200.0, 20.0])
    print(f"    np.cumsum route  {np.maximum(naive, 0.0).tolist()}")
    print(f"    the honest loop  {balances[1:]}")
    print("  Different answers, and the loop's is the right one. Sequential")
    print("  dependence is the clearest signal that the loop should stay.")
    assert balances == [100.0, 70.0, 120.0, 0.0, 20.0]
    assert np.maximum(naive, 0.0).tolist() != balances[1:]

    # -- 6. When NOT to vectorise: memory ------------------------------------
    print()
    print("6. When not to vectorise, case three: it will not fit")
    print("-" * 70)
    for n in (1_000, 10_000, 30_000, 100_000):
        gb = n * n * 8 / 1e9
        print(f"    all pairs of {n:>7,} points, float64 : {gb:>10.2f} GB")
    print()
    print("  The one-line distance matrix from section 2 of script 06 is")
    print("  `x[:, None] - x[None, :]`, and it allocates n squared elements")
    print("  whether you need them all or not. At a hundred thousand points")
    print("  that is 80 GB, and the elegant line is the reason the process")
    print("  died. The loop that processes a thousand at a time is slower and")
    print("  finishes.")
    print()
    print("  This is not hypothetical arithmetic -- here is the real allocation")
    print("  for a size that does fit:")
    x = np.arange(2000, dtype=np.float64)
    pairwise = x[:, None] - x[None, :]
    print(f"    2,000 points -> shape {pairwise.shape}, {pairwise.nbytes / 1e6:.1f} MB")
    print(f"    the input was {x.nbytes / 1e3:.0f} kB. The output is {pairwise.nbytes / x.nbytes:,.0f}x bigger.")
    assert pairwise.shape == (2000, 2000)
    assert pairwise.nbytes == 32_000_000
    del pairwise

    # -- 7. Same value, different operation -----------------------------------
    print()
    print("7. And a last honesty note about 'the same computation'")
    print("-" * 70)
    values = dataset.big_values()
    as_list = values.tolist()
    by_math = np.array([math.sqrt(x) for x in as_list])
    by_pow = np.array([x ** 0.5 for x in as_list])
    vec = roots_vec(values)
    math_matches = int(np.count_nonzero(by_math != vec))
    pow_matches = int(np.count_nonzero(by_pow != vec))
    print(f"  over {values.size:,} values, compared with np.sqrt:")
    print(f"    math.sqrt(x)  disagrees on {math_matches:>6,} of them")
    print(f"    x ** 0.5      disagrees on {pow_matches:>6,} of them")
    first = int(np.nonzero(by_pow != vec)[0][0])
    print()
    print(f"  the first disagreement, at index {first}:")
    print(f"    x           {as_list[first]!r}")
    print(f"    x ** 0.5    {float(by_pow[first])!r}")
    print(f"    np.sqrt(x)  {float(vec[first])!r}")
    print(f"    difference  {abs(by_pow[first] - vec[first]):.3e}")
    print()
    print("  One unit in the last place, on about one value in seven hundred.")
    print("  IEEE-754 requires square root to be correctly rounded and both")
    print("  math.sqrt and np.sqrt use the instruction that obeys that.")
    print("  pow(x, 0.5) is a general power routine and makes no such promise.")
    print()
    print("  So 'the vectorised version gives the same answer' is a claim about")
    print("  the OPERATION, not about anything that agrees in exact arithmetic.")
    print("  When a rewrite needs a tolerance it did not need before, that is")
    print("  worth reading rather than widening.")
    assert math_matches == 0
    assert pow_matches > 0
    assert abs(float(by_pow[first]) - float(vec[first])) < 1e-15

    print()
    print("=" * 70)
    print("07_nan_and_when_not_to_vectorise.py: every assertion held.")


if __name__ == "__main__":
    main()

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

Both `examples/` and `starter/` contain modules called `vectorize` and
`dataset`, 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 `vectorize` 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.

So: put this directory first on the import path, and drop any already-imported
`vectorize`, `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 ("vectorize", "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 data this lab measures, and the tolerances it compares with.

Three kinds of data live here, and each is here for a different reason.

1. A **big** array of a million numbers, drawn from a seeded generator. It
   exists so the memory and speed measurements have something to bite on. A
   thousand elements would make the loop look fine.

2. A **small** array of twenty integers, drawn from the same seeded generator.
   It exists so every boolean mask in this lab can be checked by eye. You can
   read the twenty numbers, count the ones above fifty, and compare with what
   the code says.

3. The **article catalogue from Day 99 and Day 103**, unchanged. Six invented
   articles described by four hand-counted features. Day 103 ranked them by
   cosine similarity; today the ranking is done with `argsort`, which is the
   part of that day that transfers directly to model code.

Nothing here is real. The articles do not exist, the sensor readings are not
sensor readings, and the counts were chosen so that every number in this lab
can be re-derived with a pen.

The seed is fixed at 104 and every number below follows from it, so two runs of
this lab on two machines produce the same values. That is a deliberate choice
rather than a convenience: a lab that asserts on random numbers is a lab that
fails for the reader and not for the author.
"""

from __future__ import annotations

import numpy as np

# ---------------------------------------------------------------------------
# Seeds and sizes
# ---------------------------------------------------------------------------

#: The one seed this lab uses. Passed to numpy.random.default_rng.
SEED = 104

#: How many elements the timing and memory comparisons use. One million is
#: large enough that the loop is unmistakably slower and small enough that the
#: array is 8 MB rather than 8 GB.
N_BIG = 1_000_000

#: How many elements the by-eye mask exercises use.
N_SMALL = 20

#: Float comparison tolerance, used wherever two routes to the same number are
#: allowed to differ in the last bits. Most comparisons in this lab do NOT use
#: it: the whole point of the from-scratch section is that the loop and the
#: vectorised version agree EXACTLY, and asserting that with a tolerance would
#: hide the very fact being demonstrated.
TOL = 1e-12


def big_values() -> np.ndarray:
    """A million float64 values in [0, 1), from the fixed seed.

    Returns a fresh array each call so that a test which mutates it cannot
    quietly change the answer of a later one.
    """
    return np.random.default_rng(SEED).random(N_BIG)


def small_readings() -> np.ndarray:
    """Twenty integers in [0, 100), from the fixed seed.

    Drawn from a generator seeded independently of `big_values`, so the two
    are reproducible on their own.
    """
    return np.random.default_rng(SEED).integers(0, 100, size=N_SMALL)


# The literal values `small_readings()` produces on the fixed seed, written out
# so you can check a mask against them without running anything. The reference
# tests assert that this list and the generator still agree; if a future NumPy
# ever changed the generator's output, the suite would say so rather than
# quietly rewriting the lesson.
SMALL_READINGS_EXPECTED = [
    70, 83, 34, 69, 26, 21, 18, 12, 65, 37,
    17, 75, 30, 73, 37, 41, 97, 64, 21, 82,
]

# ---------------------------------------------------------------------------
# The article catalogue, carried unchanged from Day 99 and Day 103
# ---------------------------------------------------------------------------

FEATURES = ("cooking", "running", "money", "weather")

ARTICLE_NAMES = (
    "roast-chicken",
    "slow-cooker-stew",
    "marathon-plan",
    "race-day-nutrition",
    "household-budget",
    "storm-bulletin",
)

#: One row per article, one column per feature, in the order above. Day 103
#: held these as six separate lists; today they are one 6 by 4 array, and that
#: change is the day's subject rather than a tidy-up.
CATALOGUE = np.array(
    [
        [9, 0, 1, 0],   # roast-chicken
        [8, 0, 2, 0],   # slow-cooker-stew
        [0, 9, 1, 2],   # marathon-plan
        [4, 6, 3, 0],   # race-day-nutrition
        [1, 0, 9, 0],   # household-budget
        [0, 1, 0, 9],   # storm-bulletin
    ],
    dtype=np.float64,
)

#: "training for a race and what to eat", written as feature counts the same
#: way the articles were. Deliberately a close call between two articles.
QUERY = np.array([2, 5, 0, 0], dtype=np.float64)

#: How many results the search returns.
TOP_K = 3

# ---------------------------------------------------------------------------
# The three operations implemented twice
# ---------------------------------------------------------------------------
#
# The multiplier, offset and clip bounds are all exactly representable in
# binary floating point (2.5 is 10.1 in binary, 1.25 is 1.01, 0.25 is 0.01,
# 0.75 is 0.11). That matters: it means the loop and the vectorised version
# perform the identical IEEE-754 operation on the identical bits, so they can
# be compared with `==` rather than with a tolerance.

SCALE_M = 2.5
SCALE_C = 1.25
CLIP_LO = 0.25
CLIP_HI = 0.75

# ---------------------------------------------------------------------------
# The dtype demonstrations
# ---------------------------------------------------------------------------

#: The largest value an int8 can hold. Adding 1 to this wraps to INT8_MIN.
INT8_MAX = 127
INT8_MIN = -128

#: Three int8 values doubled in section 2. Two of the three wrap.
INT8_DOUBLING_INPUT = [120, 125, 127]

#: A value float32 cannot tell apart from its successor: 2 ** 24. A float32
#: has 24 bits of significand, so at this magnitude the gap between
#: representable numbers is exactly 1, and adding 1 changes nothing.
FLOAT32_BLIND_SPOT = 16777216.0

# ---------------------------------------------------------------------------
# The array with a hole in it
# ---------------------------------------------------------------------------

#: Four readings, one of which is missing. `mean` on this returns nan; the
#: whole point of section 7 is that the nan is contagious and that saying so
#: loudly is better than silently dropping it.
WITH_A_HOLE = np.array([1.0, 2.0, np.nan, 4.0])

"""The reference suite: real values, real exceptions, real measurements.

Run from the LAB DIRECTORY:

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

Nothing here reads source code or checks that a file exists. Every test runs
something and looks at what came back.

Two rules this file follows deliberately:

  * Timings are asserted as SHAPES, never as figures. `test_speedup_is_large`
    asserts at least 20x. The authoring machine measured over 100x. A test
    that asserted 100x would fail on a slower laptop and teach the reader that
    the suite is unreliable rather than that their laptop is slower.

  * Where the loop and the vectorised version do the same operation, they are
    compared with `==` over a million elements rather than with a tolerance,
    because exactness is the claim under test.
"""

import math
import sys
import warnings

import numpy as np
import pytest

import dataset
from vectorize import (
    array_bytes,
    clip_loop,
    clip_vec,
    cosine_similarities,
    count_above,
    describe,
    list_bytes,
    mask_between,
    median_seconds,
    nan_aware_mean,
    roots_loop,
    roots_vec,
    scale_and_offset_loop,
    scale_and_offset_vec,
    select,
    speedup,
    time_call,
    top_k_indices,
    wrap_int8,
)

# One million-element array, built once for the whole module. Every test that
# uses it treats it as read-only; the two that mutate take a copy first.
BIG = dataset.big_values()
BIG_LIST = BIG.tolist()


# ===========================================================================
# The environment
# ===========================================================================


def test_numpy_is_version_two_or_later():
    assert int(np.__version__.split(".")[0]) >= 2


def test_the_seeded_generator_still_produces_the_documented_values():
    """If a future NumPy changed the generator, this says so rather than
    letting every hand-checked number in the lesson quietly become wrong."""
    assert dataset.small_readings().tolist() == dataset.SMALL_READINGS_EXPECTED


def test_the_big_array_is_the_documented_size_and_dtype():
    assert BIG.size == dataset.N_BIG
    assert BIG.dtype == np.float64
    assert BIG.nbytes == 8_000_000


def test_two_generators_with_the_same_seed_agree():
    first = np.random.default_rng(dataset.SEED).random(5)
    second = np.random.default_rng(dataset.SEED).random(5)
    assert np.array_equal(first, second)


def test_two_generators_with_different_seeds_do_not():
    first = np.random.default_rng(dataset.SEED).random(5)
    other = np.random.default_rng(dataset.SEED + 1).random(5)
    assert not np.array_equal(first, other)


# ===========================================================================
# Memory
# ===========================================================================


def test_the_naive_size_comparison_is_misleading():
    """sys.getsizeof on the list is within one percent of the array's nbytes,
    which is why this lab does not use it as the measurement."""
    values = list(range(dataset.N_BIG))
    array = np.arange(dataset.N_BIG, dtype=np.int64)
    ratio = sys.getsizeof(values) / array.nbytes
    assert 0.99 < ratio < 1.01


def test_the_honest_list_total_is_four_and_a_half_times_the_array():
    values = list(range(dataset.N_BIG))
    array = np.arange(dataset.N_BIG, dtype=np.int64)
    assert list_bytes(values) == 36_000_056
    assert array_bytes(array) == 8_000_000
    assert list_bytes(values) / array_bytes(array) == pytest.approx(4.5, abs=0.01)


def test_a_python_int_is_twenty_eight_bytes_here():
    assert sys.getsizeof(1_000_000) == 28


def test_an_int64_element_is_eight_bytes():
    assert np.arange(3, dtype=np.int64).itemsize == 8


def test_list_bytes_counts_a_shared_integer_once():
    """CPython caches -5 to 256, so a thousand references to 100 are one
    object and must be charged once."""
    hundreds = [int("100")] * 1000
    assert list_bytes(hundreds) == sys.getsizeof(hundreds) + sys.getsizeof(100)


def test_dtype_decides_the_bill():
    n = 1000
    assert np.zeros(n, dtype=np.int8).nbytes == 1000
    assert np.zeros(n, dtype=np.int16).nbytes == 2000
    assert np.zeros(n, dtype=np.float32).nbytes == 4000
    assert np.zeros(n, dtype=np.float64).nbytes == 8000


# ===========================================================================
# Shape, strides and contiguity
# ===========================================================================


def test_shape_dtype_and_strides_of_a_three_by_four():
    grid = np.arange(12).reshape(3, 4)
    assert grid.shape == (3, 4)
    assert grid.dtype == np.int64
    assert grid.strides == (32, 8)
    assert grid.flags["C_CONTIGUOUS"]


def test_describe_reports_all_six_facts():
    text = describe(np.arange(12).reshape(3, 4))
    for fragment in (
        "shape=(3, 4)",
        "dtype=int64",
        "itemsize=8",
        "nbytes=96",
        "strides=(32, 8)",
        "c_contiguous=True",
    ):
        assert fragment in text


def test_a_transpose_copies_nothing():
    grid = np.arange(12).reshape(3, 4)
    assert np.shares_memory(grid, grid.T)
    assert grid.T.strides == (8, 32)
    assert not grid.T.flags["C_CONTIGUOUS"]
    assert grid.T.flags["F_CONTIGUOUS"]


# ===========================================================================
# dtypes and overflow
# ===========================================================================


def test_int8_wraps_from_127_to_minus_128():
    assert wrap_int8(dataset.INT8_MAX, 1) == dataset.INT8_MIN


def test_the_wrap_is_silent_on_this_numpy():
    """Measured, not assumed. If a future NumPy started warning, this fails and
    the lesson gets corrected rather than left stale."""
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        wrap_int8(dataset.INT8_MAX, 1)
    assert [w.category.__name__ for w in caught] == []


def test_doubling_three_int8_values_wraps_two_of_them():
    doubled = np.array(dataset.INT8_DOUBLING_INPUT, dtype=np.int8) * np.int8(2)
    assert doubled.tolist() == [-16, -6, -2]
    assert doubled.dtype == np.int8


def test_a_python_scalar_does_not_widen_the_array():
    result = np.array([127], dtype=np.int8) + 1
    assert result.dtype == np.int8
    assert int(result[0]) == -128


def test_astype_widens_and_the_answer_is_right():
    result = np.array([127], dtype=np.int8).astype(np.int16) + 1
    assert result.dtype == np.int16
    assert int(result[0]) == 128


def test_int8_range_is_what_iinfo_says():
    info = np.iinfo(np.int8)
    assert (info.min, info.max) == (dataset.INT8_MIN, dataset.INT8_MAX)


def test_float32_cannot_tell_two_million_apart_at_its_blind_spot():
    blind = np.float32(dataset.FLOAT32_BLIND_SPOT)
    assert bool(blind + np.float32(1.0) == blind)


def test_float64_can():
    assert np.float64(dataset.FLOAT32_BLIND_SPOT) + 1.0 == 16777217.0


def test_float32_stores_a_worse_one_tenth_than_float64():
    assert float(np.float32(0.1)) != 0.1
    assert abs(float(np.float32(0.1)) - 0.1) > abs(0.1 - 0.1)


# ===========================================================================
# The three operations, twice each
# ===========================================================================


def test_scale_and_offset_agrees_exactly_over_a_million_elements():
    loop = np.array(scale_and_offset_loop(BIG_LIST, dataset.SCALE_M, dataset.SCALE_C))
    vec = scale_and_offset_vec(BIG, dataset.SCALE_M, dataset.SCALE_C)
    assert np.array_equal(loop, vec)


def test_roots_agree_exactly_over_a_million_elements():
    assert np.array_equal(np.array(roots_loop(BIG_LIST)), roots_vec(BIG))


def test_clip_agrees_exactly_over_a_million_elements():
    loop = np.array(clip_loop(BIG_LIST, dataset.CLIP_LO, dataset.CLIP_HI))
    vec = clip_vec(BIG, dataset.CLIP_LO, dataset.CLIP_HI)
    assert np.array_equal(loop, vec)


def test_the_hand_worked_examples_from_the_docstrings():
    assert scale_and_offset_loop([0.0, 1.0, 2.0], 2.5, 1.25) == [1.25, 3.75, 6.25]
    assert scale_and_offset_vec(np.array([0.0, 1.0, 2.0]), 2.5, 1.25).tolist() == [
        1.25,
        3.75,
        6.25,
    ]
    assert roots_loop([0.0, 1.0, 4.0]) == [0.0, 1.0, 2.0]
    assert roots_vec(np.array([0.0, 1.0, 4.0])).tolist() == [0.0, 1.0, 2.0]
    assert clip_loop([0.0, 0.5, 1.0], 0.25, 0.75) == [0.25, 0.5, 0.75]
    assert clip_vec(np.array([0.0, 0.5, 1.0]), 0.25, 0.75).tolist() == [0.25, 0.5, 0.75]


def test_clip_actually_clips():
    clipped = clip_vec(BIG, dataset.CLIP_LO, dataset.CLIP_HI)
    assert float(clipped.min()) == dataset.CLIP_LO
    assert float(clipped.max()) == dataset.CLIP_HI


def test_speedup_is_large_but_the_figure_is_not_asserted():
    """The SHAPE of the gap, not the figure. 20x survives a slow machine; the
    authoring machine measured over 100x on all three operations."""
    loop = time_call(
        lambda: scale_and_offset_loop(BIG_LIST, dataset.SCALE_M, dataset.SCALE_C), 3
    )
    vec = time_call(
        lambda: scale_and_offset_vec(BIG, dataset.SCALE_M, dataset.SCALE_C), 3
    )
    assert speedup(loop, vec) > 20.0


def test_median_seconds_is_the_middle_value():
    assert median_seconds([3.0, 1.0, 2.0]) == 2.0


def test_time_call_returns_one_timing_per_repeat():
    times = time_call(lambda: None, 4)
    assert len(times) == 4
    assert all(t >= 0.0 for t in times)


def test_x_to_the_half_is_not_the_same_operation_as_sqrt():
    """A measured disagreement, kept rather than tidied away: 1390 of a million
    values differ by one unit in the last place."""
    by_pow = np.array([x ** 0.5 for x in BIG_LIST])
    differing = int(np.count_nonzero(by_pow != roots_vec(BIG)))
    assert differing > 0
    assert differing < BIG.size // 100
    worst = float(np.max(np.abs(by_pow - roots_vec(BIG))))
    assert worst < 1e-15


def test_math_sqrt_is_the_same_operation_as_sqrt():
    by_math = np.array([math.sqrt(x) for x in BIG_LIST])
    assert int(np.count_nonzero(by_math != roots_vec(BIG))) == 0


# ===========================================================================
# Universal functions and creation
# ===========================================================================


def test_the_constructors_produce_the_documented_values():
    assert np.zeros(4).tolist() == [0.0] * 4
    assert np.ones((2, 3)).shape == (2, 3)
    assert np.full(3, 7).tolist() == [7, 7, 7]
    assert np.arange(0, 10, 2).tolist() == [0, 2, 4, 6, 8]
    assert np.linspace(0.0, 1.0, 5).tolist() == [0.0, 0.25, 0.5, 0.75, 1.0]
    assert np.eye(3).tolist() == [
        [1.0, 0.0, 0.0],
        [0.0, 1.0, 0.0],
        [0.0, 0.0, 1.0],
    ]


def test_arange_excludes_the_stop_and_linspace_includes_it():
    assert 10 not in np.arange(0, 10, 2).tolist()
    assert np.linspace(0.0, 1.0, 5)[-1] == 1.0


def test_full_takes_its_dtype_from_the_fill_value():
    assert np.full(3, 7).dtype == np.int64
    assert np.full(3, 7.0).dtype == np.float64


def test_a_ufunc_matches_the_comprehension():
    a = np.array([0.0, 1.0, 4.0, 9.0, 16.0])
    assert np.array_equal(np.sqrt(a), [math.sqrt(x) for x in a])


def test_math_sqrt_refuses_an_array():
    with pytest.raises(TypeError):
        math.sqrt(np.array([1.0, 2.0]))


def test_maximum_is_elementwise_and_max_is_a_reduction():
    small = np.array([-2.0, -0.5, 0.0, 0.5, 2.0])
    assert np.maximum(small, 0.0).tolist() == [0.0, 0.0, 0.0, 0.5, 2.0]
    assert float(np.max(small)) == 2.0


def test_star_is_elementwise_and_at_is_the_dot_product():
    left = np.array([1.0, 2.0, 3.0])
    right = np.array([10.0, 20.0, 30.0])
    assert (left * right).tolist() == [10.0, 40.0, 90.0]
    assert float(left @ right) == 140.0


def test_mismatched_shapes_raise_with_a_readable_message():
    with pytest.raises(ValueError) as excinfo:
        np.array([1.0, 2.0, 3.0]) + np.array([1.0, 2.0])
    assert "broadcast" in str(excinfo.value)


def test_a_scalar_broadcasts():
    assert (np.array([1.0, 2.0, 3.0]) + 100.0).tolist() == [101.0, 102.0, 103.0]


# ===========================================================================
# Masking and selection
# ===========================================================================


def test_a_comparison_returns_one_boolean_per_element():
    readings = dataset.small_readings()
    mask = readings > 50
    assert mask.dtype == np.bool_
    assert mask.shape == (dataset.N_SMALL,)


def test_counting_above_fifty():
    assert count_above(dataset.small_readings(), 50) == 9


def test_the_selected_values_are_the_documented_nine():
    readings = dataset.small_readings()
    assert select(readings, readings > 50).tolist() == [
        70, 83, 69, 65, 75, 73, 97, 64, 82
    ]


def test_the_selection_length_equals_the_count():
    readings = dataset.small_readings()
    mask = readings > 50
    assert select(readings, mask).size == int(mask.sum())


def test_any_all_and_mean_of_a_boolean_array():
    mask = dataset.small_readings() > 50
    assert bool(mask.any()) is True
    assert bool(mask.all()) is False
    assert float(mask.mean()) == 0.45


def test_mask_between_uses_ampersand_and_gives_the_documented_seven():
    readings = dataset.small_readings()
    mask = mask_between(readings, 30, 70)
    assert int(mask.sum()) == 7
    assert readings[mask].tolist() == [34, 69, 65, 37, 37, 41, 64]


def test_or_gives_the_documented_one():
    readings = dataset.small_readings()
    mask = (readings < 10) | (readings > 90)
    assert readings[mask].tolist() == [97]


def test_negating_a_mask():
    readings = dataset.small_readings()
    mask = mask_between(readings, 30, 70)
    assert int((~mask).sum()) == dataset.N_SMALL - int(mask.sum())


def test_and_raises_valueerror_with_the_ambiguous_truth_value_message():
    readings = dataset.small_readings()
    with pytest.raises(ValueError) as excinfo:
        (readings > 30) and (readings < 70)
    assert "truth value of an array with more than one element is ambiguous" in str(
        excinfo.value
    )


def test_or_the_keyword_raises_too():
    readings = dataset.small_readings()
    with pytest.raises(ValueError):
        (readings > 30) or (readings < 70)


def test_bool_of_a_multi_element_array_raises():
    with pytest.raises(ValueError):
        bool(dataset.small_readings() > 30)


def test_bool_of_a_one_element_array_does_not():
    """The message says "more than one element" and it means it."""
    assert bool(np.array([True])) is True


def test_the_missing_parentheses_raise_rather_than_giving_a_wrong_answer():
    readings = dataset.small_readings()
    with pytest.raises(ValueError):
        readings > 30 & readings < 70  # noqa: B015 - the exception is the point


def test_where_chooses_between_two_values():
    readings = dataset.small_readings()
    labels = np.where(readings > 50, 1, 0)
    assert labels.tolist() == (readings > 50).astype(int).tolist()
    assert int(labels.sum()) == 9


def test_where_accepts_an_array_on_either_branch():
    readings = dataset.small_readings()
    capped = np.where(readings > 50, 50, readings)
    assert int(capped.max()) == 50
    assert int(capped.min()) == int(readings.min())


def test_nonzero_gives_positions_rather_than_values():
    readings = dataset.small_readings()
    assert np.nonzero(readings > 50)[0].tolist() == [0, 1, 3, 8, 11, 13, 16, 17, 19]


def test_assigning_through_a_mask_changes_only_the_marked_elements():
    readings = dataset.small_readings()
    working = readings.copy()
    working[working > 90] = 90
    assert int(working.max()) == 90
    assert int(readings.max()) == 97
    assert int((working != readings).sum()) == 1


def test_fancy_indexing_takes_the_shape_of_the_index_and_allows_repeats():
    readings = dataset.small_readings()
    wanted = np.array([0, 5, 19, 5])
    assert readings[wanted].tolist() == [70, 21, 82, 21]
    assert readings[wanted].shape == wanted.shape


# ===========================================================================
# Axes and shapes
# ===========================================================================


def test_the_axis_you_name_disappears():
    grid = np.arange(12).reshape(3, 4)
    assert grid.sum(axis=0).shape == (4,)
    assert grid.sum(axis=1).shape == (3,)


def test_the_aggregate_values():
    grid = np.arange(12).reshape(3, 4)
    assert int(grid.sum()) == 66
    assert grid.sum(axis=0).tolist() == [12, 15, 18, 21]
    assert grid.sum(axis=1).tolist() == [6, 22, 38]
    assert grid.max(axis=1).tolist() == [3, 7, 11]
    assert grid.mean(axis=0).tolist() == [4.0, 5.0, 6.0, 7.0]


def test_keepdims_holds_the_shape_open():
    grid = np.arange(12).reshape(3, 4)
    assert grid.sum(axis=1, keepdims=True).shape == (3, 1)


def test_keepdims_is_what_lets_the_result_broadcast_back():
    grid = np.arange(1, 13, dtype=np.float64).reshape(3, 4)
    normalised = grid / grid.sum(axis=1, keepdims=True)
    assert np.allclose(normalised.sum(axis=1), 1.0, atol=dataset.TOL)


def test_newaxis_makes_a_column_and_a_row():
    v = np.array([1.0, 2.0, 3.0])
    assert v[:, np.newaxis].shape == (3, 1)
    assert v[np.newaxis, :].shape == (1, 3)


def test_a_column_against_a_row_broadcasts_to_every_pairing():
    v = np.array([1.0, 2.0, 3.0])
    table = v[:, np.newaxis] - v[np.newaxis, :]
    assert table.shape == (3, 3)
    assert table.tolist() == [[0.0, -1.0, -2.0], [1.0, 0.0, -1.0], [2.0, 1.0, 0.0]]


def test_reshape_minus_one_works_it_out():
    assert np.array([1.0, 2.0, 3.0]).reshape(-1, 1).shape == (3, 1)


def test_reshape_refuses_an_impossible_shape():
    with pytest.raises(ValueError):
        np.arange(12).reshape(5, 3)


# ===========================================================================
# Views and copies
# ===========================================================================


def test_a_row_slice_is_a_view_and_writing_to_it_writes_through():
    grid = np.arange(12).reshape(3, 4)
    row = grid[1]
    assert np.shares_memory(grid, row)
    row[0] = 999
    assert int(grid[1, 0]) == 999


def test_a_copy_breaks_the_link():
    grid = np.arange(12).reshape(3, 4)
    detached = grid[2].copy()
    detached[0] = -1
    assert int(grid[2, 0]) == 8
    assert not np.shares_memory(grid, detached)


def test_a_view_knows_its_owner_and_a_copy_does_not():
    grid = np.arange(12).reshape(3, 4)
    assert grid[1].base is not None
    assert grid[1].copy().base is None


@pytest.mark.parametrize(
    "make, expect_view",
    [
        (lambda a: a[1], True),
        (lambda a: a[:, 1], True),
        (lambda a: a[0:2, 1:3], True),
        (lambda a: a.T, True),
        (lambda a: a.reshape(4, 3), True),
        (lambda a: a.ravel(), True),
        (lambda a: a[a > 5], False),
        (lambda a: a[[0, 2]], False),
        (lambda a: a.copy(), False),
        (lambda a: a + 0, False),
        (lambda a: a.flatten(), False),
    ],
)
def test_which_operations_return_a_view(make, expect_view):
    grid = np.arange(12).reshape(3, 4)
    assert np.shares_memory(grid, make(grid)) is expect_view


def test_ravel_is_a_view_and_flatten_is_always_a_copy():
    """The pair that catches people: near-identical names, opposite behaviour."""
    grid = np.arange(12).reshape(3, 4)
    assert np.shares_memory(grid, grid.ravel())
    assert not np.shares_memory(grid, grid.flatten())


# ===========================================================================
# Sorting and ranking
# ===========================================================================


def test_np_sort_returns_a_new_array_and_leaves_the_original_alone():
    scores = np.array([5.0, 1.0, 9.0, 3.0])
    assert np.sort(scores).tolist() == [1.0, 3.0, 5.0, 9.0]
    assert scores.tolist() == [5.0, 1.0, 9.0, 3.0]


def test_the_sort_method_is_in_place_and_returns_none():
    scores = np.array([5.0, 1.0, 9.0, 3.0])
    assert scores.sort() is None
    assert scores.tolist() == [1.0, 3.0, 5.0, 9.0]


def test_argsort_gives_positions_and_reconstructs_the_sorted_values():
    scores = np.array([5.0, 1.0, 9.0, 3.0])
    order = np.argsort(scores)
    assert order.tolist() == [1, 3, 0, 2]
    assert scores[order].tolist() == np.sort(scores).tolist()


def test_top_k_indices_gives_the_best_first():
    assert top_k_indices(np.array([0.1, 0.9, 0.5]), 2).tolist() == [1, 2]


def test_top_k_asks_for_no_more_than_it_needs():
    assert top_k_indices(np.array([0.1, 0.9, 0.5]), 1).tolist() == [1]
    assert top_k_indices(np.array([0.1, 0.9, 0.5]), 3).tolist() == [1, 2, 0]


# ===========================================================================
# The Day 103 search, vectorised
# ===========================================================================


def test_the_catalogue_is_a_six_by_four_array():
    assert dataset.CATALOGUE.shape == (6, 4)
    assert len(dataset.ARTICLE_NAMES) == 6
    assert len(dataset.FEATURES) == 4


def test_cosine_similarities_returns_one_score_per_row():
    sims = cosine_similarities(dataset.CATALOGUE, dataset.QUERY)
    assert sims.shape == (6,)


def test_cosine_similarities_matches_the_one_row_at_a_time_version():
    """Day 103's loop, kept here as the thing the vectorised version replaced."""
    sims = cosine_similarities(dataset.CATALOGUE, dataset.QUERY)
    by_hand = []
    for row in dataset.CATALOGUE:
        dot = sum(float(x) * float(y) for x, y in zip(row, dataset.QUERY))
        left = math.sqrt(sum(float(x) * float(x) for x in row))
        right = math.sqrt(sum(float(y) * float(y) for y in dataset.QUERY))
        by_hand.append(dot / (left * right))
    assert np.allclose(sims, by_hand, atol=dataset.TOL)


def test_a_perfectly_aligned_row_scores_one():
    identity = np.array([[1.0, 0.0], [0.0, 1.0]])
    sims = cosine_similarities(identity, np.array([1.0, 0.0]))
    assert sims.tolist() == [1.0, 0.0]


def test_the_top_three_are_the_documented_articles_in_order():
    sims = cosine_similarities(dataset.CATALOGUE, dataset.QUERY)
    top = top_k_indices(sims, dataset.TOP_K)
    assert top.tolist() == [3, 2, 0]
    assert [dataset.ARTICLE_NAMES[i] for i in top] == [
        "race-day-nutrition",
        "marathon-plan",
        "roast-chicken",
    ]


def test_the_full_ranking_worst_first():
    sims = cosine_similarities(dataset.CATALOGUE, dataset.QUERY)
    assert np.argsort(sims).tolist() == [4, 5, 1, 0, 2, 3]


def test_the_margin_between_first_and_second_is_reported_as_small():
    sims = cosine_similarities(dataset.CATALOGUE, dataset.QUERY)
    top = top_k_indices(sims, 2)
    margin = float(sims[top[0]] - sims[top[1]])
    assert 0.0 < margin < 0.01


def test_argpartition_finds_the_same_top_three():
    sims = cosine_similarities(dataset.CATALOGUE, dataset.QUERY)
    partitioned = np.argpartition(-sims, dataset.TOP_K)[: dataset.TOP_K]
    ordered = partitioned[np.argsort(-sims[partitioned])]
    assert ordered.tolist() == top_k_indices(sims, dataset.TOP_K).tolist()


# ===========================================================================
# nan
# ===========================================================================


def test_nan_is_not_equal_to_itself():
    assert (np.nan == np.nan) is False
    assert (np.nan != np.nan) is True


def test_comparing_an_array_to_nan_finds_nothing():
    assert not (dataset.WITH_A_HOLE == np.nan).any()


def test_isnan_finds_it():
    assert np.isnan(dataset.WITH_A_HOLE).tolist() == [False, False, True, False]
    assert int(np.isnan(dataset.WITH_A_HOLE).sum()) == 1


def test_nan_propagates_through_every_plain_aggregation():
    holed = dataset.WITH_A_HOLE
    assert math.isnan(float(holed.sum()))
    assert math.isnan(float(holed.mean()))
    assert math.isnan(float(holed.max()))
    assert math.isnan(float(holed.min()))


def test_the_nan_aware_versions_skip_it():
    holed = dataset.WITH_A_HOLE
    assert float(np.nansum(holed)) == 7.0
    assert float(np.nanmax(holed)) == 4.0
    assert nan_aware_mean(holed) == 7.0 / 3.0


def test_nan_aware_mean_divides_by_the_count_that_exists():
    """Three readings, not four. Stated as arithmetic so the choice is visible."""
    assert nan_aware_mean(dataset.WITH_A_HOLE) == pytest.approx((1.0 + 2.0 + 4.0) / 3)


def test_where_a_nan_comes_from():
    with np.errstate(invalid="ignore", divide="ignore"):
        assert math.isnan(float(np.float64(0.0) / np.float64(0.0)))
        assert math.isnan(float(np.sqrt(np.array([-1.0]))[0]))
        assert math.isnan(float(np.float64(np.inf) - np.float64(np.inf)))


def test_nan_survives_a_sort_and_goes_to_the_end():
    """Worth knowing before an argsort-based top-k meets missing data."""
    holed = np.array([3.0, np.nan, 1.0, 2.0])
    order = np.argsort(holed)
    assert order.tolist()[-1] == 1
    assert math.isnan(float(holed[order][-1]))


# ===========================================================================
# When not to vectorise
# ===========================================================================


def test_on_four_elements_the_comprehension_wins():
    """A measurement, not a belief. If this ever stops being true on some
    machine the suite will say so."""
    small = [1.0, 2.0, 3.0, 4.0]
    loop = time_call(lambda: [math.sqrt(x) for x in small], 2000)
    vec = time_call(lambda: np.sqrt(np.array(small)), 2000)
    assert median_seconds(vec) > median_seconds(loop)


def test_a_sequential_dependence_has_no_one_line_equivalent():
    changes = [-30.0, 50.0, -200.0, 20.0]
    balances = [100.0]
    for change in changes:
        balances.append(max(balances[-1] + change, 0.0))
    naive = np.maximum(100.0 + np.cumsum(changes), 0.0)
    assert balances[1:] == [70.0, 120.0, 0.0, 20.0]
    assert naive.tolist() != balances[1:]


def test_the_pairwise_table_grows_with_the_square():
    x = np.arange(2000, dtype=np.float64)
    pairwise = x[:, None] - x[None, :]
    assert pairwise.shape == (2000, 2000)
    assert pairwise.nbytes == 32_000_000
    assert pairwise.nbytes == x.nbytes * 2000

"""The reference implementation: ten functions, each written twice or once.

This is the finished version of `starter/vectorize.py`. Read it after you have
tried, not before.

The organising idea of the whole module: **a vectorised expression is the same
computation as the loop, not a different one.** Where that is true, the two
results are compared with `==` and they are exactly equal. Where it is not
true, this module says so rather than reaching for a tolerance to paper over
the difference — see `roots_loop`, whose docstring records a real measured
disagreement between `x ** 0.5` and `numpy.sqrt`.

Three functions are implemented twice, once as an explicit Python loop and
once as a NumPy expression. Every one of the three pairs agrees bit for bit,
which is the claim `03_same_answer_faster.py` and the reference tests check.
"""

from __future__ import annotations

import math
import statistics
import sys
import time
from typing import Callable, Iterable, Sequence

import numpy as np

# ===========================================================================
# 1. The three operations, each implemented twice
# ===========================================================================


def scale_and_offset_loop(values: Sequence[float], m: float, c: float) -> list[float]:
    """`m * x + c` for every element, written as an explicit Python loop.

    This is the shape of code you have been writing for the last hundred days,
    and there is nothing wrong with it except its speed. It allocates the
    output list up front rather than appending, which is the fastest honest
    version of the loop -- comparing a slow loop against a fast array would be
    rigging the measurement.

    >>> scale_and_offset_loop([0.0, 1.0, 2.0], 2.5, 1.25)
    [1.25, 3.75, 6.25]
    """
    out = [0.0] * len(values)
    for i, x in enumerate(values):
        out[i] = m * x + c
    return out


def scale_and_offset_vec(a: np.ndarray, m: float, c: float) -> np.ndarray:
    """`m * x + c` for every element, written as one expression on the array.

    There is still a loop. It happens inside NumPy, compiled, over a
    contiguous block of float64 with no Python object in sight. What you have
    given up is the ability to put a `print` inside it.

    >>> scale_and_offset_vec(np.array([0.0, 1.0, 2.0]), 2.5, 1.25).tolist()
    [1.25, 3.75, 6.25]
    """
    return m * a + c


def roots_loop(values: Sequence[float]) -> list[float]:
    """The square root of every element, as a loop, using `math.sqrt`.

    `math.sqrt` and not `x ** 0.5`, and the difference is not pedantry. On
    this machine, with numpy 2.5.2 and CPython 3.14.0, `math.sqrt` agrees with
    `numpy.sqrt` on all one million values in this lab, while `x ** 0.5`
    disagrees on 1390 of them -- always by one unit in the last place. IEEE-754
    requires square root to be correctly rounded, and both `math.sqrt` and
    `numpy.sqrt` use the hardware instruction that obeys that requirement.
    `pow(x, 0.5)` is a general power routine with no such guarantee.

    So "the vectorised version gives the same answer" is true of the operation,
    not of anything that happens to compute the same value in exact
    arithmetic. `07_when_not_to_vectorise.py` measures the 1390 and prints one
    of them.

    >>> roots_loop([0.0, 1.0, 4.0])
    [0.0, 1.0, 2.0]
    """
    out = [0.0] * len(values)
    for i, x in enumerate(values):
        out[i] = math.sqrt(x)
    return out


def roots_vec(a: np.ndarray) -> np.ndarray:
    """The square root of every element, as one call to a universal function.

    `numpy.sqrt` is a *ufunc*: it applies elementwise to an array of any shape
    and returns an array of the same shape.

    >>> roots_vec(np.array([0.0, 1.0, 4.0])).tolist()
    [0.0, 1.0, 2.0]
    """
    return np.sqrt(a)


def clip_loop(values: Sequence[float], lo: float, hi: float) -> list[float]:
    """Pull every element back inside `[lo, hi]`, as a loop.

    >>> clip_loop([0.0, 0.5, 1.0], 0.25, 0.75)
    [0.25, 0.5, 0.75]
    """
    out = [0.0] * len(values)
    for i, x in enumerate(values):
        if x < lo:
            out[i] = lo
        elif x > hi:
            out[i] = hi
        else:
            out[i] = x
    return out


def clip_vec(a: np.ndarray, lo: float, hi: float) -> np.ndarray:
    """Pull every element back inside `[lo, hi]`, as one call.

    The branch has not disappeared, it has moved: `numpy.clip` evaluates the
    same two comparisons per element in C. This is the pattern that replaces
    `if` inside a loop, and it is why `numpy.where` matters so much.

    >>> clip_vec(np.array([0.0, 0.5, 1.0]), 0.25, 0.75).tolist()
    [0.25, 0.5, 0.75]
    """
    return np.clip(a, lo, hi)


# ===========================================================================
# 2. Masking, selection and ranking
# ===========================================================================


def count_above(a: np.ndarray, threshold: float) -> int:
    """How many elements are strictly greater than `threshold`.

    `a > threshold` is a boolean array, and `True` counts as 1 when summed.
    That is the whole trick, and it replaces a counter variable.

    >>> count_above(np.array([1, 5, 9]), 4)
    2
    """
    return int((a > threshold).sum())


def mask_between(a: np.ndarray, lo: float, hi: float) -> np.ndarray:
    """A boolean array that is True where `lo < x < hi`.

    The parentheses are not optional. `&` binds tighter than `<` in Python, so
    `a > lo & a < hi` parses as `a > (lo & a) < hi` and raises. Every NumPy
    user has written that line once.

    And it must be `&`, not `and`. `and` asks the left operand whether it is
    truthy, which for an array of more than one element raises
    ValueError: The truth value of an array with more than one element is
    ambiguous.

    >>> mask_between(np.array([1, 5, 9]), 2, 8).tolist()
    [False, True, False]
    """
    return (a > lo) & (a < hi)


def select(a: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """The elements where `mask` is True, as a new array.

    Indexing with a boolean array always returns a **copy**, never a view --
    it has to, because the selected elements are not evenly spaced and no
    stride can describe them.

    >>> select(np.array([1, 5, 9]), np.array([True, False, True])).tolist()
    [1, 9]
    """
    return a[mask]


def top_k_indices(scores: np.ndarray, k: int) -> np.ndarray:
    """The indices of the `k` largest scores, best first.

    `numpy.argsort` returns the indices that would sort the array, ascending.
    Reverse them and take the first `k`.

    The indices are the point. `numpy.sort` would give you the scores and lose
    which row each one came from, which is exactly the thing a search needs.

    >>> top_k_indices(np.array([0.1, 0.9, 0.5]), 2).tolist()
    [1, 2]
    """
    return np.argsort(scores)[::-1][:k]


def cosine_similarities(matrix: np.ndarray, query: np.ndarray) -> np.ndarray:
    """Cosine similarity between `query` and every row of `matrix`, at once.

    Day 103 computed these one row at a time. This is the same arithmetic with
    the loop moved into NumPy: one matrix-vector product for all the dot
    products, and `axis=1` to take one norm per row.

    `axis=1` means "collapse the columns, leaving one number per row". The rule
    worth memorising: **the axis you name is the one that disappears.** A 6 by
    4 array summed with `axis=1` gives 6 numbers, not 4.

    >>> m = np.array([[1.0, 0.0], [0.0, 1.0]])
    >>> cosine_similarities(m, np.array([1.0, 0.0])).tolist()
    [1.0, 0.0]
    """
    row_norms = np.linalg.norm(matrix, axis=1)
    return (matrix @ query) / (row_norms * np.linalg.norm(query))


# ===========================================================================
# 3. Memory, dtypes and missing values
# ===========================================================================


def list_bytes(values: list[int]) -> int:
    """An honest total for what a Python list of integers costs in memory.

    `sys.getsizeof(values)` on its own is **not** the answer, and reporting it
    as one is the mistake this function exists to avoid. It measures the list
    object: a header plus one 8-byte pointer per element. It does not measure
    the integers those pointers point at, because they are separate objects
    that the list does not own.

    So this adds the payload: every distinct integer object reachable from the
    list, counted once. CPython caches the small integers from -5 to 256, so
    those really are shared and counting them once is right rather than
    generous.

    On this machine `sys.getsizeof(list(range(1_000_000)))` is 8,000,056 bytes
    -- almost exactly the 8,000,000 an int64 array needs, which would make the
    two look equal. The truth is 36,000,056, because each of the million
    integers is a 28-byte object.

    >>> list_bytes([1, 2, 3]) > sys.getsizeof([1, 2, 3])
    True
    """
    by_identity = {id(x): x for x in values}
    payload = sum(sys.getsizeof(x) for x in by_identity.values())
    return sys.getsizeof(values) + payload


def array_bytes(a: np.ndarray) -> int:
    """What the array's data block costs: `a.nbytes`.

    This is `a.size * a.itemsize` and it is the whole story for the numbers.
    The ndarray object itself adds a small fixed header -- 112 bytes here --
    which `sys.getsizeof` includes and which stops mattering above a few
    hundred elements.

    >>> array_bytes(np.arange(3, dtype=np.int64))
    24
    """
    return int(a.nbytes)


def wrap_int8(value: int, added: int) -> int:
    """Add `added` to `value` inside an int8 array, and report what came out.

    An int8 holds -128 to 127. Adding 1 to 127 does not raise, does not
    promote to a larger type, and on numpy 2.5.2 does not even warn: it wraps
    to -128, silently. That is the single most surprising thing about dtypes,
    and the reason to state a dtype deliberately rather than let one be
    guessed.

    The addition is done with an int8 array on both sides so that NumPy 2's
    promotion rules cannot rescue it. Adding a plain Python `1` to an int8
    array wraps too, which is checked in `02_dtypes_and_overflow.py`.

    >>> wrap_int8(127, 1)
    -128
    """
    a = np.array([value], dtype=np.int8)
    b = np.array([added], dtype=np.int8)
    return int((a + b)[0])


def nan_aware_mean(a: np.ndarray) -> float:
    """The mean of the non-missing entries.

    `a.mean()` on an array containing nan returns nan, because nan propagates
    through every arithmetic operation it touches. That is correct and it is
    useful: it tells you loudly that a value is missing rather than quietly
    averaging over a hole.

    `numpy.nanmean` is the version that skips them, and choosing it should be
    a decision you make rather than a default you inherit.

    >>> nan_aware_mean(np.array([1.0, 2.0, np.nan, 4.0]))
    2.3333333333333335
    """
    return float(np.nanmean(a))


# ===========================================================================
# Helpers -- written for you. Read them; the tests use them.
# ===========================================================================


def time_call(fn: Callable[[], object], repeats: int = 5) -> list[float]:
    """Run `fn` `repeats` times and return every elapsed time in seconds.

    Every time, not the best one and not the average. A single timing is
    noise; a list of them lets the caller take a median and lets a reader see
    the spread. `time.perf_counter` is the right clock: monotonic and the
    highest resolution available.
    """
    times: list[float] = []
    for _ in range(repeats):
        start = time.perf_counter()
        fn()
        times.append(time.perf_counter() - start)
    return times


def median_seconds(times: Iterable[float]) -> float:
    """The median of a list of timings.

    The median rather than the mean, because one unlucky run in which the
    operating system decided to do something else will drag a mean around and
    leave a median alone.
    """
    return statistics.median(times)


def speedup(loop_times: Iterable[float], vec_times: Iterable[float]) -> float:
    """How many times faster the vectorised version was, by median.

    Reported to be read, never asserted on exactly. The lab's tests assert
    that this is at least 20, which is a claim about the SHAPE of the gap and
    survives a slower machine. Asserting the actual figure would make the
    suite fail for everyone whose hardware is not this one.
    """
    return median_seconds(loop_times) / median_seconds(vec_times)


def describe(a: np.ndarray) -> str:
    """The four facts that distinguish an ndarray from a list, in one line."""
    return (
        f"shape={a.shape} dtype={a.dtype} itemsize={a.itemsize} "
        f"nbytes={a.nbytes} strides={a.strides} "
        f"c_contiguous={a.flags['C_CONTIGUOUS']}"
    )

lesson_id: D104
day: 104
kind: guided-build
languages: [python, bash]
setup_commands:
  - cd labs/sections/math-statistics-and-data/day-104-numpy-arrays-and-vectorized-thinking
  - 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_list_versus_array.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 02_dtypes_and_overflow.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 03_same_answer_faster.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 04_creating_and_ufuncs.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 05_masks_and_selection.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 06_axes_views_and_ranking.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 07_nan_and_when_not_to_vectorise.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 . -path ./.venv -prune -o -type d -name '__pycache__' -print -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-17'
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 -> 80 checks, 0 failure(s), exit 0; pytest examples -> 107 passed; pytest starter -> 1 passed, 70 skipped on an untouched checkout, and 71 passed against a fully solved copy of starter/ kept outside the lab. All seven reference scripts exit 0 with every internal assertion holding. Everything was run through a real lab-local .venv created by the documented setup commands, not through an authoring environment. 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 one expectation deliberately swapped for the belief that the nan-aware mean of [1, 2, nan, 4] is 2.5 rather than 7/3, 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. Measured on this run: the three loop-versus-vectorised pairs agreed on all 1,000,000 elements under == in every case, and the vectorised versions ran 106.2x, 134.0x and 123.2x faster by median of five timings; the tests assert only that the ratio exceeds 20, never a duration. Three numbers in the captured output are real and are asserted rather than hidden: sys.getsizeof on a list of a million ints is 8,000,056 bytes against the array''s 8,000,000, which makes the naive comparison say a list is as compact as an array, while the honest total counting the 28-byte integer objects is 36,000,056; np.array([127], dtype=np.int8) + np.array([1], dtype=np.int8) is -128 with no exception and no warning at all on numpy 2.5.2, and the absence of the warning is itself asserted; and x ** 0.5 disagrees with numpy.sqrt on 1,390 of the 1,000,000 values, always by one unit in the last place, while math.sqrt agrees on all of them.'

# What this lab installs, and what it costs you

Two packages, both free, both open source, no account and no key.

| Package | Version pinned here | Licence | Why it is here |
| --- | --- | --- | --- |
| `numpy` | 2.5.2 | BSD 3-Clause | The subject of the day. Today it stops being the answer key beside a from-scratch implementation and becomes the thing being taught. |
| `pytest` | 9.1.1 | MIT | Runs both suites: the reference tests in `examples/` and your running score in `starter/`. |

Both versions are pinned exactly. Section 1 of `tests/run_tests.sh` reads the
installed numpy and compares it against this file rather than trusting it, so a
mismatch is reported rather than discovered later as a puzzling number.

## Why a version this specific

Two things in this lab depend on the NumPy major version and are asserted
rather than described:

- **Integer overflow is silent.** On numpy 2.5.2, `np.array([127],
  dtype=np.int8) + np.array([1], dtype=np.int8)` returns `-128` with no
  exception and no warning. A reference test records the absence of the warning
  so that if a future NumPy started emitting one, the suite would say so and the
  lesson could be corrected instead of quietly going stale.
- **A Python scalar does not widen an array.** Since NumPy 2, `np.array([127],
  dtype=np.int8) + 1` stays `int8` and wraps. Under NumPy 1 the promotion rules
  differed. The harness checks the major version is 2 or later for this reason.

## The network

Installing these two packages is the only thing in this lab that touches the
network. Nothing here opens a socket, reads a URL or needs an API key, and
section 7 of the test harness greps every source file in `examples/` and
`starter/` to prove it.

## If you cannot install anything at all

You cannot do most of this lab, and it would be dishonest to pretend otherwise:
NumPy is the subject.

What you can still do on a bare `python3` with only the standard library:

- write `scale_and_offset_loop`, `roots_loop` and `clip_loop`, which use
  nothing but `math`;
- reproduce the memory measurement, which needs only `sys.getsizeof` — the
  helper `list_bytes` in `starter/vectorize.py` runs without NumPy, and the
  28-bytes-per-integer figure is the whole point of it;
- reason through every prediction in `starter/answers.py`, which is where most
  of the thinking is anyway.

What you lose is the entire right-hand column: every vectorised version, the
speed comparison that motivates the day, boolean masking, `argsort`, views
versus copies, and `nan`.

## Disk

Roughly 60 MB for the virtual environment, almost all of it NumPy. `rm -rf
.venv` from the lab directory is a complete undo.

numpy==2.5.2
pytest==9.1.1

# Day 104 lab — Stop Writing the Loop

One idea holds this whole lab together:

> **A vectorised expression is the SAME computation as the loop, not a
> different one — you have moved the loop, not removed it.**

Everything below is a consequence of that sentence, including the parts where
the trade turns out to be a bad one.

Work in order; each exercise uses the one before it.

Check yourself at any point, from the **lab directory** (the one above this
file):

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

Unattempted work is **skipped**, not failed. On an untouched checkout you will
see `1 passed, 70 skipped`. When it says `71 passed`, you are finished.

---

## Exercise 1 — `vectorize.py` (ten functions)

Write the ten functions marked `raise NotImplementedError`. Each docstring
gives the derivation and a worked example you can check on paper.

| Step | Function | What it must do |
| --- | --- | --- |
| 1.1 | `scale_and_offset_loop` | `m * x + c`, as an explicit Python loop. No NumPy in this one. |
| 1.2 | `scale_and_offset_vec` | The same, as one expression on the whole array. |
| 1.3 | `roots_loop` | Square roots, as a loop. Use `math.sqrt`, and read why. |
| 1.4 | `roots_vec` | The same, as one call to a universal function. |
| 1.5 | `clip_loop` | Hold every value inside `[lo, hi]`, as a loop. |
| 1.6 | `clip_vec` | The same, in one call. |
| 1.7 | `count_above` | How many elements beat a threshold — no counter variable. |
| 1.8 | `mask_between` | A boolean array, True where `lo < x < hi`. |
| 1.9 | `top_k_indices` | The positions of the `k` largest, best first. |
| 1.10 | `cosine_similarities` | Day 103's search, done to every row at once. |

Nine helpers are written for you at the bottom of the file — `list_bytes`,
`array_bytes`, `wrap_int8`, `nan_aware_mean`, `select`, `time_call`,
`median_seconds`, `speedup` and `describe`. Read them; the tests use them, and
`list_bytes` in particular carries a fact you will need in exercise 2.

**The gotcha in 1.3.** Use `math.sqrt(x)`, not `x ** 0.5`. They give different
answers on about one value in seven hundred, and one of the tests compares a
million elements with `==`. Exercise 6.4 is about why, and it is one of the
more interesting things in the day.

**The gotcha in 1.8.** Two of them, actually. It must be `&` and not `and`,
and each comparison needs its own brackets. The docstring explains both, and
exercise 4.6 and 4.7 make you say why.

**The gotcha in 1.10.** `np.linalg.norm(matrix, axis=1)` — think about which
number comes out six times and which comes out four times before you write it.

---

## Exercise 2 — what an array actually is (`answers.py`)

Six predictions about memory and layout.

2.4 is the one worth slowing down for. `sys.getsizeof(list(range(1_000_000)))`
is 8,000,056 bytes and the equivalent int64 array is 8,000,000 — so the naive
measurement says a list is *exactly as compact as an array*, which is the
opposite of everything you have been told. One of those two numbers is not
measuring what you think it is measuring. Work out which before you answer.

2.5 asks for strides. A stride is how many **bytes** you skip to move one step
along an axis. For a 3 by 4 array of int64, moving one column across is one
element, and moving one row down is four.

---

## Exercise 3 — dtypes

Seven predictions, and 3.1 is the famous one. An int8 holds -128 to 127. What
is 127 + 1?

3.2 and 3.3 are the questions that make it matter: does it raise, and does it
warn? Answer both before you run anything. Whatever you expect, the answer is
worth having been wrong about once.

3.5 is subtler. Adding a plain Python `1` — not an int8 — to an int8 array:
which dtype wins? NumPy 2 changed this rule from NumPy 1, so anything you
remember from an older tutorial may be out of date.

---

## Exercise 4 — masking

The twenty readings are printed in `dataset.py` as `SMALL_READINGS_EXPECTED`.
Count by eye; every answer in this exercise can be got without running
anything.

4.6 and 4.7 are the pair the whole day turns on. `(a > 30) and (a < 70)` does
not work. Name the exception class, then say *why* — and note that the reason
is about Python, not about NumPy. NumPy could not fix this if it wanted to.

4.8 is a small one with a large idea inside it: the mean of a boolean array is
the fraction that are True, because `True` is 1 and `False` is 0. Once you see
that, a whole family of "what proportion of..." questions becomes one line.

---

## Exercise 5 — axes, views and copies

5.1 to 5.4 are the axis rule. **The axis you name is the one that
disappears.** If you find yourself reading it as "the axis I want to keep", you
will be off by exactly one every time.

5.5 and 5.6 are the bug. `row = grid[1]` does not give you a copy of row one.
It gives you a window onto row one, and writing through the window writes
through to the array. This is the hardest NumPy bug to find, because the code
that breaks is nowhere near the code that caused it.

5.7 asks which operations give a view. Predict first, then look for the
pattern: it has to do with whether the elements you asked for are **evenly
spaced**.

5.8 is the trap inside the trap: `ravel` and `flatten` do the same job and have
opposite behaviour.

---

## Exercise 6 — sorting, ranking and speed

6.1 to 6.3 are about `argsort`, and about the fact that `np.sort(a)` and
`a.sort()` behave differently — one returns a new array, the other mutates.

6.4 is the honest finding of the lab. Three ways to take a square root, two of
them identical to the last bit and one of them not. Predict which is the odd
one out, then read the reason in `expected-output/FIELDS.md`.

6.5 asks for the order of magnitude of the speedup on a million elements, and
6.6 asks the same question on **four** elements — where the answer reverses.
Both are measurements this lab actually takes.

---

## Exercise 7 — nan

7.1 is the rule everybody meets once: `np.nan == np.nan` is False.

7.4 and 7.6 are the pair that matters. `a.mean()` on an array containing a nan
returns nan. Is that a bug? Think about what the alternative would be: a mean
that silently averaged three readings while you believed it had averaged four,
and never told you.

---

## When you are done

Read the reference. Each script prints its working and asserts every claim it
makes, so nothing in it is decoration:

```bash
cd examples
../.venv/bin/python3 01_list_versus_array.py
../.venv/bin/python3 02_dtypes_and_overflow.py
../.venv/bin/python3 03_same_answer_faster.py
../.venv/bin/python3 04_creating_and_ufuncs.py
../.venv/bin/python3 05_masks_and_selection.py
../.venv/bin/python3 06_axes_views_and_ranking.py
../.venv/bin/python3 07_nan_and_when_not_to_vectorise.py
cd ..
```

Script 07 is the one worth reading even if you stop everything else. The first
half is `nan`. The second half is the honest case against the habit the rest of
the lab has spent an hour teaching you: three situations where the loop is the
better code, all three measured rather than asserted, including one where the
elegant one-line version allocates 80 GB and the ugly loop finishes.

A tool you can only argue for is a tool you do not understand yet.

"""Exercises 2 to 7 -- your predictions. Work them out BEFORE running anything.

Nearly all of these can be done on paper or in your head. That is deliberate:
a lab about NumPy whose answers you cannot check by hand is a lab that teaches
you to trust output.

Replace each `None` with your answer. Anything still `None` is SKIPPED by the
test suite rather than failed, so your score only ever counts work you actually
attempted.

Check yourself from the LAB DIRECTORY:

    .venv/bin/pytest starter -q
"""

# Imported for you: exercise 5.1 asks for an exception CLASS, and several
# answers are dtypes or NumPy values.
import numpy


# =============================================================================
# Exercise 2 -- what an array actually is
# =============================================================================

# 2.1 A list of one million Python integers and an int64 array of the same
#     million numbers. Roughly how many times more memory does the LIST use,
#     counting the integer objects it points at as well as its own pointers?
#     Answer with a float. The lab measures 4.5 on this machine; you are being
#     asked whether you expect roughly 1, roughly 4.5, or roughly 100.
LIST_TO_ARRAY_MEMORY_RATIO = None

# 2.2 How many bytes is one Python int on CPython 3.14, as sys.getsizeof
#     reports it? An integer.
BYTES_PER_PYTHON_INT = None

# 2.3 How many bytes is one element of an int64 array? An integer.
BYTES_PER_INT64_ELEMENT = None

# 2.4 sys.getsizeof(list(range(1_000_000))) comes out at 8,000,056 -- almost
#     exactly the same as the array's 8,000,000. Why does that number NOT show
#     that a list is as compact as an array?
#     Answer with one of these strings:
#       "the list is genuinely as compact"
#       "getsizeof measures only the pointers, not the integers"
#       "getsizeof is inaccurate for large objects"
WHY_GETSIZEOF_MISLEADS = None

# 2.5 np.arange(12).reshape(3, 4) has int64 elements. What are its strides, as
#     a tuple of two integers? Think: how many BYTES do you skip to move one
#     row down, and how many to move one column across?
STRIDES_OF_A_THREE_BY_FOUR = None

# 2.6 Does transposing that array copy any data? True or False.
TRANSPOSE_COPIES_DATA = None


# =============================================================================
# Exercise 3 -- dtypes
# =============================================================================

# 3.1 np.array([127], dtype=np.int8) + np.array([1], dtype=np.int8)
#     What is the single value that comes out? An integer.
INT8_127_PLUS_1 = None

# 3.2 Does that raise an exception? True or False.
INT8_OVERFLOW_RAISES = None

# 3.3 Does it emit a warning on numpy 2.5.2? True or False.
INT8_OVERFLOW_WARNS = None

# 3.4 np.array([120, 125, 127], dtype=np.int8) * np.int8(2)
#     Give all three values, as a list of three integers.
INT8_DOUBLED = None

# 3.5 np.array([127], dtype=np.int8) + 1  -- note the plain Python 1 this time.
#     What DTYPE does the result have? Give the numpy dtype itself, for
#     example numpy.int16, not the string "int16".
DTYPE_OF_INT8_PLUS_PYTHON_INT = None

# 3.6 A float32 has 24 bits of significand, so above 2 ** 24 the gap between
#     representable values is at least 1. Is np.float32(16777216.0) + 1 equal
#     to np.float32(16777216.0)? True or False.
FLOAT32_CANNOT_ADD_ONE = None

# 3.7 np.full(3, 7) -- what dtype? np.full(3, 7.0) -- what dtype?
#     A tuple of two numpy dtypes, for example (numpy.int8, numpy.float32).
DTYPES_OF_FULL = None


# =============================================================================
# Exercise 4 -- masking
# =============================================================================
#
# The twenty readings are in dataset.py, both as a seeded generator and
# written out as SMALL_READINGS_EXPECTED so you can count by eye:
#
#   [70, 83, 34, 69, 26, 21, 18, 12, 65, 37,
#    17, 75, 30, 73, 37, 41, 97, 64, 21, 82]

# 4.1 How many are strictly greater than 50? An integer.
COUNT_ABOVE_50 = None

# 4.2 Which ones? A list of integers, in the order they appear in the data.
VALUES_ABOVE_50 = None

# 4.3 (readings > 30) & (readings < 70) -- how many? An integer.
COUNT_BETWEEN_30_AND_70 = None

# 4.4 What is the SHAPE of `readings > 50`? A tuple.
SHAPE_OF_A_MASK = None

# 4.5 What is the DTYPE of `readings > 50`? A numpy dtype, e.g. numpy.int64.
DTYPE_OF_A_MASK = None

# 4.6 (readings > 30) and (readings < 70) -- with the keyword rather than the
#     operator. Name the exception CLASS this raises. Give the class itself,
#     not a string. It is a builtin.
EXCEPTION_FROM_KEYWORD_AND = None

# 4.7 Why does `and` fail where `&` works? One of these strings:
#       "NumPy forgot to implement it"
#       "and is a keyword, so it asks the array for a single True or False"
#       "and only works on lists"
WHY_AND_FAILS = None

# 4.8 mask.mean() on `readings > 50`, where mask has 9 Trues out of 20.
#     A float.
MEAN_OF_THE_MASK = None

# 4.9 readings[[0, 5, 19, 5]] -- fancy indexing, and note the repeat.
#     A list of four integers.
FANCY_INDEX_RESULT = None


# =============================================================================
# Exercise 5 -- axes, views and copies
# =============================================================================

# 5.1 np.arange(12).reshape(3, 4).sum(axis=0) -- what SHAPE comes out?
#     A tuple. The rule: the axis you name is the one that disappears.
SHAPE_AFTER_SUM_AXIS_0 = None

# 5.2 And .sum(axis=1)? A tuple.
SHAPE_AFTER_SUM_AXIS_1 = None

# 5.3 .sum(axis=1) on that array -- the three values, as a list of integers.
VALUES_OF_SUM_AXIS_1 = None

# 5.4 v = np.array([1.0, 2.0, 3.0]). What shape is v[:, np.newaxis]? A tuple.
SHAPE_OF_A_COLUMN = None

# 5.5 grid = np.arange(12).reshape(3, 4); row = grid[1]; row[0] = 999
#     What is grid[1, 0] afterwards? An integer.
GRID_AFTER_WRITING_THROUGH_A_SLICE = None

# 5.6 Same again, but with `row = grid[1].copy()`. An integer.
GRID_AFTER_WRITING_THROUGH_A_COPY = None

# 5.7 Which of these share memory with the array they came from? Answer with a
#     list of the labels that DO, in this order, as strings:
#       "row slice", "column slice", "transpose", "reshape",
#       "boolean mask", "fancy index"
#     For example: ["row slice", "reshape"]
WHICH_ARE_VIEWS = None

# 5.8 grid.ravel() and grid.flatten() both give you a flat array. One is a
#     view when it can be and one is always a copy. Which is always a copy?
#     The string "ravel" or the string "flatten".
ALWAYS_A_COPY = None


# =============================================================================
# Exercise 6 -- sorting, ranking and speed
# =============================================================================

# 6.1 scores = np.array([5.0, 1.0, 9.0, 3.0]). What is np.argsort(scores)?
#     A list of four integers.
ARGSORT_OF_THE_SCORES = None

# 6.2 Does `np.sort(scores)` change `scores` itself? True or False.
NP_SORT_MUTATES = None

# 6.3 Does the METHOD `scores.sort()` change `scores` itself? True or False.
SORT_METHOD_MUTATES = None

# 6.4 The lab computes the square root of a million values three ways:
#     math.sqrt in a loop, x ** 0.5 in a loop, and np.sqrt on the array.
#     Two of the three agree bit for bit. Which one is the odd one out?
#     One of these strings: "math.sqrt", "x ** 0.5", "np.sqrt"
THE_ODD_ONE_OUT = None

# 6.5 Roughly how much faster is the vectorised version than the loop, on a
#     million elements? Not the exact figure -- the ORDER OF MAGNITUDE.
#     One of these strings: "about the same", "about 2x", "about 100x"
ROUGH_SPEEDUP = None

# 6.6 On an array of FOUR elements, which is faster: the list comprehension or
#     the NumPy call? One of these strings: "comprehension", "numpy"
FASTER_ON_FOUR_ELEMENTS = None


# =============================================================================
# Exercise 7 -- nan
# =============================================================================

# 7.1 np.nan == np.nan -- True or False.
NAN_EQUALS_ITSELF = None

# 7.2 a = np.array([1.0, 2.0, np.nan, 4.0]). How many elements does the mask
#     `a == np.nan` mark as True? An integer.
COUNT_FROM_COMPARING_TO_NAN = None

# 7.3 And np.isnan(a).sum()? An integer.
COUNT_FROM_ISNAN = None

# 7.4 What does a.mean() return for that array? Answer with the string "nan"
#     if it returns nan, or with the float if it returns a number.
MEAN_OF_THE_HOLED_ARRAY = None

# 7.5 And np.nanmean(a)? A float. Work out the arithmetic first: it is the
#     mean of the values that are actually there.
NANMEAN_OF_THE_HOLED_ARRAY = None

# 7.6 Is a.mean() returning nan a bug in NumPy? True or False, and think about
#     what the alternative would mean for a reading that was never taken.
NAN_PROPAGATION_IS_A_BUG = None

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

Both `examples/` and `starter/` contain modules called `vectorize` and
`dataset`, 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 `vectorize` was seen first and then reuse it for the
other suite -- so these 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.

So: put this directory first on the import path, and drop any already-imported
`vectorize`, `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 ("vectorize", "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 data this lab measures, and the tolerances it compares with.

Three kinds of data live here, and each is here for a different reason.

1. A **big** array of a million numbers, drawn from a seeded generator. It
   exists so the memory and speed measurements have something to bite on. A
   thousand elements would make the loop look fine.

2. A **small** array of twenty integers, drawn from the same seeded generator.
   It exists so every boolean mask in this lab can be checked by eye. You can
   read the twenty numbers, count the ones above fifty, and compare with what
   the code says.

3. The **article catalogue from Day 99 and Day 103**, unchanged. Six invented
   articles described by four hand-counted features. Day 103 ranked them by
   cosine similarity; today the ranking is done with `argsort`, which is the
   part of that day that transfers directly to model code.

Nothing here is real. The articles do not exist, the sensor readings are not
sensor readings, and the counts were chosen so that every number in this lab
can be re-derived with a pen.

The seed is fixed at 104 and every number below follows from it, so two runs of
this lab on two machines produce the same values. That is a deliberate choice
rather than a convenience: a lab that asserts on random numbers is a lab that
fails for the reader and not for the author.
"""

from __future__ import annotations

import numpy as np

# ---------------------------------------------------------------------------
# Seeds and sizes
# ---------------------------------------------------------------------------

#: The one seed this lab uses. Passed to numpy.random.default_rng.
SEED = 104

#: How many elements the timing and memory comparisons use. One million is
#: large enough that the loop is unmistakably slower and small enough that the
#: array is 8 MB rather than 8 GB.
N_BIG = 1_000_000

#: How many elements the by-eye mask exercises use.
N_SMALL = 20

#: Float comparison tolerance, used wherever two routes to the same number are
#: allowed to differ in the last bits. Most comparisons in this lab do NOT use
#: it: the whole point of the from-scratch section is that the loop and the
#: vectorised version agree EXACTLY, and asserting that with a tolerance would
#: hide the very fact being demonstrated.
TOL = 1e-12


def big_values() -> np.ndarray:
    """A million float64 values in [0, 1), from the fixed seed.

    Returns a fresh array each call so that a test which mutates it cannot
    quietly change the answer of a later one.
    """
    return np.random.default_rng(SEED).random(N_BIG)


def small_readings() -> np.ndarray:
    """Twenty integers in [0, 100), from the fixed seed.

    Drawn from a generator seeded independently of `big_values`, so the two
    are reproducible on their own.
    """
    return np.random.default_rng(SEED).integers(0, 100, size=N_SMALL)


# The literal values `small_readings()` produces on the fixed seed, written out
# so you can check a mask against them without running anything. The reference
# tests assert that this list and the generator still agree; if a future NumPy
# ever changed the generator's output, the suite would say so rather than
# quietly rewriting the lesson.
SMALL_READINGS_EXPECTED = [
    70, 83, 34, 69, 26, 21, 18, 12, 65, 37,
    17, 75, 30, 73, 37, 41, 97, 64, 21, 82,
]

# ---------------------------------------------------------------------------
# The article catalogue, carried unchanged from Day 99 and Day 103
# ---------------------------------------------------------------------------

FEATURES = ("cooking", "running", "money", "weather")

ARTICLE_NAMES = (
    "roast-chicken",
    "slow-cooker-stew",
    "marathon-plan",
    "race-day-nutrition",
    "household-budget",
    "storm-bulletin",
)

#: One row per article, one column per feature, in the order above. Day 103
#: held these as six separate lists; today they are one 6 by 4 array, and that
#: change is the day's subject rather than a tidy-up.
CATALOGUE = np.array(
    [
        [9, 0, 1, 0],   # roast-chicken
        [8, 0, 2, 0],   # slow-cooker-stew
        [0, 9, 1, 2],   # marathon-plan
        [4, 6, 3, 0],   # race-day-nutrition
        [1, 0, 9, 0],   # household-budget
        [0, 1, 0, 9],   # storm-bulletin
    ],
    dtype=np.float64,
)

#: "training for a race and what to eat", written as feature counts the same
#: way the articles were. Deliberately a close call between two articles.
QUERY = np.array([2, 5, 0, 0], dtype=np.float64)

#: How many results the search returns.
TOP_K = 3

# ---------------------------------------------------------------------------
# The three operations implemented twice
# ---------------------------------------------------------------------------
#
# The multiplier, offset and clip bounds are all exactly representable in
# binary floating point (2.5 is 10.1 in binary, 1.25 is 1.01, 0.25 is 0.01,
# 0.75 is 0.11). That matters: it means the loop and the vectorised version
# perform the identical IEEE-754 operation on the identical bits, so they can
# be compared with `==` rather than with a tolerance.

SCALE_M = 2.5
SCALE_C = 1.25
CLIP_LO = 0.25
CLIP_HI = 0.75

# ---------------------------------------------------------------------------
# The dtype demonstrations
# ---------------------------------------------------------------------------

#: The largest value an int8 can hold. Adding 1 to this wraps to INT8_MIN.
INT8_MAX = 127
INT8_MIN = -128

#: Three int8 values doubled in section 2. Two of the three wrap.
INT8_DOUBLING_INPUT = [120, 125, 127]

#: A value float32 cannot tell apart from its successor: 2 ** 24. A float32
#: has 24 bits of significand, so at this magnitude the gap between
#: representable numbers is exactly 1, and adding 1 changes nothing.
FLOAT32_BLIND_SPOT = 16777216.0

# ---------------------------------------------------------------------------
# The array with a hole in it
# ---------------------------------------------------------------------------

#: Four readings, one of which is missing. `mean` on this returns nan; the
#: whole point of section 7 is that the nan is contagious and that saying so
#: loudly is better than silently dropping it.
WITH_A_HOLE = np.array([1.0, 2.0, np.nan, 4.0])

"""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.

Every test that exercises your code runs its whole body inside `written(...)`,
so a test is skipped if ANY function it needs is still unwritten -- not just
the first one. Python evaluates arguments before the call, so gating on one
function while calling another inside the arguments would let a
NotImplementedError escape and be reported as a failure. It would say
"attempted and wrong" about work you had not attempted, which is precisely the
lie this suite exists to avoid.
"""

import math
import sys

import numpy as np
import pytest

import answers
import dataset
from vectorize import (
    array_bytes,
    clip_loop,
    clip_vec,
    cosine_similarities,
    count_above,
    list_bytes,
    mask_between,
    median_seconds,
    nan_aware_mean,
    roots_loop,
    roots_vec,
    scale_and_offset_loop,
    scale_and_offset_vec,
    select,
    speedup,
    time_call,
    top_k_indices,
    wrap_int8,
)

TOL = dataset.TOL


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 int(np.__version__.split(".")[0]) >= 2, "numpy 2 or later is importable"
    assert dataset.small_readings().tolist() == dataset.SMALL_READINGS_EXPECTED, (
        "the seeded generator produces the documented twenty readings"
    )
    assert array_bytes(np.arange(3, dtype=np.int64)) == 24, "the helpers load"


# -- Exercise 1: your vectorize.py -------------------------------------------


def test_1_1_scale_and_offset_loop():
    got = written(lambda: scale_and_offset_loop([0.0, 1.0, 2.0], 2.5, 1.25))
    assert got == [1.25, 3.75, 6.25]


def test_1_1_scale_and_offset_loop_returns_a_plain_list():
    got = written(lambda: scale_and_offset_loop([0.0, 1.0], 2.0, 0.0))
    assert isinstance(got, list), "return a list, not an array"


def test_1_2_scale_and_offset_vec():
    got = written(lambda: scale_and_offset_vec(np.array([0.0, 1.0, 2.0]), 2.5, 1.25))
    assert got.tolist() == [1.25, 3.75, 6.25]


def test_1_2_scale_and_offset_vec_returns_an_array():
    got = written(lambda: scale_and_offset_vec(np.array([0.0, 1.0]), 2.0, 0.0))
    assert isinstance(got, np.ndarray), "return an array, not a list"


def test_1_1_and_1_2_agree_exactly_on_a_million_elements():
    """The claim the whole lab rests on: the SAME computation, not a similar
    one. Compared with ==, elementwise, over a million values."""
    values = dataset.big_values()

    def run():
        loop = scale_and_offset_loop(values.tolist(), dataset.SCALE_M, dataset.SCALE_C)
        vec = scale_and_offset_vec(values, dataset.SCALE_M, dataset.SCALE_C)
        return np.array(loop), vec

    loop, vec = written(run)
    assert np.array_equal(loop, vec), (
        "every one of the million elements must match bit for bit"
    )


def test_1_3_roots_loop():
    assert written(lambda: roots_loop([0.0, 1.0, 4.0])) == [0.0, 1.0, 2.0]


def test_1_4_roots_vec():
    assert written(lambda: roots_vec(np.array([0.0, 1.0, 4.0]))).tolist() == [
        0.0,
        1.0,
        2.0,
    ]


def test_1_3_and_1_4_agree_exactly_on_a_million_elements():
    """If this fails while 1.3 passes, check that roots_loop uses math.sqrt
    and not x ** 0.5. They are not the same operation; exercise 6.4 is about
    exactly this."""
    values = dataset.big_values()

    def run():
        return np.array(roots_loop(values.tolist())), roots_vec(values)

    loop, vec = written(run)
    differing = int(np.count_nonzero(loop != vec))
    assert differing == 0, (
        f"{differing} of {values.size} elements differ; math.sqrt agrees with "
        "numpy.sqrt bit for bit, and x ** 0.5 does not"
    )


def test_1_5_clip_loop():
    assert written(lambda: clip_loop([0.0, 0.5, 1.0], 0.25, 0.75)) == [0.25, 0.5, 0.75]


def test_1_5_clip_loop_leaves_interior_values_untouched():
    got = written(lambda: clip_loop([0.3, 0.4, 0.5], 0.25, 0.75))
    assert got == [0.3, 0.4, 0.5]


def test_1_6_clip_vec():
    got = written(lambda: clip_vec(np.array([0.0, 0.5, 1.0]), 0.25, 0.75))
    assert got.tolist() == [0.25, 0.5, 0.75]


def test_1_5_and_1_6_agree_exactly_on_a_million_elements():
    values = dataset.big_values()

    def run():
        loop = clip_loop(values.tolist(), dataset.CLIP_LO, dataset.CLIP_HI)
        return np.array(loop), clip_vec(values, dataset.CLIP_LO, dataset.CLIP_HI)

    loop, vec = written(run)
    assert np.array_equal(loop, vec)


def test_the_vectorised_versions_are_much_faster():
    """A claim about the SHAPE of the gap, never the figure. Twenty times
    survives a slow machine; the authoring machine measured over a hundred."""
    values = dataset.big_values()
    as_list = values.tolist()

    def run():
        scale_and_offset_loop(as_list, dataset.SCALE_M, dataset.SCALE_C)
        scale_and_offset_vec(values, dataset.SCALE_M, dataset.SCALE_C)
        loop = time_call(
            lambda: scale_and_offset_loop(as_list, dataset.SCALE_M, dataset.SCALE_C), 3
        )
        vec = time_call(
            lambda: scale_and_offset_vec(values, dataset.SCALE_M, dataset.SCALE_C), 3
        )
        return speedup(loop, vec)

    factor = written(run)
    assert factor > 20.0, f"measured {factor:.1f}x, which is below the 20x floor"


def test_1_7_count_above():
    assert written(lambda: count_above(np.array([1, 5, 9]), 4)) == 2


def test_1_7_count_above_returns_a_plain_int():
    got = written(lambda: count_above(np.array([1, 5, 9]), 4))
    assert type(got) is int, "wrap the sum in int()"


def test_1_7_count_above_on_the_twenty_readings():
    readings = dataset.small_readings()
    assert written(lambda: count_above(readings, 50)) == 9


def test_1_8_mask_between():
    got = written(lambda: mask_between(np.array([1, 5, 9]), 2, 8))
    assert got.tolist() == [False, True, False]


def test_1_8_mask_between_returns_a_boolean_array():
    got = written(lambda: mask_between(np.array([1, 5, 9]), 2, 8))
    assert got.dtype == np.bool_, "a mask is boolean, not integer"


def test_1_8_mask_between_is_strict_at_both_ends():
    got = written(lambda: mask_between(np.array([2, 5, 8]), 2, 8))
    assert got.tolist() == [False, True, False], "lo < x < hi, both strict"


def test_1_8_mask_between_on_the_twenty_readings():
    readings = dataset.small_readings()
    mask = written(lambda: mask_between(readings, 30, 70))
    assert readings[mask].tolist() == [34, 69, 65, 37, 37, 41, 64]


def test_1_9_top_k_indices():
    got = written(lambda: top_k_indices(np.array([0.1, 0.9, 0.5]), 2))
    assert list(got) == [1, 2]


def test_1_9_top_k_indices_returns_exactly_k():
    got = written(lambda: top_k_indices(np.array([0.1, 0.9, 0.5, 0.7]), 3))
    assert len(got) == 3


def test_1_9_top_k_indices_gives_positions_not_values():
    got = written(lambda: top_k_indices(np.array([10.0, 30.0, 20.0]), 1))
    assert list(got) == [1], "index 1, not the value 30.0"


def test_1_10_cosine_similarities():
    identity = np.array([[1.0, 0.0], [0.0, 1.0]])
    got = written(lambda: cosine_similarities(identity, np.array([1.0, 0.0])))
    assert got.tolist() == [1.0, 0.0]


def test_1_10_cosine_similarities_gives_one_score_per_row():
    got = written(lambda: cosine_similarities(dataset.CATALOGUE, dataset.QUERY))
    assert got.shape == (6,), "axis=1 collapses the columns, leaving one per row"


def test_1_10_cosine_similarities_matches_day_103_row_by_row():
    got = written(lambda: cosine_similarities(dataset.CATALOGUE, dataset.QUERY))
    by_hand = []
    for row in dataset.CATALOGUE:
        dot = sum(float(x) * float(y) for x, y in zip(row, dataset.QUERY))
        left = math.sqrt(sum(float(x) * float(x) for x in row))
        right = math.sqrt(sum(float(y) * float(y) for y in dataset.QUERY))
        by_hand.append(dot / (left * right))
    assert np.allclose(got, by_hand, atol=TOL)


def test_1_9_and_1_10_together_rank_the_catalogue():
    def run():
        sims = cosine_similarities(dataset.CATALOGUE, dataset.QUERY)
        return top_k_indices(sims, dataset.TOP_K)

    top = written(run)
    assert [dataset.ARTICLE_NAMES[i] for i in top] == [
        "race-day-nutrition",
        "marathon-plan",
        "roast-chicken",
    ]


# -- Exercise 2: what an array actually is ------------------------------------


def test_2_1_memory_ratio():
    guess = predicted("LIST_TO_ARRAY_MEMORY_RATIO")
    values = list(range(dataset.N_BIG))
    real = list_bytes(values) / np.arange(dataset.N_BIG, dtype=np.int64).nbytes
    assert abs(guess - real) < 1.0, f"measured {real:.2f}"


def test_2_2_bytes_per_python_int():
    assert predicted("BYTES_PER_PYTHON_INT") == sys.getsizeof(1_000_000)


def test_2_3_bytes_per_int64_element():
    assert predicted("BYTES_PER_INT64_ELEMENT") == np.zeros(1, dtype=np.int64).itemsize


def test_2_4_why_getsizeof_misleads():
    assert (
        predicted("WHY_GETSIZEOF_MISLEADS")
        == "getsizeof measures only the pointers, not the integers"
    )


def test_2_5_strides():
    assert tuple(predicted("STRIDES_OF_A_THREE_BY_FOUR")) == np.arange(12).reshape(
        3, 4
    ).strides


def test_2_6_transpose_copies_nothing():
    grid = np.arange(12).reshape(3, 4)
    shares = bool(np.shares_memory(grid, grid.T))
    assert predicted("TRANSPOSE_COPIES_DATA") is not shares
    assert shares is True, "a transpose is a view; nothing was copied"


# -- Exercise 3: dtypes -------------------------------------------------------


def test_3_1_int8_overflow_value():
    assert predicted("INT8_127_PLUS_1") == wrap_int8(127, 1)


def test_3_2_int8_overflow_does_not_raise():
    assert predicted("INT8_OVERFLOW_RAISES") is False


def test_3_3_int8_overflow_does_not_warn_on_this_numpy():
    import warnings

    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        wrap_int8(127, 1)
    assert predicted("INT8_OVERFLOW_WARNS") is (len(caught) > 0)


def test_3_4_int8_doubled():
    real = (np.array(dataset.INT8_DOUBLING_INPUT, dtype=np.int8) * np.int8(2)).tolist()
    assert list(predicted("INT8_DOUBLED")) == real


def test_3_5_dtype_of_int8_plus_python_int():
    real = (np.array([127], dtype=np.int8) + 1).dtype
    assert np.dtype(predicted("DTYPE_OF_INT8_PLUS_PYTHON_INT")) == real


def test_3_6_float32_blind_spot():
    blind = np.float32(dataset.FLOAT32_BLIND_SPOT)
    assert predicted("FLOAT32_CANNOT_ADD_ONE") is bool(blind + np.float32(1.0) == blind)


def test_3_7_dtypes_of_full():
    guessed = predicted("DTYPES_OF_FULL")
    assert np.dtype(guessed[0]) == np.full(3, 7).dtype
    assert np.dtype(guessed[1]) == np.full(3, 7.0).dtype


# -- Exercise 4: masking ------------------------------------------------------


def test_4_1_count_above_50():
    readings = dataset.small_readings()
    assert predicted("COUNT_ABOVE_50") == int((readings > 50).sum())


def test_4_2_values_above_50():
    readings = dataset.small_readings()
    assert list(predicted("VALUES_ABOVE_50")) == readings[readings > 50].tolist()


def test_4_3_count_between_30_and_70():
    readings = dataset.small_readings()
    real = int(((readings > 30) & (readings < 70)).sum())
    assert predicted("COUNT_BETWEEN_30_AND_70") == real


def test_4_4_shape_of_a_mask():
    readings = dataset.small_readings()
    assert tuple(predicted("SHAPE_OF_A_MASK")) == (readings > 50).shape


def test_4_5_dtype_of_a_mask():
    readings = dataset.small_readings()
    assert np.dtype(predicted("DTYPE_OF_A_MASK")) == (readings > 50).dtype


def test_4_6_exception_from_keyword_and():
    readings = dataset.small_readings()
    guess = predicted("EXCEPTION_FROM_KEYWORD_AND")
    with pytest.raises(guess):
        (readings > 30) and (readings < 70)


def test_4_7_why_and_fails():
    assert (
        predicted("WHY_AND_FAILS")
        == "and is a keyword, so it asks the array for a single True or False"
    )


def test_4_8_mean_of_the_mask():
    readings = dataset.small_readings()
    assert predicted("MEAN_OF_THE_MASK") == float((readings > 50).mean())


def test_4_9_fancy_index_result():
    readings = dataset.small_readings()
    real = readings[np.array([0, 5, 19, 5])].tolist()
    assert list(predicted("FANCY_INDEX_RESULT")) == real


# -- Exercise 5: axes, views and copies ---------------------------------------


def test_5_1_shape_after_sum_axis_0():
    assert tuple(predicted("SHAPE_AFTER_SUM_AXIS_0")) == np.arange(12).reshape(
        3, 4
    ).sum(axis=0).shape


def test_5_2_shape_after_sum_axis_1():
    assert tuple(predicted("SHAPE_AFTER_SUM_AXIS_1")) == np.arange(12).reshape(
        3, 4
    ).sum(axis=1).shape


def test_5_3_values_of_sum_axis_1():
    real = np.arange(12).reshape(3, 4).sum(axis=1).tolist()
    assert list(predicted("VALUES_OF_SUM_AXIS_1")) == real


def test_5_4_shape_of_a_column():
    v = np.array([1.0, 2.0, 3.0])
    assert tuple(predicted("SHAPE_OF_A_COLUMN")) == v[:, np.newaxis].shape


def test_5_5_writing_through_a_slice():
    grid = np.arange(12).reshape(3, 4)
    row = grid[1]
    row[0] = 999
    assert predicted("GRID_AFTER_WRITING_THROUGH_A_SLICE") == int(grid[1, 0])


def test_5_6_writing_through_a_copy():
    grid = np.arange(12).reshape(3, 4)
    row = grid[1].copy()
    row[0] = 999
    assert predicted("GRID_AFTER_WRITING_THROUGH_A_COPY") == int(grid[1, 0])


def test_5_7_which_are_views():
    grid = np.arange(12).reshape(3, 4)
    real = [
        label
        for label, result in (
            ("row slice", grid[1]),
            ("column slice", grid[:, 1]),
            ("transpose", grid.T),
            ("reshape", grid.reshape(4, 3)),
            ("boolean mask", grid[grid > 5]),
            ("fancy index", grid[[0, 2]]),
        )
        if np.shares_memory(grid, result)
    ]
    assert sorted(predicted("WHICH_ARE_VIEWS")) == sorted(real)


def test_5_8_always_a_copy():
    grid = np.arange(12).reshape(3, 4)
    assert not np.shares_memory(grid, grid.flatten())
    assert np.shares_memory(grid, grid.ravel())
    assert predicted("ALWAYS_A_COPY") == "flatten"


# -- Exercise 6: sorting, ranking and speed -----------------------------------


def test_6_1_argsort_of_the_scores():
    scores = np.array([5.0, 1.0, 9.0, 3.0])
    assert list(predicted("ARGSORT_OF_THE_SCORES")) == np.argsort(scores).tolist()


def test_6_2_np_sort_does_not_mutate():
    scores = np.array([5.0, 1.0, 9.0, 3.0])
    before = scores.tolist()
    np.sort(scores)
    assert predicted("NP_SORT_MUTATES") is (scores.tolist() != before)


def test_6_3_the_sort_method_does_mutate():
    scores = np.array([5.0, 1.0, 9.0, 3.0])
    before = scores.tolist()
    scores.sort()
    assert predicted("SORT_METHOD_MUTATES") is (scores.tolist() != before)


def test_6_4_the_odd_one_out():
    """Measured on a hundred values rather than a million, so this test stays
    quick. The effect is the same one script 07 measures at full size."""
    sample = dataset.big_values()[:100_000]
    by_math = np.array([math.sqrt(x) for x in sample.tolist()])
    by_pow = np.array([x ** 0.5 for x in sample.tolist()])
    vec = np.sqrt(sample)
    assert np.array_equal(by_math, vec)
    assert not np.array_equal(by_pow, vec)
    assert predicted("THE_ODD_ONE_OUT") == "x ** 0.5"


def test_6_5_rough_speedup():
    assert predicted("ROUGH_SPEEDUP") == "about 100x"


def test_6_6_faster_on_four_elements():
    small = [1.0, 2.0, 3.0, 4.0]
    loop = median_seconds(time_call(lambda: [math.sqrt(x) for x in small], 2000))
    vec = median_seconds(time_call(lambda: np.sqrt(np.array(small)), 2000))
    real = "comprehension" if loop < vec else "numpy"
    assert predicted("FASTER_ON_FOUR_ELEMENTS") == real


# -- Exercise 7: nan ----------------------------------------------------------


def test_7_1_nan_equals_itself():
    assert predicted("NAN_EQUALS_ITSELF") is (np.nan == np.nan)


def test_7_2_comparing_to_nan_finds_nothing():
    real = int((dataset.WITH_A_HOLE == np.nan).sum())
    assert predicted("COUNT_FROM_COMPARING_TO_NAN") == real


def test_7_3_isnan_finds_it():
    assert predicted("COUNT_FROM_ISNAN") == int(np.isnan(dataset.WITH_A_HOLE).sum())


def test_7_4_mean_of_the_holed_array():
    assert predicted("MEAN_OF_THE_HOLED_ARRAY") == "nan"
    assert math.isnan(float(dataset.WITH_A_HOLE.mean()))


def test_7_5_nanmean_of_the_holed_array():
    guess = predicted("NANMEAN_OF_THE_HOLED_ARRAY")
    assert abs(guess - nan_aware_mean(dataset.WITH_A_HOLE)) <= TOL


def test_7_6_nan_propagation_is_not_a_bug():
    assert predicted("NAN_PROPAGATION_IS_A_BUG") is False


# -- A final check that uses your work end to end ------------------------------


def test_the_whole_pipeline_on_the_twenty_readings():
    """Mask, count, select, rank -- all four of your functions in one go."""
    readings = dataset.small_readings()

    def run():
        mask = mask_between(readings, 30, 70)
        chosen = select(readings, mask)
        return count_above(readings, 50), chosen, top_k_indices(chosen.astype(float), 3)

    count, chosen, top = written(run)
    assert count == 9
    assert chosen.tolist() == [34, 69, 65, 37, 37, 41, 64]
    assert chosen[top].tolist() == [69, 65, 64]

"""Exercise 1 -- ten functions to write. Your work goes here.

Each function raises NotImplementedError until you write it, and the test
suite SKIPS anything still unwritten rather than failing it. So your score
only ever counts work you actually attempted.

Check yourself from the LAB DIRECTORY (the one above this file):

    .venv/bin/pytest starter -q

Read each docstring before you write the body. Every one of them gives the
derivation and a worked example small enough to check on paper.

Three of these are implemented TWICE -- once as an explicit Python loop and
once as a NumPy expression. That is the spine of the whole lab: a vectorised
expression is the SAME computation as the loop, not a different one, and the
tests prove it by comparing a million elements with `==` rather than with a
tolerance.

The helpers at the bottom of this file are written for you. Read them; the
tests use them.
"""

from __future__ import annotations

import math
import statistics
import sys
import time
from typing import Callable, Iterable, Sequence

import numpy as np

# ===========================================================================
# 1. The three operations, each implemented twice
# ===========================================================================


def scale_and_offset_loop(values: Sequence[float], m: float, c: float) -> list[float]:
    """1.1 -- `m * x + c` for every element, as an explicit Python loop.

    Return a NEW list. Do not use NumPy in this function; the whole point is
    that this is the code you have been writing for a hundred days.

    Allocate the output up front with `[0.0] * len(values)` and assign into it
    rather than appending. Appending would be a slower loop, and comparing a
    slow loop against a fast array would be rigging the measurement.

    >>> scale_and_offset_loop([0.0, 1.0, 2.0], 2.5, 1.25)
    [1.25, 3.75, 6.25]
    """
    raise NotImplementedError("scale_and_offset_loop")


def scale_and_offset_vec(a: np.ndarray, m: float, c: float) -> np.ndarray:
    """1.2 -- the same thing as one expression on the whole array.

    One line. No loop, no comprehension, no `for` anywhere.

    >>> scale_and_offset_vec(np.array([0.0, 1.0, 2.0]), 2.5, 1.25).tolist()
    [1.25, 3.75, 6.25]
    """
    raise NotImplementedError("scale_and_offset_vec")


def roots_loop(values: Sequence[float]) -> list[float]:
    """1.3 -- the square root of every element, as a loop.

    Use `math.sqrt(x)` and NOT `x ** 0.5`, and the difference is not pedantry.
    IEEE-754 requires square root to be correctly rounded, and `math.sqrt`
    uses the hardware instruction that obeys that requirement -- so it agrees
    with `numpy.sqrt` bit for bit. `pow(x, 0.5)` is a general power routine
    with no such guarantee, and on this machine it disagrees on 1390 of the
    lab's million values. Exercise 6.4 asks you to predict that.

    >>> roots_loop([0.0, 1.0, 4.0])
    [0.0, 1.0, 2.0]
    """
    raise NotImplementedError("roots_loop")


def roots_vec(a: np.ndarray) -> np.ndarray:
    """1.4 -- the same thing as one call to a universal function.

    >>> roots_vec(np.array([0.0, 1.0, 4.0])).tolist()
    [0.0, 1.0, 2.0]
    """
    raise NotImplementedError("roots_vec")


def clip_loop(values: Sequence[float], lo: float, hi: float) -> list[float]:
    """1.5 -- pull every element back inside `[lo, hi]`, as a loop.

    Anything below `lo` becomes `lo`, anything above `hi` becomes `hi`, and
    everything else is left exactly as it was.

    >>> clip_loop([0.0, 0.5, 1.0], 0.25, 0.75)
    [0.25, 0.5, 0.75]
    """
    raise NotImplementedError("clip_loop")


def clip_vec(a: np.ndarray, lo: float, hi: float) -> np.ndarray:
    """1.6 -- the same thing in one call.

    There is a NumPy function named after exactly this operation. The branch
    has not disappeared, it has moved into C.

    >>> clip_vec(np.array([0.0, 0.5, 1.0]), 0.25, 0.75).tolist()
    [0.25, 0.5, 0.75]
    """
    raise NotImplementedError("clip_vec")


# ===========================================================================
# 2. Masking, selection and ranking
# ===========================================================================


def count_above(a: np.ndarray, threshold: float) -> int:
    """1.7 -- how many elements are strictly greater than `threshold`.

    No loop and no counter variable. `a > threshold` gives you one boolean per
    element, and `True` counts as 1 when a boolean array is summed. Wrap the
    result in `int()` so the return type is a plain Python integer.

    >>> count_above(np.array([1, 5, 9]), 4)
    2
    """
    raise NotImplementedError("count_above")


def mask_between(a: np.ndarray, lo: float, hi: float) -> np.ndarray:
    """1.8 -- a boolean array that is True where `lo < x < hi`.

    Two things will bite you here and both are worth meeting now.

    `and` does not work. It is a control-flow keyword, not an operator, so
    NumPy cannot redefine it; Python asks the left operand "are you true?" and
    an array of twenty answers raises
    `ValueError: The truth value of an array with more than one element is
    ambiguous`. Use `&`, which IS an operator and which NumPy defines to mean
    elementwise-and.

    The parentheses around each comparison are not optional either. `&` binds
    tighter than `<`, so `a > lo & a < hi` parses as `a > (lo & a) < hi`.

    >>> mask_between(np.array([1, 5, 9]), 2, 8).tolist()
    [False, True, False]
    """
    raise NotImplementedError("mask_between")


def top_k_indices(scores: np.ndarray, k: int) -> np.ndarray:
    """1.9 -- the INDICES of the `k` largest scores, best first.

    `numpy.argsort` returns the indices that would sort the array, smallest
    first. You want the other end, and you want only `k` of them.

    The indices are the point. `numpy.sort` would hand you the scores and lose
    which row each one came from, which is exactly the thing a search needs to
    know. This is the function that turns Day 103's similarity scores into an
    answer.

    >>> top_k_indices(np.array([0.1, 0.9, 0.5]), 2).tolist()
    [1, 2]
    """
    raise NotImplementedError("top_k_indices")


def cosine_similarities(matrix: np.ndarray, query: np.ndarray) -> np.ndarray:
    """1.10 -- cosine similarity between `query` and EVERY row of `matrix`.

    Day 103 did this one row at a time. Do it in one expression, with no loop
    over rows.

    Two pieces:

      * all the dot products at once. `matrix @ query` gives one number per
        row -- a 6 by 4 matrix times a 4-vector is a 6-vector.
      * one length per row. `numpy.linalg.norm(matrix, axis=1)` gives six
        numbers, because **the axis you name is the one that disappears**: a
        (6, 4) array reduced along axis 1 leaves shape (6,).

    Then divide, elementwise, by those row lengths times the query's length.

    >>> m = np.array([[1.0, 0.0], [0.0, 1.0]])
    >>> cosine_similarities(m, np.array([1.0, 0.0])).tolist()
    [1.0, 0.0]
    """
    raise NotImplementedError("cosine_similarities")


# ===========================================================================
# Helpers -- written for you. Read them; the tests use them.
# ===========================================================================


def list_bytes(values: list[int]) -> int:
    """An honest total for what a Python list of integers costs in memory.

    `sys.getsizeof(values)` on its own is NOT the answer. It measures the list
    object -- a header plus one 8-byte pointer per element -- and not the
    integers those pointers point at, because the list does not own them.

    On this machine `sys.getsizeof(list(range(1_000_000)))` is 8,000,056
    bytes, almost exactly the 8,000,000 an int64 array needs, which would make
    the two look equal. The truth is 36,000,056: each of the million integers
    is a separate 28-byte object.

    So this adds the payload, counting each DISTINCT integer object once --
    CPython caches -5 through 256, so those really are shared.
    """
    by_identity = {id(x): x for x in values}
    payload = sum(sys.getsizeof(x) for x in by_identity.values())
    return sys.getsizeof(values) + payload


def array_bytes(a: np.ndarray) -> int:
    """What the array's data block costs: `a.nbytes`, which is size x itemsize."""
    return int(a.nbytes)


def wrap_int8(value: int, added: int) -> int:
    """Add `added` to `value` inside an int8 array and report what came out.

    An int8 holds -128 to 127. Adding 1 to 127 does not raise, does not
    promote, and on numpy 2.5.2 does not warn. Exercise 3 asks you what it
    gives instead.
    """
    a = np.array([value], dtype=np.int8)
    b = np.array([added], dtype=np.int8)
    return int((a + b)[0])


def nan_aware_mean(a: np.ndarray) -> float:
    """The mean of the non-missing entries, via numpy.nanmean."""
    return float(np.nanmean(a))


def select(a: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """The elements where `mask` is True, as a new array.

    Boolean indexing always returns a COPY -- it has to, because the selected
    elements are not evenly spaced and no stride can describe them.
    """
    return a[mask]


def time_call(fn: Callable[[], object], repeats: int = 5) -> list[float]:
    """Run `fn` `repeats` times and return every elapsed time in seconds.

    Every time, not the best and not the average: a single timing is noise.
    """
    times: list[float] = []
    for _ in range(repeats):
        start = time.perf_counter()
        fn()
        times.append(time.perf_counter() - start)
    return times


def median_seconds(times: Iterable[float]) -> float:
    """The median of a list of timings, which one unlucky run cannot drag."""
    return statistics.median(times)


def speedup(loop_times: Iterable[float], vec_times: Iterable[float]) -> float:
    """How many times faster the vectorised version was, by median.

    To be read, never asserted on exactly. The tests assert at least 20x,
    which is a claim about the shape of the gap and survives a slower machine.
    """
    return median_seconds(loop_times) / median_seconds(vec_times)


def describe(a: np.ndarray) -> str:
    """The facts that distinguish an ndarray from a list, in one line."""
    return (
        f"shape={a.shape} dtype={a.dtype} itemsize={a.itemsize} "
        f"nbytes={a.nbytes} strides={a.strides} "
        f"c_contiguous={a.flags['C_CONTIGUOUS']}"
    )


# `math` is imported for you because exercise 1.3 needs math.sqrt.
_ = math

#!/usr/bin/env bash
# Tests for the Day 104 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:
#
#   * a list of a million integers costs 36,000,056 bytes and the equivalent
#     int64 array costs 8,000,000 -- and sys.getsizeof on its own says they are
#     the same size, which is why the lab does not use it as the measurement;
#   * an int8 wraps from 127 to -128 with no exception and no warning, and the
#     absence of the warning is asserted rather than assumed;
#   * three operations written as a loop and as a NumPy expression agree over a
#     million elements with ==, not with a tolerance;
#   * the vectorised versions are at least 20 times faster -- the SHAPE of the
#     gap, never a millisecond figure, because a timing assertion is flaky on
#     someone else's machine;
#   * `and` on two arrays raises ValueError with the ambiguous-truth-value
#     message, and `&` does the thing that was meant;
#   * argsort returns [3, 2, 0] for Day 103's query, which is
#     race-day-nutrition, marathon-plan, roast-chicken;
#   * a slice is a view, so writing through it writes through to the original,
#     and .copy() breaks the link;
#   * np.nan != np.nan, np.isnan finds it, and np.nanmean gives 7/3;
#   * nothing is left behind on disk.
#
# Everything after 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 104 — Stop Writing the Loop"
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_list_versus_array 02_dtypes_and_overflow 03_same_answer_faster \
              04_creating_and_ufuncs 05_masks_and_selection \
              06_axes_views_and_ranking 07_nan_and_when_not_to_vectorise; 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 exceptions"
# --------------------------------------------------------------------------

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 100 ]; then
  check "the reference suite ran at least 100 tests (ran ${ref_passed})" "yes"
else
  check "the reference suite ran at least 100 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. Both directories contain modules called `vectorize` and
# `dataset`, 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. Each directory's conftest.py prevents that. 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 math
import sys
import warnings

import numpy as np

import dataset
from vectorize import (
    array_bytes,
    clip_loop,
    clip_vec,
    cosine_similarities,
    count_above,
    list_bytes,
    mask_between,
    nan_aware_mean,
    roots_loop,
    roots_vec,
    scale_and_offset_loop,
    scale_and_offset_vec,
    select,
    speedup,
    time_call,
    top_k_indices,
    wrap_int8,
)

# -- memory ---------------------------------------------------------------
values = list(range(dataset.N_BIG))
array = np.arange(dataset.N_BIG, dtype=np.int64)
print("naive_ratio_is_about_one", round(sys.getsizeof(values) / array.nbytes, 2))
print("honest_list_bytes", list_bytes(values))
print("array_bytes", array_bytes(array))
print("honest_ratio", round(list_bytes(values) / array_bytes(array), 2))
print("python_int_bytes", sys.getsizeof(1_000_000))
print("int64_element_bytes", array.itemsize)
del values

grid = np.arange(12).reshape(3, 4)
print("strides", grid.strides)
print("transpose_is_a_view", bool(np.shares_memory(grid, grid.T)))

# -- dtypes ---------------------------------------------------------------
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    wrapped = wrap_int8(127, 1)
    names = [w.category.__name__ for w in caught]
print("int8_wrap", wrapped)
print("int8_wrap_warnings", ",".join(names) if names else "none")
print("int8_doubled", (np.array(dataset.INT8_DOUBLING_INPUT, dtype=np.int8) * np.int8(2)).tolist())
print("int8_plus_python_int", (np.array([127], dtype=np.int8) + 1).tolist(),
      (np.array([127], dtype=np.int8) + 1).dtype)
print("astype_int16", (np.array([127], dtype=np.int8).astype(np.int16) + 1).tolist())
blind = np.float32(dataset.FLOAT32_BLIND_SPOT)
print("float32_blind", bool(blind + np.float32(1.0) == blind))
print("float64_not_blind", float(np.float64(dataset.FLOAT32_BLIND_SPOT) + 1.0))

# -- the three operations, twice each --------------------------------------
big = dataset.big_values()
as_list = big.tolist()
print("scale_exact", bool(np.array_equal(
    np.array(scale_and_offset_loop(as_list, dataset.SCALE_M, dataset.SCALE_C)),
    scale_and_offset_vec(big, dataset.SCALE_M, dataset.SCALE_C))))
print("roots_exact", bool(np.array_equal(
    np.array(roots_loop(as_list)), roots_vec(big))))
print("clip_exact", bool(np.array_equal(
    np.array(clip_loop(as_list, dataset.CLIP_LO, dataset.CLIP_HI)),
    clip_vec(big, dataset.CLIP_LO, dataset.CLIP_HI))))
factor = speedup(
    time_call(lambda: scale_and_offset_loop(as_list, dataset.SCALE_M, dataset.SCALE_C), 3),
    time_call(lambda: scale_and_offset_vec(big, dataset.SCALE_M, dataset.SCALE_C), 3),
)
print("speedup_over_20", factor > 20.0)
print("speedup_measured", round(factor, 1))
pow_differs = int(np.count_nonzero(np.array([x ** 0.5 for x in as_list]) != roots_vec(big)))
math_differs = int(np.count_nonzero(np.array([math.sqrt(x) for x in as_list]) != roots_vec(big)))
print("pow_disagrees", pow_differs > 0)
print("math_agrees", math_differs == 0)

# -- masking ---------------------------------------------------------------
readings = dataset.small_readings()
print("readings_match_documented", readings.tolist() == dataset.SMALL_READINGS_EXPECTED)
print("count_above_50", count_above(readings, 50))
print("values_above_50", select(readings, readings > 50).tolist())
print("mask_dtype", (readings > 50).dtype)
print("between_count", int(mask_between(readings, 30, 70).sum()))
print("between_values", readings[mask_between(readings, 30, 70)].tolist())
print("mask_mean", float((readings > 50).mean()))
print("fancy", readings[np.array([0, 5, 19, 5])].tolist())
try:
    (readings > 30) and (readings < 70)
except Exception as exc:  # deliberately broad: the TYPE is what is asserted
    print("keyword_and_raises", type(exc).__name__)
else:
    print("keyword_and_raises", "NOTHING_RAISED")
try:
    readings > 30 & readings < 70
except Exception as exc:  # deliberately broad: the TYPE is what is asserted
    print("missing_brackets_raise", type(exc).__name__)
else:
    print("missing_brackets_raise", "NOTHING_RAISED")

# -- axes, views, ranking ---------------------------------------------------
print("sum_axis0_shape", grid.sum(axis=0).shape)
print("sum_axis1_shape", grid.sum(axis=1).shape)
print("sum_axis1_values", grid.sum(axis=1).tolist())
v = np.array([1.0, 2.0, 3.0])
print("newaxis_column_shape", v[:, np.newaxis].shape)
fresh = np.arange(12).reshape(3, 4)
row = fresh[1]
row[0] = 999
print("view_wrote_through", int(fresh[1, 0]))
fresh2 = np.arange(12).reshape(3, 4)
copied = fresh2[2].copy()
copied[0] = -1
print("copy_did_not", int(fresh2[2, 0]))
print("mask_is_a_copy", bool(np.shares_memory(fresh2, fresh2[fresh2 > 5])))
print("ravel_is_a_view", bool(np.shares_memory(fresh2, fresh2.ravel())))
print("flatten_is_a_copy", bool(np.shares_memory(fresh2, fresh2.flatten())))
scores = np.array([5.0, 1.0, 9.0, 3.0])
print("argsort", np.argsort(scores).tolist())
sims = cosine_similarities(dataset.CATALOGUE, dataset.QUERY)
print("sims_shape", sims.shape)
top = top_k_indices(sims, dataset.TOP_K)
print("top3_indices", top.tolist())
print("top3_names", "|".join(dataset.ARTICLE_NAMES[i] for i in top))
print("margin_is_small", 0.0 < float(sims[top[0]] - sims[top[1]]) < 0.01)

# -- nan --------------------------------------------------------------------
holed = dataset.WITH_A_HOLE
print("nan_ne_nan", np.nan != np.nan)
print("eq_nan_finds_nothing", int((holed == np.nan).sum()))
print("isnan_finds_it", int(np.isnan(holed).sum()))
print("mean_is_nan", math.isnan(float(holed.mean())))
print("nanmean", nan_aware_mean(holed))
print("nansum", float(np.nansum(holed)))
PY
)"

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

check_eq "sys.getsizeof alone makes the list look the same size as the array" \
  "1.0" "$(get naive_ratio_is_about_one)"
check_eq "the honest list total is 36,000,056 bytes" \
  "36000056" "$(get honest_list_bytes)"
check_eq "the equivalent int64 array is 8,000,000 bytes" \
  "8000000" "$(get array_bytes)"
check_eq "so the array is 4.5 times smaller" "4.5" "$(get honest_ratio)"
check_eq "one Python int is 28 bytes here" "28" "$(get python_int_bytes)"
check_eq "one int64 array element is 8 bytes" "8" "$(get int64_element_bytes)"
check_eq "a 3 by 4 int64 array has strides (32, 8)" "(32, 8)" "$(get strides)"
check_eq "a transpose copies nothing" "True" "$(get transpose_is_a_view)"

check_eq "int8 127 + 1 wraps to -128" "-128" "$(get int8_wrap)"
check_eq "and does so with no warning at all on this numpy" \
  "none" "$(get int8_wrap_warnings)"
check_eq "doubling [120, 125, 127] as int8 wraps two of the three" \
  "[-16, -6, -2]" "$(get int8_doubled)"
check_eq "a plain Python 1 does not widen the array" \
  "[-128] int8" "$(get int8_plus_python_int)"
check_eq "asking for int16 first gives the right answer" \
  "[128]" "$(get astype_int16)"
check_eq "float32 cannot tell 16777216 from 16777217" "True" "$(get float32_blind)"
check_eq "float64 can" "16777217.0" "$(get float64_not_blind)"

check_eq "loop and vectorised scale-and-offset agree EXACTLY on a million values" \
  "True" "$(get scale_exact)"
check_eq "loop and vectorised square roots agree EXACTLY" "True" "$(get roots_exact)"
check_eq "loop and vectorised clip agree EXACTLY" "True" "$(get clip_exact)"
check_eq "the vectorised version is at least 20 times faster" \
  "True" "$(get speedup_over_20)"
check_eq "x ** 0.5 is NOT the same operation as np.sqrt" "True" "$(get pow_disagrees)"
check_eq "math.sqrt IS the same operation as np.sqrt" "True" "$(get math_agrees)"
echo "  (measured speedup on this run: $(get speedup_measured)x -- reported, not asserted)"

check_eq "the seeded readings are the twenty documented values" \
  "True" "$(get readings_match_documented)"
check_eq "nine readings are above 50" "9" "$(get count_above_50)"
check_eq "and they are the nine documented values" \
  "[70, 83, 69, 65, 75, 73, 97, 64, 82]" "$(get values_above_50)"
check_eq "a comparison produces a boolean array" "bool" "$(get mask_dtype)"
check_eq "seven readings are strictly between 30 and 70" "7" "$(get between_count)"
check_eq "and they are the seven documented values" \
  "[34, 69, 65, 37, 37, 41, 64]" "$(get between_values)"
check_eq "the mask's mean is the fraction above 50" "0.45" "$(get mask_mean)"
check_eq "fancy indexing keeps the order asked for and allows a repeat" \
  "[70, 21, 82, 21]" "$(get fancy)"
check_eq "the keyword 'and' raises ValueError on two arrays" \
  "ValueError" "$(get keyword_and_raises)"
check_eq "and so does the same expression with the brackets left off" \
  "ValueError" "$(get missing_brackets_raise)"

check_eq "summing a (3, 4) along axis 0 leaves shape (4,)" \
  "(4,)" "$(get sum_axis0_shape)"
check_eq "summing it along axis 1 leaves shape (3,)" "(3,)" "$(get sum_axis1_shape)"
check_eq "and the three row totals are 6, 22, 38" "[6, 22, 38]" "$(get sum_axis1_values)"
check_eq "np.newaxis turns a length-3 row into a 3 by 1 column" \
  "(3, 1)" "$(get newaxis_column_shape)"
check_eq "writing through a row slice writes through to the original" \
  "999" "$(get view_wrote_through)"
check_eq "writing through a .copy() does not" "8" "$(get copy_did_not)"
check_eq "a boolean mask returns a copy, never a view" "False" "$(get mask_is_a_copy)"
check_eq "ravel returns a view when it can" "True" "$(get ravel_is_a_view)"
check_eq "flatten always returns a copy" "False" "$(get flatten_is_a_copy)"
check_eq "argsort returns positions, not values" "[1, 3, 0, 2]" "$(get argsort)"
check_eq "one cosine similarity per catalogue row" "(6,)" "$(get sims_shape)"
check_eq "the top 3 by argsort are indices 3, 2 and 0" "[3, 2, 0]" "$(get top3_indices)"
check_eq "which are the three documented articles, best first" \
  "race-day-nutrition|marathon-plan|roast-chicken" "$(get top3_names)"
check_eq "the winner's margin is small enough to report rather than trumpet" \
  "True" "$(get margin_is_small)"

check_eq "nan is not equal to itself" "True" "$(get nan_ne_nan)"
check_eq "comparing an array to nan finds nothing at all" \
  "0" "$(get eq_nan_finds_nothing)"
check_eq "np.isnan finds the one that is missing" "1" "$(get isnan_finds_it)"
check_eq "the plain mean is nan, loudly" "True" "$(get mean_is_nan)"
# Section 6 re-runs this script with D104_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_nanmean="2.3333333333333335"
if [ -n "${D104_SELF_TEST:-}" ]; then
  expected_nanmean="2.5"   # the naive belief that a missing value is a zero-cost skip
fi
check_eq "np.nanmean divides by the three readings that exist, not the four wanted" \
  "${expected_nanmean}" "$(get nanmean)"
check_eq "np.nansum is 7.0" "7.0" "$(get nansum)"

# --------------------------------------------------------------------------
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 one -- 2.5, which is what you would get if a nan simply did not
# count and the divisor stayed at four -- and asserts that the re-run reports
# the failure and exits non-zero. If this section passes, section 5 is not
# decorative.
if [ -z "${D104_SELF_TEST:-}" ]; then
  self_out="$(D104_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: np.nanmean divides by the three readings"*)
      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 both searches below. The virtual environment ships
# NumPy's and pytest's own precompiled bytecode -- hundreds of __pycache__
# directories that came with the packages and have nothing to do with whether
# THIS lab tidied up after itself. Searching them would report a failure the
# reader cannot fix and did not cause. Everything the lab itself writes lives
# outside `.venv`, which is exactly what these two checks look at.

if find "${lab_dir}" -name '.venv' -prune -o -type d -name '__pycache__' -print -quit 2>/dev/null | grep -q .; then
  check "no __pycache__ directory left by the lab's own code" "no"
else
  check "no __pycache__ directory left by the lab's own code" "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 ]