"""Small NumPy implementation of the smooth-RG modewise update.

The caller supplies approximate positive curvature modes q and optionally a
modewise damping estimate a. The default quadratic fallback a=q is exact for
an independently diagonal quadratic objective.
"""
import numpy as np


def smooth_regulator(q, alpha=0.7, s=1.0):
    q = np.asarray(q, dtype=float)
    if alpha <= 0 or s <= 0:
        raise ValueError("alpha and s must be positive")
    return alpha * s**2 / np.expm1(np.minimum((q / s)**2, 700.0))


class SmoothRGModewise:
    def __init__(self, q, alpha=0.7, s=1.0, lr=0.5, eps=1e-8):
        self.q = np.asarray(q, dtype=float)
        if np.any(self.q <= 0):
            raise ValueError("curvature modes must be positive")
        self.alpha, self.s, self.lr, self.eps = alpha, s, lr, eps

    def step(self, theta, grad, damping=None):
        """Return theta - lr*grad/(|a|+R_s(q)+eps).

        In a real network q and damping can be refreshed from Hessian probes
        and perturbation responses; this MVP assumes they are mode-aligned.
        """
        a = self.q if damping is None else np.asarray(damping, dtype=float)
        denom = np.abs(a) + smooth_regulator(self.q, self.alpha, self.s) + self.eps
        return np.asarray(theta) - self.lr * np.asarray(grad) / denom