#!/usr/bin/env python3
"""
W5 — Reverse-mode autodiff in 45 lines, checked against finite differences. (~45 min)

Miniature of P13. The point of this walkthrough: automatic differentiation is
BOOKKEEPING, not calculus. The per-op derivative rules are trivial; the engineering
is the graph, the topological order, and gradient ACCUMULATION.
"""
import math, random

class Value:
    """A scalar and its place in the computation graph."""
    def __init__(self, data, parents=(), op=""):
        self.data = data; self.grad = 0.0
        self._parents = parents; self._backward = lambda: None; self._op = op

    def __add__(self, o):
        o = o if isinstance(o, Value) else Value(o)
        out = Value(self.data + o.data, (self, o), "+")
        def back():
            # ACCUMULATE (+=), never assign. If a value is used twice, both
            # contributions must sum. Assigning here makes gradients too small by
            # an exact integer factor -- which looks like a learning-rate problem
            # and gets "fixed" by raising the learning rate.
            self.grad += out.grad; o.grad += out.grad
        out._backward = back; return out

    def __mul__(self, o):
        o = o if isinstance(o, Value) else Value(o)
        out = Value(self.data * o.data, (self, o), "*")
        def back():
            self.grad += o.data * out.grad; o.grad += self.data * out.grad
        out._backward = back; return out

    def tanh(self):
        t = math.tanh(self.data)
        out = Value(t, (self,), "tanh")
        def back(): self.grad += (1 - t * t) * out.grad
        out._backward = back; return out

    def __pow__(self, k):
        out = Value(self.data ** k, (self,), f"**{k}")
        def back(): self.grad += k * self.data ** (k - 1) * out.grad
        out._backward = back; return out

    def backward(self):
        # reverse topological order: a node's grad must be complete before it is used
        topo, seen = [], set()
        def build(v):
            if id(v) in seen: return
            seen.add(id(v))
            for p in v._parents: build(p)
            topo.append(v)
        build(self)
        self.grad = 1.0
        for v in reversed(topo): v._backward()

    __radd__ = __add__; __rmul__ = __mul__
    def __neg__(self): return self * -1
    def __sub__(self, o): return self + (-o if isinstance(o, Value) else Value(-o))
    def __repr__(self): return f"Value({self.data:.4f}, grad={self.grad:.4f})"


print("TEST 1 — gradients match central finite differences")
random.seed(0)
def f(xs):
    a, b, c = xs
    return ((a * b + c).tanh() * (a + c) ** 2 + b * c)

worst = 0.0
for _ in range(200):
    vals = [random.uniform(-1.5, 1.5) for _ in range(3)]
    vs = [Value(v) for v in vals]
    out = f(vs); out.backward()
    for i in range(3):
        h = 1e-6
        up = vals[:]; up[i] += h
        dn = vals[:]; dn[i] -= h
        num = (f([Value(v) for v in up]).data - f([Value(v) for v in dn]).data) / (2 * h)
        worst = max(worst, abs(num - vs[i].grad) / (abs(num) + 1e-9))
print(f"  worst relative error over 600 partials: {worst:.2e}")
assert worst < 1e-5

print("\nTEST 2 — the accumulation bug, demonstrated")
x = Value(3.0)
y = x * x                      # x used TWICE in one expression
y.backward()
print(f"  d(x*x)/dx at x=3: got {x.grad}, correct 6.0  -> {'OK' if abs(x.grad-6)<1e-9 else 'WRONG'}")
print("  With `=` instead of `+=` in __mul__ this prints 3.0: exactly half.")
print("  An integer-factor error is the signature of a missing accumulation.")

print("\nTEST 3 — train something. Fit y = sin-ish curve with a 1-hidden-layer net.")
random.seed(1)
H = 8
W1 = [Value(random.uniform(-1, 1)) for _ in range(H)]
b1 = [Value(0.0) for _ in range(H)]
W2 = [Value(random.uniform(-1, 1)) for _ in range(H)]
b2 = Value(0.0)
params = W1 + b1 + W2 + [b2]
data = [(x / 10.0, math.sin(x / 10.0)) for x in range(-25, 26)]

def net(x):
    hs = [(W1[i] * x + b1[i]).tanh() for i in range(H)]
    out = b2
    for i in range(H): out = out + W2[i] * hs[i]
    return out

lr = 0.05
for epoch in range(601):
    loss = Value(0.0)
    for x, y in data: loss = loss + (net(x) - y) ** 2
    loss = loss * (1.0 / len(data))
    for p in params: p.grad = 0.0            # zero_grad: forgetting this accumulates
    loss.backward()                          # across epochs and training silently stalls
    for p in params: p.data -= lr * p.grad
    if epoch % 150 == 0: print(f"  epoch {epoch:>4}  mse {loss.data:.6f}")

print(f"\n  final mse {loss.data:.6f} over {len(params)} parameters")
print("  Reverse mode computed all", len(params), "gradients in ONE backward pass.")
print("  Finite differences would need", len(params)+1, "forward passes for the same.")