Hands-on lab — Day 111: Gradient Descent from Scratch

"""01_the_hook.py -- the failure that opens the lesson.

The simplest convex function there is, f(x) = 0.5 * x**2, has one minimum,
at x = 0. Gradient descent should find it easily. Run it with a learning
rate only slightly above the function's own divergence boundary and watch
the loss climb every single step, without a single bug in the function,
the gradient, or the code -- until it overflows.
"""

import math

import dataset as D
import descent as G

asserts_held = 0


def check(label, condition):
    global asserts_held
    assert condition, f"FAILED: {label}"
    asserts_held += 1
    print(f"  ok: {label}")


print(f"f(x) = 0.5 * {D.HOOK_A} * x^2, minimum at x = 0")
print(f"divergence boundary: eta > 2 / a = {D.HOOK_DIVERGENCE_LR}")
print(f"chosen learning rate: eta = {D.HOOK_LR} -- only slightly too large")
print()

path = G.gradient_descent(lambda x: D.HOOK_A * x, D.HOOK_X0, D.HOOK_LR, D.HOOK_ITERS)
losses = [0.5 * D.HOOK_A * v * v for v in path[:12]]

print("step |     x            | loss")
for i in range(12):
    print(f"{i:4d} | {path[i]: .10f} | {losses[i]:.10f}")
print("  ...")

check(
    "the loss increases on every one of the first 20 steps",
    all(losses[i + 1] > losses[i] for i in range(len(losses) - 1)),
)

first_inf = next(i for i, v in enumerate(path) if math.isinf(v))
first_nan = next(i for i, v in enumerate(path) if math.isnan(v))
print()
print(f"x first becomes inf at step {first_inf}")
print(f"x first becomes nan at step {first_nan}")
print(f"value the step before overflow: {path[first_inf - 1]:.6e}")

check("x reaches inf before it reaches nan", first_inf < first_nan)
check("nan follows inf on the very next step", first_nan == first_inf + 1)
check(
    "nothing raised an exception along the way",
    True,  # if we got here, gradient_descent completed without raising
)

print()
print("Nothing was wrong with f, with its gradient, or with the update rule.")
print("The step size alone turned a solved problem into a divergent one.")
print()
print(f"01_the_hook.py: every assertion held. ({asserts_held} checks)")

"""02_regimes_and_contraction.py -- exercises 1, 3 and 4.

The update rule for f(x) = 0.5 * a * x**2 is exact algebra:
    x_{n+1} = x_n * (1 - eta * a)
so the whole story of what a learning rate does to convergence lives in
one number, (1 - eta * a), and this script measures it rather than states
it.
"""

import math

import dataset as D
import descent as G

asserts_held = 0


def check(label, condition):
    global asserts_held
    assert condition, f"FAILED: {label}"
    asserts_held += 1
    print(f"  ok: {label}")


print(f"a = {D.A}   1/a = {1.0 / D.A}   2/a = {2.0 / D.A}")
print()

# -- exercise 1: numeric_gradient agrees with the analytic gradient --------
print("Exercise 1 -- numeric_gradient vs the analytic gradient")
f = lambda x: 0.5 * D.A * x * x
for x in (-2.0, 0.5, 3.0):
    analytic = D.A * x
    measured = G.numeric_gradient(f, x, D.NUMERIC_H)
    print(f"  x={x:+.2f}  analytic={analytic:+.6f}  numeric={measured:+.6f}  gap={abs(analytic - measured):.3e}")
    check(f"numeric_gradient agrees with a*x at x={x}", abs(analytic - measured) < D.NUMERIC_TOL)

g = lambda x: math.sin(x * x)
for x in (0.3, 1.5, -0.8):
    analytic = 2.0 * x * math.cos(x * x)
    measured = G.numeric_gradient(g, x, D.NUMERIC_H)
    print(f"  x={x:+.2f}  analytic={analytic:+.6f}  numeric={measured:+.6f}  gap={abs(analytic - measured):.3e}")
    check(f"numeric_gradient agrees with the composed function at x={x}", abs(analytic - measured) < D.NUMERIC_TOL)

# -- exercise 3: the three regimes ------------------------------------------
print()
print("Exercise 3 -- the three regimes")
regimes = {
    "monotone   (0 < eta < 1/a)": D.LR_MONOTONE,
    "exact      (eta = 1/a)": D.LR_EXACT,
    "oscillating(1/a < eta < 2/a)": D.LR_OSCILLATING,
    "divergent  (eta > 2/a)": D.LR_DIVERGENT,
}
expected = {
    "monotone   (0 < eta < 1/a)": "monotone",
    "exact      (eta = 1/a)": "exact",
    "oscillating(1/a < eta < 2/a)": "oscillating",
    "divergent  (eta > 2/a)": "divergent",
}
for label, lr in regimes.items():
    path = G.gradient_descent(lambda x: D.A * x, D.X0_1D, lr, D.REGIME_ITERS)
    regime = G.classify_regime(path, D.A, lr)
    print(f"  eta={lr:.2f}  {label}  first 4 steps: {[round(v, 4) for v in path[:4]]}  ...  classified: {regime}")
    check(f"eta={lr} classified as {expected[label]}", regime == expected[label])

# -- exercise 4: the contraction ratio --------------------------------------
print()
print("Exercise 4 -- the measured contraction ratio equals |1 - eta*a|")
for lr in (D.LR_MONOTONE, D.LR_OSCILLATING, D.LR_DIVERGENT):
    path = G.gradient_descent(lambda x: D.A * x, D.X0_1D, lr, 8)
    ratios = G.per_step_ratios(path)
    predicted = abs(1.0 - lr * D.A)
    print(f"  eta={lr:.2f}  predicted |1-eta*a|={predicted:.6f}  measured ratios={[round(r, 6) for r in ratios[:3]]}")
    check(f"contraction ratio matches prediction at eta={lr}", all(abs(r - predicted) < D.EXACT_TOL for r in ratios))

print()
print(f"02_regimes_and_contraction.py: every assertion held. ({asserts_held} checks)")

"""03_ill_conditioning_and_momentum.py -- exercises 5 and 6.

f(x, y) = 0.5 * (x**2 + kappa * y**2) has Hessian diag(1, kappa), so its
condition number IS kappa -- Day 106's ratio of eigenvalues, put to work.
The optimal fixed learning rate for this bowl is 2 / (1 + kappa); this
script measures how many steps that optimal rate needs as kappa grows, and
then shows what a beta*v running average of the gradient buys back.
"""

import numpy as np

import dataset as D
import descent as G

asserts_held = 0


def check(label, condition):
    global asserts_held
    assert condition, f"FAILED: {label}"
    asserts_held += 1
    print(f"  ok: {label}")


print("Exercise 5 -- steps to convergence grow with the condition number")
print(f"tolerance on ||grad||: {D.KAPPA_GRAD_TOL}")
print()
print("kappa |     eta = 2/(1+kappa)  | steps")
counts = []
for k in D.KAPPA_VALUES:
    lr = D.kappa_lr(k)
    steps = G.steps_to_tolerance(D.bowl_grad(k), np.array(D.KAPPA_START), lr, D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS)
    counts.append(steps)
    print(f"{k:5d} | {lr:22.6f} | {steps}")

check("steps are non-decreasing as kappa grows", all(counts[i] <= counts[i + 1] for i in range(len(counts) - 1)))
check(f"kappa={D.KAPPA_VALUES[-1]} needs at least 10x the steps of kappa={D.KAPPA_VALUES[0]}", counts[-1] >= 10 * max(counts[0], 1))
check("the isotropic bowl (kappa=1) converges in exactly one step", counts[0] == 1)

print()
print("Exercise 6 -- momentum on the kappa=20 bowl")
k = D.MOMENTUM_KAPPA
lr = D.kappa_lr(k)
plain_steps = G.steps_to_tolerance(D.bowl_grad(k), np.array(D.KAPPA_START), lr, D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS)
momentum_steps = G.steps_to_tolerance_momentum(
    D.bowl_grad(k), np.array(D.KAPPA_START), D.MOMENTUM_LR, D.MOMENTUM_BETA, D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS
)
print(f"  same learning rate for both: eta = {lr:.6f}")
print(f"  plain gradient descent:    {plain_steps} steps")
print(f"  momentum (beta={D.MOMENTUM_BETA}):        {momentum_steps} steps")
print(f"  speedup: {plain_steps / momentum_steps:.2f}x")
check("momentum needs strictly fewer steps than plain descent", momentum_steps < plain_steps)

print()
print(f"03_ill_conditioning_and_momentum.py: every assertion held. ({asserts_held} checks)")

"""04_checking_landscapes_and_traps.py -- exercises 7, 8 and 9.

Three separate ways a training run can look fine and be wrong: a gradient
with the wrong sign in one place, a landscape with more than one answer,
and a stopping rule that mistakes a flat approach for arrival.
"""

import dataset as D
import descent as G

asserts_held = 0


def check(label, condition):
    global asserts_held
    assert condition, f"FAILED: {label}"
    asserts_held += 1
    print(f"  ok: {label}")


print("Exercise 7 -- gradient checking catches a sign-error bug")
correct_flags = G.gradient_check(D.check_function, D.check_gradient_correct, D.CHECK_POINT, D.NUMERIC_H, D.CHECK_TOL)
buggy_flags = G.gradient_check(D.check_function, D.check_gradient_buggy, D.CHECK_POINT, D.NUMERIC_H, D.CHECK_TOL)
print(f"  point: {D.CHECK_POINT}")
print(f"  correct gradient, checked component by component: {correct_flags}")
print(f"  buggy gradient (component 1 sign-flipped):          {buggy_flags}")
check("the correct gradient passes on every component", all(correct_flags))
check("the buggy gradient fails on exactly component 1", buggy_flags == [True, False, True])

