"""Small-jump Gaussian compensation for an Euler neural-SDE update.

The implementation uses a symmetric Levy measure (random +/- signs), so the
retained large jumps are already centered in distribution. For a one-sided
uncompensated process with alpha < 1, ``small_mean`` gives the optional exact
mean correction described in the idea.
"""
import math
import numpy as np


def small_variance(epsilon, alpha, c=1.0):
    if not (0.0 < alpha < 2.0 and epsilon > 0.0 and c > 0.0):
        raise ValueError("require 0 < alpha < 2, epsilon > 0, c > 0")
    return c * epsilon ** (2.0 - alpha) / (2.0 - alpha)


def small_mean(epsilon, alpha, c=1.0):
    """Mean of discarded positive jumps; finite only for alpha < 1."""
    if alpha >= 1.0:
        raise ValueError("positive-jump mean diverges for alpha >= 1")
    return c * epsilon ** (1.0 - alpha) / (1.0 - alpha)


def compensated_euler_step(x, drift, h, epsilon, alpha, c=1.0,
                           rng=None, one_sided=False, add_small_mean=False):
    """Advance a batch of states by one Levy-perturbed Euler step.

    Args:
      x: array [batch, ...]; drift has the same shape and is f_theta(x,t).
      h, epsilon, alpha, c: Euler time step and Levy parameters.
      one_sided: use positive jumps instead of symmetric signs.
      add_small_mean: for one-sided alpha<1, add h times the exact discarded
        small-jump mean (useful for an uncompensated positive process).
    Returns: (new_x, number_of_large_jumps_sampled).
    """
    if rng is None:
        rng = np.random.default_rng()
    x = np.asarray(x)
    drift = np.asarray(drift)
    if x.shape != drift.shape or h <= 0:
        raise ValueError("x and drift must have equal shape; h must be positive")
    batch = x.shape[0]
    rate = h * c * epsilon ** (-alpha) / alpha
    counts = rng.poisson(rate, size=batch)
    total = int(counts.sum())
    large = np.zeros_like(x, dtype=float)
    if total:
        # Conditional law of U given U >= epsilon:
        # P(U > u)=(u/epsilon)^(-alpha).
        u = epsilon * rng.random(total) ** (-1.0 / alpha)
        signs = np.ones(total) if one_sided else rng.choice([-1.0, 1.0], total)
        owner = np.repeat(np.arange(batch), counts)
        sums = np.bincount(owner, weights=signs * u, minlength=batch)
        large[:, ...] = sums.reshape((batch,) + (1,) * (x.ndim - 1))
    small = rng.normal(0.0, math.sqrt(h * small_variance(epsilon, alpha, c)),
                       size=x.shape)
    if one_sided and add_small_mean:
        small += h * small_mean(epsilon, alpha, c)
    return x + h * drift + large + small, total