# Experiments for "Neural Networks and Backpropagation": every number in the post comes from here (numpy 2.1, seeded).
# A one-hidden-layer network with hand-written backpropagation, checked against finite differences.
#   python neural_networks.py [OUT.json]
import sys, json, time, numpy as np
sig = lambda z: 1 / (1 + np.exp(-z))
out = {}

# 1. the worked exercise: 2 inputs, 3 sigmoid hidden units, 1 sigmoid output, squared error L = 1/2 (o - t)^2, one example
W = np.array([[0.1, -0.2], [0.3, -0.4], [0.5, -0.6]]); b = np.array([0.01, 0.02, 0.03]); v = np.array([0.5, -0.3, 0.2]); c = 0.1
x = np.array([0.7, 0.9]); t = 0.85; eta = 0.01
def fwd(W, b, v, c):
    z = W @ x + b; h = sig(z); s = v @ h + c; o = sig(s); return z, h, s, o, 0.5 * (o - t) ** 2
z, h, s, o, L = fwd(W, b, v, c)
print("1. forward: z =", z.round(6), " h =", h.round(6), f" s = {s:.6f}  o = {o:.6f}  L = {L:.6f}")
d_o = (o - t) * o * (1 - o)                 # dL/ds
gv, gc = d_o * h, d_o                       # output layer
d_h = d_o * v * h * (1 - h)                 # dL/dz_j
gW, gb = np.outer(d_h, x), d_h
print(f"   backward: dL/ds = {d_o:.6f}; dL/dv = {gv.round(6)}; dL/dc = {gc:.6f}")
print("   dL/dz =", d_h.round(8), " dL/dW =", gW.round(8).tolist(), " dL/db =", gb.round(8))
def fd(param, idx, eps=1e-6):
    P = [W.copy(), b.copy(), v.copy(), np.array(c)]; P[param][idx] += eps; Lp = fwd(*P[:3], float(P[3]))[4]
    P = [W.copy(), b.copy(), v.copy(), np.array(c)]; P[param][idx] -= eps; Lm = fwd(*P[:3], float(P[3]))[4]; return (Lp - Lm) / (2 * eps)
num = np.r_[[fd(0, (i, j)) for i in range(3) for j in range(2)], [fd(1, i) for i in range(3)], [fd(2, i) for i in range(3)], fd(3, ())]
ana = np.r_[gW.ravel(), gb, gv, gc]; print(f"   finite differences vs backprop: max relative error {np.max(np.abs(num - ana) / np.abs(ana)):.2e}")
W1, b1, v1, c1 = W - eta * gW, b - eta * gb, v - eta * gv, c - eta * gc; L1 = fwd(W1, b1, v1, c1)[4]
print(f"   after one step (eta = 0.01): v = {v1.round(6)}, c = {c1:.6f}, W = {W1.round(6).tolist()}; loss {L:.8f} -> {L1:.8f}")
Wk, bk, vk, ck = W.copy(), b.copy(), v.copy(), c; Ls = []
for k in range(1, 2001):
    zz = Wk @ x + bk; hh = sig(zz); oo = sig(vk @ hh + ck); do = (oo - t) * oo * (1 - oo); dh = do * vk * hh * (1 - hh)
    vk, ck, Wk, bk = vk - eta * do * hh, ck - eta * do, Wk - eta * np.outer(dh, x), bk - eta * dh
    if k in (1, 10, 100, 1000, 2000): Ls.append((k, float(oo), float(0.5 * (oo - t) ** 2)))
print("   training on this one example:", ", ".join(f"step {k}: o {oo:.4f}, L {l:.2e}" for k, oo, l in Ls))
out["ex"] = {"W": W.tolist(), "b": b.tolist(), "v": v.tolist(), "c": c, "x": x.tolist(), "t": t, "o": o, "L": L}

# a one-hidden-layer classifier (cross-entropy) for the experiments below
def moons(n, rng):
    t_ = rng.uniform(0, np.pi, n); y = rng.integers(0, 2, n)
    X = np.where(y[:, None] == 0, np.c_[np.cos(t_), np.sin(t_)], np.c_[1 - np.cos(t_), 0.5 - np.sin(t_)]) + rng.normal(0, 0.15, (n, 2)); return X, y
def init(m, rng, zero=False, d=2):
    return [np.zeros((m, d)), np.zeros(m), np.zeros(m), 0.0] if zero else [rng.normal(0, np.sqrt(2 / d), (m, d)), np.zeros(m), rng.normal(0, np.sqrt(1 / m), m), 0.0]
