Hands-on lab — Day 110: The Chain Rule

"""Rates multiply. That is the chain rule, before any calculus is involved.

Run from inside `examples/`:

    ../.venv/bin/python3 01_gears_and_rates.py
"""

import dataset as D
from chainrule import gear_ratio, product

print("=" * 74)
print("1. Two gears")
print("=" * 74)
print()
print("  Gear A turns 2 times for every 1 turn of gear B.")
print("  Gear B turns 3 times for every 1 turn of gear C.")
print()
print("  So how many times does A turn for one turn of C?")
print()
print("  Turn C once      -> B turns 3 times")
print("  B turns 3 times  -> A turns 3 x 2 = 6 times")
print()
print(f"  overall ratio = 2 x 3 = {gear_ratio(D.GEAR_RATIOS)}")
print()
print("  You did not need calculus for that, and you did not need to be")
print("  told a rule. You multiplied, because 'per' stacks by multiplying.")

assert gear_ratio(D.GEAR_RATIOS) == D.GEAR_RATIO_PRODUCT

print()
print("=" * 74)
print("2. A longer gear train")
print("=" * 74)
print()
print("  Add two more stages and nothing about the reasoning changes.")
print()
print("     stage      ratio     running product")
running = 1.0
for i, ratio in enumerate(D.GEAR_TRAIN, start=1):
    running *= ratio
    print(f"     {i}          {ratio:<9.2f} {running:.2f}")
print()
print(f"  overall ratio = {gear_ratio(D.GEAR_TRAIN)}")
print()
print("  Four stages, four numbers, one product. The chain rule for a")
print("  composition of four functions has exactly this shape, and the only")
print("  thing calculus adds is that the ratios are allowed to depend on")
print("  where you are -- a gear ratio is fixed, a derivative is not.")

assert gear_ratio(D.GEAR_TRAIN) == D.GEAR_TRAIN_PRODUCT

print()
print("=" * 74)
print("3. The same arithmetic with money")
print("=" * 74)
print()
print("  These three rates are invented for the arithmetic. They are not")
print("  quoted from any market and no real currency is named.")
print()
print("     1 unit of the first  buys 1.25 of the second")
print("     1 unit of the second buys 0.80 of the third")
print("     1 unit of the third  buys  150 of the fourth")
print()
print(f"     1.25 x 0.80 x 150 = {product(D.CURRENCY_RATES)}")
print()
print("  Notice the middle rate is BELOW one, and it drags the product down")
print("  relative to what the other two would have given alone. Hold on to")
print("  that: fifty factors slightly below one is how a gradient vanishes,")
print("  and script 07 measures it.")

assert product(D.CURRENCY_RATES) == D.CURRENCY_PRODUCT

print()
print("=" * 74)
print("4. What the notation is for")
print("=" * 74)
print()
print("  Write the gear answer the way calculus writes it:")
print()
print("      dA     dA     dB")
print("      --  =  --  x  --")
print("      dC     dB     dC")
print()
print("      6   =   2  x   3")
print()
print("  The 'dB' looks like it cancels, top and bottom. That reading is a")
print("  useful mnemonic and it is NOT a proof: dA and dB are not numbers")
print("  and there is no division happening. The reason the rule is true is")
print("  the gear reasoning above -- rates per something stack by")
print("  multiplying -- not the accident that the symbols line up.")
print()
print("  The mnemonic also stops working the moment a variable reaches the")
print("  output by more than one route. Script 04 is that case, and there")
print("  the contributions are ADDED. Nothing cancels.")

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

"""Composition first, then the chain rule, then the rule checked by measurement.

Run from inside `examples/`:

    ../.venv/bin/python3 02_composition_and_the_chain_rule.py
"""

import dataset as D
from chainrule import central_difference, chain_rule, compose

print("=" * 74)
print("1. Composition, with numbers and no derivatives at all")
print("=" * 74)
print()
print("  Two ordinary functions:")
print()
print("      g(x) = 3x + 1        the inner function, runs first")
print("      f(u) = u squared     the outer function, runs on g's answer")
print()
print("  Composing them means feeding one into the other:")
print()
print("      f(g(x)) = (3x + 1) squared")
print()
print("  At x = 2, in the order the arithmetic actually happens:")
print()
inner_value = D.line(2.0)
outer_value = D.square(inner_value)
print(f"      g(2)     = 3*2 + 1 = {inner_value}")
print(f"      f(g(2))  = {inner_value} squared = {outer_value}")
print()
print("  Read the parentheses inside out. The inner function runs first even")
print("  though it is written second, which is the one piece of bookkeeping")
print("  that trips people up before any calculus arrives.")

assert inner_value == 7.0
assert outer_value == 49.0
assert compose(D.square, D.line)(2.0) == 49.0

print()
print("=" * 74)
print("2. Now ask how fast the answer moves")
print("=" * 74)
print()
print("  Nudge x a little. Two things happen in sequence:")
print()
print("      x moves by 1 unit    ->  u = g(x) moves by 3 units")
print("      u moves by 1 unit    ->  y = f(u) moves by 2u = 14 units")
print()
print("  So x moving by 1 moves y by 3 x 14 = 42. The rates multiplied,")
print("  exactly as the gears did.")
print()
print("      dy/dx = dy/du x du/dx = 14 x 3 = 42")
print()
print("  The single most common mistake in this line is evaluating the")
print("  outer derivative at x instead of at u. f'(u) = 2u, and u is 7 here,")
print("  not 2. Using x would give 2*2 = 4 and an answer of 12, which is")
print("  wrong by more than a factor of three.")
print()
d_outer_at_u = D.d_square(inner_value)
d_outer_at_x_wrong = D.d_square(2.0)
print(f"      f'(u) at u = 7  ->  {d_outer_at_u}   correct")
print(f"      f'(x) at x = 2  ->  {d_outer_at_x_wrong}    the mistake")
print(f"      correct answer   ->  {d_outer_at_u} x 3 = {d_outer_at_u * 3.0}")
print(f"      the mistake gives -> {d_outer_at_x_wrong} x 3 = {d_outer_at_x_wrong * 3.0}")

assert d_outer_at_u == 14.0
assert d_outer_at_u * 3.0 == 42.0
assert d_outer_at_x_wrong * 3.0 == 12.0

print()
print("=" * 74)
print("3. Six compositions, each checked against a central difference")
print("=" * 74)
print()
print("  The central difference from Day 108 knows nothing about the chain")
print("  rule. It moves x by a hair and watches the output move. That")
print("  independence is what makes it a real check rather than a restatement.")
print()
print(f"  Step size h = {D.H:g}, tolerance {D.NUMERIC_TOL:g}.")
print()
print("     composition                 chain rule      measured        gap")
print("     " + "-" * 63)
for case in D.COMPOSITIONS:
    analytic = chain_rule(case.d_outer, case.inner, case.d_inner, case.x)
    measured = central_difference(compose(case.outer, case.inner), case.x, D.H)
    gap = abs(analytic - measured)
    print(
        f"     {case.name:<26s} {analytic:>13.9f} {measured:>13.9f}  {gap:.2e}"
    )
    assert abs(analytic - case.exact) < D.ANALYTIC_TOL, case.name
    assert gap < D.NUMERIC_TOL, case.name
print()
print("  Every gap is around 1e-10 or smaller, which is the central")
print("  difference's own error and not a disagreement about the rule.")

print()
print("=" * 74)
print("4. Two of those six are worth naming")
print("=" * 74)
print()
sigmoid_case = D.COMPOSITIONS[4]
sigmoid_slope = chain_rule(
    sigmoid_case.d_outer, sigmoid_case.inner, sigmoid_case.d_inner, 0.0
)
print("  'the sigmoid' is 1 / (1 + e to the minus x), which is a composition")
print("  of a reciprocal with an exponential. Its slope at 0 is")
print()
print(f"      dy/dx = (-1/4) x (-1) = {sigmoid_slope}")
print()
print("  and 0.25 is the LARGEST slope the sigmoid ever has. Every other")
print("  point is shallower. Stack fifty sigmoid layers and you are")
print("  multiplying fifty numbers no bigger than 0.25. Script 07 does that")
print("  arithmetic; this is where the answer comes from.")
print()
tanh_case = D.COMPOSITIONS[5]
tanh_slope = chain_rule(
    tanh_case.d_outer, tanh_case.inner, tanh_case.d_inner, -0.5
)
print("  'tanh of a line' is tanh(2x + 1) at x = -0.5, where the inner")
print("  function is exactly 0. tanh'(0) = 1 - tanh(0) squared = 1, so the")
print(f"  whole answer is 1 x 2 = {tanh_slope}, exactly. That exactness is")
print("  used again in script 06 to make a whole backward pass checkable")
print("  with a pen.")

assert sigmoid_slope == 0.25
assert tanh_slope == 2.0

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

"""Two functions, then three, then five. The product grows; nothing else does.

Run from inside `examples/`:

    ../.venv/bin/python3 03_deeper_chains.py
"""

import math

import dataset as D
from chainrule import (
    central_difference,
    chain_derivative,
    chain_function,
    chain_local_rates,
    chain_values,
    running_products,
)

print("=" * 74)
print("1. Depth two, three and five, side by side")
print("=" * 74)
print()
print("  The five stages, applied left to right starting from x = 1:")
print()
print("      1. double        u -> 2u")
print("      2. add three     u -> u + 3")
print("      3. square        u -> u squared")
print("      4. square root   u -> sqrt(u)")
print("      5. logarithm     u -> ln(u)")
print()
print("  Take the first two, then the first three, then all five, and watch")
print("  the derivative be the product of however many local rates there are.")
print()
print("     depth   value out        local rates                  product")
print("     " + "-" * 66)
for depth in (2, 3, 5):
    stages = D.FIVE_STAGES[:depth]
    rates = D.FIVE_RATES[:depth]
    out = chain_function(stages)(D.FIVE_START)
    local = chain_local_rates(stages, rates, D.FIVE_START)
    analytic = chain_derivative(stages, rates, D.FIVE_START)
    measured = central_difference(chain_function(stages), D.FIVE_START, D.H)
    rates_text = " x ".join(f"{r:g}" for r in local)
    print(f"     {depth}       {out:<15.10g} {rates_text:<27s} {analytic:.10g}")
    assert abs(analytic - measured) < D.NUMERIC_TOL, depth
print()
print("  Each row was also checked against a central difference of the whole")
print(f"  composed function, and every gap was below {D.NUMERIC_TOL:g}.")

print()
print("=" * 74)
print("2. The forward pass: what arrives where")
print("=" * 74)
print()
values = chain_values(D.FIVE_STAGES, D.FIVE_START)
local = chain_local_rates(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START)
print("     stage           input      output     local rate at that input")
print("     " + "-" * 62)
names = ("double", "add three", "square", "square root", "logarithm")
for i, name in enumerate(names):
    print(
        f"     {name:<15s} {values[i]:<10.6g} {values[i + 1]:<10.6g} {local[i]:.6g}"
    )
print()
print("  The column that matters is the last one. Stage 3's local rate is")
print("  2u, and the u it is evaluated at is 5 -- the value that ARRIVES at")
print("  stage 3, not the input x and not the final answer. Getting the")
print("  evaluation point wrong is the mistake that survives longest,")
print("  because the shape of the answer still looks right.")

assert values == list(D.FIVE_VALUES)
assert local == list(D.FIVE_LOCAL_RATES)

print()
print("=" * 74)
print("3. The product, and the same answer from an entirely different route")
print("=" * 74)
print()
analytic = chain_derivative(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START)
print(f"     2 x 1 x 10 x 0.1 x 0.2 = {analytic}")
print()
print("  Now collapse the five stages by hand instead. Squaring and then")
print("  taking a square root of a positive number is the identity, so")
print()
print("      ln( sqrt( (2x + 3) squared ) ) = ln(2x + 3)")
print()
closed = D.d_five_chain_closed_form(D.FIVE_START)
measured = central_difference(chain_function(D.FIVE_STAGES), D.FIVE_START, D.H)
print(f"      d/dx ln(2x + 3) = 2 / (2x + 3),  at x = 1 that is {closed}")
print(f"      central difference of the five-stage chain:      {measured:.12f}")
print()
print("  Three routes -- five local rates multiplied, one collapsed formula,")
print("  and a measurement that knows about none of it -- agree.")

assert analytic == D.FIVE_DERIVATIVE
assert abs(closed - D.FIVE_DERIVATIVE) < D.ANALYTIC_TOL
assert abs(measured - D.FIVE_DERIVATIVE) < D.NUMERIC_TOL

print()
print("=" * 74)
print("4. What a backward pass is actually carrying")
print("=" * 74)
print()
carried = running_products(local)
print("  Walk the chain from the output end, multiplying as you go. After")
print("  k steps the number in your hand is the product of the last k local")
print("  rates -- which is exactly the gradient of the output with respect")
print("  to the value arriving at that stage.")
print()
print("     stage           value in    d(output)/d(that value)")
print("     " + "-" * 55)
for i, name in enumerate(names):
    print(f"     {name:<15s} {values[i]:<11.6g} {carried[i]:.10g}")
print()
print("  Read the last row upwards. The logarithm stage sees 0.2; the square")
print("  root stage sees 0.02; and by the time the walk reaches the input")
print(f"  the number is {carried[0]:.6g}, which is the answer.")
print()
print("  That is the entire backward pass. One walk, one multiplication per")
print("  stage, and every intermediate gradient computed on the way past --")
print("  not one walk per stage. Script 07 measures what that saves.")

assert abs(carried[0] - D.FIVE_DERIVATIVE) < D.ANALYTIC_TOL
assert carried[-1] == D.FIVE_LOCAL_RATES[-1]

print()
print("=" * 74)
print("5. A note on the order of the multiplications")
print("=" * 74)
print()
left_to_right = analytic
right_to_left = carried[0]
print(f"     multiplied forwards:  {left_to_right!r}")
print(f"     multiplied backwards: {right_to_left!r}")
print(f"     they differ by:       {abs(left_to_right - right_to_left):.3e}")
print()
print("  Float64 multiplication is not associative, so the two orders can")
print("  land on different bit patterns. The gap here is one unit in the")
print("  last place and it is reported rather than hidden, because a lab")
print(f"  that compared these with == would be lying. The tolerance used is")
print(f"  {D.ANALYTIC_TOL:g}, which is for two analytic routes to the same")
print("  number, and it is a thousand times tighter than the tolerance used")
print("  against a measured derivative.")

assert abs(left_to_right - right_to_left) < D.ANALYTIC_TOL
assert math.isfinite(left_to_right)

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

"""When a variable reaches the output twice, the contributions ADD.

This is the section to slow down in. Everything before it multiplies;
everything after it depends on knowing when to add.

Run from inside `examples/`:

    ../.venv/bin/python3 04_two_paths_add.py
"""

import dataset as D
from chainrule import (
    central_difference,
    partial_difference,
    path_contributions,
    total_derivative,
    wrong_single_path_derivative,
)

print("=" * 74)
print("1. One input, two routes to the output")
print("=" * 74)
print()
print("  Build a small graph in which x is used twice:")
print()
print("      u = x squared")
print("      v = 3x")
print("      f = u x v")
print()
print("  Draw it and x has two arrows leaving it, one into u and one into v.")
print("  Both of them end up at f.")
print()
x = D.TWO_PATH_X
u = x * x
v = 3.0 * x
f = u * v
print(f"      at x = {x}:   u = {u},  v = {v},  f = {f}")

assert u == D.TWO_PATH_U
assert v == D.TWO_PATH_V
assert f == D.TWO_PATH_OUTPUT

print()
print("=" * 74)
print("2. The two path products")
print("=" * 74)
print()
print("  Along each path, multiply the local rates as usual:")
print()
print("      path through u:   df/du x du/dx  =  v   x 2x  = 6 x 4 = 24")
print("      path through v:   df/dv x dv/dx  =  u   x 3   = 4 x 3 = 12")
print()
paths = [[v, 2.0 * x], [u, 3.0]]
contributions = path_contributions(paths)
print(f"      contributions = {contributions}")
print()
print("  Now the question the whole day turns on: is the answer 24, is it 12,")
print("  is it 24 x 12 = 288, or is it 24 + 12 = 36?")

assert contributions == list(D.TWO_PATH_CONTRIBUTIONS)

print()
print("=" * 74)
print("3. Ask the measurement, which has no opinion about rules")
print("=" * 74)
print()
measured = central_difference(D.two_path_direct, x, D.H)
summed = total_derivative(paths)
wrong_a = wrong_single_path_derivative(paths, 0)
wrong_b = wrong_single_path_derivative(paths, 1)
multiplied = contributions[0] * contributions[1]
print("  Substitute the intermediates away and the function is just")
print("  f = x squared x 3x = 3 x cubed, which we can nudge directly.")
print()
print(f"      central difference of 3x cubed at x = 2:   {measured:.9f}")
print()
print("     candidate answer            value      matches the measurement?")
print("     " + "-" * 60)
for label, value in (
    ("sum of the paths, 24 + 12", summed),
    ("path through u alone", wrong_a),
    ("path through v alone", wrong_b),
    ("product of the paths", multiplied),
):
    verdict = "YES" if abs(value - measured) < D.NUMERIC_TOL else "no"
    print(f"     {label:<27s} {value:<10.6g} {verdict}")
print()
print("  Only the sum survives. And it is not a convention: changing x moves")
print("  the output through u AND through v, both movements are real, and")
print("  both happen at once. Adding them is what 'both happen' means.")
print()
print("  Check it against the closed form too. f = 3x cubed, so df/dx = 9x")
print(f"  squared, which at x = 2 is {D.d_two_path_direct(x)}.")

assert summed == D.TWO_PATH_DERIVATIVE
assert abs(summed - measured) < D.NUMERIC_TOL
assert abs(D.d_two_path_direct(x) - summed) < D.ANALYTIC_TOL
# The instructive failures, asserted as failures so the suite would notice if
# a future edit made the single-path answer accidentally correct.
assert abs(wrong_a - measured) > 1.0
assert abs(wrong_b - measured) > 1.0
assert abs(multiplied - measured) > 1.0

print()
print("=" * 74)
print("4. Why the cancelling mnemonic breaks here")
print("=" * 74)
print()
print("  The 'du cancels' reading of dy/dx = dy/du x du/dx has nothing to")
print("  say about this graph, because there are two different intermediates")
print("  and no single symbol to cancel. Written honestly the rule is:")
print()
print("      df      df   du     df   dv")
print("      --  =   -- x --  +  -- x --")
print("      dx      du   dx     dv   dx")
print()
print("  A product for each route, a sum across routes. Every backward pass")
print("  in every framework is that formula applied node by node, and the")
print("  '+' is why a gradient is accumulated with += rather than assigned.")

print()
print("=" * 74)
print("5. The full multivariable case: two inputs, two intermediates")
print("=" * 74)
print()
print("      z = u squared + v squared,   u = s x t,   v = s - t")
print()
s, t = D.SURFACE_POINT
u2 = s * t
v2 = s - t
z = u2 * u2 + v2 * v2
print(f"      at (s, t) = ({s:g}, {t:g}):  u = {u2:g}, v = {v2:g}, z = {z:g}")
print()
print("      dz/du = 2u = 12          dz/dv = 2v = -2")
print("      du/ds = t  = 3           du/dt = s  = 2")
print("      dv/ds = 1                dv/dt = -1")
print()
dz_ds = D.SURFACE_DZ_DU * t + D.SURFACE_DZ_DV * 1.0
dz_dt = D.SURFACE_DZ_DU * s + D.SURFACE_DZ_DV * -1.0
print(f"      dz/ds = 12 x 3 + (-2) x 1    = {dz_ds:g}")
print(f"      dz/dt = 12 x 2 + (-2) x (-1) = {dz_dt:g}")
print()
num_ds = partial_difference(D.surface, D.SURFACE_POINT, 0, D.H)
num_dt = partial_difference(D.surface, D.SURFACE_POINT, 1, D.H)
print("     quantity   chain rule    partial difference     gap")
print("     " + "-" * 52)
print(f"     dz/ds      {dz_ds:<13g} {num_ds:<22.9f} {abs(dz_ds - num_ds):.2e}")
print(f"     dz/dt      {dz_dt:<13g} {num_dt:<22.9f} {abs(dz_dt - num_dt):.2e}")
print()
print("  Two inputs, two intermediates, four paths in total, and every")
print("  gradient is a sum of two products. That is the entire multivariable")
print("  chain rule, and it is the reason a branching computation graph is")
print("  no harder to differentiate than a straight line -- only longer.")

assert z == D.SURFACE_Z
assert (dz_ds, dz_dt) == D.SURFACE_GRADIENT
assert abs(dz_ds - num_ds) < D.NUMERIC_TOL
assert abs(dz_dt - num_dt) < D.NUMERIC_TOL

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

"""The reverse-mode engine, exercised and checked against measurements.

Run from inside `examples/`:

    ../.venv/bin/python3 05_the_value_engine.py
"""

import math

import dataset as D
from autodiff import (
    Value,
    graph_size,
    numeric_gradient,
    reverse_mode_gradient,
    topological_order,
)
from chainrule import central_difference

print("=" * 74)
print("1. The smallest possible graph")
print("=" * 74)
print()
print("      a = 3,  b = 4,  c = a x b")
print()
a = Value(3.0, label="a")
b = Value(4.0, label="b")
c = a * b
c.backward()
print(f"      c.data = {c.data}")
print(f"      a.grad = {a.grad}   <- dc/da, which is b")
print(f"      b.grad = {b.grad}   <- dc/db, which is a")
print()
print("  For a product, each input's local rate is the other input's value.")
print("  Nudge a by one and c moves by b. That is the whole `__mul__`")
print("  backward step, and it took two lines to write.")

assert c.data == 12.0
assert a.grad == 4.0
assert b.grad == 3.0

print()
print("=" * 74)
print("2. A value used twice: where the += earns its keep")
print("=" * 74)
print()
print("      x = 3,  y = x + x")
print()
x = Value(3.0, label="x")
y = x + x
y.backward()
print(f"      y.data = {y.data}")
print(f"      x.grad = {x.grad}")
print()
print("  x receives a contribution from each of its two uses, and they add.")
print("  If the engine assigned the gradient instead of accumulating it, the")
print("  second contribution would overwrite the first and this would print")
print("  1.0 -- a plausible-looking, confidently wrong number. That is the")
print("  same sum-over-paths fact as script 04, now as one character of code.")
print()
print("  The same thing with a multiplication, where the answer is less")
print("  obvious:")
print()
print("      x = 3,  y = x x x")
x2 = Value(3.0, label="x")
y2 = x2 * x2
y2.backward()
print(f"      y.data = {y2.data},  x.grad = {x2.grad}   <- 2x, the power rule")
print()
print("  The engine has never heard of the power rule. It applied the product")
print("  rule and added the two contributions, and the power rule fell out.")

assert y.data == 6.0
assert x.grad == 2.0
assert y2.data == 9.0
assert x2.grad == 6.0

print()
print("=" * 74)
print("3. The topological order, and why it cannot be skipped")
print("=" * 74)
print()
p = Value(2.0, label="p")
q = Value(-3.0, label="q")
r = p * q
s = r + p
out = s.tanh()
order = topological_order(out)
print("      p = 2,  q = -3,  r = p x q,  s = r + p,  out = tanh(s)")
print()
print("     position   node")
print("     " + "-" * 46)
for i, node in enumerate(order):
    name = node.label or node._op or "const"
    print(f"     {i}          {name:<10s} data = {node.data:.6g}")
print()
print("  Every node sits after everything it was computed from, so walking")
print("  the list backwards guarantees a node has received all of its")
print("  gradient before it passes any of it on. p is used twice -- once by")
print("  r and once by s -- and if s handed its gradient to p before r had")
print("  finished, p's total would be short by one path.")
print()
out.backward()
print(f"      out.data = {out.data:.12f}")
print(f"      p.grad   = {p.grad:.12f}")
print(f"      q.grad   = {q.grad:.12f}")
print()


def scalar_out(pv: float) -> float:
    """The same expression in plain floats, for a central difference."""
    return math.tanh(pv * -3.0 + pv)


measured_p = central_difference(scalar_out, 2.0, D.H)
print(f"      central difference for p: {measured_p:.12f}")
print(f"      gap: {abs(p.grad - measured_p):.3e}")

assert graph_size(out) == len(order)
assert abs(p.grad - measured_p) < D.NUMERIC_TOL

print()
print("=" * 74)
print("4. Four expressions, each differentiated two independent ways")
print("=" * 74)
print()
print("  The engine only has +, x and tanh, so it cannot build a sine or a")
print("  logarithm. What it CAN build is any polynomial and any tanh network,")
print("  which is enough to check it hard. Every gradient below is produced")
print("  by the engine and then measured with a central difference that")
print("  knows nothing about the graph.")
print()


def expr1(vals):
    (xv,) = vals
    return (3.0 * xv + 1.0) * (3.0 * xv + 1.0)


def expr2(vals):
    (xv,) = vals
    return (xv * xv * xv) + (-2.0) * xv


def expr4(vals):
    xv, yv = vals
    return (xv * yv + xv) * (yv + 3.0)


def expr5(vals):
    xv, yv, zv = vals
    return ((xv * yv).tanh() * zv + xv * zv) * (1.0 + yv)


CASES = (
    ("(3x + 1) squared            ", expr1, [2.0]),
    ("x cubed - 2x                ", expr2, [1.5]),
    ("(xy + x)(y + 3)             ", expr4, [2.0, -1.0]),
    ("tanh(xy)z + xz, times (1+y) ", expr5, [0.7, 0.4, -1.3]),
)

print("     expression                     input   engine grad     measured")
print("     " + "-" * 68)
for label, build, point in CASES:

    def plain(vals, build=build):
        node = build([Value(v) for v in vals])
        return node.data

    grads, passes = reverse_mode_gradient(build, point)
    numeric, num_passes = numeric_gradient(plain, point, D.H)
    for i, (g, n) in enumerate(zip(grads, numeric)):
        shown = label if i == 0 else " " * len(label)
        print(f"     {shown} {point[i]:>7.4g} {g:>14.9f} {n:>14.9f}")
        assert abs(g - n) < D.NUMERIC_TOL + D.NUMERIC_REL_TOL * abs(n), label
    assert passes == 1, label
    assert num_passes == 2 * len(point), label
print()
print("  Every engine gradient agrees with a central difference of the same")
print("  expression to within the numerical rule's own error. The engine had")
print("  no formula for any of these functions -- it composed +, x and tanh")
print("  and let the chain rule do the rest.")
print()
print("  Note the pass counts, which the assertions above also check: the")
print("  engine used ONE forward-and-backward sweep per expression no matter")
print("  how many inputs it had, and the central difference needed two")
print("  evaluations per input. Script 07 makes that gap the point.")

print()
print("=" * 74)
print("5. A tanh identity the engine reproduces without being told")
print("=" * 74)
print()
z = Value(0.6, label="z")
t = z.tanh()
t.backward()
print("      z = 0.6,  t = tanh(z)")
print(f"      t.data = {t.data:.12f}")
print(f"      z.grad = {z.grad:.12f}")
print(f"      1 - t squared = {1.0 - t.data * t.data:.12f}")
print()
print("  The derivative of tanh is 1 - tanh squared, and the engine's")
print("  backward step for tanh is literally that expression -- reusing the")
print("  forward value rather than recomputing anything. Every framework does")
print("  this, and it is why a backward pass needs the forward pass's values")
print("  kept in memory. That memory cost is the price reverse mode pays for")
print("  its speed, and on a large model it is the dominant one.")

assert abs(z.grad - (1.0 - t.data * t.data)) < D.ANALYTIC_TOL
assert abs(t.data - math.tanh(0.6)) < D.ANALYTIC_TOL

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

"""Backpropagation through a tiny network, by hand, then by the engine.

Every number here is exact in float64, so you can check the whole backward
pass with a pen and nothing will be off in the twelfth decimal place.

Run from inside `examples/`:

    ../.venv/bin/python3 06_backprop_by_hand.py
"""

import dataset as D
import network as N