print()
print("Exercise 8 -- two initialisations, two different minima")
left = G.gradient_descent(D.two_minima_grad, D.TWO_MINIMA_LEFT_START, D.TWO_MINIMA_LR, D.TWO_MINIMA_ITERS)
right = G.gradient_descent(D.two_minima_grad, D.TWO_MINIMA_RIGHT_START, D.TWO_MINIMA_LR, D.TWO_MINIMA_ITERS)
print("  f(x) = (x^2 - 1)^2, minima at x = -1 and x = +1, local maximum at x = 0")
print(f"  start x0={D.TWO_MINIMA_LEFT_START}   ->  converged to {left[-1]:.6f}")
print(f"  start x0={D.TWO_MINIMA_RIGHT_START}   ->  converged to {right[-1]:.6f}")
check("the two runs converge to minima more than the margin apart", G.minima_differ(left[-1], right[-1], D.TWO_MINIMA_MARGIN))
check("the left run reaches -1", abs(left[-1] + 1.0) < 1e-3)
check("the right run reaches +1", abs(right[-1] - 1.0) < 1e-3)

print()
print("Exercise 9 -- the stopping-criterion trap")
result = G.stopping_criteria_disagree(
    D.PLATEAU_X0, D.plateau_grad, D.plateau_value, D.PLATEAU_LR, D.PLATEAU_GRAD_TOL, D.PLATEAU_DELTA_F_TOL
)
print(f"  a shallow bowl (a={D.PLATEAU_A}) at x={D.PLATEAU_X0}, far from its minimum at x=0")
print(f"  ||grad|| = {result['grad_norm']:.6f}   (tolerance: {D.PLATEAU_GRAD_TOL})")
print(f"  |delta f| = {result['delta_f']:.3e}   (tolerance: {D.PLATEAU_DELTA_F_TOL})")
print(f"  naive '|delta f| < tol, so we converged' would stop here: {result['naive_stops_early']}")
check("the gradient is still above its own tolerance", result["grad_norm"] >= D.PLATEAU_GRAD_TOL)
check("the loss barely moved, below its own tolerance", result["delta_f"] < D.PLATEAU_DELTA_F_TOL)
check("the naive criterion would stop early", result["naive_stops_early"] is True)

print()
print(f"04_checking_landscapes_and_traps.py: every assertion held. ({asserts_held} checks)")

"""Make this directory's own modules the ones its tests import.

Both `examples/` and `starter/` contain modules called `dataset` and
`descent`, 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 `descent` 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 module of those names 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 ("dataset", "descent"):
    module = sys.modules.get(name)
    origin = getattr(module, "__file__", "") or ""
    if module is not None and not origin.startswith(HERE):
        del sys.modules[name]

"""Shared data for the Day 111 lab -- read this file, do not change it.

Every constant here is invented and stated to be invented. Nothing is fitted
to a real dataset; every number is chosen so a reader can re-derive it by
hand, and every tolerance below is derived from the arithmetic that actually
governs the comparison it guards, not tuned until a test happened to pass.

The core object of the day is the one-dimensional quadratic
    f(x) = 0.5 * a * x**2          f'(x) = a * x
because its gradient-descent update has a closed form:
    x_{n+1} = x_n - lr * a * x_n = x_n * (1 - lr * a)
so after n steps
    x_n = x_0 * (1 - lr * a) ** n