act = {"relu": (lambda z: np.maximum(z, 0), lambda z: (z > 0).astype(float)), "identity": (lambda z: z, lambda z: np.ones_like(z)), "tanh": (np.tanh, lambda z: 1 - np.tanh(z) ** 2)}
def grads(P, X, y, a="relu"):
    W, b, v, c = P; f, df = act[a]; Z = X @ W.T + b; H = f(Z); p = sig(H @ v + c); n = len(y)
    loss = -np.mean(y * np.log(np.clip(p, 1e-12, 1)) + (1 - y) * np.log(np.clip(1 - p, 1e-12, 1))); d = (p - y) / n
    return loss, [((d[:, None] * v) * df(Z)).T @ X, ((d[:, None] * v) * df(Z)).sum(0), H.T @ d, d.sum()]
def predict(P, X, a="relu"): W, b, v, c = P; return (act[a][0](X @ W.T + b) @ v + c > 0).astype(int)
def train(P, X, y, lr, steps, a="relu", batch=None, rng=None):
    hist = []
    for s_ in range(steps):
        if batch: i = rng.integers(0, len(y), batch); loss, g = grads(P, X[i], y[i], a)
        else: loss, g = grads(P, X, y, a)
        P = [p - lr * gg for p, gg in zip(P, g)]; hist.append(loss)
    return P, hist
rng = np.random.default_rng(0)
Xtr, ytr = moons(400, rng); Xte, yte = moons(4000, rng)
# gradient check of the classifier
P0 = init(8, rng); l0, g0 = grads(P0, Xtr[:50], ytr[:50]); eps = 1e-6; errs = []
for k in range(4):
    arr = np.asarray(P0[k], dtype=float)
    for idx in [tuple(rng.integers(0, s_) for s_ in arr.shape) for _ in range(3)] if arr.ndim else [()]:
        Pp = [np.array(p, dtype=float, copy=True) for p in P0]; Pm = [np.array(p, dtype=float, copy=True) for p in P0]; Pp[k][idx] += eps; Pm[k][idx] -= eps
        num_ = (grads(Pp, Xtr[:50], ytr[:50])[0] - grads(Pm, Xtr[:50], ytr[:50])[0]) / (2 * eps); errs.append(abs(num_ - np.asarray(g0[k])[idx]) / max(abs(num_), 1e-12))
print(f"2. classifier gradient check: max relative error {max(errs):.2e}")

# 3. why the activation: linear vs network without activation vs ReLU network (two moons)
lin = train([np.zeros((1, 2)), np.zeros(1), np.ones(1), 0.0], Xtr, ytr, 0.5, 3000, "identity")[0]
res = {"logistic regression": np.mean(predict(lin, Xte, "identity") == yte)}
for a in ("identity", "relu", "tanh"):
    P, h_ = train(init(16, np.random.default_rng(1)), Xtr, ytr, 0.5, 5000, a); res[f"16 hidden units, {a}"] = np.mean(predict(P, Xte, a) == yte)
print("3. test accuracy on two moons:", {k: round(float(v_), 4) for k, v_ in res.items()})

# 4. symmetric initialisation: every hidden unit starts identical
same = [np.full((16, 2), 0.1), np.zeros(16), np.full(16, 0.1), 0.0]
Ps, _ = train(same, Xtr, ytr, 0.5, 5000, "tanh"); Pr, _ = train(init(16, np.random.default_rng(1)), Xtr, ytr, 0.5, 5000, "tanh")
P1, _ = train(init(1, np.random.default_rng(1)), Xtr, ytr, 0.5, 5000, "tanh")
Ws_ = np.c_[Ps[0], Ps[1], Ps[2]]
print(f"4. 16 tanh units, all weights started at 0.1: distinct units after 5,000 steps {len(np.unique(Ws_.round(9), axis=0))} of 16, accuracy {np.mean(predict(Ps, Xte, 'tanh') == yte):.4f}; "
      f"a network with 1 unit {np.mean(predict(P1, Xte, 'tanh') == yte):.4f}; random start {np.mean(predict(Pr, Xte, 'tanh') == yte):.4f}")