print("=" * 74)
print("1. The network")
print("=" * 74)
print()
print("      x1, x2  ->  two tanh hidden units  ->  one linear output  ->  loss")
print()
print("      a_pre = wA1*x1 + wA2*x2 + bA        a = tanh(a_pre)")
print("      b_pre = wB1*x1 + wB2*x2 + bB        b = tanh(b_pre)")
print("      out   = vA*a + vB*b + c")
print("      loss  = (out - target) squared")
print()
print("     name     value      name     value")
print("     " + "-" * 42)
print(f"     x1       {D.NET_X1:<10g} x2       {D.NET_X2:g}")
print(f"     wA1      {D.NET_WA1:<10g} wB1      {D.NET_WB1:g}")
print(f"     wA2      {D.NET_WA2:<10g} wB2      {D.NET_WB2:g}")
print(f"     bA       {D.NET_BA:<10g} bB       {D.NET_BB:.10f}")
print(f"     vA       {D.NET_VA:<10g} vB       {D.NET_VB:g}")
print(f"     c        {D.NET_C:<10g} target   {D.NET_TARGET:g}")
print()
print("  bB is half the natural logarithm of 3. That is not a magic number:")
print("  it is chosen so that tanh(bB) is exactly 0.5 in float64 and its")
print("  slope is exactly 0.75, which makes every line below checkable with")
print("  a pen. Nothing about backpropagation depends on the choice.")

assert D.NET_BB == D.HALF_LN3

print()
print("=" * 74)
print("2. The forward pass")
print("=" * 74)
print()
fw = N.forward(D.NET_X1, D.NET_X2, N.default_parameter_values())
print(f"      a_pre = 1*1 + (-0.5)*2 + 0        = {fw['a_pre']:g}")
print(f"      a     = tanh(0)                   = {fw['a']:g}")
print(f"      b_pre = (-0.5)*1 + 0.25*2 + bB    = {fw['b_pre']:.10f}")
print(f"      b     = tanh(bB)                  = {fw['b']:g}")
print(f"      out   = 2*0 + (-3)*0.5 + 1        = {fw['out']:g}")
print(f"      loss  = (-0.5 - 1) squared        = {fw['loss']:g}")

assert fw["a_pre"] == D.NET_A_PRE
assert fw["a"] == D.NET_A
assert fw["b"] == D.NET_B
assert fw["out"] == D.NET_OUT
assert fw["loss"] == D.NET_LOSS
# The two exactness claims the whole section rests on, asserted not assumed.
assert D.NET_B == 0.5
assert (1.0 - D.NET_B * D.NET_B) == 0.75

print()
print("=" * 74)
print("3. The backward pass, one local rate at a time")
print("=" * 74)
print()
print("  Start at the end. d(loss)/d(loss) is 1 -- the derivative of anything")
print("  with respect to itself. Everything else follows by multiplying.")
print()
hand = N.hand_gradients()
print("     step                        local rate            gradient")
print("     " + "-" * 62)
rows = (
    ("d loss / d out", "2 x (out - target) = 2 x (-1.5)", "out"),
    ("d loss / d c", "x 1", "c"),
    ("d loss / d vA", "x a = x 0", "vA"),
    ("d loss / d vB", "x b = x 0.5", "vB"),
    ("d loss / d a", "x vA = x 2", "a"),
    ("d loss / d b", "x vB = x (-3)", "b"),
    ("d loss / d a_pre", "x (1 - a^2) = x 1", "a_pre"),
    ("d loss / d b_pre", "x (1 - b^2) = x 0.75", "b_pre"),
    ("d loss / d wA1", "x x1 = x 1", "wA1"),
    ("d loss / d wA2", "x x2 = x 2", "wA2"),
    ("d loss / d bA", "x 1", "bA"),
    ("d loss / d wB1", "x x1 = x 1", "wB1"),
    ("d loss / d wB2", "x x2 = x 2", "wB2"),
    ("d loss / d bB", "x 1", "bB"),
)
for label, rate, key in rows:
    print(f"     {label:<26s}  {rate:<22s} {hand[key]:>9g}")
print()
print("  Two of those deserve a second look.")
print()
print(f"      d loss / d vA = {hand['vA']:g}")
print()
print("  Not small -- exactly zero. vA multiplies a, and a is exactly 0, so")
print("  nudging vA does not move the output at all. A weight feeding a unit")
print("  whose activation is zero receives no gradient and does not learn on")
print("  this step. That is not a bug in the arithmetic; it is the arithmetic")
print("  telling you something true about the network.")
print()
print(f"      d loss / d b_pre = {hand['b']:g} x 0.75 = {hand['b_pre']:g}")
print()
print("  The 0.75 is tanh's slope, and it is below 1. Every tanh unit a")
print("  gradient passes through multiplies it by a number in (0, 1]. Stack")
print("  fifty of them and script 07 shows what happens.")

print()
print("=" * 74)
print("4. The inputs, where two paths meet")
print("=" * 74)
print()
print("  x1 is used by BOTH hidden units, so it reaches the loss twice.")
print()
first, second = D.NET_X1_CONTRIBUTIONS
print(f"      through unit A:  d loss/d a_pre x wA1 = {hand['a_pre']:g} x {D.NET_WA1:g} = {first:g}")
print(f"      through unit B:  d loss/d b_pre x wB1 = {hand['b_pre']:g} x {D.NET_WB1:g} = {second:g}")
print(f"      total:           {first:g} + {second:g} = {hand['x1']:g}")
print()
print(f"      and for x2:      {hand['a_pre']:g} x {D.NET_WA2:g} + {hand['b_pre']:g} x {D.NET_WB2:g} = {hand['x2']:g}")
print()
print("  A product-only chain rule would report -6 or -3.375 here and look")
print("  entirely reasonable doing it. The measurement in section 5 settles")
print("  it, exactly as it did for the two-path graph in script 04.")

assert hand["x1"] == first + second
assert hand["x1"] == D.NET_GRADIENTS["x1"]
assert hand["x2"] == D.NET_GRADIENTS["x2"]

print()
print("=" * 74)
print("5. Three independent routes to the same sixteen numbers")
print("=" * 74)
print()
engine = N.engine_gradients()
numeric = N.numeric_parameter_gradients(D.H)
numeric.update(N.numeric_input_gradients(D.H))
print("     quantity   by hand      engine       central difference   gap")
print("     " + "-" * 64)
for key in ("wA1", "wA2", "bA", "wB1", "wB2", "bB", "vA", "vB", "c", "x1", "x2"):
    gap = abs(hand[key] - numeric[key])
    print(
        f"     {key:<10s} {hand[key]:>10g} {engine[key]:>12g} {numeric[key]:>20.9f}  {gap:.1e}"
    )
    assert hand[key] == engine[key], key
    assert abs(hand[key] - D.NET_GRADIENTS[key]) < D.ANALYTIC_TOL, key
    assert gap < D.NUMERIC_TOL, key
print()
print("  The hand column and the engine column are equal bit for bit, because")
print("  they perform the same multiplications in the same order on the same")
print("  exact values. The central-difference column is close but not equal,")
print("  and it never will be -- it is an approximation with its own error,")
print("  which is why it is compared with a tolerance a million times looser")
print("  than the one used between the first two columns.")
print()
print(f"      analytic vs analytic tolerance:  {D.ANALYTIC_TOL:g}")
print(f"      analytic vs measured tolerance:  {D.NUMERIC_TOL:g}")

print()
print("=" * 74)
print("6. What one backward pass cost")
print("=" * 74)
print()
forward_grads, forward_passes = N.forward_mode_parameter_gradients()
print(f"     reverse mode:  1 forward pass + 1 backward pass  -> all "
      f"{len(D.NET_PARAMETERS)} parameter gradients")
print(f"     forward mode:  {forward_passes} complete runs of the network "
      f"-> the same {len(D.NET_PARAMETERS)} gradients")
print(f"     central diff:  {2 * len(D.NET_PARAMETERS)} complete runs "
      "-> the same gradients, approximately")
print()
print("  Nine parameters is a toy. A model with a hundred million parameters")
print("  and one loss makes that ratio a hundred million to one, and that")
print("  single asymmetry is why training a large model is possible at all.")
for key, value in forward_grads.items():
    assert abs(value - hand[key]) < D.ANALYTIC_TOL, key
assert forward_passes == len(D.NET_PARAMETERS)

print()
print("06_backprop_by_hand.py: every assertion held.")

"""Fifty factors below one collapse. Fifty above one blow up. Same rule.

Run from inside `examples/`:

    ../.venv/bin/python3 07_vanishing_and_exploding.py
"""

import dataset as D
from autodiff import (
    Value,
    forward_mode_gradient,
    numeric_gradient,
    reverse_mode_gradient,
)
from chainrule import order_of_magnitude, product_trace, repeated_product

print("=" * 74)
print("1. A gradient walking back through fifty layers")
print("=" * 74)
print()
print("  The chain rule says the gradient at the far end of a chain is the")
print("  product of every local rate on the way. If each layer contributes a")
print("  rate of 0.9, the product is 0.9 to the fiftieth power. If each")
print("  contributes 1.1, it is 1.1 to the fiftieth.")
print()
print("  Neither factor looks alarming. Watch anyway.")
print()
decay = repeated_product(D.DECAY_FACTOR, D.CHAIN_LENGTH)
growth = repeated_product(D.GROWTH_FACTOR, D.CHAIN_LENGTH)
decay_trace = product_trace(D.DECAY_FACTOR, D.CHAIN_LENGTH)
growth_trace = product_trace(D.GROWTH_FACTOR, D.CHAIN_LENGTH)
print("     layers      x 0.9 each        x 1.1 each")
print("     " + "-" * 48)
for n in (1, 5, 10, 20, 30, 40, 50):
    print(f"     {n:<11d} {decay_trace[n - 1]:<17.6e} {growth_trace[n - 1]:.6e}")
print()
print(f"  After {D.CHAIN_LENGTH} layers the shrinking chain has lost more than")
print("  two orders of magnitude and the growing one has gained more than two.")
print("  The early layers of the shrinking network receive a gradient a few")
print("  thousandths the size of the one the last layer receives, so they")
print("  learn a few thousandths as fast. They are not broken. They are slow")
print("  by a factor nobody budgeted for.")
print()
print("  Asserted as orders of magnitude, not as digits -- the scale is the")
print("  lesson and the digits are float64 rounding:")
print()
print(f"      0.9 ** 50 = {decay:.6e}   order {order_of_magnitude(decay)}")
print(f"      1.1 ** 50 = {growth:.6e}   order {order_of_magnitude(growth)}")

assert order_of_magnitude(decay) == D.DECAY_ORDER
assert order_of_magnitude(growth) == D.GROWTH_ORDER
assert decay < 1e-2
assert growth > 1e2

print()
print("=" * 74)
print("2. Deeper, and harsher")
print("=" * 74)
print()
print("     factor   layers    product            order    lost in a weight of 1?")
print("     " + "-" * 68)
weight = 1.0
for factor, count in (
    (D.DECAY_FACTOR, D.LONG_CHAIN_LENGTH),
    (D.GROWTH_FACTOR, D.LONG_CHAIN_LENGTH),
    (D.MILD_DECAY, D.CHAIN_LENGTH),
    (D.SHARP_DECAY, D.CHAIN_LENGTH),
    (D.SHARP_GROWTH, D.CHAIN_LENGTH),
):
    value = repeated_product(factor, count)
    lost = "yes" if weight + value == weight else "no"
    print(
        f"     {factor:<8g} {count:<9d} {value:<18.6e} "
        f"{order_of_magnitude(value):<8d} {lost}"
    )
print()
mild = repeated_product(D.MILD_DECAY, D.CHAIN_LENGTH)
sharp = repeated_product(D.SHARP_DECAY, D.CHAIN_LENGTH)
print("  The last column asks a blunt question: if this were the gradient")
print("  and you added it to a weight of about 1, would the weight change at")
print("  all, or would the update disappear into rounding?")
print()
print("  The 0.5 row is worth pausing on, because the obvious guess about it")
print("  is wrong. Fifty halvings give")
print()
print(f"      0.5 ** 50 = {mild:.6e},  which is {mild / D.EPSILON:g} x EPSILON")
print()
print(f"  EPSILON here is {D.EPSILON:.6e}, the gap between 1.0 and the next")
print("  float64 above it. So this gradient is four of those gaps wide and it")
print("  DOES still move a weight of 1. It takes about three more halvings to")
print("  disappear -- 0.5 to the 53rd is half an EPSILON, and half a gap")
print("  rounds away to nothing.")
print()
print("  The row underneath is the one that actually vanishes. 0.25 is the")
print("  largest slope the sigmoid ever has, measured back in script 02, so a")
print("  stack of sigmoid layers is multiplying numbers no bigger than this:")
print()
print(f"      0.25 ** 50 = {sharp:.6e}")
print(f"      1.0 + {sharp:.6e} == 1.0  ->  {weight + sharp == weight}")
print()
print("  That is not a metaphor for a vanishing gradient. It is one, and the")
print("  factor being used is the sigmoid at its most generous rather than at")
print("  anything like a typical value.")

assert order_of_magnitude(repeated_product(D.DECAY_FACTOR, D.LONG_CHAIN_LENGTH)) == -10
assert order_of_magnitude(repeated_product(D.GROWTH_FACTOR, D.LONG_CHAIN_LENGTH)) == 8
# Measured, and it contradicts the obvious guess: four EPSILONs still counts.
assert mild == 4.0 * D.EPSILON
assert weight + mild != weight
assert repeated_product(D.MILD_DECAY, 53) == 0.5 * D.EPSILON
assert weight + repeated_product(D.MILD_DECAY, 53) == weight
# And the sigmoid's best case, which does not.
assert sharp < D.EPSILON
assert weight + sharp == weight

print()
print("=" * 74)
print("3. The same collapse, through the real engine")
print("=" * 74)
print()
print("  Nothing above needed the autodiff engine, so it could be accused of")
print("  being a story about exponents. Build the chain for real instead:")
print("  fifty multiplications by 0.9, differentiated by one backward pass.")
print()


def deep_chain(vals):
    (xv,) = vals
    node = xv
    for _ in range(D.CHAIN_LENGTH):
        node = node * D.DECAY_FACTOR
    return node


grads, passes = reverse_mode_gradient(deep_chain, [1.0])
print(f"      engine gradient after {D.CHAIN_LENGTH} layers: {grads[0]:.6e}")
print(f"      0.9 ** {D.CHAIN_LENGTH} for comparison:        {decay:.6e}")
print(f"      backward passes used:                {passes}")
print()
print("  The engine reproduces the collapse exactly, because the collapse IS")
print("  the chain rule. Nothing has gone wrong; the arithmetic is correct")
print("  and the answer is useless. Those are different complaints, and")
print("  every fix the course reaches later -- careful initialisation,")
print("  residual connections, normalisation, gradient clipping, ReLU in")
print("  place of a saturating non-linearity -- is an attempt to keep this")
print("  product near 1 rather than to make the chain rule behave otherwise.")

assert abs(grads[0] - decay) < D.ANALYTIC_TOL
assert passes == 1

print()
print("=" * 74)
print("4. Forward mode against reverse mode, counted")
print("=" * 74)
print()
print("  Both modes apply the same chain rule. They differ in which end they")
print("  start from, and that decides the cost.")
print()


def many_inputs(vals):
    """One output built from every input, with a non-linearity in the way."""
    total = vals[0] * 1.0
    for value in vals[1:]:
        total = total + value * value
    return (total * 0.1).tanh()


print("     inputs   reverse passes   forward passes   central-diff passes")
print("     " + "-" * 66)
for n in (1, 2, 5, 10, 25):
    point = [0.1 * (i + 1) for i in range(n)]

    def plain(vals):
        return many_inputs([Value(v) for v in vals]).data

    r_grads, r_passes = reverse_mode_gradient(many_inputs, point)
    f_grads, f_passes = forward_mode_gradient(many_inputs, point)
    n_grads, n_passes = numeric_gradient(plain, point, D.H)
    print(f"     {n:<8d} {r_passes:<16d} {f_passes:<16d} {n_passes}")
    for i in range(n):
        assert abs(r_grads[i] - f_grads[i]) < D.ANALYTIC_TOL, (n, i)
        assert abs(r_grads[i] - n_grads[i]) < D.NUMERIC_TOL, (n, i)
    assert r_passes == 1
    assert f_passes == n
    assert n_passes == 2 * n
print()
print("  All three columns produce the same gradients -- the first two")
print("  agree to the last bit, and the third to about a part in a billion.")
print("  Only the cost differs, and it differs by a factor that grows with")
print("  the number of inputs.")
print()
print("  So the rule of thumb is not about elegance, it is about shape:")
print()
print("      many inputs, one output   ->  reverse mode  (training a model)")
print("      one input, many outputs   ->  forward mode  (sensitivity to a")
print("                                    single parameter)")
print()
print("  Training a model is the first shape: millions of parameters in, one")
print("  scalar loss out. Reverse mode gets every gradient for roughly the")
print("  cost of two forward passes, and that is the entire economic basis")
print("  of modern machine learning.")

print()
print("=" * 74)
print("5. What reverse mode pays for it")
print("=" * 74)
print()
print("  Reverse mode has to keep the forward pass's intermediate values")
print("  alive until the backward pass consumes them, because the local")
print("  derivatives are written in terms of those values -- tanh's backward")
print("  step needs the tanh output, and a product's needs both inputs.")
print()
node = Value(1.0)
for _ in range(D.CHAIN_LENGTH):
    node = node * D.DECAY_FACTOR
from autodiff import graph_size  # noqa: E402  (imported here to keep it visible)

print(f"      a {D.CHAIN_LENGTH}-layer chain holds {graph_size(node)} nodes alive")
print()
print("  Forward mode holds almost nothing, which is its one real advantage.")
print("  On a large model the stored activations dominate memory use, and")
print("  that is why techniques for trading recomputation against memory")
print("  exist at all. The chain rule is free; remembering where you have")
print("  been is not.")

assert graph_size(node) > D.CHAIN_LENGTH

print()
print("=" * 74)
print("6. A measurement that corrects sections 1 and 2")
print("=" * 74)
print()
print("  Everything above multiplied a CONSTANT factor. Real layers do not")
print("  work that way: a local rate depends on where it is evaluated, and")
print("  the forward pass moves that point. So stack tanh on tanh on tanh and")
print("  measure what actually happens.")
print()


def stacked_tanh(depth):
    def deep(vals):
        node = vals[0]
        for _ in range(depth):
            node = node.tanh()
        return node

    return deep


print("     depth    gradient at x = 0.9     ratio to the row above")
print("     " + "-" * 58)
previous = None
measured = {}
for depth in (1, 5, 10, 20, 40, 80, 160):
    value = reverse_mode_gradient(stacked_tanh(depth), [0.9])[0][0]
    measured[depth] = value
    ratio = "-" if previous is None else f"{previous / value:.3f}"
    print(f"     {depth:<8d} {value:<23.6e} {ratio}")
    previous = value
print()
single = measured[1]
naive = single**40
print("  Now the naive prediction. tanh's slope at 0.9 is about")
print(f"  {single:.6f}, so forty tanh layers 'should' multiply the gradient")
print(f"  by that forty times over:")
print()
print(f"      {single:.6f} ** 40 = {naive:.6e}     the prediction")
print(f"      measured at depth 40 = {measured[40]:.6e}     the measurement")
print(f"      the measurement is larger by a factor of {measured[40] / naive:.3e}")
print()
print("  Ten orders of magnitude. The prediction is not slightly off, it is")
print("  wrong in kind, and the reason is worth more than the number: each")
print("  tanh pulls its input closer to 0, and tanh's slope AT 0 is 1. The")
print("  deeper the stack goes, the closer every local rate creeps back")
print("  towards 1, so the product decays like a power of the depth rather")
print("  than exponentially.")
print()
print("  A product of constants is the wrong model for a product of rates")
print("  that depend on where they are evaluated. Sections 1 and 2 are still")
print("  the right picture of what a chain of fixed factors does -- and a")
print("  weight matrix that is too small or too large really does behave that")
print("  way -- but 'tanh saturates, therefore gradients vanish' is a claim")
print("  that has to be measured on the network in front of you rather than")
print("  assumed from the shape of the curve.")

assert measured[40] > 1e9 * naive
assert naive < 1e-12
assert list(measured.values()) == sorted(measured.values(), reverse=True)

print()
print("07_vanishing_and_exploding.py: every assertion held.")

"""A reverse-mode automatic differentiation engine, in about seventy lines.

This is the core of what every deep-learning framework does. It is not a
simplified illustration of the idea -- it *is* the idea, with the engineering
removed: no tensors, no GPU kernels, no fused operations, no memory planning.
A `Value` holds one number and one gradient, remembers which values it came
from, and knows how to hand its own gradient back to them. `backward()` walks
the graph once in reverse, applying the chain rule at every node.

Two things are worth watching as you read.

**The gradient is accumulated with `+=`, never assigned.** That single choice
is the multivariable chain rule. A value used in two places receives a
contribution from each use, and both are real, so they add. Change either
`+=` below to `=` and the engine will still run, still look sensible, and be
quietly wrong on any graph where something is used twice.

**One backward pass produces every gradient.** Not one pass per parameter --
one pass, total. The forward-mode engine at the bottom of this file does the
opposite, and the two are compared head to head in
`07_vanishing_and_exploding.py`. That asymmetry is the reason training a model
with a hundred million parameters and one loss is affordable at all.
"""

import math
from typing import Callable, Iterable, Sequence


class Value:
    """One number in a computation graph, with a gradient and a history."""

    __slots__ = ("data", "grad", "label", "_backward", "_children", "_op")

    def __init__(
        self,
        data: float,
        children: tuple["Value", ...] = (),
        op: str = "",
        label: str = "",
    ) -> None:
        self.data: float = float(data)
        #: d(final output) / d(this value). Zero until a backward pass fills
        #: it in, and accumulated rather than overwritten.
        self.grad: float = 0.0
        self.label: str = label
        #: Hands this node's gradient back to its children. The identity
        #: function for a leaf, which has no children to hand anything to.
        self._backward: Callable[[], None] = _do_nothing
        self._children: tuple["Value", ...] = children
        self._op: str = op

    # -- the two arithmetic operations -------------------------------------

    def __add__(self, other: "Value | float") -> "Value":
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")

        def backward() -> None:
            # Addition passes the gradient through untouched: nudge either
            # input by d and the sum moves by d, so the local rate is 1.
            self.grad += out.grad
            other.grad += out.grad

        out._backward = backward
        return out

    def __mul__(self, other: "Value | float") -> "Value":
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), "*")

        def backward() -> None:
            # For a product, each input's local rate is the OTHER input's
            # value. Nudge self by d and the product moves by d * other.
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad

        out._backward = backward
        return out

    # -- the one non-linearity ---------------------------------------------

    def tanh(self) -> "Value":
        """The hyperbolic tangent, and the only non-linear operation here.

        One non-linearity is enough to make the network in this lab a genuine
        network rather than a stack of matrix multiplications that collapses
        into one. Its derivative is 1 - tanh squared, which is convenient
        because the forward pass has already computed tanh.
        """
        t = math.tanh(self.data)
        out = Value(t, (self,), "tanh")

        def backward() -> None:
            self.grad += (1.0 - t * t) * out.grad

        out._backward = backward
        return out

    # -- conveniences built from the two operations above ------------------

    def __neg__(self) -> "Value":
        return self * -1.0

    def __sub__(self, other: "Value | float") -> "Value":
        return self + (-(other if isinstance(other, Value) else Value(other)))

    def __radd__(self, other: "Value | float") -> "Value":
        return self + other

    def __rmul__(self, other: "Value | float") -> "Value":
        return self * other

    def __rsub__(self, other: "Value | float") -> "Value":
        return (-self) + other

    def __repr__(self) -> str:
        name = f"{self.label}=" if self.label else ""
        return f"Value({name}{self.data:.6g}, grad={self.grad:.6g})"

    # -- the backward pass --------------------------------------------------

    def backward(self) -> None:
        """Fill in `.grad` on every value this one was computed from.

        Three steps, and none of them is subtle:

        1. Order the graph so that every node comes after everything it was
           computed from. That is a topological sort, and it matters because a
           node must not hand its gradient onwards until it has received every
           contribution owed to it.
        2. Seed this node's own gradient with 1.0. The derivative of the
           output with respect to itself is 1 -- that is the base case the
           whole chain rule hangs from.
        3. Walk the order backwards, letting each node push its gradient to
           its children by multiplying by the local derivative.
        """
        order = topological_order(self)
        for node in order:
            node.grad = 0.0
        self.grad = 1.0
        for node in reversed(order):
            node._backward()


def _do_nothing() -> None:
    """The backward step of a leaf: it has nobody to pass anything to."""
    return None


def topological_order(root: Value) -> list[Value]:
    """Every value `root` depends on, parents always after their children.

    Iterative rather than recursive, so a chain of ten thousand operations
    does not exhaust the interpreter's stack -- which a deep network's graph
    genuinely would.
    """
    order: list[Value] = []
    visited: set[int] = set()
    # Each stack entry is (node, children_already_expanded).
    stack: list[tuple[Value, bool]] = [(root, False)]
    while stack:
        node, expanded = stack.pop()
        if expanded:
            order.append(node)
            continue
        if id(node) in visited:
            continue
        visited.add(id(node))
        stack.append((node, True))
        for child in node._children:
            if id(child) not in visited:
                stack.append((child, False))
    return order


def graph_size(root: Value) -> int:
    """How many nodes a backward pass will visit."""
    return len(topological_order(root))


# --------------------------------------------------------------------------
# Forward mode, for comparison: the same chain rule, run the other way
# --------------------------------------------------------------------------


class Dual:
    """A number carried alongside its derivative with respect to ONE input.

    Forward mode is the chain rule applied left to right. Every operation
    computes both the value and the rate at which that value moves when the
    chosen input moves. It is simpler than reverse mode -- there is no graph
    and no second pass -- and that simplicity costs it the thing that matters:
    it answers about one input at a time, so a function of n inputs needs n
    separate runs.
    """

    __slots__ = ("value", "dot")

    def __init__(self, value: float, dot: float = 0.0) -> None:
        self.value: float = float(value)
        #: The derivative of this quantity with respect to the seeded input.
        self.dot: float = float(dot)

    def __add__(self, other: "Dual | float") -> "Dual":
        other = other if isinstance(other, Dual) else Dual(other)
        return Dual(self.value + other.value, self.dot + other.dot)

    def __mul__(self, other: "Dual | float") -> "Dual":
        other = other if isinstance(other, Dual) else Dual(other)
        # The product rule, which is the chain rule's travelling companion.
        return Dual(
            self.value * other.value,
            self.dot * other.value + self.value * other.dot,
        )

    def tanh(self) -> "Dual":
        t = math.tanh(self.value)
        return Dual(t, (1.0 - t * t) * self.dot)

    def __neg__(self) -> "Dual":
        return self * -1.0

    def __sub__(self, other: "Dual | float") -> "Dual":
        return self + (-(other if isinstance(other, Dual) else Dual(other)))

    def __radd__(self, other: "Dual | float") -> "Dual":
        return self + other

    def __rmul__(self, other: "Dual | float") -> "Dual":
        return self * other

    def __rsub__(self, other: "Dual | float") -> "Dual":
        return (-self) + other

    def __repr__(self) -> str:
        return f"Dual({self.value:.6g}, dot={self.dot:.6g})"


# --------------------------------------------------------------------------
# The two modes, measured against each other
# --------------------------------------------------------------------------


def reverse_mode_gradient(
    build: Callable[[Sequence[Value]], Value], xs: Sequence[float]
) -> tuple[list[float], int]:
    """Every partial derivative of `build`, and the number of passes used.

    `build` receives one `Value` per input and returns the single output.
    The count returned is the number of forward-and-backward sweeps needed to
    obtain ALL the gradients, and it is 1 no matter how many inputs there are.
    """
    inputs = [Value(x, label=f"x{i}") for i, x in enumerate(xs)]
    out = build(inputs)
    out.backward()
    return [node.grad for node in inputs], 1


def forward_mode_gradient(
    build: Callable[[Sequence[Dual]], Dual], xs: Sequence[float]
) -> tuple[list[float], int]:
    """The same gradients by forward mode, and the number of passes used.

    One pass per input, because each pass can only carry the derivative with
    respect to whichever input was seeded with a 1. The count returned is
    therefore `len(xs)`, and that is the entire argument for reverse mode.
    """
    grads: list[float] = []
    passes = 0
    for seed in range(len(xs)):
        inputs = [
            Dual(x, 1.0 if i == seed else 0.0) for i, x in enumerate(xs)
        ]
        grads.append(build(inputs).dot)
        passes += 1
    return grads, passes


