def noisy_quadratic_grad(w, rng, noise_scale=0.5):
    # loss = 0.5 * (5*w0^2 + 0.1*w1^2) — very different curvature
    # along the two axes, a classic case where momentum helps.
    true_grad = np.array([5.0 * w[0], 0.1 * w[1]])
    return true_grad + rng.normal(0, noise_scale, size=2)


def run_sgd_momentum(w0, steps, lr=0.05, beta=0.9, seed=0):
    rng = np.random.default_rng(seed)
    w, v = w0.copy(), np.zeros_like(w0)
    trajectory = [w.copy()]
    for _ in range(steps):
        g = noisy_quadratic_grad(w, rng)
        v = beta * v + g
        w = w - lr * v
        trajectory.append(w.copy())
    return np.array(trajectory)


def run_adam(w0, steps, lr=0.1, seed=0):
    rng = np.random.default_rng(seed)
    w = w0.copy()
    opt = Adam([w.shape], lr=lr)
    trajectory = [w.copy()]
    for _ in range(steps):
        g = noisy_quadratic_grad(w, rng)
        w = opt.step([w], [g])[0]
        trajectory.append(w.copy())
    return np.array(trajectory)


if __name__ == "__main__":
    w0 = np.array([1.0, 1.0])
    steps = 200

    traj_momentum = run_sgd_momentum(w0, steps)
    traj_adam = run_adam(w0, steps)

    final_loss_momentum = 0.5 * (
        5 * traj_momentum[-1, 0] ** 2 + 0.1 * traj_momentum[-1, 1] ** 2
    )
    final_loss_adam = 0.5 * (
        5 * traj_adam[-1, 0] ** 2 + 0.1 * traj_adam[-1, 1] ** 2
    )
    print(f"Final loss (SGD+momentum): {final_loss_momentum:.6f}")
    print(f"Final loss (Adam):         {final_loss_adam:.6f}")