and the whole regime structure of the lesson -- monotone, exact, oscillating,
divergent -- falls out of the single number (1 - lr * a).
"""

from __future__ import annotations

import numpy as np

EPSILON = float(np.finfo(np.float64).eps)

# ---------------------------------------------------------------------------
# The 1-D quadratic that the first half of the lab is built on.
# ---------------------------------------------------------------------------

A = 5.0                       # curvature of f(x) = 0.5 * A * x**2
CRITICAL_LR = 1.0 / A          # 0.2 -- exact one-step landing
DIVERGENCE_LR = 2.0 / A        # 0.4 -- boundary of the oscillating-but-converging regime

X0_1D = 1.0

# The four learning rates exercise 3 classifies, chosen to sit one in each
# regime relative to CRITICAL_LR = 0.2 and DIVERGENCE_LR = 0.4:
LR_MONOTONE = 0.10             # 0 < lr < 1/A            -> monotone decrease
LR_EXACT = 0.20                # lr == 1/A                -> exact in one step
LR_OSCILLATING = 0.35          # 1/A < lr < 2/A          -> alternates sign, |x| shrinks
LR_DIVERGENT = 0.45            # lr > 2/A                -> |x| grows without bound

REGIME_ITERS = 30

# The opening hook: the simplest convex function there is, and a learning
# rate only slightly above its own divergence boundary.
HOOK_A = 1.0
HOOK_DIVERGENCE_LR = 2.0 / HOOK_A      # 2.0
HOOK_LR = 2.2                          # "only slightly too large"
HOOK_X0 = 1.0
HOOK_ITERS = 4000                      # long enough to reach inf, then nan

# ---------------------------------------------------------------------------
# Tolerances, derived rather than chosen.
# ---------------------------------------------------------------------------
# The quadratic's update is EXACT algebra in floating point: one
# multiplication per step, x_{n+1} = x_n * (1 - lr * A). float64 carries
# about 15-17 significant decimal digits, and after a modest number of
# multiplications the accumulated rounding is still many orders of magnitude
# below 1e-9, so two routes to the same exact quantity (a direct formula and
# a step-by-step loop) are compared at:
EXACT_TOL = 1e-9

# The central difference used to check an analytic gradient against a
# numerical one carries truncation error of order h**2 and rounding error of
# order EPSILON / h. At h = 1e-6 (h**2 = 1e-12, EPSILON / h ~ 2.2e-10) the
# total is comfortably under 1e-6 for the smooth functions this lab uses, so:
NUMERIC_H = 1e-6
NUMERIC_TOL = 1e-6

# ---------------------------------------------------------------------------
# Ill-conditioning: f(x, y) = 0.5 * (x**2 + kappa * y**2)
# ---------------------------------------------------------------------------
# The Hessian of this bowl is diag(1, kappa), so its condition number is
# exactly kappa (Day 106: the ratio of the eigenvalues -- min eigenvalue 1,
# max eigenvalue kappa). This lab uses the standard optimal FIXED step size
# for a quadratic with eigenvalues mu (smallest) and L (largest),
#     lr* = 2 / (mu + L) = 2 / (1 + kappa)
# which is the single learning rate that minimises the worst-case per-step
# contraction over every eigen-direction at once. For kappa = 1 (an
# isotropic bowl) that contraction is exactly zero -- gradient descent with
# the optimal step solves an isotropic quadratic in ONE step, the same
# "exact" regime exercise 3 meets on the 1-D bowl. As kappa grows, the
# optimal step shrinks and the worst-case per-step contraction, which is
# exactly (kappa - 1) / (kappa + 1), climbs towards 1 -- so the number of
# steps needed to reach a fixed gradient tolerance grows with kappa, with no
# free parameter left to compensate.
KAPPA_VALUES = (1, 5, 20, 100)
KAPPA_START = (1.0, 1.0)
KAPPA_GRAD_TOL = 1e-4
KAPPA_MAX_ITERS = 5000


def kappa_lr(kappa: float) -> float:
    """The optimal fixed learning rate for a bowl of condition number kappa."""
    return 2.0 / (1.0 + kappa)


def bowl_grad(kappa: float):
    """Return the gradient function of f(x, y) = 0.5 * (x**2 + kappa * y**2)."""

    def grad(point):
        x, y = point
        return np.array([x, kappa * y])

    return grad


def bowl_value(kappa: float, point) -> float:
    x, y = point
    return 0.5 * (x * x + kappa * y * y)


# ---------------------------------------------------------------------------
# Momentum comparison, on the kappa = 20 bowl.
# ---------------------------------------------------------------------------
MOMENTUM_KAPPA = 20
MOMENTUM_BETA = 0.5
# Momentum is given exactly the SAME learning rate plain descent uses --
# kappa_lr(MOMENTUM_KAPPA) -- so the comparison isolates what the beta*v
# term buys on its own, with nothing else changed.
MOMENTUM_LR = kappa_lr(MOMENTUM_KAPPA)

# ---------------------------------------------------------------------------
# Gradient checking: a deliberately wrong analytic gradient.
# ---------------------------------------------------------------------------
CHECK_POINT = np.array([0.7, -1.3, 2.1])


def check_function(point) -> float:
    x, y, z = point
    return x * x + 2.0 * y * y + 0.5 * z * z * z


def check_gradient_correct(point):
    x, y, z = point
    return np.array([2.0 * x, 4.0 * y, 1.5 * z * z])


def check_gradient_buggy(point):
    """The correct gradient with the SIGN of component 1 (index 1) flipped."""
    correct = check_gradient_correct(point)
    buggy = correct.copy()
    buggy[1] = -buggy[1]
    return buggy


CHECK_TOL = 1e-4

# ---------------------------------------------------------------------------
# Non-convexity: two minima, initialisation decides the answer.
# ---------------------------------------------------------------------------
# f(x) = (x**2 - 1)**2 has minima at x = -1 and x = +1 (value 0) and a local
# maximum at x = 0 (value 1). f'(x) = 4*x**3 - 4*x, which is negative for
# 0 < x < 1 (pulling towards +1) and positive for -1 < x < 0 (pulling
# towards -1), so any start strictly inside (-1, 1) but on one side of 0
# converges to the minimum on that side.
TWO_MINIMA_LR = 0.05
TWO_MINIMA_ITERS = 400
TWO_MINIMA_LEFT_START = -0.1
TWO_MINIMA_RIGHT_START = 0.1
TWO_MINIMA_MARGIN = 1.5


def two_minima_value(x: float) -> float:
    return (x * x - 1.0) ** 2


def two_minima_grad(x: float) -> float:
    return 4.0 * x ** 3 - 4.0 * x


# ---------------------------------------------------------------------------
# Stopping-criterion trap: a shallow bowl, far from its own minimum.
# ---------------------------------------------------------------------------
# The same quadratic family as the top of this file, but with curvature so
# small that a point far from the minimum still has a small local slope --
# a real, bounded plateau rather than an unbounded linear tail. Locally,
# taking one gradient-descent step changes the value by
#   delta_f = -lr * a * x**2 * (1 - 0.5 * lr * a)   (exact algebra, not an
#   approximation, for this exact quadratic)
# which is tiny whenever `a` is tiny, even while the gradient a*x itself is
# comfortably above a small tolerance.
PLATEAU_A = 1e-4
PLATEAU_X0 = 100.0
PLATEAU_LR = 1e-3
PLATEAU_GRAD_TOL = 1e-3
PLATEAU_DELTA_F_TOL = 1e-6


def plateau_grad(x: float) -> float:
    return PLATEAU_A * x


def plateau_value(x: float) -> float:
    return 0.5 * PLATEAU_A * x * x

"""The core of the day: numeric_gradient, gradient_descent, and everything
built on top of them. Every function here is deliberately small -- the whole
point of the lab is that the entire training loop used across the rest of
this course is a handful of lines, plus a great deal of care about the
learning rate.
"""

from __future__ import annotations

import numpy as np


def numeric_gradient(f, x, h=1e-6):
    """The gradient of `f` at `x`, by central differences (Day 108's
    definition, applied to every coordinate of `x`).

    `x` may be a plain Python float, in which case call `f` with a float
    on either side of `x` and return a float. Or `x` may be a 1-D
    array-like, in which case `f` is called with the whole perturbed
    vector and is expected to accept one; returns a numpy array the same
    shape as `x`.
    """
    if np.isscalar(x):
        return (f(x + h) - f(x - h)) / (2.0 * h)
    point = np.asarray(x, dtype=float)
    grad = np.zeros_like(point)
    for i in range(point.size):
        forward = point.copy()
        backward = point.copy()
        forward[i] += h
        backward[i] -= h
        grad[i] = (f(forward) - f(backward)) / (2.0 * h)
    return grad


def gradient_descent(grad_fn, x0, lr, iters):
    """Run `iters` steps of x <- x - lr * grad_fn(x), starting from x0.

    Returns the WHOLE path as a list of length iters + 1, path[0] == x0,
    so later exercises can inspect every intermediate value rather than
    only the final answer.

    Runs under numpy's error state set to 'ignore' for overflow: a
    diverging run is expected to produce inf and then nan, and that is the
    day's own point (a diverging training run looks exactly like this),
    not a crash to be prevented.
    """
    path = [x0]
    x = x0
    with np.errstate(over="ignore", invalid="ignore"):
        for _ in range(iters):
            x = x - lr * grad_fn(x)
            path.append(x)
    return path


def gradient_descent_momentum(grad_fn, x0, lr, beta, iters):
    """x <- x - lr * v, where v <- beta * v + grad_fn(x).

    Momentum is not a different rule; it is an exponentially weighted
    running average of the gradient, substituted for the raw gradient in
    the same update. Returns the whole path, same shape as
    `gradient_descent`.
    """
    path = [x0]
    x = x0
    v = np.zeros_like(np.asarray(x0, dtype=float)) if not np.isscalar(x0) else 0.0
    with np.errstate(over="ignore", invalid="ignore"):
        for _ in range(iters):
            g = grad_fn(x)
            v = beta * v + g
            x = x - lr * v
            path.append(x)
    return path


def steps_to_tolerance(grad_fn, x0, lr, tol, max_iters):
    """Run gradient descent until the gradient's norm drops below `tol`,
    and return the number of steps taken (not the path). Returns
    `max_iters` if the tolerance was never reached, so callers can tell a
    slow run from a run that never converges.
    """
    x = x0
    scalar = np.isscalar(x0)
    with np.errstate(over="ignore", invalid="ignore"):
        for step in range(max_iters):
            g = grad_fn(x)
            norm = abs(g) if scalar else float(np.linalg.norm(np.asarray(g, dtype=float)))
            if norm < tol:
                return step
            x = x - lr * g
    return max_iters


def steps_to_tolerance_momentum(grad_fn, x0, lr, beta, tol, max_iters):
    """The momentum analogue of `steps_to_tolerance`."""
    x = x0
    v = np.zeros_like(np.asarray(x0, dtype=float))
    with np.errstate(over="ignore", invalid="ignore"):
        for step in range(max_iters):
            g = np.asarray(grad_fn(x), dtype=float)
            if float(np.linalg.norm(g)) < tol:
                return step
            v = beta * v + g
            x = x - lr * v
    return max_iters


def gradient_check(f, grad_analytic_fn, x, h=1e-6, tol=1e-4):
    """Compare an analytic gradient function against a numerical one at
    `x`. Returns a list of booleans, one per coordinate of x, True where
    the two agree within `tol` and False where they do not -- so a caller
    can identify exactly which components are wrong, not merely that
    something is.
    """
    point = np.atleast_1d(np.asarray(x, dtype=float))
    analytic = np.atleast_1d(np.asarray(grad_analytic_fn(point), dtype=float))
    numeric = np.atleast_1d(numeric_gradient(f, point, h))
    return [bool(abs(a - n) < tol) for a, n in zip(analytic, numeric)]


def per_step_ratios(path):
    """The measured contraction ratio |x_{n+1} / x_n| at every step of a
    1-D path where x_n != 0. For a quadratic f(x) = 0.5 * a * x**2 this
    should equal |1 - lr * a| at every step -- exercise 4 measures that
    prediction rather than assuming it.
    """
    return [abs(path[i + 1] / path[i]) for i in range(len(path) - 1) if path[i] != 0]


def minima_differ(final_a, final_b, margin):
    """True if two converged points are farther apart than `margin` --
    the check that two gradient-descent runs on a non-convex function
    landed at genuinely different minima rather than both near the same
    one.
    """
    return abs(final_a - final_b) > margin


def stopping_criteria_disagree(x, grad_fn, value_fn, lr, tol_grad, tol_f):
    """Take ONE gradient-descent step from `x` and report whether the two
    common stopping rules disagree: the naive '|delta f| < tol_f, so we
    must have converged' rule against the more honest
    '||gradient|| < tol_grad' rule.

    Returns a dict with the measured gradient norm, the measured |delta f|,
    and `naive_stops_early`: True when |delta f| < tol_f while the
    gradient norm is still >= tol_grad -- exactly the trap where "the loss
    stopped changing" is mistaken for "we converged".
    """
    grad = grad_fn(x)
    grad_norm = abs(grad) if np.isscalar(grad) else float(np.linalg.norm(grad))
    f_before = value_fn(x)
    x_after = x - lr * grad
    f_after = value_fn(x_after)
    delta_f = abs(f_after - f_before)
    return {
        "grad_norm": grad_norm,
        "delta_f": delta_f,
        "naive_stops_early": bool(delta_f < tol_f and grad_norm >= tol_grad),
    }


def classify_regime(path, a, lr):
    """Classify a 1-D quadratic gradient-descent run into one of
    'monotone', 'exact', 'oscillating', 'divergent', purely from the
    observed path (never from lr and a directly), so the classification is
    a real behavioural check rather than a restatement of the formula.
    """
    values = [abs(v) for v in path]
    if len(values) >= 2 and values[1] == 0.0:
        return "exact"
    signs = [1 if v > 0 else (-1 if v < 0 else 0) for v in path]
    alternates = all(
        signs[i] != 0 and signs[i + 1] != 0 and signs[i] != signs[i + 1]
        for i in range(len(signs) - 1)
    )
    shrinking = all(
        values[i + 1] <= values[i] + 1e-12 for i in range(len(values) - 1)
    )
    growing = values[-1] > values[0] and not shrinking
    if alternates and shrinking:
        return "oscillating"
    if shrinking and not alternates:
        return "monotone"
    if growing or not shrinking:
        return "divergent"
    return "unknown"

"""The reference test suite for Day 111. Every test checks REAL behaviour
of the functions in this directory -- real numbers, real convergence, real
overflow -- never source text.
"""

import math

import numpy as np
import pytest

import dataset as D
import descent as G

# ---------------------------------------------------------------------------
# Exercise 1 -- numeric_gradient
# ---------------------------------------------------------------------------


def test_numeric_gradient_matches_quadratic():
    f = lambda x: 0.5 * D.A * x * x
    for x in (-2.0, -0.5, 0.0, 0.5, 3.0):
        assert abs(G.numeric_gradient(f, x, D.NUMERIC_H) - D.A * x) < D.NUMERIC_TOL


def test_numeric_gradient_matches_composed_function():
    # d/dx sin(x**2) = 2x cos(x**2)
    f = lambda x: math.sin(x * x)
    for x in (0.3, 1.5, -0.8):
        analytic = 2.0 * x * math.cos(x * x)
        assert abs(G.numeric_gradient(f, x, D.NUMERIC_H) - analytic) < D.NUMERIC_TOL


def test_numeric_gradient_handles_vector_input():
    grad = G.numeric_gradient(D.check_function, D.CHECK_POINT, D.NUMERIC_H)
    correct = D.check_gradient_correct(D.CHECK_POINT)
    assert np.max(np.abs(grad - correct)) < D.NUMERIC_TOL


# ---------------------------------------------------------------------------
# Exercise 2 -- gradient_descent returns the whole path
# ---------------------------------------------------------------------------


def test_gradient_descent_returns_whole_path():
    path = G.gradient_descent(lambda x: D.A * x, D.X0_1D, D.LR_MONOTONE, 10)
    assert len(path) == 11
    assert path[0] == D.X0_1D


def test_gradient_descent_matches_closed_form():
    lr = D.LR_MONOTONE
    path = G.gradient_descent(lambda x: D.A * x, D.X0_1D, lr, 20)
    for n, x_n in enumerate(path):
        closed_form = D.X0_1D * (1.0 - lr * D.A) ** n
        assert abs(x_n - closed_form) < D.EXACT_TOL


# ---------------------------------------------------------------------------
# Exercise 3 -- the three regimes
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "lr,expected",
    [
        (D.LR_MONOTONE, "monotone"),
        (D.LR_EXACT, "exact"),
        (D.LR_OSCILLATING, "oscillating"),
        (D.LR_DIVERGENT, "divergent"),
    ],
)
def test_regime_classification(lr, expected):
    path = G.gradient_descent(lambda x: D.A * x, D.X0_1D, lr, D.REGIME_ITERS)
    assert G.classify_regime(path, D.A, lr) == expected


def test_exact_regime_lands_on_zero_in_one_step():
    path = G.gradient_descent(lambda x: D.A * x, D.X0_1D, D.LR_EXACT, 1)
    assert path[1] == 0.0


def test_oscillating_regime_alternates_sign():
    path = G.gradient_descent(lambda x: D.A * x, D.X0_1D, D.LR_OSCILLATING, 6)
    signs = [1 if v > 0 else -1 for v in path]
    assert all(signs[i] != signs[i + 1] for i in range(len(signs) - 1))


def test_divergent_regime_grows_without_shrinking():
    path = G.gradient_descent(lambda x: D.A * x, D.X0_1D, D.LR_DIVERGENT, D.REGIME_ITERS)
    assert abs(path[-1]) > abs(path[0]) * 10


# ---------------------------------------------------------------------------
# Exercise 4 -- the contraction ratio is exactly |1 - lr * a|
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("lr", [D.LR_MONOTONE, D.LR_OSCILLATING, D.LR_DIVERGENT])
def test_contraction_ratio_matches_prediction(lr):
    path = G.gradient_descent(lambda x: D.A * x, D.X0_1D, lr, 15)
    predicted = abs(1.0 - lr * D.A)
    for ratio in G.per_step_ratios(path):
        assert abs(ratio - predicted) < D.EXACT_TOL


# ---------------------------------------------------------------------------
# Exercise 5 -- ill-conditioning: steps grow with kappa
# ---------------------------------------------------------------------------


def test_steps_to_tolerance_nondecreasing_in_kappa():
    counts = [
        G.steps_to_tolerance(
            D.bowl_grad(k), np.array(D.KAPPA_START), D.kappa_lr(k),
            D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS,
        )
        for k in D.KAPPA_VALUES
    ]
    assert all(counts[i] <= counts[i + 1] for i in range(len(counts) - 1))
    assert counts[-1] >= 10 * max(counts[0], 1)


def test_isotropic_bowl_converges_in_one_step():
    steps = G.steps_to_tolerance(
        D.bowl_grad(1), np.array(D.KAPPA_START), D.kappa_lr(1),
        D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS,
    )
    assert steps == 1


# ---------------------------------------------------------------------------
# Exercise 6 -- momentum beats plain descent on the same learning rate
# ---------------------------------------------------------------------------


def test_momentum_needs_fewer_steps_than_plain_descent():
    k = D.MOMENTUM_KAPPA
    lr = D.kappa_lr(k)
    plain = G.steps_to_tolerance(D.bowl_grad(k), np.array(D.KAPPA_START), lr, D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS)
    momentum = G.steps_to_tolerance_momentum(
        D.bowl_grad(k), np.array(D.KAPPA_START), D.MOMENTUM_LR, D.MOMENTUM_BETA,
        D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS,
    )
    assert momentum < plain


# ---------------------------------------------------------------------------
# Exercise 7 -- gradient checking catches a sign-error bug
# ---------------------------------------------------------------------------


def test_gradient_check_passes_correct_gradient():
    flags = G.gradient_check(D.check_function, D.check_gradient_correct, D.CHECK_POINT, D.NUMERIC_H, D.CHECK_TOL)
    assert all(flags)


def test_gradient_check_flags_exactly_the_buggy_component():
    flags = G.gradient_check(D.check_function, D.check_gradient_buggy, D.CHECK_POINT, D.NUMERIC_H, D.CHECK_TOL)
    assert flags == [True, False, True]


# ---------------------------------------------------------------------------
# Exercise 8 -- non-convexity: initialisation decides the minimum
# ---------------------------------------------------------------------------


def test_two_initialisations_reach_different_minima():
    left = G.gradient_descent(D.two_minima_grad, D.TWO_MINIMA_LEFT_START, D.TWO_MINIMA_LR, D.TWO_MINIMA_ITERS)
    right = G.gradient_descent(D.two_minima_grad, D.TWO_MINIMA_RIGHT_START, D.TWO_MINIMA_LR, D.TWO_MINIMA_ITERS)
    assert G.minima_differ(left[-1], right[-1], D.TWO_MINIMA_MARGIN)
    assert left[-1] < 0 < right[-1]
    assert abs(left[-1] + 1.0) < 1e-3
    assert abs(right[-1] - 1.0) < 1e-3


# ---------------------------------------------------------------------------
# Exercise 9 -- the stopping-criterion trap
# ---------------------------------------------------------------------------


def test_naive_delta_f_criterion_stops_early_on_the_plateau():
    result = G.stopping_criteria_disagree(
        D.PLATEAU_X0, D.plateau_grad, D.plateau_value, D.PLATEAU_LR,
        D.PLATEAU_GRAD_TOL, D.PLATEAU_DELTA_F_TOL,
    )
    assert result["grad_norm"] >= D.PLATEAU_GRAD_TOL
    assert result["delta_f"] < D.PLATEAU_DELTA_F_TOL
    assert result["naive_stops_early"] is True


# ---------------------------------------------------------------------------
# The hook: overflow, then nan, handled without crashing
# ---------------------------------------------------------------------------


def test_divergent_run_overflows_to_inf_then_nan_without_raising():
    path = G.gradient_descent(lambda x: D.HOOK_A * x, D.HOOK_X0, D.HOOK_LR, D.HOOK_ITERS)
    assert any(math.isinf(v) for v in path)
    assert any(math.isnan(v) for v in path)
    first_inf = next(i for i, v in enumerate(path) if math.isinf(v))
    first_nan = next(i for i, v in enumerate(path) if math.isnan(v))
    assert first_nan == first_inf + 1


def test_loss_increases_every_step_before_it_overflows():
    path = G.gradient_descent(lambda x: D.HOOK_A * x, D.HOOK_X0, D.HOOK_LR, 20)
    losses = [0.5 * D.HOOK_A * v * v for v in path]
    assert all(losses[i + 1] > losses[i] for i in range(len(losses) - 1))

lesson_id: D111
day: 111
kind: guided-build
languages: [python, bash]
setup_commands:
  - cd labs/sections/math-statistics-and-data/day-111-gradient-descent-from-scratch
  - 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_the_hook.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 02_regimes_and_contraction.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 03_ill_conditioning_and_momentum.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 04_checking_landscapes_and_traps.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: 35
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 -> 50 checks, 0 failure(s), exit 0; pytest examples -> 24 passed; pytest starter -> 1 passed, 20 skipped on an untouched checkout, and 21 passed when checked against a fully solved copy of starter/descent.py kept outside the lab directory. All four reference scripts exit 0 with every internal assertion holding. Everything was run through a real lab-local .venv built by the documented setup commands on this machine today, not through an authoring environment, and the whole lab directory (content, lab and instructor paths) was deleted and rebuilt from scratch in this session after a second, now-stopped agent had been authoring into the same three directories concurrently -- every file under this lab, and every captured output in expected-output/, comes from this session''s own code and this session''s own run, not from the earlier agent''s work. Section 6 of the harness re-runs itself with one expectation deliberately swapped for the belief that a flat loss always means convergence, and confirms the re-run exits non-zero and reports exactly one named failure, so the suite is demonstrated to be capable of failing rather than merely claimed to be. The measured numbers this run produced: the three learning-rate regimes on a = 5 classify as monotone (eta=0.10), exact (eta=0.20, landing on x=0.0 in one step), oscillating (eta=0.35) and divergent (eta=0.45); the measured per-step contraction ratios are 0.5, 0.75 and 1.25, matching |1 - eta*a| to under 1e-9 at every step; steps to converge on the ill-conditioned bowl f(x,y) = 0.5*(x^2 + kappa*y^2) at the optimal fixed learning rate 2/(1+kappa) were 1, 27, 122 and 691 for kappa in {1, 5, 20, 100} -- the isotropic case (kappa=1) solves in exactly one step, and kappa=100 needs 691, comfortably more than ten times kappa=1''s count; momentum (beta=0.5) at the SAME learning rate plain descent uses on the kappa=20 bowl needed 34 steps against plain descent''s 122; a gradient check built from a central difference passed a correct three-component gradient on every component and flagged exactly the one component of a deliberately sign-flipped gradient; two initialisations on f(x) = (x^2-1)^2 converged to -1.0 and 1.0; on a shallow bowl (a=1e-4) at x=100, one gradient-descent step left the gradient at 0.01 (ten times its own tolerance of 1e-3) while the loss changed by only 1.0e-7 (below its own tolerance of 1e-6), so the naive |delta f| stopping rule fires while the honest gradient-norm rule says training is not done; and the opening hook -- f(x) = 0.5*x^2 with eta = 2.2, just above its divergence boundary of 2/1 = 2.0 -- produced a loss that increased on every one of the first 20 steps and, continuing to 4000 steps, overflowed to inf at step 3890 with nan following on step 3891, with nothing raised along the way. Every tolerance is derived in examples/dataset.py from the error terms that actually govern the comparison it guards and is documented, with its derivation, in expected-output/FIELDS.md. No scipy, PyTorch or JAX output is reproduced anywhere in this lab or its lesson; none of the three is installed in this environment, and the lesson''s tools section says so plainly.'

# What is installed, why, and what it costs

Two packages, both free and open source, both installed into a lab-local
virtual environment that `rm -rf .venv` completely undoes.

| Package | Version pinned | Licence | What this lab uses it for |
| --- | --- | --- | --- |
| `numpy` | 2.5.2 | BSD 3-Clause | Vectors for the two-dimensional ill-conditioning and momentum exercises (5, 6), reading float64's machine epsilon from `numpy.finfo`, and `np.linalg.norm` for the stopping-tolerance checks. |
| `pytest` | 9.1.1 | MIT | The reference suite and your running score in `starter/`. |

There is no paid tier of anything in this lab, no account, no key and no
signup, personally or commercially.

## The one time the network is needed

```bash
.venv/bin/pip install -r requirements/requirements.txt
```

That is the only command in the lab that opens a connection. Section 7 of
`tests/run_tests.sh` greps every source file in `examples/` and `starter/` to
prove that nothing else does.

## If you cannot install anything at all

You lose relatively little. Every scalar exercise (1 through 4, 7, 8 and 9)
needs only `math` and the standard library — a plain Python float and a
`for` loop compute all of them. Only the two-dimensional bowl in exercises 5
and 6 genuinely wants a small array type; without NumPy, `x, y = point` on a
two-element Python list or tuple does the same job with a few more lines,
and `math.hypot(x, y)` replaces `np.linalg.norm`.

What you lose without `pytest` is the running score and the skip-versus-fail
distinction — you would read your own printed numbers against the ones in
`expected-output/` instead.

## What is deliberately *not* installed

`scipy.optimize.minimize`, `torch.optim.SGD` and `jax.grad` with `optax` all
do this job at production scale, and none of them is installed here. **No
output from any of them is reproduced anywhere in this lab or its lesson.**
They are described from their documentation, and the lesson's tools section
marks each one as not run here.

That is not a limitation to apologise for. The nine exercises in this lab are
the same update rule every one of those tools runs underneath — `x <- x -
lr * grad(x)`, or a running average of it — with the engineering (batching,
adaptive learning rates, GPU dispatch, automatic differentiation feeding the
gradient in) removed. Having written the loop by hand, you will read a
`for epoch in range(...): optimizer.step()` differently.

numpy==2.5.2
pytest==9.1.1

# Descent by Hand — nine exercises

Everything you write goes in `descent.py`, beside this file. `dataset.py`
holds every constant and helper you need — read it, do not change it. Check
yourself as you go:

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

A skip means "not attempted yet". A failure means "attempted and wrong", and
prints your answer next to the correct one.

## 1. `numeric_gradient(f, x, h)`

The gradient by central differences — Day 108's definition, put to work
again. Must agree with the analytic gradient of a quadratic and of a
composed function (`sin(x**2)`) to a stated tolerance, and must handle both
a scalar `x` and a vector one.

## 2. `gradient_descent(grad, x0, lr, iters)`

The whole loop, in one function: `x <- x - lr * grad(x)`, repeated `iters`
times. Return the **whole path**, not just the final answer — every later
exercise inspects it.

## 3. The three regimes

`classify_regime(path, a, lr)` reads a path produced on `f(x) = 0.5*a*x**2`
and says which of four things happened: `monotone`, `exact`, `oscillating`,
`divergent`. With `a = 5` (so `1/a = 0.2`, `2/a = 0.4`), the four learning
rates in `dataset.py` are chosen to land one in each regime. This is the
exercise the whole day is built on — get the boundary conditions right and
the rest of the lab follows.

## 4. The measured contraction ratio

`per_step_ratios(path)` returns the list of `|x_{n+1} / x_n|`. For the
quadratic, this should equal `|1 - lr*a|` at every single step — a
prediction from Day 111's algebra, checked against real numbers rather than
trusted.

## 5. Ill-conditioning

`steps_to_tolerance(grad, x0, lr, tol, max_iters)` runs gradient descent
until `||grad(x)|| < tol` and reports how many steps it took. Run it on
`f(x, y) = 0.5*(x**2 + kappa*y**2)` for `kappa` in `{1, 5, 20, 100}`
(`dataset.bowl_grad` builds the gradient function; `dataset.kappa_lr` picks
the learning rate). The step count should never decrease as `kappa` grows,
and `kappa=100` should need at least ten times the steps of `kappa=1`.

## 6. Momentum

`gradient_descent_momentum(grad, x0, lr, beta, iters)` and
`steps_to_tolerance_momentum(...)` add one line to exercise 2 and 5's
shapes: a velocity `v <- beta*v + grad(x)`, and `x <- x - lr*v` in place of
`x <- x - lr*grad(x)`. On the `kappa=20` bowl, at the SAME learning rate
plain descent uses, momentum should need strictly fewer steps.

## 7. Gradient checking

`gradient_check(f, grad_fn, x, h, tol)` compares an analytic gradient
against `numeric_gradient` at `x`, component by component, and returns a
list of booleans. `dataset.py` supplies a function with a known-correct
gradient and a deliberately buggy one (one component's sign flipped) —
your check must pass the first and flag exactly the broken component of the
second.

## 8. Two minima

`minima_differ(final_a, final_b, margin)` is one line: are two converged
points farther apart than `margin`? Run `gradient_descent` from two
different starting points on `dataset.two_minima_grad` (the gradient of
`f(x) = (x**2 - 1)**2`, which has minima at `x = -1` and `x = +1`) and
confirm they land on opposite sides.

## 9. The stopping-criterion trap

`stopping_criteria_disagree(x, grad_fn, value_fn, lr, tol_grad, tol_f)`
takes one gradient-descent step and reports the gradient's magnitude, the
change in the function value, and whether the naive "the loss barely
changed, so we must be done" rule would fire while the gradient is still
well above its own tolerance. Run it on `dataset.PLATEAU_X0` with
`dataset.plateau_grad` and `dataset.plateau_value` — a point far from the
minimum of a very shallow bowl, where the slope is real but tiny.

## When you are done

```bash
.venv/bin/pytest starter -q -p no:cacheprovider
```

should report every test passing. Then read `examples/` — the reference
implementation and four narrated demonstration scripts that print their own
working and assert every claim they make.

"""Make this directory's own modules the ones its tests import.