def numeric_gradient(
    f: Callable[[Sequence[float]], float], xs: Sequence[float], h: float
) -> tuple[list[float], int]:
    """The same gradients by central differences, and the passes used.

    Two evaluations per input, so 2n passes -- worse than forward mode and far
    worse than reverse mode, and approximate into the bargain. It is the
    checking tool, not the production tool, and the counts here are why.
    """
    grads: list[float] = []
    passes = 0
    for i in range(len(xs)):
        ahead = list(xs)
        behind = list(xs)
        ahead[i] += h
        behind[i] -= h
        grads.append((f(ahead) - f(behind)) / (2.0 * h))
        passes += 2
    return grads, passes


def parameters_of(root: Value) -> list[Value]:
    """Every leaf in the graph -- the nodes with no children of their own."""
    return [node for node in topological_order(root) if not node._children]


def sum_values(values: Iterable[Value]) -> Value:
    """Add up an iterable of Values, starting from a fresh zero."""
    total = Value(0.0)
    for value in values:
        total = total + value
    return total

"""The chain rule, written out in code: local rates, products, and paths.

Nothing in this module knows anything about neural networks. It knows that
when one quantity depends on another which depends on another, the rates
multiply -- and that when a quantity reaches the output by more than one
route, the routes are added.

Every function here is checked against a central difference in
`test_reference.py`, which is the whole point of the day: the chain rule is
not a rule to be believed, it is a rule that can be measured.
"""

import math
from typing import Callable, Iterable, Sequence

Scalar = Callable[[float], float]


# --------------------------------------------------------------------------
# The measuring instrument, carried over from Day 108
# --------------------------------------------------------------------------


def central_difference(f: Scalar, x: float, h: float) -> float:
    """Estimate f'(x) by straddling x: (f(x+h) - f(x-h)) / (2h).

    This is the checking tool for the whole lab. It knows nothing about the
    chain rule -- it just moves x a little and watches the output move -- and
    that independence is what makes it a valid check.
    """
    if h <= 0.0:
        raise ValueError("h must be positive")
    return (f(x + h) - f(x - h)) / (2.0 * h)


def partial_difference(
    f: Callable[..., float], point: Sequence[float], index: int, h: float
) -> float:
    """Estimate one partial derivative of a multi-input function.

    Nudge coordinate `index` and hold every other coordinate still. This is
    Day 109's tool, and it is what the multivariable chain rule is checked
    against here.
    """
    if h <= 0.0:
        raise ValueError("h must be positive")
    ahead = list(point)
    behind = list(point)
    ahead[index] += h
    behind[index] -= h
    return (f(*ahead) - f(*behind)) / (2.0 * h)


# --------------------------------------------------------------------------
# Rates multiply: the idea with no calculus in it at all
# --------------------------------------------------------------------------


def product(factors: Iterable[float]) -> float:
    """Multiply an iterable of numbers together.

    An empty product is 1.0, which is the identity for multiplication and is
    also the right answer to "how much does x change per unit of x".
    """
    total = 1.0
    for factor in factors:
        total *= factor
    return total


def gear_ratio(ratios: Iterable[float]) -> float:
    """The overall ratio of a gear train: every stage ratio multiplied.

    If the first gear turns twice per turn of the second, and the second turns
    three times per turn of the third, the first turns six times per turn of
    the third. That sentence is the chain rule with the notation removed.
    """
    return product(ratios)


# --------------------------------------------------------------------------
# Composition, and the one-variable chain rule
# --------------------------------------------------------------------------


def compose(outer: Scalar, inner: Scalar) -> Scalar:
    """Return the function x -> outer(inner(x)).

    Read the parentheses from the inside out: `inner` runs first on x, and
    `outer` runs on whatever came back.
    """

    def composed(x: float) -> float:
        return outer(inner(x))

    return composed


def chain_rule(
    d_outer: Scalar, inner: Scalar, d_inner: Scalar, x: float
) -> float:
    """The chain rule for one variable, in one line.

    dy/dx = dy/du * du/dx, where u = inner(x). The outer derivative is
    evaluated **at the inner value**, not at x -- which is the single most
    common way to get this wrong, and the reason the argument is written out
    here instead of being tucked into a lambda.
    """
    u = inner(x)
    return d_outer(u) * d_inner(x)


# --------------------------------------------------------------------------
# Chains of any length
# --------------------------------------------------------------------------


def chain_values(stages: Sequence[Scalar], x: float) -> list[float]:
    """The forward pass: every intermediate value, starting with x itself.

    For n stages this returns n + 1 numbers: the input, then the output of
    each stage in turn. Keeping the input in the list means value[i] is
    always the input to stage i, which makes the backward pass easy to read.
    """
    values = [x]
    current = x
    for stage in stages:
        current = stage(current)
        values.append(current)
    return values


def chain_local_rates(
    stages: Sequence[Scalar], rates: Sequence[Scalar], x: float
) -> list[float]:
    """The local derivative of every stage, each evaluated at its own input.

    This is the step people skip. Stage i's derivative is evaluated at the
    value that *arrives* at stage i, which is the output of stage i-1 -- not
    at x, and not at the final answer.
    """
    if len(stages) != len(rates):
        raise ValueError("every stage needs exactly one derivative")
    values = chain_values(stages, x)
    return [rate(values[i]) for i, rate in enumerate(rates)]


def chain_derivative(
    stages: Sequence[Scalar], rates: Sequence[Scalar], x: float
) -> float:
    """The derivative of the whole chain: every local rate multiplied.

    This is the chain rule for a composition of any depth. Two functions or
    two hundred, the shape does not change -- which is exactly why a network
    with a hundred layers is trainable at all.
    """
    return product(chain_local_rates(stages, rates, x))


def chain_function(stages: Sequence[Scalar]) -> Scalar:
    """Collapse a list of stages into the single function they compose."""

    def composed(x: float) -> float:
        current = x
        for stage in stages:
            current = stage(current)
        return current

    return composed


def running_products(rates: Sequence[float]) -> list[float]:
    """The partial products of the local rates, taken from the output end.

    This is what a backward pass actually computes: after visiting k stages
    from the end, the number it is carrying is the product of the last k
    local rates. The list is returned in stage order, so entry i is the
    gradient of the output with respect to the value arriving at stage i.
    """
    out: list[float] = []
    carried = 1.0
    for rate in reversed(rates):
        carried *= rate
        out.append(carried)
    out.reverse()
    return out


# --------------------------------------------------------------------------
# More than one path: the part where contributions ADD
# --------------------------------------------------------------------------


def path_contributions(local_rates_per_path: Sequence[Sequence[float]]) -> list[float]:
    """One number per path: the product of the local rates along that path."""
    return [product(path) for path in local_rates_per_path]


def total_derivative(local_rates_per_path: Sequence[Sequence[float]]) -> float:
    """Multiply along each path, then ADD across paths.

    The addition is the half that gets dropped. If a variable influences the
    output through two routes, changing it moves the output twice, and both
    movements happen. There is no rule of nature that makes one of them the
    real one.
    """
    return sum(path_contributions(local_rates_per_path))


def wrong_single_path_derivative(
    local_rates_per_path: Sequence[Sequence[float]], path_index: int = 0
) -> float:
    """The mistake, implemented deliberately so a test can catch it.

    This takes one path's product and stops. It is here to be compared
    against `total_derivative` and against a central difference, so that the
    failure is a measurement rather than a warning in a comment.
    """
    return product(local_rates_per_path[path_index])


# --------------------------------------------------------------------------
# Products that collapse and products that blow up
# --------------------------------------------------------------------------


def repeated_product(factor: float, count: int) -> float:
    """Multiply `factor` by itself `count` times, one multiplication at a time.

    Written as a loop rather than as `factor ** count` because the loop is
    what a backward pass through `count` layers actually does, and because the
    intermediate values are the interesting part.
    """
    if count < 0:
        raise ValueError("count must not be negative")
    total = 1.0
    for _ in range(count):
        total *= factor
    return total


def product_trace(factor: float, count: int) -> list[float]:
    """Every running value of `repeated_product`, for plotting or printing."""
    trace: list[float] = []
    total = 1.0
    for _ in range(count):
        total *= factor
        trace.append(total)
    return trace


def order_of_magnitude(value: float) -> int:
    """floor(log10(|value|)) -- the exponent, ignoring the digits.

    Vanishing and exploding gradients are a statement about scale, so this is
    what the tests assert. Claiming an exact value for 0.9 to the fiftieth
    power would be asserting float64 rounding, which is not the lesson.
    """
    if value == 0.0:
        raise ValueError("zero has no order of magnitude")
    return math.floor(math.log10(abs(value)))

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

