"""Small numpy MLP (2 hidden ReLU layers, softmax cross-entropy, Adam) used
as the neural-network baseline.  Deliberately plain: the point of the
comparison is bias, not architecture search."""

from __future__ import annotations

import numpy as np


class MLP:
    def __init__(self, n_in: int, n_out: int, hidden: int = 64, seed: int = 0):
        rng = np.random.default_rng(seed)
        dims = [n_in, hidden, hidden, n_out]
        self.W = [rng.normal(0, np.sqrt(2.0 / dims[i]),
                             (dims[i], dims[i + 1])).astype(np.float32)
                  for i in range(3)]
        self.b = [np.zeros(dims[i + 1], dtype=np.float32) for i in range(3)]
        self.n_params = sum(w.size for w in self.W) + sum(b.size for b in self.b)

    def _forward(self, X):
        a1 = np.maximum(X @ self.W[0] + self.b[0], 0)
        a2 = np.maximum(a1 @ self.W[1] + self.b[1], 0)
        z = a2 @ self.W[2] + self.b[2]
        return a1, a2, z

    def predict(self, X, chunk: int = 65536):
        out = np.empty(len(X), dtype=np.int8)
        for s in range(0, len(X), chunk):
            _, _, z = self._forward(X[s:s + chunk])
            out[s:s + chunk] = z.argmax(axis=1)
        return out

    def fit(self, X, y, epochs: int = 60, batch: int = 128, lr: float = 1e-3,
            seed: int = 0, class_weight: np.ndarray | None = None,
            verbose: bool = False):
        rng = np.random.default_rng(seed)
        n = len(X)
        params = self.W + self.b
        m = [np.zeros_like(p) for p in params]
        v = [np.zeros_like(p) for p in params]
        t = 0
        for ep in range(epochs):
            order = rng.permutation(n)
            for s in range(0, n, batch):
                idx = order[s:s + batch]
                xb, yb = X[idx], y[idx]
                a1, a2, z = self._forward(xb)
                z = z - z.max(axis=1, keepdims=True)
                p = np.exp(z)
                p /= p.sum(axis=1, keepdims=True)
                dz = p
                dz[np.arange(len(yb)), yb] -= 1.0
                if class_weight is not None:
                    dz *= class_weight[yb][:, None]
                dz /= len(yb)
                gW2 = a2.T @ dz
                gb2 = dz.sum(0)
                da2 = dz @ self.W[2].T
                da2[a2 <= 0] = 0
                gW1 = a1.T @ da2
                gb1 = da2.sum(0)
                da1 = da2 @ self.W[1].T
                da1[a1 <= 0] = 0
                gW0 = xb.T @ da1
                gb0 = da1.sum(0)
                grads = [gW0, gW1, gW2, gb0, gb1, gb2]
                t += 1
                for p_, g, mi, vi in zip(params, grads, m, v):
                    mi *= 0.9
                    mi += 0.1 * g
                    vi *= 0.999
                    vi += 0.001 * g * g
                    mh = mi / (1 - 0.9 ** t)
                    vh = vi / (1 - 0.999 ** t)
                    p_ -= lr * mh / (np.sqrt(vh) + 1e-8)
            if verbose and (ep + 1) % 20 == 0:
                acc = (self.predict(X) == y).mean()
                print(f"    epoch {ep+1}: train acc {acc*100:.2f}%")
        return self