Both `examples/` and `starter/` contain modules called `dataset` and
`descent`, 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 `descent` 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 module of those names 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 ("dataset", "descent"):
    module = sys.modules.get(name)
    origin = getattr(module, "__file__", "") or ""
    if module is not None and not origin.startswith(HERE):
        del sys.modules[name]

"""Shared data for the Day 111 lab -- read this file, do not change it.

Every constant here is invented and stated to be invented. Nothing is fitted
to a real dataset; every number is chosen so a reader can re-derive it by
hand, and every tolerance below is derived from the arithmetic that actually
governs the comparison it guards, not tuned until a test happened to pass.

The core object of the day is the one-dimensional quadratic
    f(x) = 0.5 * a * x**2          f'(x) = a * x
because its gradient-descent update has a closed form:
    x_{n+1} = x_n - lr * a * x_n = x_n * (1 - lr * a)
so after n steps
    x_n = x_0 * (1 - lr * a) ** n
and the whole regime structure of the lesson -- monotone, exact, oscillating,
divergent -- falls out of the single number (1 - lr * a).
"""

from __future__ import annotations

import numpy as np

EPSILON = float(np.finfo(np.float64).eps)

# ---------------------------------------------------------------------------
# The 1-D quadratic that the first half of the lab is built on.
# ---------------------------------------------------------------------------