Both `examples/` and `starter/` contain modules called `autodiff`,
`chainrule`, `dataset` and `network`, 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 `autodiff` 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 ("autodiff", "chainrule", "dataset", "network", "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, the functions and every tolerance this lab compares against.

Read this file. Nothing here is tuned: every tolerance below is derived from
the error terms that actually govern the comparison being made, and the
arithmetic is written out beside it. A tolerance reached by running a test and
enlarging the number until it went green is a tolerance chosen by whatever bug
happened to exist at the time.

Almost every number in this lab is exact in float64. The chains, the two-path
example and the whole two-layer network were chosen so that the reader can
re-derive each one with a pen. Where a value is not exact -- a sine, a
logarithm, an exponential -- it is computed here from `math` rather than
written down as a literal, so nothing in this lab is a remembered constant.
"""

import math
from typing import Callable, NamedTuple

import numpy as np

# --------------------------------------------------------------------------
# Machine constants
# --------------------------------------------------------------------------

#: float64 machine epsilon, read from NumPy rather than trusted as a literal.
EPSILON: float = float(np.finfo(np.float64).eps)

# --------------------------------------------------------------------------
# The numerical step, and the three tolerances
# --------------------------------------------------------------------------

#: The central-difference step. Day 108 measured the bottom of the error U for
#: the central rule on e**x and found it in the 1e-7 to 1e-4 band; 1e-5 sits
#: inside it with room on both sides.
H: float = 1e-5

# The central difference (f(x+h) - f(x-h)) / (2h) carries two errors:
#
#   truncation  ~ (h**2 / 6) * |f'''(x)|   = 1.667e-11 * |f'''(x)|  at h = 1e-5
#   rounding    ~ EPSILON * |f(x)| / h     = 2.220e-11 * |f(x)|     at h = 1e-5
#
# No function differentiated in this lab has |f| or |f'''| above about 250 at
# the points used, so the bound is about 250 * (1.667e-11 + 2.220e-11), which
# is roughly 9.7e-9. The tolerance below is 1e-6, so there is about a
# hundredfold margin -- enough that ordinary float64 noise cannot trip it, and
# far too tight to hide a chain rule that dropped a factor or summed the wrong
# paths. The smallest true gradient this tolerance guards is 0.25, so a
# mistake would have to be smaller than four parts in a million to slip past.
#: Analytic gradient against a central difference. See the derivation above.
NUMERIC_TOL: float = 1e-6

#: The same comparison, stated relatively, for the handful of quantities whose
#: magnitude runs into the hundreds -- where an absolute 1e-6 would be
#: stricter than the arithmetic can honestly support.
NUMERIC_REL_TOL: float = 1e-6

# Two analytic computations of the same quantity -- for example the product of
# five local rates against the closed-form derivative of the whole chain --
# differ only in the order the multiplications happen, so they differ by a few
# units in the last place. At a magnitude of 1e2 one ulp is about 1.4e-14, and
# a chain of five products can accumulate a handful of them, so the honest
# bound is a few times 1e-13.
#: Analytic against analytic: rounding only, no truncation.
ANALYTIC_TOL: float = 1e-12

# --------------------------------------------------------------------------
# Gears: the whole idea, before any calculus
# --------------------------------------------------------------------------

#: Gear A turns twice for every turn of B; B turns three times for every turn
#: of C. So A turns six times per turn of C. The rates multiplied.
GEAR_RATIOS: tuple[float, ...] = (2.0, 3.0)
GEAR_RATIO_PRODUCT: float = 6.0

#: A longer train, to show the product keeping on multiplying.
GEAR_TRAIN: tuple[float, ...] = (2.0, 3.0, 1.5, 4.0)
GEAR_TRAIN_PRODUCT: float = 36.0

#: The same arithmetic with money instead of teeth. 1 unit of the first
#: currency buys 1.25 of the second, which buys 0.8 of the third, which buys
#: 150 of the fourth. These rates are invented for the arithmetic and are not
#: quoted from any market.
CURRENCY_RATES: tuple[float, ...] = (1.25, 0.8, 150.0)
CURRENCY_PRODUCT: float = 150.0

# --------------------------------------------------------------------------
# Functions used as the outer and inner halves of a composition
# --------------------------------------------------------------------------


def square(x: float) -> float:
    """x squared."""
    return x * x


def d_square(x: float) -> float:
    """The derivative of x squared."""
    return 2.0 * x


def line(x: float) -> float:
    """3x + 1."""
    return 3.0 * x + 1.0


def d_line(x: float) -> float:
    """The derivative of 3x + 1: a constant 3."""
    return 3.0


def half_negative_square(x: float) -> float:
    """-x squared over 2 -- the inside of a Gaussian bump."""
    return -0.5 * x * x


def d_half_negative_square(x: float) -> float:
    """The derivative of -x squared over 2."""
    return -x


def shifted_square(x: float) -> float:
    """x squared plus 1, which is never zero, so its logarithm is safe."""
    return x * x + 1.0


def reciprocal(x: float) -> float:
    """1 / x."""
    return 1.0 / x


def d_reciprocal(x: float) -> float:
    """The derivative of 1 / x."""
    return -1.0 / (x * x)


def one_plus_exp_negative(x: float) -> float:
    """1 + e to the minus x -- the denominator of the sigmoid."""
    return 1.0 + math.exp(-x)


def d_one_plus_exp_negative(x: float) -> float:
    """The derivative of 1 + e to the minus x."""
    return -math.exp(-x)


def double_plus_one(x: float) -> float:
    """2x + 1."""
    return 2.0 * x + 1.0


def d_double_plus_one(x: float) -> float:
    """The derivative of 2x + 1: a constant 2."""
    return 2.0


def d_tanh(x: float) -> float:
    """The derivative of tanh, written in terms of x rather than of tanh(x)."""
    t = math.tanh(x)
    return 1.0 - t * t


def d_ln(x: float) -> float:
    """The derivative of the natural logarithm."""
    return 1.0 / x


# --------------------------------------------------------------------------
# The one-variable chain rule: six compositions, each checked numerically
# --------------------------------------------------------------------------


class Composition(NamedTuple):
    """One composed function f(g(x)), with both halves and both derivatives.

    `exact` is the closed-form derivative at `x`, computed from `math` rather
    than written down as a literal, so nothing here is a remembered constant.
    """

    name: str
    outer: Callable[[float], float]
    d_outer: Callable[[float], float]
    inner: Callable[[float], float]
    d_inner: Callable[[float], float]
    x: float
    exact: float


COMPOSITIONS: tuple[Composition, ...] = (
    # (3x + 1) squared at x = 2. Inner is 7, outer rate is 2*7 = 14, inner
    # rate is 3, so the answer is 42 -- exact, and checkable in one line.
    Composition("square of a line", square, d_square, line, d_line, 2.0, 42.0),
    # sin(x squared) at x = 1.5. Rate is cos(2.25) * 3.
    Composition(
        "sine of a square",
        math.sin,
        math.cos,
        square,
        d_square,
        1.5,
        math.cos(2.25) * 3.0,
    ),
    # A Gaussian bump, e to the minus x squared over 2, at x = 0.8.
    Composition(
        "gaussian bump",
        math.exp,
        math.exp,
        half_negative_square,
        d_half_negative_square,
        0.8,
        math.exp(-0.32) * -0.8,
    ),
    # ln(x squared + 1) at x = 2. Inner is 5, outer rate is 1/5, inner rate
    # is 4, so the answer is 0.8 -- exact.
    Composition(
        "log of a shifted square",
        math.log,
        d_ln,
        shifted_square,
        d_square,
        2.0,
        0.8,
    ),
    # The sigmoid, written as 1 / (1 + e to the minus x), at x = 0. Inner is
    # 2, outer rate is -1/4, inner rate is -1, so the answer is 0.25 -- which
    # is the largest value the sigmoid's slope ever takes.
    Composition(
        "the sigmoid",
        reciprocal,
        d_reciprocal,
        one_plus_exp_negative,
        d_one_plus_exp_negative,
        0.0,
        0.25,
    ),
    # tanh(2x + 1) at x = -0.5, where the inner function is exactly 0 and
    # tanh's slope is exactly 1, so the answer is exactly 2.
    Composition(
        "tanh of a line",
        math.tanh,
        d_tanh,
        double_plus_one,
        d_double_plus_one,
        -0.5,
        2.0,
    ),
)

# --------------------------------------------------------------------------
# A chain of five functions, built so every local rate is a round number
# --------------------------------------------------------------------------

#: Applied in order, left to right: double, add 3, square, square root,
#: natural logarithm. Starting from x = 1 the values are 1, 2, 5, 25, 5,
#: ln 5, and the local rates are 2, 1, 10, 0.1, 0.2. Their product is 0.4.
#:
#: The whole chain collapses to ln(2x + 3), whose derivative is 2 / (2x + 3),
#: which at x = 1 is 2/5 = 0.4. Two routes, the same number.
FIVE_STAGES: tuple[Callable[[float], float], ...] = (
    lambda u: 2.0 * u,
    lambda u: u + 3.0,
    lambda u: u * u,
    math.sqrt,
    math.log,
)

FIVE_RATES: tuple[Callable[[float], float], ...] = (
    lambda u: 2.0,
    lambda u: 1.0,
    lambda u: 2.0 * u,
    lambda u: 0.5 / math.sqrt(u),
    lambda u: 1.0 / u,
)

FIVE_START: float = 1.0
FIVE_VALUES: tuple[float, ...] = (1.0, 2.0, 5.0, 25.0, 5.0, math.log(5.0))
FIVE_LOCAL_RATES: tuple[float, ...] = (2.0, 1.0, 10.0, 0.1, 0.2)
FIVE_DERIVATIVE: float = 0.4


def five_chain_closed_form(x: float) -> float:
    """The five stages collapsed by hand: ln(2x + 3)."""
    return math.log(2.0 * x + 3.0)


def d_five_chain_closed_form(x: float) -> float:
    """Its derivative: 2 / (2x + 3)."""
    return 2.0 / (2.0 * x + 3.0)


# --------------------------------------------------------------------------
# Two paths into one output: the case where contributions ADD
# --------------------------------------------------------------------------

#: x reaches the output twice: once through u = x squared, once through
#: v = 3x. The output is f = u * v.
#:
#:   df/dx = (df/du)(du/dx) + (df/dv)(dv/dx)
#:         = v * 2x         + u * 3
#:         = 3x * 2x        + x squared * 3
#:         = 6 x squared    + 3 x squared     = 9 x squared
#:
#: At x = 2 that is 36. The two path contributions are 24 and 12. Neither one
#: alone is the answer, and this is exactly the mistake the lab is built to
#: catch: taking the product along one path and stopping there.
TWO_PATH_X: float = 2.0
TWO_PATH_U: float = 4.0
TWO_PATH_V: float = 6.0
TWO_PATH_OUTPUT: float = 24.0
TWO_PATH_CONTRIBUTIONS: tuple[float, float] = (24.0, 12.0)
TWO_PATH_DERIVATIVE: float = 36.0


def two_path_direct(x: float) -> float:
    """The same function with the composition already done: 3 x cubed."""
    return 3.0 * x * x * x


def d_two_path_direct(x: float) -> float:
    """Its derivative: 9 x squared."""
    return 9.0 * x * x


# --------------------------------------------------------------------------
# The full multivariable chain rule: two inputs, two intermediates
# --------------------------------------------------------------------------

#: z = u squared + v squared, with u = s*t and v = s - t, at (s, t) = (2, 3).
#:
#:   u = 6,  v = -1,  z = 37
#:   dz/du = 12,  dz/dv = -2
#:   du/ds = t = 3,   du/dt = s = 2
#:   dv/ds = 1,       dv/dt = -1
#:
#:   dz/ds = 12*3 + (-2)*1  = 34
#:   dz/dt = 12*2 + (-2)*-1 = 26
#:
#: Every one of those is an integer, so the whole thing is exact in float64.
SURFACE_POINT: tuple[float, float] = (2.0, 3.0)
SURFACE_U: float = 6.0
SURFACE_V: float = -1.0
SURFACE_Z: float = 37.0
SURFACE_DZ_DU: float = 12.0
SURFACE_DZ_DV: float = -2.0
SURFACE_GRADIENT: tuple[float, float] = (34.0, 26.0)


def surface(s: float, t: float) -> float:
    """z as a function of s and t, with the intermediates substituted in."""
    u = s * t
    v = s - t
    return u * u + v * v


# --------------------------------------------------------------------------
# The tiny two-layer network, backpropagated by hand
# --------------------------------------------------------------------------

#: Half the natural logarithm of 3. tanh of this number is exactly 0.5 in
#: float64, and 1 - tanh squared is then exactly 0.75. Both facts are asserted
#: by the test suite rather than assumed.
#:
#: The bias of the second hidden unit is set to this value on purpose, so that
#: every number in the hand-worked backward pass is exact and can be checked
#: with a pen. Nothing about the chain rule depends on the choice; it is a
#: convenience for the reader, and it is declared rather than hidden.
HALF_LN3: float = 0.5 * math.log(3.0)

#: The two inputs.
NET_X1: float = 1.0
NET_X2: float = 2.0

#: Hidden unit A: weights and bias. Pre-activation is 1*1 + (-0.5)*2 + 0 = 0.
NET_WA1: float = 1.0
NET_WA2: float = -0.5
NET_BA: float = 0.0

#: Hidden unit B: pre-activation is -0.5*1 + 0.25*2 + HALF_LN3 = HALF_LN3.
NET_WB1: float = -0.5
NET_WB2: float = 0.25
NET_BB: float = HALF_LN3

#: The output layer: a linear combination of the two hidden activations.
NET_VA: float = 2.0
NET_VB: float = -3.0
NET_C: float = 1.0

#: The target the loss is measured against.
NET_TARGET: float = 1.0

#: The forward pass, every value exact.
NET_A_PRE: float = 0.0
NET_A: float = 0.0
NET_B_PRE: float = HALF_LN3
NET_B: float = 0.5
NET_OUT: float = -0.5
NET_LOSS: float = 2.25

#: The backward pass, every value exact. Worked out in full in
#: `06_backprop_by_hand.py` and asserted against the engine in the tests.
NET_GRADIENTS: dict[str, float] = {
    "out": -3.0,
    "c": -3.0,
    "vA": 0.0,  # zero, because hidden unit A output exactly 0
    "vB": -1.5,
    "a": -6.0,
    "b": 9.0,
    "a_pre": -6.0,
    "b_pre": 6.75,
    "wA1": -6.0,
    "wA2": -12.0,
    "bA": -6.0,
    "wB1": 6.75,
    "wB2": 13.5,
    "bB": 6.75,
    # x1 and x2 each reach the loss through BOTH hidden units, so each of
    # these is a sum over two paths, not a single product.
    "x1": -9.375,
    "x2": 4.6875,
}

#: The two per-path contributions to dL/dx1, which must be added.
NET_X1_CONTRIBUTIONS: tuple[float, float] = (-6.0, -3.375)

#: The names of the twelve parameters, in the order the scripts print them.
NET_PARAMETERS: tuple[str, ...] = (
    "wA1",
    "wA2",
    "bA",
    "wB1",
    "wB2",
    "bB",
    "vA",
    "vB",
    "c",
)

# --------------------------------------------------------------------------
# Products that collapse and products that blow up
# --------------------------------------------------------------------------

#: A gradient that passes through 50 layers, each contributing a local rate
#: slightly below or slightly above 1.
DECAY_FACTOR: float = 0.9
GROWTH_FACTOR: float = 1.1
CHAIN_LENGTH: int = 50
LONG_CHAIN_LENGTH: int = 200

#: Asserted as orders of magnitude rather than as exact values, because the
#: point is the scale and not the digits.
DECAY_ORDER: int = -3  # 0.9**50 is about 5.15e-3
GROWTH_ORDER: int = 2  # 1.1**50 is about 1.17e+2

#: A harsher pair, to show how quickly the arithmetic leaves the useful range.
#: 0.25 is not an arbitrary choice: it is the LARGEST slope the sigmoid ever
#: has, measured in `02_composition_and_the_chain_rule.py`. A stack of sigmoid
#: layers is multiplying numbers no bigger than this one.
SHARP_DECAY: float = 0.25
SHARP_GROWTH: float = 2.0

#: A middling decay, kept because its behaviour is genuinely surprising and
#: contradicts the obvious guess. See `07_vanishing_and_exploding.py`.
MILD_DECAY: float = 0.5

"""The tiny two-layer network the lab backpropagates by hand.

Two inputs, two hidden units with a tanh non-linearity, one linear output, and
a squared-error loss. Nine parameters. That is small enough to differentiate
on paper and large enough to contain every structural feature of a real
network -- including the one that matters most today: each input reaches the
loss through **both** hidden units, so its gradient is a sum over two paths
rather than a single product.

The numbers were chosen so that every quantity in both passes is exact in
float64. Hidden unit A sits at a pre-activation of exactly 0, where tanh is 0
and its slope is 1. Hidden unit B sits at exactly half the natural logarithm
of 3, where tanh is exactly 0.5 and its slope is exactly 0.75. Those two
facts are asserted by the test suite rather than assumed, and the choice is a
convenience for the reader rather than anything the chain rule depends on.
"""

from typing import Sequence

import dataset as D
from autodiff import Dual, Value


def forward(
    x1: float, x2: float, params: Sequence[float]
) -> dict[str, float]:
    """Run the network in plain floats and return every intermediate value.

    `params` is in the order given by `dataset.NET_PARAMETERS`:
    wA1, wA2, bA, wB1, wB2, bB, vA, vB, c.
    """
    import math

    wA1, wA2, bA, wB1, wB2, bB, vA, vB, c = params
    a_pre = wA1 * x1 + wA2 * x2 + bA
    a = math.tanh(a_pre)
    b_pre = wB1 * x1 + wB2 * x2 + bB
    b = math.tanh(b_pre)
    out = vA * a + vB * b + c
    loss = (out - D.NET_TARGET) * (out - D.NET_TARGET)
    return {
        "a_pre": a_pre,
        "a": a,
        "b_pre": b_pre,
        "b": b,
        "out": out,
        "loss": loss,
    }


def loss_only(x1: float, x2: float, params: Sequence[float]) -> float:
    """Just the loss, for feeding to a central difference."""
    return forward(x1, x2, params)["loss"]


def build_graph(
    x1: Value, x2: Value, params: Sequence[Value]
) -> dict[str, Value]:
    """Build the same network out of `Value` nodes and return them all.

    The returned dictionary is keyed the same way as `dataset.NET_GRADIENTS`,
    so a test can walk the two side by side.
    """
    wA1, wA2, bA, wB1, wB2, bB, vA, vB, c = params
    a_pre = wA1 * x1 + wA2 * x2 + bA
    a = a_pre.tanh()
    b_pre = wB1 * x1 + wB2 * x2 + bB
    b = b_pre.tanh()
    out = vA * a + vB * b + c
    diff = out - D.NET_TARGET
    loss = diff * diff
    return {
        "x1": x1,
        "x2": x2,
        "wA1": wA1,
        "wA2": wA2,
        "bA": bA,
        "wB1": wB1,
        "wB2": wB2,
        "bB": bB,
        "vA": vA,
        "vB": vB,
        "c": c,
        "a_pre": a_pre,
        "a": a,
        "b_pre": b_pre,
        "b": b,
        "out": out,
        "loss": loss,
    }


def default_parameter_values() -> list[float]:
    """The nine parameters from `dataset.py`, in the documented order."""
    return [
        D.NET_WA1,
        D.NET_WA2,
        D.NET_BA,
        D.NET_WB1,
        D.NET_WB2,
        D.NET_BB,
        D.NET_VA,
        D.NET_VB,
        D.NET_C,
    ]


def engine_gradients() -> dict[str, float]:
    """Every gradient in the network, from one backward pass of the engine."""
    x1 = Value(D.NET_X1, label="x1")
    x2 = Value(D.NET_X2, label="x2")
    params = [
        Value(v, label=name)
        for name, v in zip(D.NET_PARAMETERS, default_parameter_values())
    ]
    nodes = build_graph(x1, x2, params)
    nodes["loss"].backward()
    return {name: node.grad for name, node in nodes.items()}


def hand_gradients() -> dict[str, float]:
    """The same gradients, written out as the arithmetic of the backward pass.

    This function deliberately repeats by hand what `engine_gradients` gets
    from the graph walk. Every line is one application of the chain rule, and
    the two lines marked SUM are where two paths meet and are added.
    """
    x1, x2 = D.NET_X1, D.NET_X2
    a, b = D.NET_A, D.NET_B
    out = D.NET_OUT

    # The seed: d(loss)/d(loss) = 1, and loss = (out - target) squared.
    d_out = 2.0 * (out - D.NET_TARGET)

    # Output layer: out = vA*a + vB*b + c.
    d_c = d_out * 1.0
    d_vA = d_out * a
    d_vB = d_out * b
    d_a = d_out * D.NET_VA
    d_b = d_out * D.NET_VB

    # Through the non-linearity: a = tanh(a_pre), so da/da_pre = 1 - a^2.
    d_a_pre = d_a * (1.0 - a * a)
    d_b_pre = d_b * (1.0 - b * b)

    # First layer: a_pre = wA1*x1 + wA2*x2 + bA.
    d_wA1 = d_a_pre * x1
    d_wA2 = d_a_pre * x2
    d_bA = d_a_pre * 1.0
    d_wB1 = d_b_pre * x1
    d_wB2 = d_b_pre * x2
    d_bB = d_b_pre * 1.0

    # The inputs reach the loss through BOTH hidden units. SUM the paths.
    d_x1 = d_a_pre * D.NET_WA1 + d_b_pre * D.NET_WB1  # SUM over two paths
    d_x2 = d_a_pre * D.NET_WA2 + d_b_pre * D.NET_WB2  # SUM over two paths

    return {
        "out": d_out,
        "c": d_c,
        "vA": d_vA,
        "vB": d_vB,
        "a": d_a,
        "b": d_b,
        "a_pre": d_a_pre,
        "b_pre": d_b_pre,
        "wA1": d_wA1,
        "wA2": d_wA2,
        "bA": d_bA,
        "wB1": d_wB1,
        "wB2": d_wB2,
        "bB": d_bB,
        "x1": d_x1,
        "x2": d_x2,
    }


def numeric_parameter_gradients(h: float) -> dict[str, float]:
    """Every parameter gradient by central difference, for cross-checking."""
    base = default_parameter_values()
    grads: dict[str, float] = {}
    for i, name in enumerate(D.NET_PARAMETERS):
        ahead = list(base)
        behind = list(base)
        ahead[i] += h
        behind[i] -= h
        grads[name] = (
            loss_only(D.NET_X1, D.NET_X2, ahead)
            - loss_only(D.NET_X1, D.NET_X2, behind)
        ) / (2.0 * h)
    return grads


def numeric_input_gradients(h: float) -> dict[str, float]:
    """The two input gradients by central difference.

    These are the two that are sums over paths, so they are the two worth
    checking hardest: a product-only chain rule gets them visibly wrong.
    """
    params = default_parameter_values()
    grads: dict[str, float] = {}
    for name, i in (("x1", 0), ("x2", 1)):
        point = [D.NET_X1, D.NET_X2]
        ahead = list(point)
        behind = list(point)
        ahead[i] += h
        behind[i] -= h
        grads[name] = (
            loss_only(ahead[0], ahead[1], params)
            - loss_only(behind[0], behind[1], params)
        ) / (2.0 * h)
    return grads


def forward_mode_parameter_gradients() -> tuple[dict[str, float], int]:
    """Every parameter gradient by forward mode, and the passes it needed.

    One complete run of the network per parameter. With nine parameters that
    is nine runs to reverse mode's one, and the ratio is the whole reason
    training uses reverse mode.
    """
    base = default_parameter_values()
    grads: dict[str, float] = {}
    passes = 0
    for seed, name in enumerate(D.NET_PARAMETERS):
        x1 = Dual(D.NET_X1, 0.0)
        x2 = Dual(D.NET_X2, 0.0)
        params = [
            Dual(v, 1.0 if i == seed else 0.0) for i, v in enumerate(base)
        ]
        wA1, wA2, bA, wB1, wB2, bB, vA, vB, c = params
        a = (wA1 * x1 + wA2 * x2 + bA).tanh()
        b = (wB1 * x1 + wB2 * x2 + bB).tanh()
        out = vA * a + vB * b + c
        diff = out - D.NET_TARGET
        grads[name] = (diff * diff).dot
        passes += 1
    return grads, passes

"""The reference suite: every claim this lab makes, checked against a value.

Run from the lab directory:

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

Nothing here asserts that a function exists or that a file is present. Every
test runs code and compares a number, and every float comparison names the
tolerance it uses and why that tolerance rather than another.
"""

import math

import numpy as np
import pytest

import dataset as D
import network as N
from autodiff import (
    Dual,
    Value,
    forward_mode_gradient,
    graph_size,
    numeric_gradient,
    parameters_of,
    reverse_mode_gradient,
    sum_values,
    topological_order,
)
from chainrule import (
    central_difference,
    chain_derivative,
    chain_function,
    chain_local_rates,
    chain_rule,
    chain_values,
    compose,
    gear_ratio,
    order_of_magnitude,
    partial_difference,
    path_contributions,
    product,
    product_trace,
    repeated_product,
    running_products,
    total_derivative,
    wrong_single_path_derivative,
)

# --------------------------------------------------------------------------
# The tolerances themselves
# --------------------------------------------------------------------------


def test_epsilon_is_numpy_float64_epsilon():
    assert D.EPSILON == float(np.finfo(np.float64).eps)


def test_this_interpreter_uses_ieee754_doubles():
    import sys

    assert sys.float_info.mant_dig == 53


def test_step_size_sits_inside_day_108s_measured_band():
    # Day 108 found the central rule's best step for e**x in 1e-7 to 1e-4.
    assert 1e-7 <= D.H <= 1e-4


@pytest.mark.parametrize(
    "name,value,ceiling",
    [
        ("NUMERIC_TOL", D.NUMERIC_TOL, 1e-5),
        ("NUMERIC_REL_TOL", D.NUMERIC_REL_TOL, 1e-5),
        ("ANALYTIC_TOL", D.ANALYTIC_TOL, 1e-10),
    ],
)
def test_no_tolerance_is_loose_enough_to_be_meaningless(name, value, ceiling):
    # A tolerance large enough to pass anything is not a test.
    assert 0.0 < value < ceiling, name


def test_the_analytic_tolerance_is_far_tighter_than_the_numeric_one():
    # Comparing two analytic routes must be much stricter than comparing an
    # analytic route against a measurement. A thousandfold, here.
    assert D.ANALYTIC_TOL * 1000.0 < D.NUMERIC_TOL


def test_the_numeric_tolerance_has_headroom_over_the_error_bound():
    # truncation ~ (h^2/6)|f'''| and rounding ~ EPSILON|f|/h, at |f|,|f'''| <= 250
    bound = 250.0 * (D.H * D.H / 6.0 + D.EPSILON / D.H)
    assert bound < D.NUMERIC_TOL / 50.0


# --------------------------------------------------------------------------
# Rates multiply, with no calculus in sight
# --------------------------------------------------------------------------


def test_two_gears_multiply_to_six():
    assert gear_ratio(D.GEAR_RATIOS) == D.GEAR_RATIO_PRODUCT


def test_a_four_stage_gear_train_multiplies_to_thirty_six():
    assert gear_ratio(D.GEAR_TRAIN) == D.GEAR_TRAIN_PRODUCT


def test_three_currency_rates_multiply_to_one_hundred_and_fifty():
    assert product(D.CURRENCY_RATES) == D.CURRENCY_PRODUCT


def test_an_empty_product_is_one():
    # The identity for multiplication, and the right answer to "how much does
    # x change per unit of x".
    assert product([]) == 1.0


def test_the_order_of_the_gear_stages_does_not_change_the_ratio():
    assert gear_ratio(D.GEAR_TRAIN) == gear_ratio(tuple(reversed(D.GEAR_TRAIN)))


# --------------------------------------------------------------------------
# Composition
# --------------------------------------------------------------------------


def test_composition_runs_the_inner_function_first():
    assert compose(D.square, D.line)(2.0) == 49.0
    # The other order gives a different answer, which is why order matters.
    assert compose(D.line, D.square)(2.0) == 13.0


def test_composition_is_not_commutative():
    assert compose(D.square, D.line)(2.0) != compose(D.line, D.square)(2.0)


@pytest.mark.parametrize("case", D.COMPOSITIONS, ids=lambda c: c.name)
def test_chain_rule_matches_the_closed_form(case):
    got = chain_rule(case.d_outer, case.inner, case.d_inner, case.x)
    assert abs(got - case.exact) < D.ANALYTIC_TOL


@pytest.mark.parametrize("case", D.COMPOSITIONS, ids=lambda c: c.name)
def test_chain_rule_matches_a_central_difference(case):
    got = chain_rule(case.d_outer, case.inner, case.d_inner, case.x)
    measured = central_difference(compose(case.outer, case.inner), case.x, D.H)
    assert abs(got - measured) < D.NUMERIC_TOL


@pytest.mark.parametrize("case", D.COMPOSITIONS, ids=lambda c: c.name)
def test_the_composed_function_is_finite_at_the_test_point(case):
    assert math.isfinite(compose(case.outer, case.inner)(case.x))


def test_evaluating_the_outer_derivative_at_x_gives_the_wrong_answer():
    # The single most common chain-rule mistake, asserted as a mistake so a
    # future edit cannot make it accidentally right.
    correct = D.d_square(D.line(2.0)) * D.d_line(2.0)
    mistake = D.d_square(2.0) * D.d_line(2.0)
    assert correct == 42.0
    assert mistake == 12.0
    assert abs(correct - mistake) > 1.0


def test_the_sigmoid_slope_at_zero_is_exactly_a_quarter():
    case = D.COMPOSITIONS[4]
    assert chain_rule(case.d_outer, case.inner, case.d_inner, 0.0) == 0.25


def test_a_quarter_is_the_sigmoids_largest_slope_anywhere():
    case = D.COMPOSITIONS[4]
    for x in (-4.0, -2.0, -0.5, 0.5, 2.0, 4.0):
        assert chain_rule(case.d_outer, case.inner, case.d_inner, x) < 0.25


def test_tanh_of_a_line_at_minus_a_half_is_exactly_two():
    case = D.COMPOSITIONS[5]
    assert chain_rule(case.d_outer, case.inner, case.d_inner, -0.5) == 2.0


# --------------------------------------------------------------------------
# Chains of any depth
# --------------------------------------------------------------------------


def test_the_forward_pass_returns_one_more_value_than_there_are_stages():
    values = chain_values(D.FIVE_STAGES, D.FIVE_START)
    assert len(values) == len(D.FIVE_STAGES) + 1
    assert values[0] == D.FIVE_START


def test_the_five_stage_forward_values_are_the_documented_ones():
    assert chain_values(D.FIVE_STAGES, D.FIVE_START) == list(D.FIVE_VALUES)


def test_the_five_local_rates_are_the_documented_ones():
    got = chain_local_rates(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START)
    assert got == list(D.FIVE_LOCAL_RATES)


def test_the_five_rates_multiply_to_the_derivative():
    assert chain_derivative(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START) == 0.4


def test_the_collapsed_formula_agrees_with_the_product_of_local_rates():
    closed = D.d_five_chain_closed_form(D.FIVE_START)
    chained = chain_derivative(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START)
    assert abs(closed - chained) < D.ANALYTIC_TOL


def test_the_five_stage_chain_collapses_to_the_logarithm_of_two_x_plus_three():
    composed = chain_function(D.FIVE_STAGES)
    for x in (0.5, 1.0, 2.0, 7.5):
        assert abs(composed(x) - D.five_chain_closed_form(x)) < D.ANALYTIC_TOL


@pytest.mark.parametrize("x", [0.5, 1.0, 2.0, 7.5])
def test_the_chain_derivative_matches_a_measurement_at_several_points(x):
    analytic = chain_derivative(D.FIVE_STAGES, D.FIVE_RATES, x)
    measured = central_difference(chain_function(D.FIVE_STAGES), x, D.H)
    assert abs(analytic - measured) < D.NUMERIC_TOL


@pytest.mark.parametrize("depth", [1, 2, 3, 4, 5])
def test_every_prefix_of_the_chain_agrees_with_a_measurement(depth):
    stages = D.FIVE_STAGES[:depth]
    rates = D.FIVE_RATES[:depth]
    analytic = chain_derivative(stages, rates, D.FIVE_START)
    measured = central_difference(chain_function(stages), D.FIVE_START, D.H)
    assert abs(analytic - measured) < D.NUMERIC_TOL


def test_a_chain_of_zero_stages_has_derivative_one():
    assert chain_derivative((), (), 3.0) == 1.0


def test_mismatched_stages_and_rates_are_refused_rather_than_guessed():
    with pytest.raises(ValueError):
        chain_local_rates(D.FIVE_STAGES, D.FIVE_RATES[:2], 1.0)


def test_the_running_products_end_at_the_last_local_rate():
    rates = chain_local_rates(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START)
    assert running_products(rates)[-1] == rates[-1]


def test_the_running_products_start_at_the_whole_derivative():
    rates = chain_local_rates(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START)
    assert abs(running_products(rates)[0] - 0.4) < D.ANALYTIC_TOL


def test_the_two_multiplication_orders_differ_only_by_rounding():
    rates = chain_local_rates(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START)
    forwards = product(rates)
    backwards = running_products(rates)[0]
    # They are NOT equal -- float64 multiplication is not associative -- and
    # this test asserts both halves of that: close, but not identical.
    assert abs(forwards - backwards) < D.ANALYTIC_TOL
    assert abs(forwards - backwards) < 4.0 * D.EPSILON


@pytest.mark.parametrize("i", range(5))
def test_each_running_product_is_the_derivative_of_the_rest_of_the_chain(i):
    rates = chain_local_rates(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START)
    carried = running_products(rates)
    assert abs(carried[i] - product(rates[i:])) < D.ANALYTIC_TOL


def test_central_difference_refuses_a_non_positive_step():
    with pytest.raises(ValueError):
        central_difference(D.square, 1.0, 0.0)
    with pytest.raises(ValueError):
        central_difference(D.square, 1.0, -1e-5)


def test_partial_difference_refuses_a_non_positive_step():
    with pytest.raises(ValueError):
        partial_difference(D.surface, (1.0, 1.0), 0, 0.0)


# --------------------------------------------------------------------------
# Two paths, and the sum
# --------------------------------------------------------------------------

TWO_PATHS = [[D.TWO_PATH_V, 2.0 * D.TWO_PATH_X], [D.TWO_PATH_U, 3.0]]


def test_the_two_path_contributions_are_twenty_four_and_twelve():
    assert path_contributions(TWO_PATHS) == list(D.TWO_PATH_CONTRIBUTIONS)


def test_the_contributions_are_added_not_multiplied():
    assert total_derivative(TWO_PATHS) == 36.0
    assert total_derivative(TWO_PATHS) != 24.0 * 12.0


def test_the_sum_of_the_paths_matches_a_central_difference():
    measured = central_difference(D.two_path_direct, D.TWO_PATH_X, D.H)
    assert abs(total_derivative(TWO_PATHS) - measured) < D.NUMERIC_TOL


def test_the_sum_of_the_paths_matches_the_closed_form():
    closed = D.d_two_path_direct(D.TWO_PATH_X)
    assert abs(total_derivative(TWO_PATHS) - closed) < D.ANALYTIC_TOL


@pytest.mark.parametrize("path_index", [0, 1])
def test_neither_single_path_is_the_answer(path_index):
    # The instructive failure, asserted AS a failure. If a future edit made
    # one path accidentally correct, this test would notice.
    measured = central_difference(D.two_path_direct, D.TWO_PATH_X, D.H)
    wrong = wrong_single_path_derivative(TWO_PATHS, path_index)
    assert abs(wrong - measured) > 1.0


@pytest.mark.parametrize("x", [0.5, 1.0, 2.0, 3.5])
def test_the_two_path_rule_holds_away_from_the_documented_point(x):
    paths = [[3.0 * x, 2.0 * x], [x * x, 3.0]]
    measured = central_difference(D.two_path_direct, x, D.H)
    assert abs(total_derivative(paths) - measured) < D.NUMERIC_TOL
    assert abs(9.0 * x * x - measured) < D.NUMERIC_TOL


def test_the_surface_value_at_the_documented_point():
    assert D.surface(*D.SURFACE_POINT) == D.SURFACE_Z


def test_the_surface_intermediates_are_the_documented_ones():
    s, t = D.SURFACE_POINT
    assert s * t == D.SURFACE_U
    assert s - t == D.SURFACE_V


@pytest.mark.parametrize("index,expected", [(0, 34.0), (1, 26.0)])
def test_the_multivariable_chain_rule_matches_a_partial_difference(index, expected):
    measured = partial_difference(D.surface, D.SURFACE_POINT, index, D.H)
    assert abs(expected - measured) < D.NUMERIC_TOL


def test_the_multivariable_gradient_is_built_from_two_products_each():
    s, t = D.SURFACE_POINT
    dz_ds = D.SURFACE_DZ_DU * t + D.SURFACE_DZ_DV * 1.0
    dz_dt = D.SURFACE_DZ_DU * s + D.SURFACE_DZ_DV * -1.0
    assert (dz_ds, dz_dt) == D.SURFACE_GRADIENT


def test_dropping_the_second_path_of_the_surface_gets_it_wrong():
    s, t = D.SURFACE_POINT
    only_u = D.SURFACE_DZ_DU * t
    assert only_u != D.SURFACE_GRADIENT[0]
    assert abs(only_u - D.SURFACE_GRADIENT[0]) == 2.0


# --------------------------------------------------------------------------
# The Value engine
# --------------------------------------------------------------------------


def test_a_leaf_starts_with_a_zero_gradient():
    assert Value(3.0).grad == 0.0


def test_addition_computes_the_right_value_and_gradients():
    a, b = Value(3.0), Value(4.0)
    c = a + b
    c.backward()
    assert c.data == 7.0
    assert a.grad == 1.0
    assert b.grad == 1.0


def test_multiplication_computes_the_right_value_and_gradients():
    a, b = Value(3.0), Value(4.0)
    c = a * b
    c.backward()
    assert c.data == 12.0
    assert a.grad == 4.0
    assert b.grad == 3.0


def test_a_value_added_to_itself_accumulates_both_contributions():
    # If `+=` were `=` this would be 1.0. That one character is the entire
    # multivariable chain rule.
    x = Value(3.0)
    y = x + x
    y.backward()
    assert y.data == 6.0
    assert x.grad == 2.0


def test_a_value_multiplied_by_itself_reproduces_the_power_rule():
    x = Value(3.0)
    y = x * x
    y.backward()
    assert x.grad == 6.0


def test_a_value_cubed_reproduces_the_power_rule_too():
    x = Value(2.0)
    y = x * x * x
    y.backward()
    assert x.grad == 12.0  # 3 x squared


def test_the_output_seeds_its_own_gradient_with_one():
    x = Value(3.0)
    y = x * 2.0
    y.backward()
    assert y.grad == 1.0


def test_backward_resets_gradients_so_it_can_be_run_twice():
    x = Value(3.0)
    y = x * x
    y.backward()
    first = x.grad
    y.backward()
    assert x.grad == first


def test_scalars_can_be_mixed_in_from_either_side():
    x = Value(2.0)
    assert (x + 3.0).data == 5.0
    assert (3.0 + x).data == 5.0
    assert (x * 3.0).data == 6.0
    assert (3.0 * x).data == 6.0
    assert (x - 1.0).data == 1.0
    assert (1.0 - x).data == -1.0
    assert (-x).data == -2.0


def test_subtraction_gradients_have_the_right_signs():
    a, b = Value(5.0), Value(3.0)
    c = a - b
    c.backward()
    assert c.data == 2.0
    assert a.grad == 1.0
    assert b.grad == -1.0


def test_tanh_computes_the_right_value():
    assert abs(Value(0.6).tanh().data - math.tanh(0.6)) < D.ANALYTIC_TOL


def test_tanh_gradient_is_one_minus_tanh_squared():
    z = Value(0.6)
    t = z.tanh()
    t.backward()
    assert abs(z.grad - (1.0 - t.data * t.data)) < D.ANALYTIC_TOL


def test_tanh_gradient_at_zero_is_exactly_one():
    z = Value(0.0)
    z.tanh().backward()
    assert z.grad == 1.0


def test_tanh_at_half_ln_three_is_exactly_a_half():
    # The fact the whole hand-worked network rests on. Asserted, not assumed.
    assert math.tanh(D.HALF_LN3) == 0.5


def test_tanh_slope_at_half_ln_three_is_exactly_three_quarters():
    z = Value(D.HALF_LN3)
    z.tanh().backward()
    assert z.grad == 0.75


def test_tanh_saturates_and_its_slope_goes_to_nearly_nothing():
    z = Value(5.0)
    z.tanh().backward()
    assert z.grad < 1e-3


def test_the_topological_order_puts_children_before_parents():
    p, q = Value(2.0), Value(-3.0)
    r = p * q
    out = r + p
    order = topological_order(out)
    position = {id(node): i for i, node in enumerate(order)}
    for node in order:
        for child in node._children:
            assert position[id(child)] < position[id(node)]


def test_a_node_used_twice_appears_once_in_the_order():
    p = Value(2.0)
    out = p * p + p
    order = topological_order(out)
    assert sum(1 for node in order if node is p) == 1


def test_graph_size_counts_every_node_once():
    p = Value(2.0)
    out = p * p + p
    assert graph_size(out) == len(topological_order(out))


def test_the_engine_handles_a_deep_chain_without_recursion_limits():
    # Ten thousand operations. A recursive topological sort would fail here.
    node = Value(1.0)
    for _ in range(10_000):
        node = node * 1.0
    node.backward()
    assert graph_size(node) > 10_000


def test_parameters_of_finds_exactly_the_leaves():
    a, b = Value(1.0), Value(2.0)
    out = a * b + a
    leaves = parameters_of(out)
    assert a in leaves
    assert b in leaves
    assert out not in leaves


def test_sum_values_adds_a_list_of_values():
    total = sum_values([Value(1.0), Value(2.0), Value(3.5)])
    total.backward()
    assert total.data == 6.5


def test_repr_shows_the_data_and_the_gradient():
    x = Value(2.0, label="x")
    text = repr(x)
    assert "x=" in text
    assert "grad" in text


ENGINE_CASES = [
    ("square_of_line", lambda v: (3.0 * v[0] + 1.0) * (3.0 * v[0] + 1.0), [2.0]),
    ("cubic", lambda v: v[0] * v[0] * v[0] + (-2.0) * v[0], [1.5]),
    ("product_pair", lambda v: (v[0] * v[1] + v[0]) * (v[1] + 3.0), [2.0, -1.0]),
    ("tanh_mix", lambda v: ((v[0] * v[1]).tanh() * v[2] + v[0] * v[2]) * (1.0 + v[1]),
     [0.7, 0.4, -1.3]),
    ("shared_input", lambda v: (v[0] * v[0]) * (3.0 * v[0]), [2.0]),
    ("deep_tanh", lambda v: ((v[0] * 0.5).tanh() * 2.0 + v[0]).tanh(), [1.1]),
]


@pytest.mark.parametrize("name,build,point", ENGINE_CASES, ids=[c[0] for c in ENGINE_CASES])
def test_reverse_mode_matches_a_central_difference(name, build, point):
    def plain(vals):
        return build([Value(v) for v in vals]).data

    grads, _ = reverse_mode_gradient(build, point)
    numeric, _ = numeric_gradient(plain, point, D.H)
    for got, want in zip(grads, numeric):
        assert abs(got - want) < D.NUMERIC_TOL + D.NUMERIC_REL_TOL * abs(want)


@pytest.mark.parametrize("name,build,point", ENGINE_CASES, ids=[c[0] for c in ENGINE_CASES])
def test_reverse_mode_matches_forward_mode_to_the_last_bits(name, build, point):
    reverse, _ = reverse_mode_gradient(build, point)
    forward, _ = forward_mode_gradient(build, point)
    for got, want in zip(reverse, forward):
        assert abs(got - want) < D.ANALYTIC_TOL


@pytest.mark.parametrize("name,build,point", ENGINE_CASES, ids=[c[0] for c in ENGINE_CASES])
def test_reverse_mode_always_uses_exactly_one_pass(name, build, point):
    _, passes = reverse_mode_gradient(build, point)
    assert passes == 1


@pytest.mark.parametrize("name,build,point", ENGINE_CASES, ids=[c[0] for c in ENGINE_CASES])
def test_forward_mode_uses_one_pass_per_input(name, build, point):
    _, passes = forward_mode_gradient(build, point)
    assert passes == len(point)


@pytest.mark.parametrize("name,build,point", ENGINE_CASES, ids=[c[0] for c in ENGINE_CASES])
def test_a_central_difference_uses_two_evaluations_per_input(name, build, point):
    def plain(vals):
        return build([Value(v) for v in vals]).data

    _, passes = numeric_gradient(plain, point, D.H)
    assert passes == 2 * len(point)


def test_the_shared_input_case_is_the_sum_over_paths_again():
    # (x*x)*(3x) is 3 x cubed, and x reaches the output three times.
    grads, _ = reverse_mode_gradient(lambda v: (v[0] * v[0]) * (3.0 * v[0]), [2.0])
    assert abs(grads[0] - 9.0 * 4.0) < D.ANALYTIC_TOL


# --------------------------------------------------------------------------
# Dual numbers (forward mode)
# --------------------------------------------------------------------------


def test_a_dual_carries_a_value_and_a_derivative():
    d = Dual(3.0, 1.0)
    assert d.value == 3.0
    assert d.dot == 1.0


def test_dual_addition_adds_both_parts():
    d = Dual(3.0, 1.0) + Dual(4.0, 0.0)
    assert (d.value, d.dot) == (7.0, 1.0)


def test_dual_multiplication_uses_the_product_rule():
    d = Dual(3.0, 1.0) * Dual(4.0, 0.0)
    assert (d.value, d.dot) == (12.0, 4.0)


def test_dual_tanh_uses_one_minus_tanh_squared():
    d = Dual(0.6, 1.0).tanh()
    assert abs(d.dot - (1.0 - math.tanh(0.6) ** 2)) < D.ANALYTIC_TOL


def test_dual_scalars_work_from_either_side():
    d = Dual(2.0, 1.0)
    assert (3.0 + d).value == 5.0
    assert (3.0 * d).dot == 3.0
    assert (1.0 - d).dot == -1.0
    assert repr(d).startswith("Dual(")


def test_an_unseeded_dual_reports_a_zero_derivative():
    # The reason forward mode needs one pass per input: an input that was not
    # seeded contributes nothing to this pass.
    d = Dual(3.0, 0.0) * Dual(4.0, 0.0)
    assert d.dot == 0.0


# --------------------------------------------------------------------------
# The tiny network
# --------------------------------------------------------------------------

FORWARD = N.forward(D.NET_X1, D.NET_X2, N.default_parameter_values())


@pytest.mark.parametrize(
    "key,expected",
    [
        ("a_pre", D.NET_A_PRE),
        ("a", D.NET_A),
        ("b_pre", D.NET_B_PRE),
        ("b", D.NET_B),
        ("out", D.NET_OUT),
        ("loss", D.NET_LOSS),
    ],
)
def test_the_forward_pass_is_exact(key, expected):
    assert FORWARD[key] == expected


def test_the_bias_of_unit_b_is_half_the_log_of_three():
    assert D.NET_BB == 0.5 * math.log(3.0)


def test_unit_b_lands_on_a_tanh_value_that_is_exactly_a_half():
    assert FORWARD["b"] == 0.5
    assert 1.0 - FORWARD["b"] ** 2 == 0.75


def test_unit_a_lands_on_a_tanh_value_that_is_exactly_zero():
    assert FORWARD["a"] == 0.0
    assert 1.0 - FORWARD["a"] ** 2 == 1.0


HAND = N.hand_gradients()
ENGINE = N.engine_gradients()
NUMERIC = {
    **N.numeric_parameter_gradients(D.H),
    **N.numeric_input_gradients(D.H),
}


@pytest.mark.parametrize("key", sorted(D.NET_GRADIENTS))
def test_the_hand_worked_gradient_matches_the_table(key):
    assert abs(HAND[key] - D.NET_GRADIENTS[key]) < D.ANALYTIC_TOL


@pytest.mark.parametrize("key", sorted(D.NET_GRADIENTS))
def test_the_engine_matches_the_hand_computation_bit_for_bit(key):
    # Same multiplications, same order, same exact values -- so equality is
    # the honest comparison here, not a tolerance.
    assert ENGINE[key] == HAND[key]


@pytest.mark.parametrize("key", sorted(NUMERIC))
def test_the_engine_matches_a_central_difference(key):
    assert abs(ENGINE[key] - NUMERIC[key]) < D.NUMERIC_TOL


def test_the_gradient_of_a_weight_feeding_a_dead_unit_is_exactly_zero():
    # vA multiplies an activation of exactly 0, so nudging it moves nothing.
    assert ENGINE["vA"] == 0.0


def test_the_input_gradients_are_sums_over_two_paths():
    first, second = D.NET_X1_CONTRIBUTIONS
    assert HAND["x1"] == first + second
    assert first != 0.0
    assert second != 0.0


def test_taking_only_one_path_for_x1_is_visibly_wrong():
    first, second = D.NET_X1_CONTRIBUTIONS
    measured = NUMERIC["x1"]
    assert abs(first - measured) > 1.0
    assert abs(second - measured) > 1.0
    assert abs((first + second) - measured) < D.NUMERIC_TOL


def test_the_loss_gradient_chain_starts_at_two_times_the_residual():
    assert HAND["out"] == 2.0 * (D.NET_OUT - D.NET_TARGET)


def test_the_gradient_through_unit_b_is_scaled_by_three_quarters():
    assert HAND["b_pre"] == HAND["b"] * 0.75


def test_forward_mode_reproduces_every_parameter_gradient():
    grads, _ = N.forward_mode_parameter_gradients()
    for key, value in grads.items():
        assert abs(value - HAND[key]) < D.ANALYTIC_TOL


def test_forward_mode_needs_one_pass_per_parameter():
    _, passes = N.forward_mode_parameter_gradients()
    assert passes == len(D.NET_PARAMETERS)
    assert passes == 9


def test_reverse_mode_needs_one_pass_for_all_nine():
    # Built directly rather than via the helper, so the claim is about the
    # network and not about a convenience wrapper.
    x1, x2 = Value(D.NET_X1), Value(D.NET_X2)
    params = [Value(v) for v in N.default_parameter_values()]
    nodes = N.build_graph(x1, x2, params)
    nodes["loss"].backward()
    assert all(p.grad != 0.0 or name == "vA"
               for name, p in zip(D.NET_PARAMETERS, params))


def test_a_gradient_step_reduces_the_loss():
    # The point of having gradients at all, and the whole of Day 111.
    base = N.default_parameter_values()
    grads = N.numeric_parameter_gradients(D.H)
    step = 0.01
    moved = [
        v - step * grads[name] for name, v in zip(D.NET_PARAMETERS, base)
    ]
    before = N.loss_only(D.NET_X1, D.NET_X2, base)
    after = N.loss_only(D.NET_X1, D.NET_X2, moved)
    assert after < before


def test_the_network_has_nine_parameters():
    assert len(D.NET_PARAMETERS) == 9
    assert len(N.default_parameter_values()) == 9


# --------------------------------------------------------------------------
# Products that collapse and products that blow up
# --------------------------------------------------------------------------


def test_repeated_product_of_zero_factors_is_one():
    assert repeated_product(0.9, 0) == 1.0


def test_repeated_product_refuses_a_negative_count():
    with pytest.raises(ValueError):
        repeated_product(0.9, -1)


def test_a_product_trace_has_one_entry_per_factor():
    assert len(product_trace(0.9, D.CHAIN_LENGTH)) == D.CHAIN_LENGTH


def test_the_trace_ends_where_the_product_ends():
    trace = product_trace(0.9, D.CHAIN_LENGTH)
    assert trace[-1] == repeated_product(0.9, D.CHAIN_LENGTH)


def test_fifty_factors_of_nine_tenths_collapse_by_two_orders():
    value = repeated_product(D.DECAY_FACTOR, D.CHAIN_LENGTH)
    assert order_of_magnitude(value) == D.DECAY_ORDER
    assert value < 1e-2


def test_fifty_factors_of_eleven_tenths_grow_by_two_orders():
    value = repeated_product(D.GROWTH_FACTOR, D.CHAIN_LENGTH)
    assert order_of_magnitude(value) == D.GROWTH_ORDER
    assert value > 1e2


@pytest.mark.parametrize(
    "factor,count,expected_order",
    [
        (D.DECAY_FACTOR, D.LONG_CHAIN_LENGTH, -10),
        (D.GROWTH_FACTOR, D.LONG_CHAIN_LENGTH, 8),
        (D.MILD_DECAY, D.CHAIN_LENGTH, -16),
        (D.SHARP_DECAY, D.CHAIN_LENGTH, -31),
        (D.SHARP_GROWTH, D.CHAIN_LENGTH, 15),
    ],
)
def test_the_documented_orders_of_magnitude(factor, count, expected_order):
    assert order_of_magnitude(repeated_product(factor, count)) == expected_order


def test_half_to_the_fiftieth_is_four_epsilons_and_still_counts():
    # This contradicts the obvious guess, so it is asserted in both halves.
    value = repeated_product(D.MILD_DECAY, D.CHAIN_LENGTH)
    assert value == 4.0 * D.EPSILON
    assert 1.0 + value != 1.0


def test_three_more_halvings_do_make_it_disappear():
    value = repeated_product(D.MILD_DECAY, 53)
    assert value == 0.5 * D.EPSILON
    assert 1.0 + value == 1.0


def test_the_sigmoids_best_case_vanishes_completely_in_fifty_layers():
    value = repeated_product(D.SHARP_DECAY, D.CHAIN_LENGTH)
    assert value < D.EPSILON
    assert 1.0 + value == 1.0


def test_a_factor_of_exactly_one_neither_vanishes_nor_explodes():
    assert repeated_product(1.0, 10_000) == 1.0


def test_order_of_magnitude_refuses_zero():
    with pytest.raises(ValueError):
        order_of_magnitude(0.0)


@pytest.mark.parametrize(
    "value,expected", [(1.0, 0), (9.99, 0), (10.0, 1), (0.1, -1), (-500.0, 2)]
)
def test_order_of_magnitude_reads_the_exponent(value, expected):
    assert order_of_magnitude(value) == expected


def test_the_engine_reproduces_the_collapse_through_a_real_graph():
    def deep(vals):
        node = vals[0]
        for _ in range(D.CHAIN_LENGTH):
            node = node * D.DECAY_FACTOR
        return node

    grads, passes = reverse_mode_gradient(deep, [1.0])
    assert abs(grads[0] - repeated_product(D.DECAY_FACTOR, D.CHAIN_LENGTH)) < D.ANALYTIC_TOL
    assert passes == 1


def test_the_engine_reproduces_the_explosion_too():
    def deep(vals):
        node = vals[0]
        for _ in range(D.CHAIN_LENGTH):
            node = node * D.GROWTH_FACTOR
        return node

    grads, _ = reverse_mode_gradient(deep, [1.0])
    assert order_of_magnitude(grads[0]) == D.GROWTH_ORDER


def _stacked_tanh(depth):
    def deep(vals):
        node = vals[0]
        for _ in range(depth):
            node = node.tanh()
        return node

    return deep


@pytest.mark.parametrize("depth", [1, 5, 10, 20, 40, 80, 160])
def test_a_stacked_tanh_gradient_stays_positive_and_below_one(depth):
    grads, _ = reverse_mode_gradient(_stacked_tanh(depth), [0.9])
    assert 0.0 < grads[0] < 1.0


def test_a_stacked_tanh_gradient_falls_monotonically_with_depth():
    values = [
        reverse_mode_gradient(_stacked_tanh(d), [0.9])[0][0]
        for d in (1, 5, 10, 20, 40, 80, 160)
    ]
    assert values == sorted(values, reverse=True)


def test_stacked_tanh_decays_far_more_slowly_than_a_constant_factor_predicts():
    """A measured correction to the naive vanishing-gradient story.

    The obvious reasoning says: tanh's slope at the input is about 0.487, so
    forty tanh layers should multiply the gradient by 0.487 forty times and
    land near 1e-13. The measurement says otherwise, and by ten orders of
    magnitude.

    The reason is that the local rates are not constant. Each tanh pulls its
    input closer to 0, and tanh's slope at 0 is 1 -- so the deeper the stack
    goes, the closer each local rate creeps back towards 1. A product of
    constants is the wrong model for a product of rates that depend on where
    they are evaluated, and this test asserts the gap rather than glossing it.
    """
    single = reverse_mode_gradient(_stacked_tanh(1), [0.9])[0][0]
    deep = reverse_mode_gradient(_stacked_tanh(40), [0.9])[0][0]
    naive = single**40
    assert 1e-3 < deep < 1e-1
    assert naive < 1e-12
    assert deep > 1e9 * naive


def test_a_stack_of_constant_factors_does_vanish_geometrically():
    # The contrast: when the local rate really is a constant, the naive
    # reasoning is correct and the collapse is geometric.
    def deep(vals):
        node = vals[0]
        for _ in range(40):
            node = node * 0.487
        return node

    grads, _ = reverse_mode_gradient(deep, [0.9])
    assert grads[0] < 1e-12

lesson_id: D110
day: 110
kind: guided-build
languages: [python, bash]
setup_commands:
  - cd labs/sections/math-statistics-and-data/day-110-the-chain-rule
  - 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_gears_and_rates.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 02_composition_and_the_chain_rule.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 03_deeper_chains.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 04_two_paths_add.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 05_the_value_engine.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 06_backprop_by_hand.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 07_vanishing_and_exploding.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 -> 120 checks, 0 failure(s), exit 0; pytest examples -> 235 passed; pytest starter -> 2 passed, 163 skipped on an untouched checkout, and 165 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. The harness was additionally confirmed to exit 0 in three further configurations: with no .venv present at all and PYTEST pointing at an interpreter elsewhere; with a fake .venv present containing nothing but a stray __pycache__ and .pytest_cache, which the run left untouched; and, in all of those, with the README''s own documented `pytest starter -q` run immediately beforehand so the tree is already dirty when the harness starts. That last case is the one that matters: `pytest starter -q` legitimately writes starter/__pycache__ and .pytest_cache, and an earlier version of this harness would have reported them at the end as litter — failing the reader for following the instructions in the README. The harness now clears both at the START of its run, pruning .venv, so the check at the end measures what THIS run left rather than what a previous command left. With that block removed the dirty-tree scenario produces three failures, which was verified rather than assumed, so the block is load-bearing and not decorative. .venv itself is never treated as a stray file and nothing inside it is ever deleted. Network is needed once to install numpy and pytest; nothing else in the lab opens a socket, and section 7 greps every source file in examples/ and starter/ to prove it. Section 6 re-runs the harness with one expectation deliberately swapped for the belief that d loss/d x1 is -6.0 — the value you get by following only the first of the two paths from x1 to the loss — 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. The two-layer network was constructed so that every quantity in both passes is exact in float64: one hidden unit sits at a pre-activation of exactly 0 where tanh is 0 and its slope is 1, and the other at exactly half the natural logarithm of 3, where math.tanh returns exactly 0.5 and 1 - tanh squared is exactly 0.75 on this platform. Both exactness claims are asserted by the reference suite rather than assumed, and the choice is declared in dataset.py as a convenience for the reader rather than anything the chain rule depends on. Consequently the hand-worked backward pass and the from-scratch engine agree bit for bit on all sixteen gradients and are compared with == rather than a tolerance; central differences agree with both to within 4.463e-09, measured and reported rather than asserted to a value. Three findings were discovered while building the lab rather than planned, and all three are kept. FIRST and most valuable: a stack of forty real tanh operations does NOT vanish geometrically. The naive prediction — take tanh''s slope at the input, 0.486917, and raise it to the fortieth — gives 3.149274e-13, while the measurement gives 8.397332e-03, larger by a factor of 2.666e+10. The cause is that each tanh pulls its input closer to zero where tanh''s slope is 1, so the local rates climb back towards 1 as the stack deepens and the product decays like a power of the depth rather than exponentially; the measured ratio settles near 2.8 per doubling of depth. The suite therefore asserts the monotonic fall, the nine-orders-of-magnitude gap against the prediction, and the contrast case where a genuinely constant factor of 0.487 does collapse below 1e-12 in the same forty steps — but never the value. This is a real correction to the standard "tanh saturates, therefore gradients vanish" story and it is stated as such in the lesson rather than smoothed over. SECOND: 0.5 to the fiftieth power is exactly four times float64 epsilon, so 1.0 + 0.5**50 != 1.0 — it still moves a weight of 1, which contradicts the obvious guess. It takes three more halvings, to 0.5**53 = half an epsilon, before the addition rounds away to nothing; both halves are asserted. The example that genuinely vanishes at fifty layers is 0.25**50 = 7.888609e-31, and 0.25 was chosen because it is the LARGEST slope the sigmoid ever has, measured in script 02 rather than recalled. THIRD: the hand route reaches d loss/d vA as -3.0 x 0.0, which IEEE-754 makes negative zero, while the engine reaches the same gradient by a different route and gets +0.0; they compare equal and behave identically, so the harness asserts gradient == 0.0 and reports the sign rather than pretending both printed the same characters. Two further honest notes. Multiplying the five local rates forwards gives 0.4 exactly while multiplying them backwards gives 0.4000000000000001, because float64 multiplication is not associative; the lab measures the gap as under four epsilons and asserts both that they are close and that they are NOT identical, rather than comparing them with ==. And every tolerance is derived in examples/dataset.py from the error terms that govern its comparison with the arithmetic written out — ANALYTIC_TOL 1e-12 for two analytic routes, NUMERIC_TOL 1e-6 for an analytic route against a central difference at h = 1e-5, where the truncation and rounding bound is about 9.7e-9 and the margin is therefore about a hundredfold — and a reference test asserts that neither is loose enough to be meaningless and that the analytic tolerance is at least a thousand times tighter than the numeric one. No deep-learning framework is installed here and no output from PyTorch, JAX, TensorFlow or SymPy is reproduced anywhere in this lab or its lesson.'

# 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 | Exactly one thing: reading float64's machine epsilon from `numpy.finfo` rather than writing the literal `2.220446049250313e-16` and hoping. The vanishing-gradient section compares products against that epsilon, so it matters that the number is read from the platform rather than remembered. |
| `pytest` | 9.1.1 | MIT | The reference suite (235 tests) 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 can still do almost all of this lab, which is unusual and worth saying
plainly.

The autodiff engine — the most valuable thing here — needs `math` and nothing
else. So do all fourteen functions in `starter/chainrule.py`, the whole
two-layer network, the hand-worked backward pass, every composition, the
five-stage chain, the two-path example and the vanishing and exploding
products. Run the reference scripts directly with any Python 3.10 or later and
they will print their working and assert every claim they make:

```bash
cd examples && python3 01_gears_and_rates.py
```

What you lose is:

- `pytest`, so no running score and no skip-versus-fail distinction — you would
  read the numbers yourself instead;
- the epsilon cross-check, since `dataset.py` imports `numpy` at the top. If
  you want to run without NumPy, replace that import and the `EPSILON` line
  with `EPSILON = sys.float_info.epsilon`, which is the same value from the
  standard library. The lab does not ship it that way because reading it from
  two independent sources and comparing them is a better habit than reading it
  from one.

## What is deliberately *not* installed

PyTorch, JAX, TensorFlow and SymPy all do this job, 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
Alternatives section marks them as not run here.

That is not a limitation to apologise for. The engine you write in
`starter/autodiff.py` is the same idea as `torch.autograd`: a graph of
operations, a local derivative at each node, and one reverse walk applying the
chain rule. The difference between the two is engineering — tensors instead of
scalars, fused kernels, GPU dispatch, memory planning — and not concept. Having
written the seventy-line version, you will read the documentation of the real
one differently.

numpy==2.5.2
pytest==9.1.1

# Day 110 lab — the brief

Nine exercises, in order. Work top to bottom; each one leans on the one before.

Check yourself at any point:

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

On an untouched checkout that prints `2 passed, 163 skipped`. A **skip** means
"not attempted". A **failure** means "attempted and wrong", and it prints your
answer beside the real one. When it prints `165 passed`, you are finished.

Predict before you run. Exercises 5 to 9 are predictions on purpose, and two of
them are traps that only catch you if you commit to an answer first.

---

## Exercise 1 — `chainrule.py`, fourteen functions

The plumbing. None of it is hard and all of it is used later.

| Function | The thing to get right |
| --- | --- |
| `product` | The empty product is `1.0`, not `0.0` |
| `gear_ratio` | One line, calling `product` |
| `central_difference` | Divide by `2h`, not `h`. Raise `ValueError` on a non-positive step |
| `partial_difference` | Nudge one coordinate, hold the rest still |
| `compose` | Returns a **function**, not a number |
| `chain_rule` | Evaluate the outer derivative at `u = inner(x)`, not at `x` |
| `chain_values` | `n` stages give `n + 1` values, starting with `x` itself |
| `chain_local_rates` | Rate `i` is evaluated at the value **arriving** at stage `i` |
| `chain_derivative` | One line, calling `product` |
| `chain_function` | Returns a function, like `compose` |
| `running_products` | Entry `i` is the product of `rates[i:]`, in stage order |
| `path_contributions` | One product per path |
| `total_derivative` | **Add** across paths. This is the whole day |
| `repeated_product` | A loop, not `**`. `ValueError` on a negative count |
| `order_of_magnitude` | `floor(log10(abs(v)))`. `ValueError` on zero |

## Exercise 2 — `autodiff.py`, the engine

The most valuable thing in the calculus arc. Roughly seventy lines when you are
done, and it is the core of what every deep-learning framework does.

- **2a `__add__`** — the shape of all three operations. The approach block in
  the docstring is written out in full, because it is worth having one to copy
  from.
- **2b `__mul__`** — each input's local rate is the *other* input's value.
- **2c `tanh`** — slope is `1 - tanh²`. Compute the tanh once and reuse it.
- **2d `topological_order`** — iterative, not recursive. A ten-thousand-node
  graph is a test.
- **2e `backward`** — zero the gradients, seed the output with `1.0`, walk the
  order in reverse.
- **2f `Dual`** — forward mode, for the cost comparison in exercise 3.

**Use `+=` in every backward step, never `=`.** That one character is the
multivariable chain rule. The engine will run either way, look sensible either
way, and be wrong on every graph where anything is used twice.

## Exercise 3 — the two modes and their cost

`reverse_mode_gradient`, `forward_mode_gradient` and `numeric_gradient`. Each
returns its gradients **and the number of passes it needed**. The counts are
the point: 1, `n`, and `2n` for a function of `n` inputs. Training a model is
`n` in the millions and one scalar out.

## Exercise 4 — `network.py`, backpropagation by hand

Two inputs, two tanh hidden units, one linear output, squared-error loss, nine
parameters. Every quantity in both passes is exact in float64, so you can check
the whole thing with a pen.

- **4a `hand_gradients`** — sixteen gradients, by your own arithmetic. Work
  backwards from `d(loss)/d(loss) = 1`.
- **4b `engine_gradients`** — the same sixteen from one backward pass. If 2 is
  right these match 4a *exactly*, not approximately.
- **4c `numeric_parameter_gradients`** — central differences, which will not
  match exactly and should not.

Two of the sixteen deserve thought before you write them. `vA` multiplies an
activation of exactly zero. And `x1` and `x2` each reach the loss through
**both** hidden units, so each of those gradients is a sum of two products.

## Exercises 5 to 9 — `answers.py`, forty-two predictions

Fill in every `None`. Grouped as:

- **5** rates multiply (4 predictions)
- **6** composition and the one-variable chain rule (7)
- **7** depth, and the sum over paths (12)
- **8** the engine, and the network (12)
- **9** cost, collapse, and one honest surprise (7)

Two of these are designed to catch you:

**7.8** asks for the correct `df/dx` when a variable reaches the output twice.
The two contributions are 24 and 12. It is not 24, not 12, and not 288.

**9.6** asks whether `1.0 + 0.5**50 == 1.0`. `0.5**50` is about `8.88e-16` and
float64's epsilon is about `2.22e-16`, so the obvious guess is `True`. Work out
how many epsilons `8.88e-16` actually is before you answer.

---

## When you are done

Read the reference. Every script prints its working and asserts every claim it
makes:

```bash
cd examples
../.venv/bin/python3 01_gears_and_rates.py
../.venv/bin/python3 02_composition_and_the_chain_rule.py
../.venv/bin/python3 03_deeper_chains.py
../.venv/bin/python3 04_two_paths_add.py
../.venv/bin/python3 05_the_value_engine.py
../.venv/bin/python3 06_backprop_by_hand.py
../.venv/bin/python3 07_vanishing_and_exploding.py
cd ..
```

Section 6 of script 07 is the one to read even if you read nothing else: it
measures a case where the standard story about vanishing gradients is wrong by
ten orders of magnitude, and explains why.

"""Exercises 5 to 9 -- forty predictions.

Replace each `None` with the value you think is correct. A `None` is a skip,
not a failure: `pytest starter -q` counts only what you have attempted. When
you are wrong it prints both your answer and the real one, so a wrong guess
is worth more than a blank.

Predict BEFORE you run anything. The two exercises that catch almost everyone
are 7 (where paths meet) and 9 (where a plausible product turns out to be
ten orders of magnitude off), and they only catch you if you commit first.

Every answer is either an exact number, an integer, or a Python bool.
"""

ANSWERS: dict[str, object] = {
    # ----------------------------------------------------------------------
    # Exercise 5 -- rates multiply
    # ----------------------------------------------------------------------
    # 5.1 Gear A turns 2x per turn of B; B turns 3x per turn of C.
    #     How many times does A turn per turn of C?
    "gears_two_stage": None,
    # 5.2 A four-stage train with ratios 2, 3, 1.5, 4. Overall ratio?
    "gears_four_stage": None,
    # 5.3 What does `product([])` return -- the empty product?
    "empty_product": None,
    # 5.4 Does reversing the order of the gear stages change the overall
    #     ratio? True or False.
    "gear_order_matters": None,
    # ----------------------------------------------------------------------
    # Exercise 6 -- composition and the one-variable chain rule
    # ----------------------------------------------------------------------
    # 6.1 f(u) = u squared, g(x) = 3x + 1. What is f(g(2))?
    "composed_value_at_two": None,
    # 6.2 And g(f(2)) -- the other order?
    "composed_other_order_at_two": None,
    # 6.3 d/dx of (3x + 1) squared at x = 2, done correctly.
    "chain_rule_at_two": None,
    # 6.4 The same calculation with the outer derivative wrongly evaluated at
    #     x instead of at u. What number does that mistake produce?
    "chain_rule_mistake_at_two": None,
    # 6.5 The slope of the sigmoid 1/(1 + e**-x) at x = 0.
    "sigmoid_slope_at_zero": None,
    # 6.6 Is that value the sigmoid's LARGEST slope anywhere? True or False.
    "sigmoid_slope_is_maximum": None,
    # 6.7 d/dx of tanh(2x + 1) at x = -0.5. (The inner function is 0 there,
    #     and tanh has slope exactly 1 at 0.)
    "tanh_of_line_slope": None,
    # ----------------------------------------------------------------------
    # Exercise 7 -- depth, and the sum over paths
    # ----------------------------------------------------------------------
    # 7.1 The five stages from dataset.py, starting at x = 1. How many
    #     numbers does `chain_values` return?
    "five_chain_value_count": None,
    # 7.2 The third of the five local rates (the `square` stage). Careful:
    #     it is 2u, and u is the value ARRIVING at that stage.
    "five_chain_third_rate": None,
    # 7.3 The product of all five local rates.
    "five_chain_derivative": None,
    # 7.4 The five stages collapse to ln(2x + 3). Its derivative at x = 1?
    "five_chain_closed_form_derivative": None,
    # 7.5 In `running_products`, what is the LAST entry equal to?
    #     Give the number for the five-stage chain.
    "running_products_last": None,
    # 7.6 Now the two-path graph: u = x squared, v = 3x, f = u x v, at x = 2.
    #     The contribution through the u path.
    "two_path_u_contribution": None,
    # 7.7 The contribution through the v path.
    "two_path_v_contribution": None,
    # 7.8 The correct df/dx. (Not one of the two above, and not their
    #     product.)
    "two_path_total": None,
    # 7.9 The same function written directly is f = 3x cubed. Its derivative
    #     at x = 2 by the power rule?
    "two_path_closed_form": None,
    # 7.10 z = u squared + v squared, u = st, v = s - t, at (s, t) = (2, 3).
    #      What is z?
    "surface_value": None,
    # 7.11 dz/ds at that point. Two paths again: through u and through v.
    "surface_dz_ds": None,
    # 7.12 dz/dt at that point.
    "surface_dz_dt": None,
    # ----------------------------------------------------------------------
    # Exercise 8 -- the engine, and the network
    # ----------------------------------------------------------------------
    # 8.1 x = Value(3.0); y = x + x; y.backward(). What is x.grad?
    "engine_x_plus_x_grad": None,
    # 8.2 x = Value(3.0); y = x * x; y.backward(). What is x.grad?
    "engine_x_times_x_grad": None,
    # 8.3 x = Value(2.0); y = x * x * x; y.backward(). What is x.grad?
    "engine_x_cubed_grad": None,
    # 8.4 What is the gradient of tanh at 0, exactly?
    "tanh_slope_at_zero": None,
    # 8.5 tanh at half the natural logarithm of 3 -- exactly.
    "tanh_at_half_ln_three": None,
    # 8.6 And its slope there -- exactly.
    "tanh_slope_at_half_ln_three": None,
    # 8.7 The network's forward pass: the value of the loss.
    "network_loss": None,
    # 8.8 d loss / d out.
    "network_d_out": None,
    # 8.9 d loss / d vA. Look at what vA multiplies before you answer.
    "network_d_vA": None,
    # 8.10 d loss / d b_pre. (d loss/d b is 9.0, and tanh's slope there
    #      is 0.75.)
    "network_d_b_pre": None,
    # 8.11 d loss / d wB2.
    "network_d_wB2": None,
    # 8.12 d loss / d x1. This one is a SUM over two paths.
    "network_d_x1": None,
    # ----------------------------------------------------------------------
    # Exercise 9 -- cost, collapse, and one honest surprise
    # ----------------------------------------------------------------------
    # 9.1 Reverse mode on a function of 25 inputs and 1 output. How many
    #     forward-and-backward passes to get ALL 25 gradients?
    "reverse_passes_for_25_inputs": None,
    # 9.2 Forward mode, same function. How many passes?
    "forward_passes_for_25_inputs": None,
    # 9.3 Central differences, same function. How many evaluations?
    "numeric_passes_for_25_inputs": None,
    # 9.4 order_of_magnitude(0.9 ** 50) -- the exponent, not the digits.
    "decay_order": None,
    # 9.5 order_of_magnitude(1.1 ** 50).
    "growth_order": None,
    # 9.6 0.5 ** 50 is about 8.88e-16, and float64's epsilon is about
    #     2.22e-16. Does `1.0 + 0.5**50 == 1.0`? True or False.
    #     Think before you answer; the obvious guess is wrong.
    "half_to_the_fiftieth_vanishes": None,
    # 9.7 0.25 is the sigmoid's largest slope. Does `1.0 + 0.25**50 == 1.0`?
    "quarter_to_the_fiftieth_vanishes": None,
}

"""Exercise 2 -- build the reverse-mode autodiff engine.

This is the most valuable thing in the calculus arc. When it works you will
have written the core of what every deep-learning framework does, and checked
it against a numerical derivative that knows nothing about your graph.

Work in this order, checking with `pytest starter -q` after each step:

    2a  __add__          the value, and both local rates of 1
    2b  __mul__          the value, and each local rate being the OTHER input
    2c  tanh             the value, and the local rate 1 - tanh squared
    2d  topological_order   children always before parents
    2e  backward         seed 1.0, then walk the order in reverse
    2f  Dual             forward mode, for the cost comparison

The single most important line in the whole file is the one that accumulates
a gradient with `+=` rather than assigning it with `=`. That one character is
the multivariable chain rule: a value used in two places receives a
contribution from each use, and both are real, so they add. Get it wrong and
the engine will still run, still look sensible, and be quietly wrong on every
graph where anything is used twice.
"""

import math
from typing import Callable, Sequence


class Value:
    """One number in a computation graph, with a gradient and a history."""

    __slots__ = ("data", "grad", "label", "_backward", "_children", "_op")

    def __init__(
        self,
        data: float,
        children: tuple["Value", ...] = (),
        op: str = "",
        label: str = "",
    ) -> None:
        self.data: float = float(data)
        self.grad: float = 0.0
        self.label: str = label
        self._backward: Callable[[], None] = _do_nothing
        self._children: tuple["Value", ...] = children
        self._op: str = op

    # -- 2a ----------------------------------------------------------------

    def __add__(self, other: "Value | float") -> "Value":
        """Return a new Value holding self.data + other.data.

        The new Value's children are (self, other) and its op is "+".

        Then give it a `_backward` function that adds `out.grad` to BOTH
        children's gradients -- addition passes the gradient through
        untouched, because nudging either input by d moves the sum by d.

        Approach:

            other = other if isinstance(other, Value) else Value(other)
            out = Value(self.data + other.data, (self, other), "+")
            def backward():
                self.grad += out.grad
                other.grad += out.grad
            out._backward = backward
            return out

        That approach block is the answer, written out, because the shape of
        this method is the shape of all three and it is worth having one to
        copy from. The next two are yours.
        """
        return None

    # -- 2b ----------------------------------------------------------------

    def __mul__(self, other: "Value | float") -> "Value":
        """Return a new Value holding self.data * other.data, op "*".

        For a product, each input's local rate is the OTHER input's value:
        nudge self by d and the product moves by d x other.data. So the
        backward step adds `other.data * out.grad` to self.grad, and
        `self.data * out.grad` to other.grad.

        Use `+=` for both. Check it afterwards with `x * x`, where x is used
        twice: if the answer for x = 3 is 6.0 you have it right, and if it is
        3.0 you assigned where you should have accumulated.
        """
        return None

    # -- 2c ----------------------------------------------------------------

    def tanh(self) -> "Value":
        """Return a new Value holding tanh(self.data), op "tanh".

        The derivative of tanh is 1 - tanh squared. Compute the tanh once,
        before building the output, and reuse it in the backward step rather
        than recomputing it -- that reuse is exactly what a real framework
        does, and it is why a backward pass needs the forward pass's values
        kept in memory.

        Check: tanh at 0 has slope exactly 1, and at half the natural log of
        3 it has value exactly 0.5 and slope exactly 0.75.
        """
        return None

    # -- conveniences, already written for you -----------------------------

    def __neg__(self) -> "Value":
        return self * -1.0

    def __sub__(self, other: "Value | float") -> "Value":
        return self + (-(other if isinstance(other, Value) else Value(other)))

    def __radd__(self, other: "Value | float") -> "Value":
        return self + other

    def __rmul__(self, other: "Value | float") -> "Value":
        return self * other

    def __rsub__(self, other: "Value | float") -> "Value":
        return (-self) + other

    def __repr__(self) -> str:
        name = f"{self.label}=" if self.label else ""
        return f"Value({name}{self.data:.6g}, grad={self.grad:.6g})"

    # -- 2e ----------------------------------------------------------------

    def backward(self) -> None:
        """Fill in `.grad` on every value this one was computed from.

        Four steps:

        1. Get the topological order of the graph (exercise 2d).
        2. Zero every gradient, so calling backward twice gives the same
           answer as calling it once.
        3. Set this node's own grad to 1.0 -- the derivative of the output
           with respect to itself, which is the base case everything else
           hangs from.
        4. Walk the order BACKWARDS, calling each node's `_backward()`.

        Step 4 must be in reverse topological order, not any order that
        happens to work on a straight chain. A node must have received every
        contribution owed to it before it passes any of them on, and on a
        branching graph an arbitrary order will silently lose one.

        This function returns None on purpose -- it fills in gradients as a
        side effect. The test suite decides you have attempted it by checking
        whether a gradient actually changed.
        """
        return None


def _do_nothing() -> None:
    """The backward step of a leaf: it has nobody to pass anything to."""
    return None


# -- 2d --------------------------------------------------------------------


def topological_order(root: Value) -> list[Value]:
    """Every value `root` depends on, parents always AFTER their children.

    Write it iteratively rather than recursively. A chain of ten thousand
    operations is an ordinary size for a computation graph and would exhaust
    the interpreter's stack; one of the reference tests builds exactly that.

    A node reached twice must appear exactly once in the result.

    Approach: a stack of (node, already_expanded) pairs. Pop one; if it is
    already expanded, append it to the output. Otherwise mark it visited,
    push it back as expanded, and push its unvisited children. Track visited
    nodes by `id(node)`, because Value has no meaningful equality.
    """
    return None


def graph_size(root: Value) -> int:
    """How many nodes a backward pass will visit.

    Approach: one line, once `topological_order` works.
    """
    return None


# -- 2f --------------------------------------------------------------------


class Dual:
    """A number carried alongside its derivative with respect to ONE input.

    Forward mode: the same chain rule, applied left to right. Each operation
    computes the value AND the rate at which that value moves when the
    seeded input moves. There is no graph and no second pass, and the price
    is that one run answers about one input only.
    """

    __slots__ = ("value", "dot")

    def __init__(self, value: float, dot: float = 0.0) -> None:
        self.value: float = float(value)
        self.dot: float = float(dot)

    def __add__(self, other: "Dual | float") -> "Dual":
        """Add both parts: values add, and derivatives add.

        Remember to accept a plain float on the right, as `Value` does.
        """
        return None

    def __mul__(self, other: "Dual | float") -> "Dual":
        """The product rule.

        The value is self.value x other.value. The derivative is

            self.dot x other.value  +  self.value x other.dot

        which is the product rule, and which is the chain rule's constant
        travelling companion.
        """
        return None

    def tanh(self) -> "Dual":
        """tanh of the value, with the derivative scaled by 1 - tanh squared."""
        return None

    # -- conveniences, already written for you -----------------------------

    def __neg__(self) -> "Dual":
        return self * -1.0

    def __sub__(self, other: "Dual | float") -> "Dual":
        return self + (-(other if isinstance(other, Dual) else Dual(other)))

    def __radd__(self, other: "Dual | float") -> "Dual":
        return self + other

    def __rmul__(self, other: "Dual | float") -> "Dual":
        return self * other

    def __rsub__(self, other: "Dual | float") -> "Dual":
        return (-self) + other

    def __repr__(self) -> str:
        return f"Dual({self.value:.6g}, dot={self.dot:.6g})"


# --------------------------------------------------------------------------
# Exercise 3 -- the two modes, and their cost
# --------------------------------------------------------------------------


def reverse_mode_gradient(
    build: Callable[[Sequence[Value]], Value], xs: Sequence[float]
) -> tuple[list[float], int]:
    """Every partial derivative of `build`, and the number of passes used.

    `build` receives one `Value` per entry of `xs` and returns one output
    Value. Build the inputs, call `build`, call `.backward()` on the output,
    and return the list of input gradients together with the pass count.

    The pass count is 1, always, no matter how many inputs there are. That
    is the entire point of reverse mode, and the test suite checks it.
    """
    return None


def forward_mode_gradient(
    build: Callable[[Sequence[Dual]], Dual], xs: Sequence[float]
) -> tuple[list[float], int]:
    """The same gradients by forward mode, and the number of passes used.

    Run `build` once per input. On run number k, seed input k with a dot of
    1.0 and every other input with 0.0; the output's `.dot` is then the
    partial derivative with respect to input k.

    The pass count is `len(xs)`. Comparing that against the 1 above, on a
    model with a hundred million parameters, is the whole argument.
    """
    return None


def numeric_gradient(
    f: Callable[[Sequence[float]], float], xs: Sequence[float], h: float
) -> tuple[list[float], int]:
    """The same gradients by central differences, and the passes used.

    Two evaluations per input, so the count is 2 x len(xs) -- worse than
    forward mode and far worse than reverse mode, and approximate as well.
    This is the checking tool, not the production tool.

    `f` takes a whole list of coordinates, not one at a time.
    """
    return None

"""Exercise 1 -- fourteen functions to write.

Every function below 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 you have not attempted it yet: `pytest starter -q` will
SKIP an unattempted function rather than fail it, so your score only ever
counts work you have actually done.

Check yourself as you go:

    .venv/bin/pytest starter -q

Nothing here needs NumPy and nothing here needs the network. `math` is the
only import you should need.
"""

import math
from typing import Callable, Iterable, Sequence

Scalar = Callable[[float], float]


# --------------------------------------------------------------------------
# 1.1 -- the two building blocks
# --------------------------------------------------------------------------


def product(factors: Iterable[float]) -> float:
    """Multiply an iterable of numbers together and return the result.

    An empty iterable must give 1.0, not 0.0. One is the identity for
    multiplication, and it is also the honest answer to "how much does x
    change per unit of x" when nothing happens in between.

    >>> product([2.0, 3.0])
    6.0

    Approach: start a running total at 1.0 and multiply each factor into it.
    """
    return None


def gear_ratio(ratios: Iterable[float]) -> float:
    """The overall ratio of a gear train: every stage ratio multiplied.

    If gear A turns twice per turn of B, and B turns three times per turn of
    C, then A turns six times per turn of C.

    Approach: this is one line, and it calls `product`.
    """
    return None


# --------------------------------------------------------------------------
# 1.2 -- the measuring instruments, carried over from Days 108 and 109
# --------------------------------------------------------------------------


def central_difference(f: Scalar, x: float, h: float) -> float:
    """Estimate f'(x) as (f(x + h) - f(x - h)) / (2h).

    Raise `ValueError` if h is zero or negative -- a step of zero would
    divide by zero, and a negative step is almost certainly a typo rather
    than an intention.

    This is the tool that checks everything else in the lab, and it must know
    nothing about the chain rule for that check to mean anything.

    Approach: guard the step, then one subtraction over 2h. Note the
    denominator is 2h and not h; dividing by h is the most common way to get
    this exactly half right.
    """
    return None


def partial_difference(
    f: Callable[..., float], point: Sequence[float], index: int, h: float
) -> float:
    """Estimate one partial derivative of a function of several inputs.

    Nudge coordinate `index` by +h and by -h, hold every other coordinate
    still, and divide the difference by 2h. Raise `ValueError` on a
    non-positive h, as above.

    `f` is called as `f(*coordinates)`, so a two-input function is called
    `f(s, t)`.

    Approach: make two copies of `point` as lists, change one entry in each,
    and call f with each copy unpacked.
    """
    return None


# --------------------------------------------------------------------------
# 1.3 -- composition and the one-variable chain rule
# --------------------------------------------------------------------------


def compose(outer: Scalar, inner: Scalar) -> Scalar:
    """Return the FUNCTION x -> outer(inner(x)).

    Note that this returns a function, not a number. `inner` runs first even
    though it is written second.

    Approach: define a small function inside this one and return it.
    """
    return None


def chain_rule(
    d_outer: Scalar, inner: Scalar, d_inner: Scalar, x: float
) -> float:
    """The chain rule for one variable: dy/dx = dy/du x du/dx.

    The trap is in one word: the outer derivative is evaluated **at the inner
    value**, not at x. Compute u = inner(x) first, then multiply d_outer(u)
    by d_inner(x).

    With outer f(u) = u squared, inner g(x) = 3x + 1 and x = 2, the answer is
    2 x 7 x 3 = 42. If you get 12, you evaluated the outer derivative at x.

    Approach: two lines. Do not try to make it one.
    """
    return None


# --------------------------------------------------------------------------
# 1.4 -- chains of any depth
# --------------------------------------------------------------------------


def chain_values(stages: Sequence[Scalar], x: float) -> list[float]:
    """The forward pass: the input, then the output of each stage in turn.

    For n stages this returns n + 1 numbers, starting with x itself. Keeping
    x in the list means entry i is always the value that ARRIVES at stage i,
    which is what the next function needs.

    With the five stages in `dataset.FIVE_STAGES` starting from 1.0 the
    answer begins 1.0, 2.0, 5.0, 25.0, 5.0, ...

    Approach: a list starting with [x], then a loop that applies each stage
    to the running value and appends it.
    """
    return None


def chain_local_rates(
    stages: Sequence[Scalar], rates: Sequence[Scalar], x: float
) -> list[float]:
    """The local derivative of every stage, each evaluated at its own input.

    Raise `ValueError` if `stages` and `rates` are different lengths -- a
    chain with a missing derivative should be refused rather than silently
    truncated.

    For the five stages this gives 2.0, 1.0, 10.0, 0.1, 0.2. Stage 3's
    derivative is 2u and the u it sees is 5, so its rate is 10 -- not 2, and
    not anything computed at x.

    Approach: call `chain_values` first, then evaluate rate i at values[i].
    """
    return None


def chain_derivative(
    stages: Sequence[Scalar], rates: Sequence[Scalar], x: float
) -> float:
    """The derivative of the whole chain: every local rate multiplied.

    For the five stages starting at 1.0 this is 2 x 1 x 10 x 0.1 x 0.2 = 0.4,
    and the same chain collapses by hand to ln(2x + 3), whose derivative at
    x = 1 is 2/5. Two routes, one number.

    Approach: one line, calling `product` and `chain_local_rates`.
    """
    return None


def chain_function(stages: Sequence[Scalar]) -> Scalar:
    """Collapse a list of stages into the single function they compose.

    Returns a function, like `compose` does. This is what you feed to
    `central_difference` to check `chain_derivative`.

    Approach: define an inner function that loops the stages over a running
    value, and return it.
    """
    return None


def running_products(rates: Sequence[float]) -> list[float]:
    """The partial products of the local rates, taken from the OUTPUT end.

    Entry i must be the product of rates[i:], so entry 0 is the whole
    derivative and the last entry is just the final local rate. Return the
    list in stage order.

    This is what a backward pass is actually carrying as it walks: after k
    steps from the end, the number in hand is the product of the last k local
    rates.

    For 2, 1, 10, 0.1, 0.2 the answer is about 0.4, 0.2, 0.2, 0.02, 0.2.

    Approach: walk `rates` in reverse with a running total, appending as you
    go, then reverse the list you built.
    """
    return None


# --------------------------------------------------------------------------
# 1.5 -- more than one path
# --------------------------------------------------------------------------


def path_contributions(
    local_rates_per_path: Sequence[Sequence[float]],
) -> list[float]:
    """One number per path: the product of the local rates along that path.

    Approach: one product per path.
    """
    return None


def total_derivative(local_rates_per_path: Sequence[Sequence[float]]) -> float:
    """Multiply along each path, then ADD across paths.

    This is the half of the chain rule that gets dropped. If a variable
    reaches the output by two routes, changing it moves the output twice and
    both movements are real, so they add. For the two paths in the lab the
    contributions are 24 and 12 and the answer is 36 -- not 24, not 12, and
    not 288.

    Approach: sum the contributions.
    """
    return None


# --------------------------------------------------------------------------
# 1.6 -- products that collapse and products that blow up
# --------------------------------------------------------------------------


def repeated_product(factor: float, count: int) -> float:
    """Multiply `factor` by itself `count` times, one multiplication at a time.

    Raise `ValueError` if `count` is negative. A count of zero gives 1.0.

    Write it as a loop rather than as `factor ** count`: the loop is what a
    backward pass through `count` layers actually does.

    Approach: a running total and a `for _ in range(count)` loop.
    """
    return None


def order_of_magnitude(value: float) -> int:
    """floor(log10(|value|)) -- the exponent, ignoring the digits.

    Raise `ValueError` on zero, which has no order of magnitude.

    0.9 to the fiftieth is about 5.15e-3, so its order is -3. Asserting the
    order rather than the digits is the honest way to state a claim about a
    gradient vanishing: the scale is the lesson and the digits are float64
    rounding.

    Approach: `math.floor(math.log10(abs(value)))`, with the guard first.
    """
    return None

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

Both `examples/` and `starter/` contain modules called `autodiff`,
`chainrule`, `dataset` and `network`, 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 `autodiff` 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 ("autodiff", "chainrule", "dataset", "network", "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, the functions and every tolerance this lab compares against.

Read this file. Nothing here is tuned: every tolerance below is derived from
the error terms that actually govern the comparison being made, and the
arithmetic is written out beside it. A tolerance reached by running a test and
enlarging the number until it went green is a tolerance chosen by whatever bug
happened to exist at the time.

Almost every number in this lab is exact in float64. The chains, the two-path
example and the whole two-layer network were chosen so that the reader can
re-derive each one with a pen. Where a value is not exact -- a sine, a
logarithm, an exponential -- it is computed here from `math` rather than
written down as a literal, so nothing in this lab is a remembered constant.
"""

import math
from typing import Callable, NamedTuple

import numpy as np

# --------------------------------------------------------------------------
# Machine constants
# --------------------------------------------------------------------------

#: float64 machine epsilon, read from NumPy rather than trusted as a literal.
EPSILON: float = float(np.finfo(np.float64).eps)

# --------------------------------------------------------------------------
# The numerical step, and the three tolerances
# --------------------------------------------------------------------------

#: The central-difference step. Day 108 measured the bottom of the error U for
#: the central rule on e**x and found it in the 1e-7 to 1e-4 band; 1e-5 sits
#: inside it with room on both sides.
H: float = 1e-5

# The central difference (f(x+h) - f(x-h)) / (2h) carries two errors:
#
#   truncation  ~ (h**2 / 6) * |f'''(x)|   = 1.667e-11 * |f'''(x)|  at h = 1e-5
#   rounding    ~ EPSILON * |f(x)| / h     = 2.220e-11 * |f(x)|     at h = 1e-5
#
# No function differentiated in this lab has |f| or |f'''| above about 250 at
# the points used, so the bound is about 250 * (1.667e-11 + 2.220e-11), which
# is roughly 9.7e-9. The tolerance below is 1e-6, so there is about a
# hundredfold margin -- enough that ordinary float64 noise cannot trip it, and
# far too tight to hide a chain rule that dropped a factor or summed the wrong
# paths. The smallest true gradient this tolerance guards is 0.25, so a
# mistake would have to be smaller than four parts in a million to slip past.
#: Analytic gradient against a central difference. See the derivation above.
NUMERIC_TOL: float = 1e-6

#: The same comparison, stated relatively, for the handful of quantities whose
#: magnitude runs into the hundreds -- where an absolute 1e-6 would be
#: stricter than the arithmetic can honestly support.
NUMERIC_REL_TOL: float = 1e-6

# Two analytic computations of the same quantity -- for example the product of
# five local rates against the closed-form derivative of the whole chain --
# differ only in the order the multiplications happen, so they differ by a few
# units in the last place. At a magnitude of 1e2 one ulp is about 1.4e-14, and
# a chain of five products can accumulate a handful of them, so the honest
# bound is a few times 1e-13.
#: Analytic against analytic: rounding only, no truncation.
ANALYTIC_TOL: float = 1e-12

# --------------------------------------------------------------------------
# Gears: the whole idea, before any calculus
# --------------------------------------------------------------------------

#: Gear A turns twice for every turn of B; B turns three times for every turn
#: of C. So A turns six times per turn of C. The rates multiplied.
GEAR_RATIOS: tuple[float, ...] = (2.0, 3.0)
GEAR_RATIO_PRODUCT: float = 6.0

#: A longer train, to show the product keeping on multiplying.
GEAR_TRAIN: tuple[float, ...] = (2.0, 3.0, 1.5, 4.0)
GEAR_TRAIN_PRODUCT: float = 36.0

#: The same arithmetic with money instead of teeth. 1 unit of the first
#: currency buys 1.25 of the second, which buys 0.8 of the third, which buys
#: 150 of the fourth. These rates are invented for the arithmetic and are not
#: quoted from any market.
CURRENCY_RATES: tuple[float, ...] = (1.25, 0.8, 150.0)
CURRENCY_PRODUCT: float = 150.0

# --------------------------------------------------------------------------
# Functions used as the outer and inner halves of a composition
# --------------------------------------------------------------------------


def square(x: float) -> float:
    """x squared."""
    return x * x


def d_square(x: float) -> float:
    """The derivative of x squared."""
    return 2.0 * x


def line(x: float) -> float:
    """3x + 1."""
    return 3.0 * x + 1.0


def d_line(x: float) -> float:
    """The derivative of 3x + 1: a constant 3."""
    return 3.0


def half_negative_square(x: float) -> float:
    """-x squared over 2 -- the inside of a Gaussian bump."""
    return -0.5 * x * x


def d_half_negative_square(x: float) -> float:
    """The derivative of -x squared over 2."""
    return -x


def shifted_square(x: float) -> float:
    """x squared plus 1, which is never zero, so its logarithm is safe."""
    return x * x + 1.0


def reciprocal(x: float) -> float:
    """1 / x."""
    return 1.0 / x


def d_reciprocal(x: float) -> float:
    """The derivative of 1 / x."""
    return -1.0 / (x * x)


def one_plus_exp_negative(x: float) -> float:
    """1 + e to the minus x -- the denominator of the sigmoid."""
    return 1.0 + math.exp(-x)


def d_one_plus_exp_negative(x: float) -> float:
    """The derivative of 1 + e to the minus x."""
    return -math.exp(-x)


def double_plus_one(x: float) -> float:
    """2x + 1."""
    return 2.0 * x + 1.0


def d_double_plus_one(x: float) -> float:
    """The derivative of 2x + 1: a constant 2."""
    return 2.0


def d_tanh(x: float) -> float:
    """The derivative of tanh, written in terms of x rather than of tanh(x)."""
    t = math.tanh(x)
    return 1.0 - t * t


def d_ln(x: float) -> float:
    """The derivative of the natural logarithm."""
    return 1.0 / x


# --------------------------------------------------------------------------
# The one-variable chain rule: six compositions, each checked numerically
# --------------------------------------------------------------------------


class Composition(NamedTuple):
    """One composed function f(g(x)), with both halves and both derivatives.

    `exact` is the closed-form derivative at `x`, computed from `math` rather
    than written down as a literal, so nothing here is a remembered constant.
    """

    name: str
    outer: Callable[[float], float]
    d_outer: Callable[[float], float]
    inner: Callable[[float], float]
    d_inner: Callable[[float], float]
    x: float
    exact: float


COMPOSITIONS: tuple[Composition, ...] = (
    # (3x + 1) squared at x = 2. Inner is 7, outer rate is 2*7 = 14, inner
    # rate is 3, so the answer is 42 -- exact, and checkable in one line.
    Composition("square of a line", square, d_square, line, d_line, 2.0, 42.0),
    # sin(x squared) at x = 1.5. Rate is cos(2.25) * 3.
    Composition(
        "sine of a square",
        math.sin,
        math.cos,
        square,
        d_square,
        1.5,
        math.cos(2.25) * 3.0,
    ),
    # A Gaussian bump, e to the minus x squared over 2, at x = 0.8.
    Composition(
        "gaussian bump",
        math.exp,
        math.exp,
        half_negative_square,
        d_half_negative_square,
        0.8,
        math.exp(-0.32) * -0.8,
    ),
    # ln(x squared + 1) at x = 2. Inner is 5, outer rate is 1/5, inner rate
    # is 4, so the answer is 0.8 -- exact.
    Composition(
        "log of a shifted square",
        math.log,
        d_ln,
        shifted_square,
        d_square,
        2.0,
        0.8,
    ),
    # The sigmoid, written as 1 / (1 + e to the minus x), at x = 0. Inner is
    # 2, outer rate is -1/4, inner rate is -1, so the answer is 0.25 -- which
    # is the largest value the sigmoid's slope ever takes.
    Composition(
        "the sigmoid",
        reciprocal,
        d_reciprocal,
        one_plus_exp_negative,
        d_one_plus_exp_negative,
        0.0,
        0.25,
    ),
    # tanh(2x + 1) at x = -0.5, where the inner function is exactly 0 and
    # tanh's slope is exactly 1, so the answer is exactly 2.
    Composition(
        "tanh of a line",
        math.tanh,
        d_tanh,
        double_plus_one,
        d_double_plus_one,
        -0.5,
        2.0,
    ),
)

# --------------------------------------------------------------------------
# A chain of five functions, built so every local rate is a round number
# --------------------------------------------------------------------------

#: Applied in order, left to right: double, add 3, square, square root,
#: natural logarithm. Starting from x = 1 the values are 1, 2, 5, 25, 5,
#: ln 5, and the local rates are 2, 1, 10, 0.1, 0.2. Their product is 0.4.
#:
#: The whole chain collapses to ln(2x + 3), whose derivative is 2 / (2x + 3),
#: which at x = 1 is 2/5 = 0.4. Two routes, the same number.
FIVE_STAGES: tuple[Callable[[float], float], ...] = (
    lambda u: 2.0 * u,
    lambda u: u + 3.0,
    lambda u: u * u,
    math.sqrt,
    math.log,
)

FIVE_RATES: tuple[Callable[[float], float], ...] = (
    lambda u: 2.0,
    lambda u: 1.0,
    lambda u: 2.0 * u,
    lambda u: 0.5 / math.sqrt(u),
    lambda u: 1.0 / u,
)

FIVE_START: float = 1.0
FIVE_VALUES: tuple[float, ...] = (1.0, 2.0, 5.0, 25.0, 5.0, math.log(5.0))
FIVE_LOCAL_RATES: tuple[float, ...] = (2.0, 1.0, 10.0, 0.1, 0.2)
FIVE_DERIVATIVE: float = 0.4


def five_chain_closed_form(x: float) -> float:
    """The five stages collapsed by hand: ln(2x + 3)."""
    return math.log(2.0 * x + 3.0)


def d_five_chain_closed_form(x: float) -> float:
    """Its derivative: 2 / (2x + 3)."""
    return 2.0 / (2.0 * x + 3.0)


# --------------------------------------------------------------------------
# Two paths into one output: the case where contributions ADD
# --------------------------------------------------------------------------

#: x reaches the output twice: once through u = x squared, once through
#: v = 3x. The output is f = u * v.
#:
#:   df/dx = (df/du)(du/dx) + (df/dv)(dv/dx)
#:         = v * 2x         + u * 3
#:         = 3x * 2x        + x squared * 3
#:         = 6 x squared    + 3 x squared     = 9 x squared
#:
#: At x = 2 that is 36. The two path contributions are 24 and 12. Neither one
#: alone is the answer, and this is exactly the mistake the lab is built to
#: catch: taking the product along one path and stopping there.
TWO_PATH_X: float = 2.0
TWO_PATH_U: float = 4.0
TWO_PATH_V: float = 6.0
TWO_PATH_OUTPUT: float = 24.0
TWO_PATH_CONTRIBUTIONS: tuple[float, float] = (24.0, 12.0)
TWO_PATH_DERIVATIVE: float = 36.0


def two_path_direct(x: float) -> float:
    """The same function with the composition already done: 3 x cubed."""
    return 3.0 * x * x * x


def d_two_path_direct(x: float) -> float:
    """Its derivative: 9 x squared."""
    return 9.0 * x * x


# --------------------------------------------------------------------------
# The full multivariable chain rule: two inputs, two intermediates
# --------------------------------------------------------------------------

#: z = u squared + v squared, with u = s*t and v = s - t, at (s, t) = (2, 3).
#:
#:   u = 6,  v = -1,  z = 37
#:   dz/du = 12,  dz/dv = -2
#:   du/ds = t = 3,   du/dt = s = 2
#:   dv/ds = 1,       dv/dt = -1
#:
#:   dz/ds = 12*3 + (-2)*1  = 34
#:   dz/dt = 12*2 + (-2)*-1 = 26
#:
#: Every one of those is an integer, so the whole thing is exact in float64.
SURFACE_POINT: tuple[float, float] = (2.0, 3.0)
SURFACE_U: float = 6.0
SURFACE_V: float = -1.0
SURFACE_Z: float = 37.0
SURFACE_DZ_DU: float = 12.0
SURFACE_DZ_DV: float = -2.0
SURFACE_GRADIENT: tuple[float, float] = (34.0, 26.0)


def surface(s: float, t: float) -> float:
    """z as a function of s and t, with the intermediates substituted in."""
    u = s * t
    v = s - t
    return u * u + v * v


# --------------------------------------------------------------------------
# The tiny two-layer network, backpropagated by hand
# --------------------------------------------------------------------------

#: Half the natural logarithm of 3. tanh of this number is exactly 0.5 in
#: float64, and 1 - tanh squared is then exactly 0.75. Both facts are asserted
#: by the test suite rather than assumed.
#:
#: The bias of the second hidden unit is set to this value on purpose, so that
#: every number in the hand-worked backward pass is exact and can be checked
#: with a pen. Nothing about the chain rule depends on the choice; it is a
#: convenience for the reader, and it is declared rather than hidden.
HALF_LN3: float = 0.5 * math.log(3.0)

#: The two inputs.
NET_X1: float = 1.0
NET_X2: float = 2.0

#: Hidden unit A: weights and bias. Pre-activation is 1*1 + (-0.5)*2 + 0 = 0.
NET_WA1: float = 1.0
NET_WA2: float = -0.5
NET_BA: float = 0.0

#: Hidden unit B: pre-activation is -0.5*1 + 0.25*2 + HALF_LN3 = HALF_LN3.
NET_WB1: float = -0.5
NET_WB2: float = 0.25
NET_BB: float = HALF_LN3

#: The output layer: a linear combination of the two hidden activations.
NET_VA: float = 2.0
NET_VB: float = -3.0
NET_C: float = 1.0

#: The target the loss is measured against.
NET_TARGET: float = 1.0

#: The forward pass, every value exact.
NET_A_PRE: float = 0.0
NET_A: float = 0.0
NET_B_PRE: float = HALF_LN3
NET_B: float = 0.5
NET_OUT: float = -0.5
NET_LOSS: float = 2.25

#: The backward pass, every value exact. Worked out in full in
#: `06_backprop_by_hand.py` and asserted against the engine in the tests.
NET_GRADIENTS: dict[str, float] = {
    "out": -3.0,
    "c": -3.0,
    "vA": 0.0,  # zero, because hidden unit A output exactly 0
    "vB": -1.5,
    "a": -6.0,
    "b": 9.0,
    "a_pre": -6.0,
    "b_pre": 6.75,
    "wA1": -6.0,
    "wA2": -12.0,
    "bA": -6.0,
    "wB1": 6.75,
    "wB2": 13.5,
    "bB": 6.75,
    # x1 and x2 each reach the loss through BOTH hidden units, so each of
    # these is a sum over two paths, not a single product.
    "x1": -9.375,
    "x2": 4.6875,
}

#: The two per-path contributions to dL/dx1, which must be added.
NET_X1_CONTRIBUTIONS: tuple[float, float] = (-6.0, -3.375)

#: The names of the twelve parameters, in the order the scripts print them.
NET_PARAMETERS: tuple[str, ...] = (
    "wA1",
    "wA2",
    "bA",
    "wB1",
    "wB2",
    "bB",
    "vA",
    "vB",
    "c",
)

# --------------------------------------------------------------------------
# Products that collapse and products that blow up
# --------------------------------------------------------------------------

#: A gradient that passes through 50 layers, each contributing a local rate
#: slightly below or slightly above 1.
DECAY_FACTOR: float = 0.9
GROWTH_FACTOR: float = 1.1
CHAIN_LENGTH: int = 50
LONG_CHAIN_LENGTH: int = 200

#: Asserted as orders of magnitude rather than as exact values, because the
#: point is the scale and not the digits.
DECAY_ORDER: int = -3  # 0.9**50 is about 5.15e-3
GROWTH_ORDER: int = 2  # 1.1**50 is about 1.17e+2

#: A harsher pair, to show how quickly the arithmetic leaves the useful range.
#: 0.25 is not an arbitrary choice: it is the LARGEST slope the sigmoid ever
#: has, measured in `02_composition_and_the_chain_rule.py`. A stack of sigmoid
#: layers is multiplying numbers no bigger than this one.
SHARP_DECAY: float = 0.25
SHARP_GROWTH: float = 2.0

#: A middling decay, kept because its behaviour is genuinely surprising and
#: contradicts the obvious guess. See `07_vanishing_and_exploding.py`.
MILD_DECAY: float = 0.5

"""Exercise 4 -- backpropagate the tiny network by hand, then by your engine.

The forward pass and the graph builder are written for you, because they are
not the exercise. The exercise is the backward pass: writing out every local
derivative and multiplying along every path, with your own arithmetic, and
then watching your engine from exercise 2 produce the identical numbers.

The network is two inputs, two tanh hidden units, one linear output, a
squared-error loss, and nine parameters. Every quantity in both passes is
exact in float64, so you can check the whole thing with a pen. Read
`dataset.py` for the values and for why the bias of unit B is half the
natural logarithm of 3.
"""

import math
from typing import Sequence

import dataset as D
from autodiff import Value


def forward(x1: float, x2: float, params: Sequence[float]) -> dict[str, float]:
    """Run the network in plain floats. Written for you."""
    wA1, wA2, bA, wB1, wB2, bB, vA, vB, c = params
    a_pre = wA1 * x1 + wA2 * x2 + bA
    a = math.tanh(a_pre)
    b_pre = wB1 * x1 + wB2 * x2 + bB
    b = math.tanh(b_pre)
    out = vA * a + vB * b + c
    loss = (out - D.NET_TARGET) * (out - D.NET_TARGET)
    return {
        "a_pre": a_pre,
        "a": a,
        "b_pre": b_pre,
        "b": b,
        "out": out,
        "loss": loss,
    }


def loss_only(x1: float, x2: float, params: Sequence[float]) -> float:
    """Just the loss, for feeding to a central difference. Written for you."""
    return forward(x1, x2, params)["loss"]


def default_parameter_values() -> list[float]:
    """The nine parameters, in the documented order. Written for you."""
    return [
        D.NET_WA1,
        D.NET_WA2,
        D.NET_BA,
        D.NET_WB1,
        D.NET_WB2,
        D.NET_BB,
        D.NET_VA,
        D.NET_VB,
        D.NET_C,
    ]


def build_graph(
    x1: Value, x2: Value, params: Sequence[Value]
) -> dict[str, Value]:
    """The same network built out of your `Value` class. Written for you.

    This will not work until exercise 2 does, which is deliberate: the graph
    is only as good as the operations it is built from.
    """
    wA1, wA2, bA, wB1, wB2, bB, vA, vB, c = params
    a_pre = wA1 * x1 + wA2 * x2 + bA
    a = a_pre.tanh()
    b_pre = wB1 * x1 + wB2 * x2 + bB
    b = b_pre.tanh()
    out = vA * a + vB * b + c
    diff = out - D.NET_TARGET
    loss = diff * diff
    return {
        "x1": x1,
        "x2": x2,
        "wA1": wA1,
        "wA2": wA2,
        "bA": bA,
        "wB1": wB1,
        "wB2": wB2,
        "bB": bB,
        "vA": vA,
        "vB": vB,
        "c": c,
        "a_pre": a_pre,
        "a": a,
        "b_pre": b_pre,
        "b": b,
        "out": out,
        "loss": loss,
    }


# --------------------------------------------------------------------------
# 4a -- the backward pass, by hand
# --------------------------------------------------------------------------


def hand_gradients() -> dict[str, float]:
    """Return every gradient in the network, computed by your own arithmetic.

    The keys must be exactly these sixteen:

        out, c, vA, vB, a, b, a_pre, b_pre,
        wA1, wA2, bA, wB1, wB2, bB, x1, x2

    Work from the end backwards. The seed is d(loss)/d(loss) = 1, and:

        loss = (out - target) squared    ->  d loss/d out = 2 x (out - target)
        out  = vA*a + vB*b + c           ->  local rates 1, a, b, vA, vB
        a    = tanh(a_pre)               ->  local rate 1 - a squared
        a_pre = wA1*x1 + wA2*x2 + bA     ->  local rates x1, x2, 1, wA1, wA2

    Two of the sixteen are worth thinking about before you write them.

    `vA` multiplies an activation of exactly 0, so ask yourself what nudging
    vA does to the output before you reach for a formula.

    `x1` and `x2` each reach the loss through BOTH hidden units. Each of
    those two gradients is therefore a SUM of two products, not a single
    product. This is the one place in the lab where getting it wrong still
    produces a completely reasonable-looking number, and the test that checks
    it compares against a central difference, which has no opinion about
    which path you meant.

    Approach: read the forward values out of `forward(...)` or straight from
    `dataset.py`, then write one line per gradient in the order above. Do not
    call your engine here -- the whole point is that the two agree.
    """
    return None


# --------------------------------------------------------------------------
# 4b -- the same thing from the engine, in one backward pass
# --------------------------------------------------------------------------


def engine_gradients() -> dict[str, float]:
    """Return every gradient in the network, from ONE backward pass.

    Build `Value` nodes for x1, x2 and the nine parameters, pass them to
    `build_graph`, call `.backward()` on the loss node, and read `.grad` off
    every node in the returned dictionary.

    If exercise 2 is correct, this will match `hand_gradients` bit for bit --
    not approximately, exactly -- because it performs the same
    multiplications in the same order on the same exact values.

    Approach: four lines, ending in a dictionary comprehension over the
    dictionary that `build_graph` returns.
    """
    return None


# --------------------------------------------------------------------------
# 4c -- the numerical cross-check
# --------------------------------------------------------------------------


def numeric_parameter_gradients(h: float) -> dict[str, float]:
    """Every parameter gradient by central difference, keyed by name.

    Nudge one parameter at a time by +h and -h, holding the rest still, and
    divide the change in the loss by 2h. The names are in
    `dataset.NET_PARAMETERS`, in the same order as
    `default_parameter_values()`.

    These will NOT match the other two exactly, and they should not. A
    central difference has its own error, which is why the test suite
    compares it with a tolerance a million times looser than the one it uses
    between your hand computation and your engine.

    Approach: a loop over `enumerate(D.NET_PARAMETERS)`, two copies of the
    parameter list per step, and `loss_only` for the evaluations.
    """
    return None

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

Run from the lab directory:

    .venv/bin/pytest starter -q

On an untouched checkout this reports one pass and everything else skipped.
A skip means "not attempted". A failure means "attempted and wrong", and the
message shows your answer next to the real one so you can see the gap rather
than guess at it.

Nothing in here checks that a function exists or that a file is present.
Every test runs your code and compares a value.
"""

import math

import pytest

import answers
import autodiff as A
import chainrule as C
import dataset as D
import network as N

# --------------------------------------------------------------------------
# The skip machinery
# --------------------------------------------------------------------------


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


def attempt(fn, what):
    """Call something that may not be written yet, and skip if it is not.

    An unwritten `__add__` returns None, so the next operation on it raises a
    TypeError or an AttributeError. That is "not attempted", not "wrong".
    """
    try:
        result = fn()
    except (TypeError, AttributeError, NotImplementedError):
        pytest.skip(f"not attempted yet: {what}")
    if result is None:
        pytest.skip(f"not attempted yet: {what}")
    return result


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():
    """One test that 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 -- the fourteen functions
# --------------------------------------------------------------------------


def test_1_product_multiplies():
    assert need(C.product([2.0, 3.0, 4.0]), "product") == 24.0


def test_1_product_of_nothing_is_one():
    assert need(C.product([]), "product") == 1.0, "the empty product must be 1.0"


def test_1_gear_ratio_of_two_stages():
    assert need(C.gear_ratio(D.GEAR_RATIOS), "gear_ratio") == 6.0


def test_1_gear_ratio_of_four_stages():
    assert need(C.gear_ratio(D.GEAR_TRAIN), "gear_ratio") == 36.0


def test_1_central_difference_on_a_parabola_is_exact():
    got = need(C.central_difference(D.square, 3.0, 0.1), "central_difference")
    close(got, 6.0, 1e-11, "central difference of x squared at 3")


def test_1_central_difference_divides_by_two_h_not_h():
    # Dividing by h instead of 2h gives exactly double. This is the most
    # common way to get this function half right.
    got = need(C.central_difference(D.square, 3.0, 0.1), "central_difference")
    assert abs(got - 12.0) > 1.0, "you divided by h instead of 2h"


def test_1_central_difference_refuses_a_zero_step():
    try:
        C.central_difference(D.square, 1.0, 0.0)
    except ValueError:
        return
    except ZeroDivisionError:
        pytest.fail("raise ValueError for a non-positive h, not ZeroDivisionError")
    pytest.skip("not attempted yet: central_difference guard")


def test_1_central_difference_refuses_a_negative_step():
    try:
        C.central_difference(D.square, 1.0, -1e-5)
    except ValueError:
        return
    pytest.skip("not attempted yet: central_difference guard")


def test_1_partial_difference_on_the_surface():
    got = need(
        C.partial_difference(D.surface, D.SURFACE_POINT, 0, D.H),
        "partial_difference",
    )
    close(got, 34.0, D.NUMERIC_TOL, "dz/ds at (2, 3)")


def test_1_partial_difference_holds_the_other_coordinate_still():
    got = need(
        C.partial_difference(D.surface, D.SURFACE_POINT, 1, D.H),
        "partial_difference",
    )
    close(got, 26.0, D.NUMERIC_TOL, "dz/dt at (2, 3)")


def test_1_compose_returns_a_function_that_runs_the_inner_one_first():
    composed = need(C.compose(D.square, D.line), "compose")
    assert composed(2.0) == 49.0, "inner runs first: (3*2 + 1) squared"


def test_1_compose_is_not_commutative():
    other = need(C.compose(D.line, D.square), "compose")
    assert other(2.0) == 13.0, "3*(2 squared) + 1"


@pytest.mark.parametrize("case", D.COMPOSITIONS, ids=lambda c: c.name)
def test_1_chain_rule_matches_the_closed_form(case):
    got = need(
        C.chain_rule(case.d_outer, case.inner, case.d_inner, case.x),
        "chain_rule",
    )
    close(got, case.exact, D.ANALYTIC_TOL, f"chain rule on {case.name}")


@pytest.mark.parametrize("case", D.COMPOSITIONS, ids=lambda c: c.name)
def test_1_chain_rule_matches_a_measurement(case):
    got = need(
        C.chain_rule(case.d_outer, case.inner, case.d_inner, case.x),
        "chain_rule",
    )
    measured = (
        case.outer(case.inner(case.x + D.H)) - case.outer(case.inner(case.x - D.H))
    ) / (2.0 * D.H)
    close(got, measured, D.NUMERIC_TOL, f"chain rule vs measurement, {case.name}")


def test_1_chain_rule_evaluates_the_outer_derivative_at_the_inner_value():
    got = need(C.chain_rule(D.d_square, D.line, D.d_line, 2.0), "chain_rule")
    assert got != 12.0, (
        "12.0 means the outer derivative was evaluated at x rather than at "
        "u = inner(x). The answer is 42.0."
    )


def test_1_chain_values_returns_one_more_than_the_stage_count():
    got = need(C.chain_values(D.FIVE_STAGES, D.FIVE_START), "chain_values")
    assert len(got) == 6, "n stages give n + 1 values, starting with x itself"


def test_1_chain_values_are_the_documented_ones():
    got = need(C.chain_values(D.FIVE_STAGES, D.FIVE_START), "chain_values")
    assert got == list(D.FIVE_VALUES)


def test_1_chain_local_rates_are_the_documented_ones():
    got = need(
        C.chain_local_rates(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START),
        "chain_local_rates",
    )
    assert got == list(D.FIVE_LOCAL_RATES), (
        "each rate is evaluated at the value ARRIVING at its stage"
    )


def test_1_chain_local_rates_refuses_a_mismatched_length():
    try:
        C.chain_local_rates(D.FIVE_STAGES, D.FIVE_RATES[:2], 1.0)
    except ValueError:
        return
    except (TypeError, IndexError):
        pytest.fail("raise ValueError when stages and rates disagree in length")
    pytest.skip("not attempted yet: chain_local_rates guard")


def test_1_chain_derivative_is_the_product_of_the_local_rates():
    got = need(
        C.chain_derivative(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START),
        "chain_derivative",
    )
    close(got, 0.4, D.ANALYTIC_TOL, "the five-stage derivative")


def test_1_chain_function_collapses_the_stages():
    composed = need(C.chain_function(D.FIVE_STAGES), "chain_function")
    close(composed(1.0), math.log(5.0), D.ANALYTIC_TOL, "the five-stage value")


def test_1_chain_derivative_matches_a_measurement():
    composed = need(C.chain_function(D.FIVE_STAGES), "chain_function")
    analytic = need(
        C.chain_derivative(D.FIVE_STAGES, D.FIVE_RATES, 2.0), "chain_derivative"
    )
    measured = (composed(2.0 + D.H) - composed(2.0 - D.H)) / (2.0 * D.H)
    close(analytic, measured, D.NUMERIC_TOL, "chain derivative at x = 2")


def test_1_running_products_starts_at_the_whole_derivative():
    got = need(C.running_products(list(D.FIVE_LOCAL_RATES)), "running_products")
    close(got[0], 0.4, D.ANALYTIC_TOL, "the first running product")


def test_1_running_products_ends_at_the_last_local_rate():
    got = need(C.running_products(list(D.FIVE_LOCAL_RATES)), "running_products")
    assert got[-1] == 0.2, "the last entry is the final local rate alone"


def test_1_running_products_has_one_entry_per_rate():
    got = need(C.running_products(list(D.FIVE_LOCAL_RATES)), "running_products")
    assert len(got) == 5


def test_1_path_contributions_multiplies_along_each_path():
    paths = [[6.0, 4.0], [4.0, 3.0]]
    got = need(C.path_contributions(paths), "path_contributions")
    assert got == [24.0, 12.0]


def test_1_total_derivative_adds_across_paths():
    paths = [[6.0, 4.0], [4.0, 3.0]]
    got = need(C.total_derivative(paths), "total_derivative")
    assert got == 36.0, (
        "24 and 12 are added, not multiplied and not chosen between"
    )


def test_1_total_derivative_matches_a_measurement():
    paths = [[6.0, 4.0], [4.0, 3.0]]
    got = need(C.total_derivative(paths), "total_derivative")
    measured = (D.two_path_direct(2.0 + D.H) - D.two_path_direct(2.0 - D.H)) / (
        2.0 * D.H
    )
    close(got, measured, D.NUMERIC_TOL, "the two-path derivative")


def test_1_repeated_product_of_fifty_nine_tenths():
    got = need(C.repeated_product(0.9, 50), "repeated_product")
    close(got, 0.9**50, 1e-18, "0.9 to the fiftieth")


def test_1_repeated_product_of_zero_factors_is_one():
    assert need(C.repeated_product(0.9, 0), "repeated_product") == 1.0


def test_1_repeated_product_refuses_a_negative_count():
    try:
        C.repeated_product(0.9, -1)
    except ValueError:
        return
    pytest.skip("not attempted yet: repeated_product guard")


@pytest.mark.parametrize(
    "value,expected", [(1.0, 0), (9.99, 0), (10.0, 1), (0.1, -1), (-500.0, 2)]
)
def test_1_order_of_magnitude_reads_the_exponent(value, expected):
    assert need(C.order_of_magnitude(value), "order_of_magnitude") == expected


def test_1_order_of_magnitude_refuses_zero():
    try:
        C.order_of_magnitude(0.0)
    except ValueError:
        return
    except ValueError:
        return
    pytest.skip("not attempted yet: order_of_magnitude guard")


# --------------------------------------------------------------------------
# Exercise 2 -- the engine
# --------------------------------------------------------------------------


def test_2a_addition_computes_the_value():
    out = attempt(lambda: A.Value(3.0) + A.Value(4.0), "Value.__add__")
    assert out.data == 7.0


def test_2a_addition_passes_the_gradient_through():
    def build():
        a, b = A.Value(3.0), A.Value(4.0)
        c = a + b
        c.backward()
        return (a.grad, b.grad) if a.grad or b.grad else None

    grads = attempt(build, "Value.__add__ and backward")
    assert grads == (1.0, 1.0)


def test_2a_a_plain_float_can_be_added():
    out = attempt(lambda: A.Value(3.0) + 4.0, "Value.__add__ with a float")
    assert out.data == 7.0


def test_2b_multiplication_computes_the_value():
    out = attempt(lambda: A.Value(3.0) * A.Value(4.0), "Value.__mul__")
    assert out.data == 12.0


def test_2b_each_factors_local_rate_is_the_other_factor():
    def build():
        a, b = A.Value(3.0), A.Value(4.0)
        c = a * b
        c.backward()
        return (a.grad, b.grad) if a.grad or b.grad else None

    grads = attempt(build, "Value.__mul__ and backward")
    assert grads == (4.0, 3.0)


def test_2b_a_value_used_twice_accumulates_both_contributions():
    def build():
        x = A.Value(3.0)
        y = x + x
        y.backward()
        return x.grad or None

    grad = attempt(build, "Value.__add__ and backward")
    assert grad == 2.0, (
        "1.0 means the gradient was assigned instead of accumulated. "
        "Use += in every backward step."
    )


def test_2b_multiplying_a_value_by_itself_gives_the_power_rule():
    def build():
        x = A.Value(3.0)
        y = x * x
        y.backward()
        return x.grad or None

    grad = attempt(build, "Value.__mul__ and backward")
    assert grad == 6.0, "d/dx of x squared at 3 is 6, not 3"


def test_2b_a_value_cubed_gives_the_power_rule_too():
    def build():
        x = A.Value(2.0)
        y = x * x * x
        y.backward()
        return x.grad or None

    grad = attempt(build, "Value.__mul__ and backward")
    assert grad == 12.0


def test_2c_tanh_computes_the_value():
    out = attempt(lambda: A.Value(0.6).tanh(), "Value.tanh")
    close(out.data, math.tanh(0.6), D.ANALYTIC_TOL, "tanh(0.6)")


def test_2c_tanh_gradient_is_one_minus_tanh_squared():
    def build():
        z = A.Value(0.6)
        t = z.tanh()
        t.backward()
        return z.grad or None

    grad = attempt(build, "Value.tanh and backward")
    close(grad, 1.0 - math.tanh(0.6) ** 2, D.ANALYTIC_TOL, "tanh's slope at 0.6")


def test_2c_tanh_slope_at_zero_is_exactly_one():
    def build():
        z = A.Value(0.0)
        z.tanh().backward()
        return z.grad or None

    assert attempt(build, "Value.tanh and backward") == 1.0


def test_2c_tanh_at_half_ln_three_is_exactly_a_half():
    out = attempt(lambda: A.Value(D.HALF_LN3).tanh(), "Value.tanh")
    assert out.data == 0.5


def test_2d_topological_order_puts_children_before_parents():
    def build():
        p, q = A.Value(2.0), A.Value(-3.0)
        r = p * q
        out = r + p
        return A.topological_order(out)

    order = attempt(build, "topological_order")
    position = {id(node): i for i, node in enumerate(order)}
    for node in order:
        for child in node._children:
            assert position[id(child)] < position[id(node)], (
                "every node must come after everything it was computed from"
            )


def test_2d_a_node_used_twice_appears_once():
    def build():
        p = A.Value(2.0)
        out = p * p + p
        order = A.topological_order(out)
        return (order, p) if order is not None else None

    order, p = attempt(build, "topological_order")
    assert sum(1 for node in order if node is p) == 1


def test_2d_graph_size_counts_the_nodes():
    def build():
        p = A.Value(2.0)
        out = p * p + p
        return A.graph_size(out)

    # p, (p*p) and (p*p + p). The reused p is counted once, not twice.
    assert attempt(build, "graph_size") == 3


def test_2d_the_order_is_iterative_and_survives_ten_thousand_nodes():
    def build():
        node = A.Value(1.0)
        for _ in range(10_000):
            node = node * 1.0
        return A.graph_size(node)

    try:
        size = attempt(build, "topological_order on a deep graph")
    except RecursionError:
        pytest.fail(
            "a recursive topological sort overflows here; write it iteratively"
        )
    assert size > 10_000


def test_2e_backward_can_be_run_twice_with_the_same_answer():
    def build():
        x = A.Value(3.0)
        y = x * x
        y.backward()
        first = x.grad
        y.backward()
        return (first, x.grad) if first else None

    first, second = attempt(build, "Value.backward")
    assert first == second, "backward must zero the gradients before it runs"


def test_2e_the_output_seeds_its_own_gradient_with_one():
    def build():
        x = A.Value(3.0)
        y = x * 2.0
        y.backward()
        return y.grad or None

    assert attempt(build, "Value.backward") == 1.0


def test_2f_dual_addition_adds_both_parts():
    d = attempt(lambda: A.Dual(3.0, 1.0) + A.Dual(4.0, 0.0), "Dual.__add__")
    assert (d.value, d.dot) == (7.0, 1.0)


def test_2f_dual_multiplication_uses_the_product_rule():
    d = attempt(lambda: A.Dual(3.0, 1.0) * A.Dual(4.0, 0.0), "Dual.__mul__")
    assert (d.value, d.dot) == (12.0, 4.0)


def test_2f_dual_tanh_scales_the_derivative():
    d = attempt(lambda: A.Dual(0.6, 1.0).tanh(), "Dual.tanh")
    close(d.dot, 1.0 - math.tanh(0.6) ** 2, D.ANALYTIC_TOL, "Dual tanh slope")


def test_2f_an_unseeded_dual_reports_nothing():
    d = attempt(lambda: A.Dual(3.0, 0.0) * A.Dual(4.0, 0.0), "Dual.__mul__")
    assert d.dot == 0.0


# --------------------------------------------------------------------------
# Exercise 3 -- the two modes and their cost
# --------------------------------------------------------------------------

CUBIC = ("x cubed - 2x", lambda v: v[0] * v[0] * v[0] + (-2.0) * v[0], [1.5])
PAIR = ("(xy + x)(y + 3)", lambda v: (v[0] * v[1] + v[0]) * (v[1] + 3.0), [2.0, -1.0])


@pytest.mark.parametrize("name,build,point", [CUBIC, PAIR], ids=["cubic", "pair"])
def test_3_reverse_mode_matches_a_measurement(name, build, point):
    result = attempt(
        lambda: A.reverse_mode_gradient(build, point), "reverse_mode_gradient"
    )
    grads, _ = result

    def plain(vals):
        return build([A.Value(v) for v in vals]).data

    for i, got in enumerate(grads):
        ahead, behind = list(point), list(point)
        ahead[i] += D.H
        behind[i] -= D.H
        measured = (plain(ahead) - plain(behind)) / (2.0 * D.H)
        close(got, measured, D.NUMERIC_TOL, f"{name}, input {i}")


@pytest.mark.parametrize("name,build,point", [CUBIC, PAIR], ids=["cubic", "pair"])
def test_3_reverse_mode_always_uses_one_pass(name, build, point):
    _, passes = attempt(
        lambda: A.reverse_mode_gradient(build, point), "reverse_mode_gradient"
    )
    assert passes == 1


@pytest.mark.parametrize("name,build,point", [CUBIC, PAIR], ids=["cubic", "pair"])
def test_3_forward_mode_agrees_with_reverse_mode(name, build, point):
    forward, _ = attempt(
        lambda: A.forward_mode_gradient(build, point), "forward_mode_gradient"
    )
    reverse, _ = attempt(
        lambda: A.reverse_mode_gradient(build, point), "reverse_mode_gradient"
    )
    for got, want in zip(forward, reverse):
        close(got, want, D.ANALYTIC_TOL, f"{name}, forward vs reverse")


@pytest.mark.parametrize("name,build,point", [CUBIC, PAIR], ids=["cubic", "pair"])
def test_3_forward_mode_uses_one_pass_per_input(name, build, point):
    _, passes = attempt(
        lambda: A.forward_mode_gradient(build, point), "forward_mode_gradient"
    )
    assert passes == len(point)


def test_3_numeric_gradient_uses_two_evaluations_per_input():
    def plain(vals):
        return vals[0] * vals[0] + vals[1]

    _, passes = attempt(
        lambda: A.numeric_gradient(plain, [1.0, 2.0], D.H), "numeric_gradient"
    )
    assert passes == 4


# --------------------------------------------------------------------------
# Exercise 4 -- the network
# --------------------------------------------------------------------------


@pytest.mark.parametrize("key", sorted(D.NET_GRADIENTS))
def test_4a_the_hand_worked_gradients(key):
    grads = need(N.hand_gradients(), "network.hand_gradients")
    assert key in grads, f"hand_gradients is missing the key {key!r}"
    close(grads[key], D.NET_GRADIENTS[key], D.ANALYTIC_TOL, f"d loss / d {key}")


@pytest.mark.parametrize("key", sorted(D.NET_GRADIENTS))
def test_4b_the_engine_matches_the_hand_computation_exactly(key):
    hand = need(N.hand_gradients(), "network.hand_gradients")
    engine = attempt(N.engine_gradients, "network.engine_gradients")
    assert engine[key] == hand[key], (
        f"d loss/d {key}: engine {engine[key]!r} vs hand {hand[key]!r}. "
        "These perform the same multiplications on the same exact values, "
        "so they should agree bit for bit."
    )


def test_4b_the_engine_gets_the_loss_right():
    engine = attempt(N.engine_gradients, "network.engine_gradients")
    assert "out" in engine


@pytest.mark.parametrize("key", D.NET_PARAMETERS)
def test_4c_the_numerical_gradients_agree_within_tolerance(key):
    numeric = need(
        N.numeric_parameter_gradients(D.H), "network.numeric_parameter_gradients"
    )
    close(numeric[key], D.NET_GRADIENTS[key], D.NUMERIC_TOL, f"numeric d/d {key}")


def test_4_the_input_gradient_is_a_sum_over_two_paths():
    grads = need(N.hand_gradients(), "network.hand_gradients")
    first, second = D.NET_X1_CONTRIBUTIONS
    assert grads["x1"] != first, (
        "-6.0 is the contribution through hidden unit A alone. x1 also "
        "reaches the loss through unit B, and the two are added."
    )
    assert grads["x1"] != second
    close(grads["x1"], first + second, D.ANALYTIC_TOL, "d loss / d x1")


# --------------------------------------------------------------------------
# Exercises 5 to 9 -- the forty predictions
# --------------------------------------------------------------------------

EXPECTED: dict[str, object] = {
    "gears_two_stage": 6.0,
    "gears_four_stage": 36.0,
    "empty_product": 1.0,
    "gear_order_matters": False,
    "composed_value_at_two": 49.0,
    "composed_other_order_at_two": 13.0,
    "chain_rule_at_two": 42.0,
    "chain_rule_mistake_at_two": 12.0,
    "sigmoid_slope_at_zero": 0.25,
    "sigmoid_slope_is_maximum": True,
    "tanh_of_line_slope": 2.0,
    "five_chain_value_count": 6,
    "five_chain_third_rate": 10.0,
    "five_chain_derivative": 0.4,
    "five_chain_closed_form_derivative": 0.4,
    "running_products_last": 0.2,
    "two_path_u_contribution": 24.0,
    "two_path_v_contribution": 12.0,
    "two_path_total": 36.0,
    "two_path_closed_form": 36.0,
    "surface_value": 37.0,
    "surface_dz_ds": 34.0,
    "surface_dz_dt": 26.0,
    "engine_x_plus_x_grad": 2.0,
    "engine_x_times_x_grad": 6.0,
    "engine_x_cubed_grad": 12.0,
    "tanh_slope_at_zero": 1.0,
    "tanh_at_half_ln_three": 0.5,
    "tanh_slope_at_half_ln_three": 0.75,
    "network_loss": 2.25,
    "network_d_out": -3.0,
    "network_d_vA": 0.0,
    "network_d_b_pre": 6.75,
    "network_d_wB2": 13.5,
    "network_d_x1": -9.375,
    "reverse_passes_for_25_inputs": 1,
    "forward_passes_for_25_inputs": 25,
    "numeric_passes_for_25_inputs": 50,
    "decay_order": -3,
    "growth_order": 2,
    "half_to_the_fiftieth_vanishes": False,
    "quarter_to_the_fiftieth_vanishes": True,
}

HINTS: dict[str, str] = {
    "chain_rule_mistake_at_two": (
        "This one asks for the WRONG answer on purpose: the value you get by "
        "evaluating the outer derivative at x rather than at u."
    ),
    "network_d_vA": (
        "vA multiplies hidden unit A's activation, which is exactly 0. Nudging "
        "vA therefore moves the output by exactly nothing."
    ),
    "network_d_x1": (
        "x1 reaches the loss through BOTH hidden units. Add the two path "
        "contributions rather than picking one."
    ),
    "half_to_the_fiftieth_vanishes": (
        "0.5 to the fiftieth is four times float64's epsilon, so it is four "
        "representable gaps wide and still shifts a weight of 1. It takes "
        "three more halvings to disappear."
    ),
    "two_path_total": (
        "The two contributions are added. Not multiplied, and not chosen "
        "between."
    ),
}


@pytest.mark.parametrize("key", sorted(EXPECTED))
def test_5_to_9_predictions(key):
    got = need(answers.ANSWERS.get(key), f"answers.ANSWERS[{key!r}]")
    want = EXPECTED[key]
    hint = HINTS.get(key, "")
    if isinstance(want, bool) or isinstance(want, int):
        assert got == want, f"{key}: your answer {got!r}, expected {want!r}. {hint}"
    else:
        assert abs(float(got) - want) < D.ANALYTIC_TOL, (
            f"{key}: your answer {got!r}, expected {want!r}. {hint}"
        )


def test_every_answer_key_is_still_present():
    missing = sorted(set(EXPECTED) - set(answers.ANSWERS))
    assert not missing, f"answers.py is missing these keys: {missing}"

#!/usr/bin/env bash
# Tests for the Day 110 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:
#
#   * rates multiply -- two gears at 2 and 3 give 6, and a four-stage train
#     gives 36 -- and the same arithmetic drives every chain in the lab;
#   * the one-variable chain rule agrees with Day 108's central difference on
#     six compositions, and evaluating the outer derivative at x instead of at
#     u is asserted to give the wrong answer rather than merely warned about;
#   * a chain of five functions has derivative 2 x 1 x 10 x 0.1 x 0.2 = 0.4,
#     which the collapsed formula ln(2x + 3) and a measurement both confirm;
#   * when a variable reaches the output by two paths the contributions ADD --
#     24 + 12 = 36 -- and the suite asserts that neither single path and not
#     their product matches the measurement;
#   * a reverse-mode autodiff engine written from scratch reproduces every
#     gradient, agrees with forward mode to the last bit, and agrees with a
#     central difference to about a part in a billion;
#   * a two-layer network is backpropagated by hand, by the engine and by
#     central differences, and all three agree on all sixteen gradients;
#   * fifty factors of 0.9 collapse and fifty of 1.1 blow up, asserted as
#     orders of magnitude rather than digits -- and one measured result that
#     contradicts the naive story is reported rather than hidden;
#   * 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 110 — Rates Multiply"
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_gears_and_rates 02_composition_and_the_chain_rule \
              03_deeper_chains 04_two_paths_add 05_the_value_engine \
              06_backprop_by_hand 07_vanishing_and_exploding; 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 200 ]; then
  check "the reference suite ran at least 200 tests (ran ${ref_passed})" "yes"