Pz, _ = train(init(16, rng, zero=True), Xtr, ytr, 0.5, 5000, "tanh")
print(f"   all weights zero: first layer still zero after 5,000 steps: {bool(np.all(Pz[0] == 0) and np.all(Pz[2] == 0))}, accuracy {np.mean(predict(Pz, Xte, 'tanh') == yte):.4f}")

# 5. learning rate (full batch, 2,000 steps)
lrs = [0.001, 0.01, 0.1, 1.0, 10.0, 100.0]; lrrow = {}
for lr in lrs:
    with np.errstate(all="ignore"): P, h_ = train(init(16, np.random.default_rng(1)), Xtr, ytr, lr, 2000, "relu")
    lrrow[lr] = (h_[0], h_[-1], float(np.mean(predict(P, Xte) == yte)) if np.all(np.isfinite(P[0])) else float("nan"))
print("5. learning rate: loss at step 1 -> step 2,000 (test accuracy):", {lr: (round(a_, 3), round(b_, 3), round(c_, 3)) for lr, (a_, b_, c_) in lrrow.items()})

# 6. full batch vs mini-batch vs single-example SGD: same number of passes over the data (30 epochs)
for bs in (400, 32, 1):
    steps = 30 * 400 // bs; t0 = time.perf_counter(); P, h_ = train(init(16, np.random.default_rng(1)), Xtr, ytr, 0.5 if bs > 1 else 0.05, steps, "relu", batch=None if bs == 400 else bs, rng=np.random.default_rng(2))
    print(f"6. batch {bs:3d}: {steps:5d} updates in 30 epochs, final full training loss {grads(P, Xtr, ytr)[0]:.4f}, test accuracy {np.mean(predict(P, Xte) == yte):.4f}, {time.perf_counter() - t0:.2f} s")
Pg = init(16, np.random.default_rng(1)); full = grads(Pg, Xtr, ytr)[1][2]
for bs in (1, 8, 32, 128):
    est = np.array([grads(Pg, Xtr[i], ytr[i])[1][2] for i in (rng.integers(0, 400, bs) for _ in range(2000))])
    print(f"   batch {bs:3d}: mean squared error of the gradient estimate {np.mean(np.sum((est - full) ** 2, 1)):.5f}")

# 7. vanishing gradients: 10 layers of 32 units at initialisation
for a, std in (("sigmoid", 1.0), ("relu", None)):
    rr = np.random.default_rng(3); Xv = rr.normal(0, 1, (256, 32)); Ws = [rr.normal(0, std if std else np.sqrt(2 / 32), (32, 32)) for _ in range(10)]
    H = [Xv]; Z = []
    for Wl in Ws: Z.append(H[-1] @ Wl.T); H.append(sig(Z[-1]) if a == "sigmoid" else np.maximum(Z[-1], 0))
    g = np.ones_like(H[-1]) / H[-1].size; norms = []
    for l in range(9, -1, -1):
        dz = g * (H[l + 1] * (1 - H[l + 1]) if a == "sigmoid" else (Z[l] > 0)); norms.append(np.linalg.norm(dz.T @ H[l])); g = dz @ Ws[l]
    print(f"7. {a}: gradient norm of layer 1 / layer 10 = {norms[-1] / norms[0]:.2e}")

# 8. why backpropagation: cost of a gradient for a 784-100-10 network
P_ = 784 * 100 + 100 + 100 * 10 + 10; Xb = rng.normal(0, 1, (64, 784)); W1_ = rng.normal(0, 0.05, (100, 784)); W2_ = rng.normal(0, 0.1, (10, 100))
def fwd8(): return np.maximum(Xb @ W1_.T, 0) @ W2_.T
def bwd8():
    Hh = np.maximum(Xb @ W1_.T, 0); O = Hh @ W2_.T; G = O / O.size; gW2 = G.T @ Hh; dH = (G @ W2_) * (Hh > 0); gW1 = dH.T @ Xb; return gW1, gW2
for fn in (fwd8, bwd8): fn()
t0 = time.perf_counter(); [fwd8() for _ in range(200)]; tf = (time.perf_counter() - t0) / 200
t0 = time.perf_counter(); [bwd8() for _ in range(200)]; tb = (time.perf_counter() - t0) / 200
print(f"8. 784-100-10 network: {P_:,} parameters; forward+backward costs {tb / tf:.1f} forward passes; central finite differences would need {2 * P_:,}")
if len(sys.argv) > 1: json.dump(out, open(sys.argv[1], "w"), separators=(",", ":"))