"""Weight perturbation (parameter-space zeroth order) baseline: MeZO / CD-RGE / weight-space forward
gradient.  All parameters are perturbed by eps*u (Rademacher), the loss is measured at +/- and the
gradient estimate is (1/n) sum_i u_i (L(theta+eps u_i) - L(theta-eps u_i)) / (2 eps).
2n forward passes of the whole network per step, no backward pass, variance ~ P/n with P = #params."""
import torch


def wp_step(model, loss_fn, x, y, n, eps, gen):
    params = [p for p in model.parameters() if p.requires_grad]
    grads = [torch.zeros_like(p) for p in params]
    with torch.no_grad():
        for i in range(n):
            us = [(torch.randint(0, 2, p.shape, generator=gen, device=p.device).to(p.dtype) * 2 - 1) for p in params]
            for p, u in zip(params, us):
                p.add_(eps * u)
            lp = loss_fn(model(x), y).item()
            for p, u in zip(params, us):
                p.sub_(2 * eps * u)
            lm = loss_fn(model(x), y).item()
            for p, u in zip(params, us):
                p.add_(eps * u)
            D = (lp - lm) / (2 * eps)
            for g, u in zip(grads, us):
                g.add_(u, alpha=D / n)
        for p, g in zip(params, grads):
            p.grad = g
        loss = loss_fn(model(x), y)
    return loss