else
  check "the reference suite ran at least 200 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 `autodiff`,
# `chainrule`, `dataset` and `network`, 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 dataset as D
import network as N
from autodiff import (
    Dual,
    Value,
    forward_mode_gradient,
    graph_size,
    numeric_gradient,
    reverse_mode_gradient,
    topological_order,
)
from chainrule import (
    central_difference,
    chain_derivative,
    chain_function,
    chain_local_rates,
    chain_rule,
    chain_values,
    compose,
    gear_ratio,
    order_of_magnitude,
    partial_difference,
    path_contributions,
    product,
    repeated_product,
    running_products,
    total_derivative,
    wrong_single_path_derivative,
)

# -- rates multiply ---------------------------------------------------------
print("gears_two", gear_ratio(D.GEAR_RATIOS))
print("gears_four", gear_ratio(D.GEAR_TRAIN))
print("currency", product(D.CURRENCY_RATES))
print("empty_product", product([]))
print("gear_order_irrelevant", gear_ratio(D.GEAR_TRAIN) == gear_ratio(tuple(reversed(D.GEAR_TRAIN))))

# -- composition ------------------------------------------------------------
print("composed_value", compose(D.square, D.line)(2.0))
print("composed_other_order", compose(D.line, D.square)(2.0))
print("chain_rule_correct", chain_rule(D.d_square, D.line, D.d_line, 2.0))
print("chain_rule_mistake", D.d_square(2.0) * D.d_line(2.0))
gaps = []
for case in D.COMPOSITIONS:
    analytic = chain_rule(case.d_outer, case.inner, case.d_inner, case.x)
    measured = central_difference(compose(case.outer, case.inner), case.x, D.H)
    gaps.append(abs(analytic - measured))
    print(f"exact_{case.name.replace(' ', '_')}", abs(analytic - case.exact) < D.ANALYTIC_TOL)
