"""Small NumPy MLPs, Adam, and a running observation normaliser.

Deliberately dependency-free. A 2x64 tanh policy is all a balancer needs, it
trains fine on CPU, and the weights drop straight into fixed-point C for an
MCU later without a framework in the way.
"""

import numpy as np


class RunningNorm:
    def __init__(self, dim, clip=10.0):
        self.mean = np.zeros(dim)
        self.var = np.ones(dim)
        self.count = 1e-4
        self.clip = clip

    def update(self, x):
        bm, bv, bc = x.mean(0), x.var(0), x.shape[0]
        delta = bm - self.mean
        tot = self.count + bc
        self.mean += delta * bc / tot
        m_a = self.var * self.count
        m_b = bv * bc
        self.var = (m_a + m_b + delta ** 2 * self.count * bc / tot) / tot
        self.count = tot

    def __call__(self, x):
        return np.clip((x - self.mean) / np.sqrt(self.var + 1e-8), -self.clip, self.clip)

    def state(self):
        return {"mean": self.mean.tolist(), "var": self.var.tolist(),
                "count": float(self.count), "clip": self.clip}

    @classmethod
    def load(cls, s):
        o = cls(len(s["mean"]), s.get("clip", 10.0))
        o.mean = np.array(s["mean"])
        o.var = np.array(s["var"])
        o.count = s["count"]
        return o


class MLP:
    """tanh hidden layers, linear output. Explicit forward/backward."""

    def __init__(self, sizes, out_gain=0.01, rng=None):
        rng = rng or np.random.default_rng(0)
        self.W, self.b = [], []
        for i in range(len(sizes) - 1):
            gain = out_gain if i == len(sizes) - 2 else np.sqrt(2.0)
            fan_in = sizes[i]
            self.W.append(rng.normal(0, gain / np.sqrt(fan_in), (sizes[i], sizes[i + 1])))
            self.b.append(np.zeros(sizes[i + 1]))

    def forward(self, x, cache=None):
        h = x
        for i in range(len(self.W)):
            z = h @ self.W[i] + self.b[i]
            if cache is not None:
                cache.append((h, z))
            h = np.tanh(z) if i < len(self.W) - 1 else z
        return h

    def backward(self, cache, dout):
        gW = [None] * len(self.W)
        gb = [None] * len(self.b)
        d = dout
        for i in reversed(range(len(self.W))):
            h, z = cache[i]
            if i < len(self.W) - 1:
                d = d * (1.0 - np.tanh(z) ** 2)
            gW[i] = h.T @ d
            gb[i] = d.sum(0)
            d = d @ self.W[i].T
        return gW, gb

    def params(self):
        return self.W + self.b

    def state(self):
        return {"W": [w.tolist() for w in self.W], "b": [b.tolist() for b in self.b]}

    @classmethod
    def load(cls, s):
        o = cls([1, 1])
        o.W = [np.array(w) for w in s["W"]]
        o.b = [np.array(b) for b in s["b"]]
        return o


class Adam:
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.999), eps=1e-8):
        self.p = params
        self.lr, self.b1, self.b2, self.eps = lr, betas[0], betas[1], eps
        self.m = [np.zeros_like(x) for x in params]
        self.v = [np.zeros_like(x) for x in params]
        self.t = 0

    def step(self, grads, max_norm=0.5):
        total = np.sqrt(sum(float(np.sum(g ** 2)) for g in grads))
        scale = min(1.0, max_norm / (total + 1e-8))
        self.t += 1
        bc1 = 1 - self.b1 ** self.t
        bc2 = 1 - self.b2 ** self.t
        for i, g in enumerate(grads):
            g = g * scale
            self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g
            self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * g * g
            self.p[i] -= self.lr * (self.m[i] / bc1) / (np.sqrt(self.v[i] / bc2) + self.eps)
        return total