"""Nonlinear noise-tightening drift utilities."""
import numpy as np


def linear_drift(x, lam):
    return -lam * np.asarray(x)


def cubic_drift(x, lam, beta):
    x = np.asarray(x)
    return -lam * x - beta * x**3


def radial_cubic_drift(x, lam, beta):
    """Vector extension: -lambda*x - beta*x*||x||^2."""
    x = np.asarray(x)
    norm2 = np.sum(x * x, axis=-1, keepdims=True)
    return -lam * x - beta * x * norm2


def scalar_pairwise_rate(x, y, lam, beta):
    """r=(v(x)-v(y))/(y-x) for x != y, extended continuously at x=y."""
    x, y = np.asarray(x), np.asarray(y)
    return lam + beta * (x*x + x*y + y*y)


def euler_maruyama(drift, x0, dt, steps, D, rng):
    x = np.array(x0, dtype=float, copy=True)
    noise_scale = np.sqrt(2.0 * D * dt)
    for _ in range(steps):
        x += drift(x) * dt + noise_scale * rng.normal(size=x.shape)
    return x