"""
stochGD.py — numerical companion to the glossary entry
'stochastic gradient descent (SGD)'.

One block per paragraph of the entry (marked [P...]): each block verifies
numerically what the corresponding statement asserts. Self-contained
(numpy/matplotlib only), fixed seed.

Blocks
------
[P-approx]    SGD replaces the full gradient — a sum of per-data-point
              gradients over the entire trainset — with the sum over a
              randomly chosen batch: the batch gradient is an unbiased
              approximation (its average over many random batches
              matches the full gradient), and one SGD step touches only
              |B| of the m data points.
[P-batchsize] The batch size trades gradient accuracy against cost: the
              approximation error of the batch gradient shrinks like
              1/sqrt(|B|) as the batch grows (variance scaling), while
              the per-step cost grows linearly in |B|.
[P-minibatch] Mini-batch SGD (|B| > 1) converges to the ERM solution on
              a least-squares problem while evaluating only a small
              fraction of the per-data-point gradients that full GD
              uses for the same number of passes.

Outputs
-------
stochGD.png : preview figure (checking only).

Data generated by pythondemos/stochGD.py.
"""

import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

rng = np.random.default_rng(42)
report = []


def check(name, ok):
    report.append((name, bool(ok)))
    print(f"  [{'ok' if ok else 'FAIL'}] {name}")


# least-squares ERM objective on m data points
m, d = 2000, 5
X = rng.normal(size=(m, d))
w_true = rng.normal(size=d)
y = X @ w_true + 0.3 * rng.normal(size=m)
full_grad = lambda w: (2 / m) * X.T @ (X @ w - y)
def batch_grad(w, B):
    idx = rng.choice(m, B, replace=False)
    return (2 / B) * X[idx].T @ (X[idx] @ w - y[idx]), B

# --------------------------------------------------------- [P-approx]
print("[P-approx] the batch gradient approximates the full-sum gradient")
w0 = np.zeros(d)
g_full = full_grad(w0)
g_avg = np.mean([batch_grad(w0, 20)[0] for _ in range(4000)], axis=0)
check("averaging batch gradients over many draws recovers the full "
      "gradient (unbiasedness)",
      np.linalg.norm(g_avg - g_full) < 0.05 * np.linalg.norm(g_full))
check("one SGD step touches |B| = 20 of the m = 2000 data points",
      batch_grad(w0, 20)[1] == 20 < m)

# ------------------------------------------------------ [P-batchsize]
print("[P-batchsize] batch size trades accuracy against cost")
errs = []
for B in (10, 100, 1000):
    errs.append(np.mean([np.linalg.norm(batch_grad(w0, B)[0] - g_full)
                         for _ in range(300)]))
print(f"    mean gradient error at |B| = 10, 100, 1000: "
      f"{errs[0]:.3f}, {errs[1]:.3f}, {errs[2]:.3f}")
check("the gradient error shrinks as the batch grows",
      errs[0] > errs[1] > errs[2])
check("error scaling is consistent with 1/sqrt(|B|) "
      "(10x batch -> ~3.2x smaller error)",
      2.0 < errs[0] / errs[1] < 5.0 and 2.0 < errs[1] / errs[2] < 5.0)

# ------------------------------------------------------ [P-minibatch]
print("[P-minibatch] mini-batch SGD reaches the ERM solution cheaply")
w_hat = np.linalg.solve(X.T @ X, X.T @ y)          # ERM solution
w = np.zeros(d)
grads_evaluated = 0
for t in range(1, 1201):
    g, B = batch_grad(w, 20)
    w -= (0.05 / np.sqrt(t)) * g
    grads_evaluated += B
gd_grads = 1200 * m                                 # full GD, same steps
check("mini-batch SGD converges near the ERM solution",
      np.linalg.norm(w - w_hat) < 0.1)
check("using 1% of the per-data-point gradient evaluations of full GD",
      grads_evaluated == 0.01 * gd_grads)

# ------------------------------------------------------------ preview
fig, ax = plt.subplots(figsize=(4.8, 3.2))
ax.loglog([10, 100, 1000], errs, "o-")
ax.set_xlabel("batch size |B|"); ax.set_ylabel("gradient error")
ax.set_title("[P-batchsize] accuracy vs batch size")
fig.tight_layout()
fig.savefig("stochGD.png", dpi=110)
print(f"\n{sum(ok for _, ok in report)}/{len(report)} checks passed")
assert all(ok for _, ok in report)