A = 5.0                       # curvature of f(x) = 0.5 * A * x**2
CRITICAL_LR = 1.0 / A          # 0.2 -- exact one-step landing
DIVERGENCE_LR = 2.0 / A        # 0.4 -- boundary of the oscillating-but-converging regime

X0_1D = 1.0

# The four learning rates exercise 3 classifies, chosen to sit one in each
# regime relative to CRITICAL_LR = 0.2 and DIVERGENCE_LR = 0.4:
LR_MONOTONE = 0.10             # 0 < lr < 1/A            -> monotone decrease
LR_EXACT = 0.20                # lr == 1/A                -> exact in one step
LR_OSCILLATING = 0.35          # 1/A < lr < 2/A          -> alternates sign, |x| shrinks
LR_DIVERGENT = 0.45            # lr > 2/A                -> |x| grows without bound

REGIME_ITERS = 30

# The opening hook: the simplest convex function there is, and a learning
# rate only slightly above its own divergence boundary.
HOOK_A = 1.0
HOOK_DIVERGENCE_LR = 2.0 / HOOK_A      # 2.0
HOOK_LR = 2.2                          # "only slightly too large"
HOOK_X0 = 1.0
HOOK_ITERS = 4000                      # long enough to reach inf, then nan

# ---------------------------------------------------------------------------
# Tolerances, derived rather than chosen.
# ---------------------------------------------------------------------------
# The quadratic's update is EXACT algebra in floating point: one
# multiplication per step, x_{n+1} = x_n * (1 - lr * A). float64 carries
# about 15-17 significant decimal digits, and after a modest number of
# multiplications the accumulated rounding is still many orders of magnitude
# below 1e-9, so two routes to the same exact quantity (a direct formula and
# a step-by-step loop) are compared at:
EXACT_TOL = 1e-9

# The central difference used to check an analytic gradient against a
# numerical one carries truncation error of order h**2 and rounding error of
# order EPSILON / h. At h = 1e-6 (h**2 = 1e-12, EPSILON / h ~ 2.2e-10) the
# total is comfortably under 1e-6 for the smooth functions this lab uses, so:
NUMERIC_H = 1e-6
NUMERIC_TOL = 1e-6

# ---------------------------------------------------------------------------
# Ill-conditioning: f(x, y) = 0.5 * (x**2 + kappa * y**2)
# ---------------------------------------------------------------------------
# The Hessian of this bowl is diag(1, kappa), so its condition number is
# exactly kappa (Day 106: the ratio of the eigenvalues -- min eigenvalue 1,
# max eigenvalue kappa). This lab uses the standard optimal FIXED step size
# for a quadratic with eigenvalues mu (smallest) and L (largest),
#     lr* = 2 / (mu + L) = 2 / (1 + kappa)
# which is the single learning rate that minimises the worst-case per-step
# contraction over every eigen-direction at once. For kappa = 1 (an
# isotropic bowl) that contraction is exactly zero -- gradient descent with
# the optimal step solves an isotropic quadratic in ONE step, the same
# "exact" regime exercise 3 meets on the 1-D bowl. As kappa grows, the
# optimal step shrinks and the worst-case per-step contraction, which is
# exactly (kappa - 1) / (kappa + 1), climbs towards 1 -- so the number of
# steps needed to reach a fixed gradient tolerance grows with kappa, with no
# free parameter left to compensate.
KAPPA_VALUES = (1, 5, 20, 100)
KAPPA_START = (1.0, 1.0)
KAPPA_GRAD_TOL = 1e-4
KAPPA_MAX_ITERS = 5000


def kappa_lr(kappa: float) -> float:
    """The optimal fixed learning rate for a bowl of condition number kappa."""
    return 2.0 / (1.0 + kappa)


def bowl_grad(kappa: float):
    """Return the gradient function of f(x, y) = 0.5 * (x**2 + kappa * y**2)."""

    def grad(point):
        x, y = point
        return np.array([x, kappa * y])

    return grad


def bowl_value(kappa: float, point) -> float:
    x, y = point
    return 0.5 * (x * x + kappa * y * y)


# ---------------------------------------------------------------------------
# Momentum comparison, on the kappa = 20 bowl.
# ---------------------------------------------------------------------------
MOMENTUM_KAPPA = 20
MOMENTUM_BETA = 0.5
# Momentum is given exactly the SAME learning rate plain descent uses --
# kappa_lr(MOMENTUM_KAPPA) -- so the comparison isolates what the beta*v
# term buys on its own, with nothing else changed.
MOMENTUM_LR = kappa_lr(MOMENTUM_KAPPA)