print("all_six_match_measurement", all(g < D.NUMERIC_TOL for g in gaps))
print("worst_composition_gap", f"{max(gaps):.3e}")
print("sigmoid_slope_at_zero", chain_rule(D.COMPOSITIONS[4].d_outer, D.COMPOSITIONS[4].inner, D.COMPOSITIONS[4].d_inner, 0.0))
print("sigmoid_quarter_is_max", all(
    chain_rule(D.COMPOSITIONS[4].d_outer, D.COMPOSITIONS[4].inner,
               D.COMPOSITIONS[4].d_inner, x) < 0.25
    for x in (-4.0, -2.0, -0.5, 0.5, 2.0, 4.0)))
print("tanh_of_line_slope", chain_rule(D.COMPOSITIONS[5].d_outer, D.COMPOSITIONS[5].inner, D.COMPOSITIONS[5].d_inner, -0.5))

# -- deep chains ------------------------------------------------------------
values = chain_values(D.FIVE_STAGES, D.FIVE_START)
rates = chain_local_rates(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START)
print("five_values", "|".join(f"{v:g}" for v in values))
print("five_rates", "|".join(f"{r:g}" for r in rates))
print("five_derivative", chain_derivative(D.FIVE_STAGES, D.FIVE_RATES, D.FIVE_START))
print("five_closed_form", D.d_five_chain_closed_form(1.0))
print("five_measured", f"{central_difference(chain_function(D.FIVE_STAGES), 1.0, D.H):.9f}")
carried = running_products(rates)
print("running_first", f"{carried[0]:.10g}")
print("running_last", carried[-1])
print("orders_differ_by_rounding", abs(product(rates) - carried[0]) < 4.0 * D.EPSILON)
print("orders_are_not_identical", product(rates) != carried[0])

