"""From-scratch multilayer perceptron with manual backpropagation.

No autograd: the forward pass caches activations, and `backward` implements the
chain rule by hand to produce the exact gradients that PyTorch would compute
automatically.  ReLU hidden units, sigmoid output, binary cross-entropy loss,
He initialisation, mini-batch SGD, optional L2 weight decay.

The point of the notebook is that backpropagation *is* the chain rule applied
layer by layer:

    output layer   delta_L = (a_L - y) / n                 (BCE + sigmoid)
    hidden layer   delta_l = (delta_{l+1} W_{l+1}^T) * ReLU'(z_l)
    gradients      dW_l = a_{l-1}^T delta_l,   db_l = sum_i delta_l

`gradient_check` verifies these analytic gradients against finite differences to
~1e-6 -- the standard proof that a from-scratch backprop implementation is
correct -- and the notebook confirms the trained network matches a PyTorch net
of the same architecture.
"""
import numpy as np


def relu(z):
    return np.maximum(0.0, z)


def relu_grad(z):
    return (z > 0).astype(float)


def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-np.clip(z, -500, 500)))


class MLP:
    """Multilayer perceptron for binary classification.

    Parameters
    ----------
    sizes : layer widths, e.g. [n_features, 32, 16, 1] (last must be 1).
    lr    : SGD learning rate.
    l2    : L2 weight-decay strength (0 = off).
    """

    def __init__(self, sizes, lr=0.05, l2=0.0, seed=0):
        self.sizes = sizes; self.lr = lr; self.l2 = l2
        rng = np.random.default_rng(seed)
        self.W = [rng.standard_normal((sizes[i], sizes[i + 1])) * np.sqrt(2.0 / sizes[i])
                  for i in range(len(sizes) - 1)]                       # He init
        self.b = [np.zeros(sizes[i + 1]) for i in range(len(sizes) - 1)]

    def forward(self, X):
        self.a = [X]; self.z = []
        A = X
        for i in range(len(self.W)):
            Z = A @ self.W[i] + self.b[i]; self.z.append(Z)
            A = sigmoid(Z) if i == len(self.W) - 1 else relu(Z)        # sigmoid out, ReLU hidden
            self.a.append(A)
        return A

    def backward(self, y):
        n = len(y); L = len(self.W)
        delta = (self.a[-1] - y.reshape(-1, 1)) / n                    # BCE+sigmoid combined gradient
        gW = [None] * L; gb = [None] * L
        for i in reversed(range(L)):
            gW[i] = self.a[i].T @ delta + self.l2 * self.W[i]
            gb[i] = delta.sum(0)
            if i > 0:
                delta = (delta @ self.W[i].T) * relu_grad(self.z[i - 1])   # chain rule into the previous layer
        return gW, gb

    @staticmethod
    def _bce(y, p):
        p = np.clip(p, 1e-7, 1 - 1e-7)
        return float(-np.mean(y * np.log(p) + (1 - y) * np.log(1 - p)))

    def fit(self, X, y, epochs=100, batch=128, seed=0, verbose=False):
        rng = np.random.default_rng(seed); n = len(y); self.loss_ = []
        for ep in range(epochs):
            idx = rng.permutation(n)
            for s in range(0, n, batch):
                bi = idx[s:s + batch]
                self.forward(X[bi]); gW, gb = self.backward(y[bi])
                for i in range(len(self.W)):
                    self.W[i] -= self.lr * gW[i]; self.b[i] -= self.lr * gb[i]
            self.loss_.append(self._bce(y, self.forward(X).ravel()))
            if verbose and ep % 10 == 0:
                print(f"epoch {ep:4d}  BCE {self.loss_[-1]:.4f}")
        return self

    def predict_proba(self, X):
        return self.forward(X).ravel()


def gradient_check(net, X, y, n_params=8, eps=1e-5, seed=0):
    """Compare analytic backprop gradients to finite-difference estimates on a
    random handful of weights.  Returns the max relative error (~1e-8 in float64).

    The differenced objective must be the SAME function `backward` differentiates.
    `backward` adds ``l2 * W`` to every weight gradient, which is the derivative of
    ``0.5 * l2 * sum(W**2)``, so that penalty belongs in the objective too -- omit it
    and the check reports a spurious failure (relative error ~0.7 at l2=1e-4) on a
    net whose gradients are perfectly correct.  Including it keeps the check valid
    for any net, not only the unregularised one.
    """
    def objective():
        loss = net._bce(y, net.forward(X).ravel())
        if net.l2:
            loss += 0.5 * net.l2 * sum(float((W ** 2).sum()) for W in net.W)
        return loss

    net.forward(X); gW, _ = net.backward(y)
    rng = np.random.default_rng(seed); errs = []
    for _ in range(n_params):
        li = rng.integers(len(net.W))
        i = rng.integers(net.W[li].shape[0]); j = rng.integers(net.W[li].shape[1])
        orig = net.W[li][i, j]
        net.W[li][i, j] = orig + eps; lp = objective()
        net.W[li][i, j] = orig - eps; lm = objective()
        net.W[li][i, j] = orig
        num = (lp - lm) / (2 * eps); ana = gW[li][i, j]
        errs.append(abs(num - ana) / max(1e-12, abs(num) + abs(ana)))
    return max(errs)