# ---------------------------------------------------------------------------
# Gradient checking: a deliberately wrong analytic gradient.
# ---------------------------------------------------------------------------
CHECK_POINT = np.array([0.7, -1.3, 2.1])


def check_function(point) -> float:
    x, y, z = point
    return x * x + 2.0 * y * y + 0.5 * z * z * z


def check_gradient_correct(point):
    x, y, z = point
    return np.array([2.0 * x, 4.0 * y, 1.5 * z * z])


def check_gradient_buggy(point):
    """The correct gradient with the SIGN of component 1 (index 1) flipped."""
    correct = check_gradient_correct(point)
    buggy = correct.copy()
    buggy[1] = -buggy[1]
    return buggy


CHECK_TOL = 1e-4

# ---------------------------------------------------------------------------
# Non-convexity: two minima, initialisation decides the answer.
# ---------------------------------------------------------------------------
# f(x) = (x**2 - 1)**2 has minima at x = -1 and x = +1 (value 0) and a local
# maximum at x = 0 (value 1). f'(x) = 4*x**3 - 4*x, which is negative for
# 0 < x < 1 (pulling towards +1) and positive for -1 < x < 0 (pulling
# towards -1), so any start strictly inside (-1, 1) but on one side of 0
# converges to the minimum on that side.
TWO_MINIMA_LR = 0.05
TWO_MINIMA_ITERS = 400
TWO_MINIMA_LEFT_START = -0.1
TWO_MINIMA_RIGHT_START = 0.1
TWO_MINIMA_MARGIN = 1.5


def two_minima_value(x: float) -> float:
    return (x * x - 1.0) ** 2


def two_minima_grad(x: float) -> float:
    return 4.0 * x ** 3 - 4.0 * x


# ---------------------------------------------------------------------------
# Stopping-criterion trap: a shallow bowl, far from its own minimum.
# ---------------------------------------------------------------------------
# The same quadratic family as the top of this file, but with curvature so
# small that a point far from the minimum still has a small local slope --
# a real, bounded plateau rather than an unbounded linear tail. Locally,
# taking one gradient-descent step changes the value by
#   delta_f = -lr * a * x**2 * (1 - 0.5 * lr * a)   (exact algebra, not an
#   approximation, for this exact quadratic)
# which is tiny whenever `a` is tiny, even while the gradient a*x itself is
# comfortably above a small tolerance.
PLATEAU_A = 1e-4
PLATEAU_X0 = 100.0
PLATEAU_LR = 1e-3
PLATEAU_GRAD_TOL = 1e-3
PLATEAU_DELTA_F_TOL = 1e-6


def plateau_grad(x: float) -> float:
    return PLATEAU_A * x


def plateau_value(x: float) -> float:
    return 0.5 * PLATEAU_A * x * x

"""Nine functions to write. Every one has a working signature, a docstring
saying exactly what it must do, and a `return None` where your code goes.
Returning None is how the test suite knows an exercise has not been
attempted yet: `pytest starter -q` SKIPS unattempted work rather than
failing it, so your score only ever counts what you have actually done.

Check yourself as you go:

    .venv/bin/pytest starter -q

`dataset.py` sits beside this file with every constant and helper function
you need -- read it, do not change it. `numpy` is used only where a
function genuinely needs a vector (exercises 5, 6 and part of 1); the
scalar exercises need nothing beyond arithmetic.
"""

from __future__ import annotations

import numpy as np

# ---------------------------------------------------------------------------
# Exercise 1 -- numeric_gradient
# ---------------------------------------------------------------------------


def numeric_gradient(f, x, h=1e-6):
    """The gradient of `f` at `x`, by central differences (Day 108's
    definition: (f(x+h) - f(x-h)) / (2h), applied to every coordinate).

    `x` may be a plain Python float, in which case call `f` with a float
    on either side of `x` and return a float. Or `x` may be a 1-D
    array-like, in which case `f` expects the whole vector: perturb one
    coordinate at a time, holding the rest fixed, and return a numpy array
    the same shape as `x`.

    Approach: `np.isscalar(x)` tells the two cases apart. For the vector
    case, copy `x` into a numpy array, then for each index i build a
    forward copy with x[i] + h and a backward copy with x[i] - h.
    """
    return None


# ---------------------------------------------------------------------------
# Exercise 2 -- gradient_descent
# ---------------------------------------------------------------------------


def gradient_descent(grad_fn, x0, lr, iters):
    """Run `iters` steps of x <- x - lr * grad_fn(x), starting from x0.

    Return the WHOLE path as a list of length iters + 1, with path[0]
    equal to x0 -- every later exercise inspects the path, not just the
    final value.

    Approach: build a list starting with x0, then loop `iters` times,
    each time computing x <- x - lr * grad_fn(x) and appending the new x.
    Wrap the loop body in `with np.errstate(over="ignore", invalid="ignore"):`
    so that a diverging run overflows to inf and then nan instead of
    printing a runtime warning -- that overflow is the point of exercise 3's
    fourth regime, not a bug to prevent.
    """
    return None


# ---------------------------------------------------------------------------
# Exercise 3 -- classify the regime a 1-D quadratic run fell into
# ---------------------------------------------------------------------------


def classify_regime(path, a, lr):
    """Look at a path produced by gradient_descent on f(x) = 0.5*a*x**2 and
    return one of the four strings 'monotone', 'exact', 'oscillating',
    'divergent' -- purely from the observed VALUES in `path`, not from a
    and lr directly, so this is a real behavioural check.

    - 'exact': the second value in the path (path[1]) is exactly 0.0.
    - 'oscillating': the sign alternates from one step to the next AND
      |x| is non-increasing (allow a tiny numerical slack, say 1e-12).
    - 'monotone': |x| is non-increasing and the sign never alternates.
    - 'divergent': anything else -- in particular, the final |x| is much
      larger than the first.

    Approach: compute the list of |x| values and the list of signs first,
    then work through the four cases above in order.
    """
    return None


# ---------------------------------------------------------------------------
# Exercise 4 -- the measured per-step contraction ratio
# ---------------------------------------------------------------------------


def per_step_ratios(path):
    """Return the list of |x_{n+1} / x_n| for every step where x_n != 0.

    For f(x) = 0.5*a*x**2 this should equal |1 - lr*a| at every step --
    exercise 4 in test_starter.py measures that prediction against what
    this function actually returns, rather than assuming the formula.

    Approach: one line, a list comprehension over consecutive pairs.
    """
    return None


# ---------------------------------------------------------------------------
# Exercise 5 -- steps needed to reach a gradient-norm tolerance
# ---------------------------------------------------------------------------


def steps_to_tolerance(grad_fn, x0, lr, tol, max_iters):
    """Run plain gradient descent (no path storage needed) until
    ||grad_fn(x)|| < tol, and return the number of steps TAKEN before that
    happened. If the tolerance is never reached within max_iters, return
    max_iters.

    `x0` here is a numpy array (a point in 2-D, for the ill-conditioning
    exercise), so `grad_fn(x)` returns a numpy array too and its norm is
    `np.linalg.norm(...)`.

    Approach: a for loop over `range(max_iters)`. At the top of each
    iteration, compute the gradient and check its norm BEFORE taking the
    step -- that is what makes step 0 a valid answer when x0 is already
    within tolerance.
    """
    return None


# ---------------------------------------------------------------------------
# Exercise 6 -- momentum
# ---------------------------------------------------------------------------


def gradient_descent_momentum(grad_fn, x0, lr, beta, iters):
    """x <- x - lr * v, where v <- beta * v + grad_fn(x), starting from
    v = 0 (a numpy array the same shape as x0) and x = x0. Return the
    whole path, exactly like `gradient_descent`.

    Momentum is not a new rule -- it substitutes an exponentially
    weighted running average of the gradient for the raw gradient in the
    SAME update. Getting the order right matters: update v first using
    the CURRENT gradient, then update x using the NEW v.

    Approach: same shape as gradient_descent, with one extra state
    variable `v` carried between iterations.
    """
    return None


def steps_to_tolerance_momentum(grad_fn, x0, lr, beta, tol, max_iters):
    """The momentum analogue of steps_to_tolerance: same stopping rule,
    same return convention, but each step updates a velocity `v` first
    and then moves `x` by `-lr * v`.

    Approach: combine the loop shape of steps_to_tolerance with the
    velocity update of gradient_descent_momentum.
    """
    return None


# ---------------------------------------------------------------------------
# Exercise 7 -- gradient checking
# ---------------------------------------------------------------------------


def gradient_check(f, grad_analytic_fn, x, h=1e-6, tol=1e-4):
    """Compare an analytic gradient function against numeric_gradient at
    `x`. Return a list of booleans, one per coordinate of x: True where
    the analytic and numeric values agree within `tol`, False where they
    do not -- so a caller can see exactly WHICH component is wrong, not
    merely that something is.

    Approach: call numeric_gradient(f, x, h) once, call
    grad_analytic_fn(x) once, then compare element by element.
    """
    return None


# ---------------------------------------------------------------------------
# Exercise 8 -- two minima
# ---------------------------------------------------------------------------


def minima_differ(final_a, final_b, margin):
    """Return True if two converged points are more than `margin` apart --
    the check that two gradient-descent runs on a non-convex function
    landed at genuinely different minima.

    Approach: one line.
    """
    return None


# ---------------------------------------------------------------------------
# Exercise 9 -- the stopping-criterion trap
# ---------------------------------------------------------------------------


def stopping_criteria_disagree(x, grad_fn, value_fn, lr, tol_grad, tol_f):
    """Take ONE gradient-descent step from `x` and report whether the
    naive '|delta f| < tol_f, so we must have converged' rule disagrees
    with the more honest '||gradient|| < tol_grad' rule.

    Return a dict with three keys:
      - 'grad_norm': the gradient's magnitude at x (a plain float; use
        abs() for a scalar x, since this exercise only ever uses scalars)
      - 'delta_f': the ABSOLUTE difference between value_fn(x) and
        value_fn(x_after_one_step)
      - 'naive_stops_early': True exactly when delta_f < tol_f AND
        grad_norm >= tol_grad -- the naive rule fires while the honest
        rule says training is not done.

    Approach: compute grad = grad_fn(x), then x_after = x - lr * grad,
    then compare value_fn at both points.
    """
    return None