# -- two paths --------------------------------------------------------------
paths = [[D.TWO_PATH_V, 2.0 * D.TWO_PATH_X], [D.TWO_PATH_U, 3.0]]
measured_two = central_difference(D.two_path_direct, D.TWO_PATH_X, D.H)
print("two_path_contributions", "|".join(f"{c:g}" for c in path_contributions(paths)))
print("two_path_sum", total_derivative(paths))
print("two_path_closed_form", D.d_two_path_direct(D.TWO_PATH_X))
print("two_path_sum_matches", abs(total_derivative(paths) - measured_two) < D.NUMERIC_TOL)
print("path_a_alone_is_wrong", abs(wrong_single_path_derivative(paths, 0) - measured_two) > 1.0)
print("path_b_alone_is_wrong", abs(wrong_single_path_derivative(paths, 1) - measured_two) > 1.0)
print("product_of_paths_is_wrong", abs(24.0 * 12.0 - measured_two) > 1.0)
print("surface_z", D.surface(*D.SURFACE_POINT))
print("surface_ds", D.SURFACE_DZ_DU * 3.0 + D.SURFACE_DZ_DV * 1.0)
print("surface_dt", D.SURFACE_DZ_DU * 2.0 + D.SURFACE_DZ_DV * -1.0)
print("surface_ds_measured", f"{partial_difference(D.surface, D.SURFACE_POINT, 0, D.H):.9f}")
print("surface_dt_measured", f"{partial_difference(D.surface, D.SURFACE_POINT, 1, D.H):.9f}")

