"""Gaussian compensation for the small jumps of a stable Levy noise."""
import math
import numpy as np
from scipy.stats import levy_stable


def small_jump_variance(eps, alpha, c=1.0):
    return c * eps ** (2.0 - alpha) / (2.0 - alpha)


def small_jump_third_abs_moment(eps, alpha, c=1.0):
    return c * eps ** (3.0 - alpha) / (3.0 - alpha)


def stable_scale(alpha, c=1.0):
    # Symmetric Levy measure with total radial intensity c*u^(-1-alpha)du.
    # Integral (cos(tx)-1)c*x^(-1-alpha)dx = -scale**alpha*|t|**alpha.
    if abs(alpha - 1.0) < 1e-12:
        return c * math.pi / 2.0
    return -c * math.gamma(-alpha) * math.cos(math.pi * alpha / 2.0)


def sample_exact_small_symmetric(eps, alpha, n, rng, c=1.0):
    """Sample the centered residual over 0<u<eps.

    It is represented as a symmetric alpha-stable variable minus its jumps
    above eps. This avoids imposing an artificial lower cutoff.
    """
    scale = stable_scale(alpha, c) ** (1.0 / alpha)
    # scipy's random_state accepts a RandomState, not Generator.
    rs = np.random.RandomState(rng.integers(0, 2**31 - 1))
    full = levy_stable.rvs(alpha, 0.0, loc=0.0, scale=scale,
                           size=n, random_state=rs)
    rate = c * eps ** (-alpha) / alpha
    k = rng.poisson(rate, size=n)
    total = int(k.sum())
    if total:
        # Conditional radial law: P(U>u)=(u/eps)^(-alpha).
        u = eps * rng.random(total) ** (-1.0 / alpha)
        signs = rng.choice(np.array([-1.0, 1.0]), size=total)
        owners = np.repeat(np.arange(n), k)
        large = np.bincount(owners, weights=signs * u, minlength=n)
    else:
        large = np.zeros(n)
    return full - large, int(total)


def approximate_samples(eps, alpha, n, rng, c=1.0):
    exact, jumps = sample_exact_small_symmetric(eps, alpha, n, rng, c)
    naive = np.zeros(n)
    gaussian = rng.normal(0.0, math.sqrt(small_jump_variance(eps, alpha, c)), n)
    return exact, naive, gaussian, jumps


def w1_1d(x, y):
    """Empirical one-dimensional W1 for equal-size samples."""
    return float(np.mean(np.abs(np.sort(x) - np.sort(y))))