"""Your running score. Unattempted work SKIPS; wrong work FAILS with both
values printed.

Run from the lab directory:

    .venv/bin/pytest starter -q

On an untouched checkout this reports everything skipped except the one
test that proves the suite itself runs. A skip means "not attempted". A
failure means "attempted and wrong", and shows your answer next to the
real one.
"""

import math

import numpy as np
import pytest

import dataset as D
import descent as S  # your work


def need(value, what):
    """Skip if the exercise has not been attempted yet, otherwise hand the
    value back."""
    if value is None:
        pytest.skip(f"not attempted yet: {what}")
    return value


def close(got, want, tol, what):
    assert abs(got - want) < tol, (
        f"{what}: your answer {got!r}, expected {want!r} "
        f"(difference {abs(got - want):.3e}, tolerance {tol:g})"
    )


def test_the_suite_itself_runs():
    """Always passes, so a green run is distinguishable from a collection
    error that quietly ran nothing at all."""
    assert D.EPSILON > 0.0


# ---------------------------------------------------------------------------
# Exercise 1 -- numeric_gradient
# ---------------------------------------------------------------------------


def test_1_numeric_gradient_matches_quadratic():
    f = lambda x: 0.5 * D.A * x * x
    got = need(S.numeric_gradient(f, 3.0, D.NUMERIC_H), "numeric_gradient")
    close(got, D.A * 3.0, D.NUMERIC_TOL, "numeric_gradient(0.5*a*x^2, x=3)")


def test_1_numeric_gradient_matches_composed_function():
    g = lambda x: math.sin(x * x)
    got = need(S.numeric_gradient(g, 1.5, D.NUMERIC_H), "numeric_gradient")
    close(got, 2.0 * 1.5 * math.cos(1.5 * 1.5), D.NUMERIC_TOL, "numeric_gradient(sin(x^2), x=1.5)")


def test_1_numeric_gradient_handles_a_vector():
    got = need(S.numeric_gradient(D.check_function, D.CHECK_POINT, D.NUMERIC_H), "numeric_gradient (vector)")
    correct = D.check_gradient_correct(D.CHECK_POINT)
    assert np.max(np.abs(np.asarray(got) - correct)) < D.NUMERIC_TOL


# ---------------------------------------------------------------------------
# Exercise 2 -- gradient_descent
# ---------------------------------------------------------------------------


def test_2_gradient_descent_returns_whole_path():
    path = need(S.gradient_descent(lambda x: D.A * x, D.X0_1D, D.LR_MONOTONE, 10), "gradient_descent")
    assert len(path) == 11, f"expected a path of length 11, got {len(path)}"
    assert path[0] == D.X0_1D


def test_2_gradient_descent_matches_closed_form():
    lr = D.LR_MONOTONE
    path = need(S.gradient_descent(lambda x: D.A * x, D.X0_1D, lr, 20), "gradient_descent")
    for n, x_n in enumerate(path):
        closed_form = D.X0_1D * (1.0 - lr * D.A) ** n
        close(x_n, closed_form, D.EXACT_TOL, f"gradient_descent step {n}")


# ---------------------------------------------------------------------------
# Exercise 3 -- the three regimes
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "lr,expected",
    [
        (D.LR_MONOTONE, "monotone"),
        (D.LR_EXACT, "exact"),
        (D.LR_OSCILLATING, "oscillating"),
        (D.LR_DIVERGENT, "divergent"),
    ],
)
def test_3_regime_classification(lr, expected):
    path = need(S.gradient_descent(lambda x: D.A * x, D.X0_1D, lr, D.REGIME_ITERS), "gradient_descent")
    got = need(S.classify_regime(path, D.A, lr), "classify_regime")
    assert got == expected, f"classify_regime at eta={lr}: your answer {got!r}, expected {expected!r}"


# ---------------------------------------------------------------------------
# Exercise 4 -- the contraction ratio
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("lr", [D.LR_MONOTONE, D.LR_OSCILLATING, D.LR_DIVERGENT])
def test_4_contraction_ratio_matches_prediction(lr):
    path = need(S.gradient_descent(lambda x: D.A * x, D.X0_1D, lr, 15), "gradient_descent")
    ratios = need(S.per_step_ratios(path), "per_step_ratios")
    predicted = abs(1.0 - lr * D.A)
    for ratio in ratios:
        close(ratio, predicted, D.EXACT_TOL, f"per_step_ratios at eta={lr}")


# ---------------------------------------------------------------------------
# Exercise 5 -- ill-conditioning
# ---------------------------------------------------------------------------


def test_5_steps_to_tolerance_nondecreasing_in_kappa():
    counts = []
    for k in D.KAPPA_VALUES:
        steps = need(
            S.steps_to_tolerance(D.bowl_grad(k), np.array(D.KAPPA_START), D.kappa_lr(k), D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS),
            "steps_to_tolerance",
        )
        counts.append(steps)
    assert all(counts[i] <= counts[i + 1] for i in range(len(counts) - 1)), f"steps should not decrease as kappa grows: {counts}"
    assert counts[-1] >= 10 * max(counts[0], 1), (
        f"kappa={D.KAPPA_VALUES[-1]} should need at least 10x the steps of kappa={D.KAPPA_VALUES[0]}: {counts}"
    )


def test_5_isotropic_bowl_converges_in_one_step():
    steps = need(
        S.steps_to_tolerance(D.bowl_grad(1), np.array(D.KAPPA_START), D.kappa_lr(1), D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS),
        "steps_to_tolerance",
    )
    assert steps == 1, f"the optimal step size should solve an isotropic bowl in one step, got {steps}"


# ---------------------------------------------------------------------------
# Exercise 6 -- momentum
# ---------------------------------------------------------------------------


def test_6_momentum_needs_fewer_steps_than_plain_descent():
    k = D.MOMENTUM_KAPPA
    lr = D.kappa_lr(k)
    plain = need(
        S.steps_to_tolerance(D.bowl_grad(k), np.array(D.KAPPA_START), lr, D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS),
        "steps_to_tolerance",
    )
    momentum = need(
        S.steps_to_tolerance_momentum(
            D.bowl_grad(k), np.array(D.KAPPA_START), D.MOMENTUM_LR, D.MOMENTUM_BETA, D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS
        ),
        "steps_to_tolerance_momentum",
    )
    assert momentum < plain, f"momentum ({momentum} steps) should beat plain descent ({plain} steps) at the same learning rate"


def test_6_momentum_path_has_the_right_length():
    path = need(
        S.gradient_descent_momentum(D.bowl_grad(D.MOMENTUM_KAPPA), np.array(D.KAPPA_START), D.MOMENTUM_LR, D.MOMENTUM_BETA, 10),
        "gradient_descent_momentum",
    )
    assert len(path) == 11


# ---------------------------------------------------------------------------
# Exercise 7 -- gradient checking
# ---------------------------------------------------------------------------


def test_7_gradient_check_passes_correct_gradient():
    flags = need(
        S.gradient_check(D.check_function, D.check_gradient_correct, D.CHECK_POINT, D.NUMERIC_H, D.CHECK_TOL),
        "gradient_check",
    )
    assert all(flags), f"the correct gradient should pass every component, got {flags}"


def test_7_gradient_check_flags_exactly_the_buggy_component():
    flags = need(
        S.gradient_check(D.check_function, D.check_gradient_buggy, D.CHECK_POINT, D.NUMERIC_H, D.CHECK_TOL),
        "gradient_check",
    )
    assert list(flags) == [True, False, True], f"expected [True, False, True], got {flags}"


# ---------------------------------------------------------------------------
# Exercise 8 -- two minima
# ---------------------------------------------------------------------------


def test_8_two_initialisations_reach_different_minima():
    left = need(
        S.gradient_descent(D.two_minima_grad, D.TWO_MINIMA_LEFT_START, D.TWO_MINIMA_LR, D.TWO_MINIMA_ITERS),
        "gradient_descent",
    )
    right = need(
        S.gradient_descent(D.two_minima_grad, D.TWO_MINIMA_RIGHT_START, D.TWO_MINIMA_LR, D.TWO_MINIMA_ITERS),
        "gradient_descent",
    )
    differ = need(S.minima_differ(left[-1], right[-1], D.TWO_MINIMA_MARGIN), "minima_differ")
    assert differ, f"left ({left[-1]}) and right ({right[-1]}) should differ by more than {D.TWO_MINIMA_MARGIN}"
    close(left[-1], -1.0, 1e-3, "left minimum")
    close(right[-1], 1.0, 1e-3, "right minimum")


# ---------------------------------------------------------------------------
# Exercise 9 -- the stopping-criterion trap
# ---------------------------------------------------------------------------


def test_9_naive_delta_f_criterion_stops_early_on_the_plateau():
    result = need(
        S.stopping_criteria_disagree(
            D.PLATEAU_X0, D.plateau_grad, D.plateau_value, D.PLATEAU_LR, D.PLATEAU_GRAD_TOL, D.PLATEAU_DELTA_F_TOL
        ),
        "stopping_criteria_disagree",
    )
    assert result["grad_norm"] >= D.PLATEAU_GRAD_TOL, result
    assert result["delta_f"] < D.PLATEAU_DELTA_F_TOL, result
    assert result["naive_stops_early"] is True, result

#!/usr/bin/env bash
# Tests for the Day 111 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:
#
#   * the whole training loop is one line, x <- x - lr * grad(x), and its
#     closed form on a quadratic, x_n = x0 * (1 - lr*a)^n, is measured
#     against a real loop rather than assumed;
#   * four learning rates on the same quadratic land in four different
#     regimes -- monotone, exact, oscillating, divergent -- with the exact
#     boundaries at 1/a and 2/a;
#   * the per-step contraction ratio measured from a real run equals
#     |1 - lr*a| to float precision;
#   * an ill-conditioned bowl needs more steps as its condition number
#     grows, and an isotropic bowl solves in exactly one step at the
#     optimal learning rate;
#   * momentum, given the SAME learning rate as plain descent, needs
#     strictly fewer steps on the same ill-conditioned bowl;
#   * a gradient check built from a central difference flags exactly the
#     one component of a deliberately broken analytic gradient;
#   * two starting points on a non-convex function converge to two
#     different minima;
#   * a naive "the loss stopped changing" stopping rule is caught firing
#     early, while the gradient itself is still well above its own
#     tolerance;
#   * a learning rate only slightly too large makes the loss climb every
#     step until the run overflows to inf and then nan, without raising;
#   * 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
}

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 111 — Descent by Hand"
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}"