# -- the engine -------------------------------------------------------------
a, b = Value(3.0), Value(4.0)
c = a * b
c.backward()
print("mul_grads", f"{a.grad:g}|{b.grad:g}")
x = Value(3.0)
(x + x).backward()
print("used_twice_grad", x.grad)
x2 = Value(3.0)
(x2 * x2).backward()
print("squared_grad", x2.grad)
z = Value(0.0)
z.tanh().backward()
print("tanh_slope_at_zero", z.grad)
z2 = Value(D.HALF_LN3)
t2 = z2.tanh()
t2.backward()
print("tanh_at_half_ln3", t2.data)
print("tanh_slope_at_half_ln3", z2.grad)
p, q = Value(2.0), Value(-3.0)
order = topological_order((p * q) + p)
position = {id(n): i for i, n in enumerate(order)}
print("topo_children_first", all(
    position[id(child)] < position[id(node)]
    for node in order for child in node._children))
deep = Value(1.0)
for _ in range(10000):
    deep = deep * 1.0
deep.backward()
print("deep_graph_nodes", graph_size(deep))
print("deep_graph_grad", deep.grad)
d = Dual(3.0, 1.0) * Dual(4.0, 0.0)
print("dual_product_rule", f"{d.value:g}|{d.dot:g}")

# -- the network ------------------------------------------------------------
fw = N.forward(D.NET_X1, D.NET_X2, N.default_parameter_values())
hand = N.hand_gradients()
engine = N.engine_gradients()
numeric = {**N.numeric_parameter_gradients(D.H), **N.numeric_input_gradients(D.H)}
print("net_a", fw["a"])
print("net_b", fw["b"])
print("net_out", fw["out"])
print("net_loss", fw["loss"])
print("net_b_slope_exact", (1.0 - fw["b"] * fw["b"]) == 0.75)
print("net_d_out", hand["out"])
print("net_d_vA_is_zero", hand["vA"] == 0.0)
print("net_d_vA_repr", repr(hand["vA"]))
print("net_d_vB", hand["vB"])
print("net_d_b_pre", hand["b_pre"])
print("net_d_wA2", hand["wA2"])
print("net_d_wB2", hand["wB2"])
print("net_d_x1", hand["x1"])
print("net_d_x2", hand["x2"])
print("net_x1_contributions", "|".join(f"{v:g}" for v in D.NET_X1_CONTRIBUTIONS))
print("net_x1_single_path_wrong", abs(D.NET_X1_CONTRIBUTIONS[0] - numeric["x1"]) > 1.0)
print("net_engine_equals_hand", all(engine[k] == hand[k] for k in D.NET_GRADIENTS))
print("net_numeric_agrees", all(abs(engine[k] - numeric[k]) < D.NUMERIC_TOL for k in numeric))
print("net_worst_numeric_gap", f"{max(abs(engine[k] - numeric[k]) for k in numeric):.3e}")
fwd_grads, fwd_passes = N.forward_mode_parameter_gradients()
print("net_forward_mode_agrees", all(abs(fwd_grads[k] - hand[k]) < D.ANALYTIC_TOL for k in fwd_grads))
print("net_forward_passes", fwd_passes)
print("net_parameter_count", len(D.NET_PARAMETERS))

# -- cost -------------------------------------------------------------------
def many(vals):
    total = vals[0] * 1.0
    for v in vals[1:]:
        total = total + v * v
    return (total * 0.1).tanh()

point = [0.1 * (i + 1) for i in range(25)]
r_grads, r_passes = reverse_mode_gradient(many, point)
f_grads, f_passes = forward_mode_gradient(many, point)
n_grads, n_passes = numeric_gradient(lambda v: many([Value(q) for q in v]).data, point, D.H)
print("reverse_passes_25", r_passes)
print("forward_passes_25", f_passes)
print("numeric_passes_25", n_passes)
print("modes_agree_exactly", all(abs(r - f) < D.ANALYTIC_TOL for r, f in zip(r_grads, f_grads)))
print("modes_agree_with_measurement", all(abs(r - n) < D.NUMERIC_TOL for r, n in zip(r_grads, n_grads)))

# -- vanishing and exploding ------------------------------------------------
decay = repeated_product(D.DECAY_FACTOR, D.CHAIN_LENGTH)
growth = repeated_product(D.GROWTH_FACTOR, D.CHAIN_LENGTH)
mild = repeated_product(D.MILD_DECAY, D.CHAIN_LENGTH)
sharp = repeated_product(D.SHARP_DECAY, D.CHAIN_LENGTH)
print("decay_50", f"{decay:.6e}")
print("growth_50", f"{growth:.6e}")
print("decay_order", order_of_magnitude(decay))
print("growth_order", order_of_magnitude(growth))
print("decay_200_order", order_of_magnitude(repeated_product(D.DECAY_FACTOR, D.LONG_CHAIN_LENGTH)))
print("growth_200_order", order_of_magnitude(repeated_product(D.GROWTH_FACTOR, D.LONG_CHAIN_LENGTH)))
print("mild_is_four_epsilons", mild == 4.0 * D.EPSILON)
print("mild_still_moves_a_weight", (1.0 + mild) != 1.0)
print("three_more_halvings_vanish", (1.0 + repeated_product(D.MILD_DECAY, 53)) == 1.0)
print("sharp_vanishes", (1.0 + sharp) == 1.0)
print("sharp_order", order_of_magnitude(sharp))

def stacked(depth):
    def f(vals):
        node = vals[0]
        for _ in range(depth):
            node = node.tanh()
        return node
    return f

single = reverse_mode_gradient(stacked(1), [0.9])[0][0]
depth40 = reverse_mode_gradient(stacked(40), [0.9])[0][0]
print("tanh_stack_single", f"{single:.6f}")
print("tanh_stack_40", f"{depth40:.6e}")
print("tanh_stack_naive_prediction", f"{single ** 40:.6e}")
print("tanh_stack_beats_prediction_by", f"{depth40 / single ** 40:.3e}")
print("tanh_stack_gap_is_over_nine_orders", depth40 > 1e9 * single ** 40)
print("tanh_stack_monotone", [
    reverse_mode_gradient(stacked(d), [0.9])[0][0] for d in (1, 5, 10, 20, 40, 80, 160)
] == sorted([
    reverse_mode_gradient(stacked(d), [0.9])[0][0] for d in (1, 5, 10, 20, 40, 80, 160)
], reverse=True))
print("constant_factor_does_vanish", repeated_product(0.487, 40) < 1e-12)
print("epsilon", D.EPSILON == float(__import__("numpy").finfo(__import__("numpy").float64).eps))
PY
)"

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

check_eq "two gears at 2 and 3 give an overall ratio of 6" "6.0" "$(get gears_two)"
check_eq "a four-stage train at 2, 3, 1.5 and 4 gives 36" "36.0" "$(get gears_four)"
check_eq "three currency rates multiply to 150" "150.0" "$(get currency)"
check_eq "the empty product is 1.0, not 0.0" "1.0" "$(get empty_product)"
check_eq "reversing the gear stages does not change the ratio" "True" "$(get gear_order_irrelevant)"

check_eq "composing square after 3x + 1 at x = 2 gives 49" "49.0" "$(get composed_value)"
check_eq "composing the other way gives 13, so order matters" "13.0" "$(get composed_other_order)"
check_eq "the chain rule gives 42 for (3x + 1) squared at x = 2" "42.0" "$(get chain_rule_correct)"
check_eq "evaluating the outer derivative at x instead gives 12, which is wrong" \
  "12.0" "$(get chain_rule_mistake)"
check_eq "all six compositions agree with a central difference" "True" "$(get all_six_match_measurement)"
for name in square_of_a_line sine_of_a_square gaussian_bump log_of_a_shifted_square the_sigmoid tanh_of_a_line; do
  check_eq "the chain rule matches the closed form for ${name//_/ }" "True" "$(get "exact_${name}")"
done
check_eq "the sigmoid's slope at zero is exactly 0.25" "0.25" "$(get sigmoid_slope_at_zero)"
check_eq "and 0.25 is the largest slope the sigmoid ever has" "True" "$(get sigmoid_quarter_is_max)"
check_eq "tanh(2x + 1) has slope exactly 2 at x = -0.5" "2.0" "$(get tanh_of_line_slope)"
echo "  (measured on this run: the worst gap between the chain rule and a central difference across the six compositions was $(get worst_composition_gap) -- reported, not asserted)"

check_eq "the five-stage forward pass is 1, 2, 5, 25, 5, ln 5" \
  "1|2|5|25|5|1.60944" "$(get five_values)"
check_eq "its five local rates are 2, 1, 10, 0.1, 0.2" "2|1|10|0.1|0.2" "$(get five_rates)"
check_eq "their product is 0.4" "0.4" "$(get five_derivative)"
check_eq "and the collapsed formula 2/(2x + 3) agrees" "0.4" "$(get five_closed_form)"
check_eq "as does a central difference of the whole chain" "0.400000000" "$(get five_measured)"
check_eq "the backward walk's first carried value is the whole derivative" \
  "0.4" "$(get running_first)"
check_eq "and its last is the final local rate alone" "0.2" "$(get running_last)"
check_eq "multiplying forwards and backwards differs by under four epsilons" \
  "True" "$(get orders_differ_by_rounding)"
check_eq "but the two orders are NOT bit-identical, which is float64, not a bug" \
  "True" "$(get orders_are_not_identical)"

check_eq "the two path contributions are 24 and 12" "24|12" "$(get two_path_contributions)"
check_eq "and the derivative is their SUM, 36" "36.0" "$(get two_path_sum)"
check_eq "which the closed form 9x squared confirms" "36.0" "$(get two_path_closed_form)"
check_eq "and a central difference confirms" "True" "$(get two_path_sum_matches)"
check_eq "the u path alone does not match the measurement" "True" "$(get path_a_alone_is_wrong)"
check_eq "the v path alone does not match the measurement" "True" "$(get path_b_alone_is_wrong)"
check_eq "and multiplying the paths does not match either" "True" "$(get product_of_paths_is_wrong)"
check_eq "the surface z at (2, 3) is 37" "37.0" "$(get surface_z)"
check_eq "dz/ds is 12x3 + (-2)x1 = 34" "34.0" "$(get surface_ds)"
check_eq "dz/dt is 12x2 + (-2)x(-1) = 26" "26.0" "$(get surface_dt)"
check_eq "and a partial difference measures dz/ds as 34" "34.000000001" "$(get surface_ds_measured)"
check_eq "and dz/dt as 26" "26.000000000" "$(get surface_dt_measured)"

check_eq "a product's two local rates are each the other input" "4|3" "$(get mul_grads)"
check_eq "a value used twice accumulates both contributions, giving 2" \
  "2.0" "$(get used_twice_grad)"
check_eq "so x times x reproduces the power rule without being told it" \
  "6.0" "$(get squared_grad)"
check_eq "tanh's slope at zero is exactly 1" "1.0" "$(get tanh_slope_at_zero)"
check_eq "tanh at half the log of 3 is exactly 0.5" "0.5" "$(get tanh_at_half_ln3)"
check_eq "and its slope there is exactly 0.75" "0.75" "$(get tanh_slope_at_half_ln3)"
check_eq "the topological order puts every child before its parent" "True" "$(get topo_children_first)"
# 20001, not 10001: each `node * 1.0` creates a Value for the constant as well
# as a Value for the product, so ten thousand operations leave twenty thousand
# nodes plus the original leaf. Counting them is the point -- this is the
# memory reverse mode pays for its speed.
check_eq "a ten-thousand-operation graph is walked without recursion limits" \
  "20001" "$(get deep_graph_nodes)"
check_eq "and its gradient is exactly 1 after ten thousand multiplications by 1" \
  "1.0" "$(get deep_graph_grad)"
check_eq "a dual number applies the product rule" "12|4" "$(get dual_product_rule)"

check_eq "hidden unit A activates at exactly 0" "0.0" "$(get net_a)"
check_eq "hidden unit B activates at exactly 0.5" "0.5" "$(get net_b)"
check_eq "so tanh's slope at B is exactly 0.75" "True" "$(get net_b_slope_exact)"
check_eq "the network output is -0.5" "-0.5" "$(get net_out)"
check_eq "and the loss is 2.25" "2.25" "$(get net_loss)"
check_eq "d loss / d out is 2 x (-1.5) = -3" "-3.0" "$(get net_d_out)"
check_eq "d loss / d vA is exactly zero, because it multiplies a dead unit" \
  "True" "$(get net_d_vA_is_zero)"
# The hand computation reaches it as -3.0 x 0.0, which in IEEE-754 is negative
# zero. It compares equal to 0.0 and behaves as zero everywhere in this lab, so
# the check above asks the question that matters; the repr is reported rather
# than asserted, because the sign of a zero is arithmetic trivia and not the
# lesson.
echo "  (measured on this run: the hand route reaches that gradient as $(get net_d_vA_repr), IEEE-754 negative zero, which compares equal to 0.0 -- reported, not asserted)"
check_eq "d loss / d vB is -1.5" "-1.5" "$(get net_d_vB)"
check_eq "d loss / d b_pre is 9 x 0.75 = 6.75" "6.75" "$(get net_d_b_pre)"
check_eq "d loss / d wA2 is -12" "-12.0" "$(get net_d_wA2)"
check_eq "d loss / d wB2 is 13.5" "13.5" "$(get net_d_wB2)"
check_eq "x1 reaches the loss twice, contributing -6 and -3.375" \
  "-6|-3.375" "$(get net_x1_contributions)"
# Section 6 re-runs this script with D110_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_d_x1="-9.375"
if [ -n "${D110_SELF_TEST:-}" ]; then
  expected_d_x1="-6.0"   # the belief that a gradient follows one path only
fi
check_eq "so d loss / d x1 is their SUM, -9.375" "${expected_d_x1}" "$(get net_d_x1)"
check_eq "and d loss / d x2 is 4.6875" "4.6875" "$(get net_d_x2)"
check_eq "taking only x1's first path is measurably wrong" "True" "$(get net_x1_single_path_wrong)"
check_eq "the engine matches the hand computation on all sixteen, bit for bit" \
  "True" "$(get net_engine_equals_hand)"
check_eq "and a central difference agrees with both within tolerance" \
  "True" "$(get net_numeric_agrees)"
check_eq "forward mode reproduces every parameter gradient" "True" "$(get net_forward_mode_agrees)"
check_eq "forward mode needed one pass per parameter" "9" "$(get net_forward_passes)"
check_eq "the network has nine parameters" "9" "$(get net_parameter_count)"
echo "  (measured on this run: the worst gap between the engine and a central difference across all sixteen network gradients was $(get net_worst_numeric_gap) -- reported, not asserted)"

check_eq "reverse mode needs 1 pass for 25 inputs" "1" "$(get reverse_passes_25)"
check_eq "forward mode needs 25" "25" "$(get forward_passes_25)"
check_eq "central differences need 50" "50" "$(get numeric_passes_25)"
check_eq "the two modes agree to the last bits" "True" "$(get modes_agree_exactly)"
check_eq "and both agree with the measurement" "True" "$(get modes_agree_with_measurement)"

check_eq "0.9 to the fiftieth is about 5.15e-3" "5.153775e-03" "$(get decay_50)"
check_eq "1.1 to the fiftieth is about 1.17e+2" "1.173909e+02" "$(get growth_50)"
check_eq "so the decayed order of magnitude is -3" "-3" "$(get decay_order)"
check_eq "and the grown one is +2" "2" "$(get growth_order)"
check_eq "at 200 layers the decay reaches order -10" "-10" "$(get decay_200_order)"
check_eq "and the growth reaches order +8" "8" "$(get growth_200_order)"
check_eq "0.5 to the fiftieth is exactly four epsilons" "True" "$(get mild_is_four_epsilons)"
check_eq "so it still moves a weight of 1, which contradicts the obvious guess" \
  "True" "$(get mild_still_moves_a_weight)"
check_eq "three more halvings do make it disappear" "True" "$(get three_more_halvings_vanish)"
check_eq "the sigmoid's best case vanishes completely in fifty layers" \
  "True" "$(get sharp_vanishes)"
check_eq "at order -31" "-31" "$(get sharp_order)"
check_eq "a stacked tanh beats the constant-factor prediction by over nine orders" \
  "True" "$(get tanh_stack_gap_is_over_nine_orders)"
check_eq "while still falling monotonically with depth" "True" "$(get tanh_stack_monotone)"
check_eq "and a genuinely constant factor does vanish geometrically" \
  "True" "$(get constant_factor_does_vanish)"
check_eq "the EPSILON in dataset.py is numpy's float64 epsilon" "True" "$(get epsilon)"
echo "  (measured on this run: 40 stacked tanh layers gave a gradient of $(get tanh_stack_40) against a naive prediction of $(get tanh_stack_naive_prediction) from a single-layer slope of $(get tanh_stack_single) -- larger by a factor of $(get tanh_stack_beats_prediction_by), reported and not asserted to a value)"

# --------------------------------------------------------------------------
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 -- -6.0, which is what you would get by following only the
# first of the two paths from x1 to the loss -- and asserts that the re-run
# reports the failure and exits non-zero. If this section passes, section 5 is
# not decorative.
if [ -z "${D110_SELF_TEST:-}" ]; then
  self_out="$(D110_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 d loss / d x1 is their SUM, -9.375"*)
      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 ]