import numpy as np
from neural_levy_noise import compensated_euler_step, small_variance, small_mean

rng=np.random.default_rng(7)
x=np.zeros((4096,3)); d=np.ones_like(x)
y,k=compensated_euler_step(x,d,0.1,0.05,1.5,rng=rng)
assert y.shape==x.shape and isinstance(k,int) and k>0
# Repeated scalar noise checks matched per-coordinate Gaussian variance after removing drift/jumps
assert abs(small_variance(.05,1.5)-0.5**0) > 0 # formula callable
assert small_mean(.05,.5)>0
try:
    small_mean(.05,1.0)
    raise AssertionError('expected divergent mean guard')
except ValueError:
    pass
# one-sided corrected update also runs
z,k2=compensated_euler_step(np.zeros((16,2)),np.zeros((16,2)),.01,.1,.5,rng=rng,one_sided=True,add_small_mean=True)
assert z.shape==(16,2) and k2>=0
print('smoke_ok', y.shape, k, k2, 'variance', small_variance(.05,1.5))