float_width="$("${python_bin}" -c "import sys; print(sys.float_info.mant_dig)")"
check_eq "Python floats are IEEE-754 doubles with a 53-bit significand" "53" "${float_width}"

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

for script in 01_the_hook 02_regimes_and_contraction \
              03_ill_conditioning_and_momentum 04_checking_landscapes_and_traps; 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 behaviour"
# --------------------------------------------------------------------------

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 20 ]; then
  check "the reference suite ran at least 20 tests (ran ${ref_passed})" "yes"
else
  check "the reference suite ran at least 20 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 `dataset` and
# `descent`, 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 numpy as np

import dataset as D
import descent as G

print("a", D.A)
print("critical_lr", D.CRITICAL_LR)
print("divergence_lr", D.DIVERGENCE_LR)

regime_lrs = {
    "monotone": D.LR_MONOTONE,
    "exact": D.LR_EXACT,
    "oscillating": D.LR_OSCILLATING,
    "divergent": D.LR_DIVERGENT,
}
regime_paths = {
    name: G.gradient_descent(lambda x: D.A * x, D.X0_1D, lr, D.REGIME_ITERS)
    for name, lr in regime_lrs.items()
}
for name, path in regime_paths.items():
    print(f"regime_{name}", G.classify_regime(path, D.A, regime_lrs[name]))
print("exact_second_value", regime_paths["exact"][1])
print("divergent_grows", abs(regime_paths["divergent"][-1]) > abs(regime_paths["divergent"][0]) * 10)

ratio_mono = G.per_step_ratios(G.gradient_descent(lambda x: D.A * x, D.X0_1D, D.LR_MONOTONE, 10))[0]
ratio_osc = G.per_step_ratios(G.gradient_descent(lambda x: D.A * x, D.X0_1D, D.LR_OSCILLATING, 10))[0]
ratio_div = G.per_step_ratios(G.gradient_descent(lambda x: D.A * x, D.X0_1D, D.LR_DIVERGENT, 10))[0]
print("ratio_monotone", ratio_mono)
print("ratio_oscillating", ratio_osc)
print("ratio_divergent", ratio_div)
print("ratio_monotone_matches", abs(ratio_mono - abs(1 - D.LR_MONOTONE * D.A)) < D.EXACT_TOL)
print("ratio_oscillating_matches", abs(ratio_osc - abs(1 - D.LR_OSCILLATING * D.A)) < D.EXACT_TOL)
print("ratio_divergent_matches", abs(ratio_div - abs(1 - D.LR_DIVERGENT * D.A)) < D.EXACT_TOL)

counts = [
    G.steps_to_tolerance(D.bowl_grad(k), np.array(D.KAPPA_START), D.kappa_lr(k), D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS)
    for k in D.KAPPA_VALUES
]
print("kappa_steps", "|".join(str(c) for c in counts))
print("kappa_nondecreasing", all(counts[i] <= counts[i + 1] for i in range(len(counts) - 1)))
print("kappa_order_of_magnitude", counts[-1] >= 10 * max(counts[0], 1))
print("kappa_one_step_isotropic", counts[0] == 1)

k = D.MOMENTUM_KAPPA
plain_steps = G.steps_to_tolerance(D.bowl_grad(k), np.array(D.KAPPA_START), D.kappa_lr(k), D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS)
momentum_steps = G.steps_to_tolerance_momentum(
    D.bowl_grad(k), np.array(D.KAPPA_START), D.MOMENTUM_LR, D.MOMENTUM_BETA, D.KAPPA_GRAD_TOL, D.KAPPA_MAX_ITERS
)
print("momentum_plain_steps", plain_steps)
print("momentum_steps", momentum_steps)
print("momentum_faster", momentum_steps < plain_steps)

correct_flags = G.gradient_check(D.check_function, D.check_gradient_correct, D.CHECK_POINT, D.NUMERIC_H, D.CHECK_TOL)
buggy_flags = G.gradient_check(D.check_function, D.check_gradient_buggy, D.CHECK_POINT, D.NUMERIC_H, D.CHECK_TOL)
print("check_correct_all_pass", all(correct_flags))
print("check_buggy_flags", "|".join(str(f) for f in buggy_flags))

left = G.gradient_descent(D.two_minima_grad, D.TWO_MINIMA_LEFT_START, D.TWO_MINIMA_LR, D.TWO_MINIMA_ITERS)
right = G.gradient_descent(D.two_minima_grad, D.TWO_MINIMA_RIGHT_START, D.TWO_MINIMA_LR, D.TWO_MINIMA_ITERS)
print("two_minima_differ", G.minima_differ(left[-1], right[-1], D.TWO_MINIMA_MARGIN))
print("two_minima_left_near_neg1", abs(left[-1] + 1.0) < 1e-3)
print("two_minima_right_near_pos1", abs(right[-1] - 1.0) < 1e-3)

r = G.stopping_criteria_disagree(D.PLATEAU_X0, D.plateau_grad, D.plateau_value, D.PLATEAU_LR, D.PLATEAU_GRAD_TOL, D.PLATEAU_DELTA_F_TOL)
print("plateau_grad_above_tol", r["grad_norm"] >= D.PLATEAU_GRAD_TOL)
print("plateau_delta_f_below_tol", r["delta_f"] < D.PLATEAU_DELTA_F_TOL)
print("plateau_naive_stops_early", r["naive_stops_early"])

hook_path = G.gradient_descent(lambda x: D.HOOK_A * x, D.HOOK_X0, D.HOOK_LR, D.HOOK_ITERS)
first_inf = next(i for i, v in enumerate(hook_path) if math.isinf(v))
first_nan = next(i for i, v in enumerate(hook_path) if math.isnan(v))
print("hook_nan_follows_inf", first_nan == first_inf + 1)
losses20 = [0.5 * D.HOOK_A * v * v for v in hook_path[:21]]
print("hook_loss_always_increases", all(losses20[i + 1] > losses20[i] for i in range(len(losses20) - 1)))
PY
)"

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

check_eq "the quadratic used throughout has a = 5" "5.0" "$(get a)"
check_eq "the exact-landing boundary is 1/a = 0.2" "0.2" "$(get critical_lr)"
check_eq "the divergence boundary is 2/a = 0.4" "0.4" "$(get divergence_lr)"

check_eq "eta=0.10 (0 < eta < 1/a) is classified monotone" "monotone" "$(get regime_monotone)"
check_eq "eta=0.20 (eta = 1/a) is classified exact" "exact" "$(get regime_exact)"
check_eq "eta=0.35 (1/a < eta < 2/a) is classified oscillating" "oscillating" "$(get regime_oscillating)"
check_eq "eta=0.45 (eta > 2/a) is classified divergent" "divergent" "$(get regime_divergent)"
check_eq "at eta = 1/a, x lands exactly on 0 after one step" "0.0" "$(get exact_second_value)"
check_eq "the divergent run grows by more than 10x over 30 steps" "True" "$(get divergent_grows)"

check_eq "the monotone ratio matches |1 - eta*a| = 0.5" "True" "$(get ratio_monotone_matches)"
check_eq "the oscillating ratio matches |1 - eta*a| = 0.75" "True" "$(get ratio_oscillating_matches)"
check_eq "the divergent ratio matches |1 - eta*a| = 1.25" "True" "$(get ratio_divergent_matches)"
echo "  (measured on this run: monotone ratio $(get ratio_monotone), oscillating ratio $(get ratio_oscillating), divergent ratio $(get ratio_divergent))"

check_eq "steps to converge are non-decreasing as kappa grows" "True" "$(get kappa_nondecreasing)"
check_eq "kappa=100 needs at least 10x the steps of kappa=1" "True" "$(get kappa_order_of_magnitude)"
check_eq "the isotropic bowl (kappa=1) converges in exactly one step" "True" "$(get kappa_one_step_isotropic)"
echo "  (measured on this run: steps for kappa in {1,5,20,100} were $(get kappa_steps))"

check_eq "momentum needs strictly fewer steps than plain descent at the same eta" "True" "$(get momentum_faster)"
echo "  (measured on this run: plain $(get momentum_plain_steps) steps, momentum $(get momentum_steps) steps)"

check_eq "the correct gradient passes every component of the check" "True" "$(get check_correct_all_pass)"
check_eq "the buggy gradient is flagged on exactly component 1" "True|False|True" "$(get check_buggy_flags)"

check_eq "the two initialisations converge to minima farther apart than the margin" "True" "$(get two_minima_differ)"
check_eq "the left run converges to -1" "True" "$(get two_minima_left_near_neg1)"
check_eq "the right run converges to +1" "True" "$(get two_minima_right_near_pos1)"

check_eq "on the plateau, the gradient stays at or above its own tolerance" "True" "$(get plateau_grad_above_tol)"
check_eq "and the loss change falls below its own tolerance" "True" "$(get plateau_delta_f_below_tol)"
# Section 6 re-runs this script with D111_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_naive_stop="True"
if [ -n "${D111_SELF_TEST:-}" ]; then
  expected_naive_stop="False"   # the belief that a flat loss always means convergence
fi
check_eq "so the naive |delta f| stopping rule fires early" "${expected_naive_stop}" "$(get plateau_naive_stops_early)"

check_eq "the diverging run's nan follows its inf on the very next step" "True" "$(get hook_nan_follows_inf)"
check_eq "the loss increases on every one of the first 20 steps before it overflows" "True" "$(get hook_loss_always_increases)"

# --------------------------------------------------------------------------
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 -- the belief that a flat loss always means convergence --
# and asserts that the re-run reports the failure and exits non-zero. If
# this section passes, section 5 is not decorative.
if [ -z "${D111_SELF_TEST:-}" ]; then
  self_out="$(D111_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: so the naive |delta f| stopping rule fires early"*)
      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.

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 ]