"""
heavytails.py -- From-scratch HEAVY-TAILED & SKEWED continuous distributions.

Backs the notebooks in  Z-Statistical Distributions-HeavyTails-Skew  (Folder 3 of
the Statistical Distributions catalog).

These are the laws built to model what the light-tailed Gaussian family misses:
fat tails (rare extremes) and asymmetry. Everything is implemented from the
definitions using only elementary math and a few special functions as primitives
(erf/erfc, the regularized incomplete gamma). No distribution objects. `scipy.stats`
(Python notebook) and R (R notebook) are used as INDEPENDENT references.

For every distribution we provide, where they are defined:
    *_pdf(x, ...)      probability density function
    *_cdf(x, ...)      cumulative distribution function
    *_rng(..., n, rng) a random-number generator (the sampling ALGORITHM matters)
    *_moments(...)     mean, variance, skewness, excess kurtosis (closed form; NaN where undefined)

Highlights (the from-scratch RNG algorithms):
    Laplace, Logistic, Cauchy, Pareto : inverse transform (closed-form quantile)
    GED (exponential power)           : signed Gamma^(1/beta)
    Skew-Normal                       : Azzalini's  delta|Z0| + sqrt(1-delta^2) Z1
    alpha-Stable                      : Chambers-Mallows-Stuck (CMS) -- the general heavy-tail family
    Levy                              : c / Normal^2  (the alpha=1/2 stable)
    Hansen skew-t                     : numerical inverse-transform (no elementary quantile)

Parameterisations match scipy where a match exists:
    laplace(loc, scale) ; logistic(loc, scale) ; cauchy(loc, scale)
    student_t_ls(loc, scale, df)          -- location-scale Student-t
    gennorm/GED(loc, scale, beta)         -- scipy.stats.gennorm(beta, loc, scale)
    skewnorm(xi, omega, alpha)            -- scipy.stats.skewnorm(alpha, loc=xi, scale=omega)
    pareto(xm, alpha)                     -- scipy.stats.pareto(alpha, scale=xm)
    levy(loc, scale)                      -- scipy.stats.levy(loc, scale)
    hansen_skewt(eta, lam)                -- standardized (mean 0, variance 1), -1<lam<1, eta>2
"""

import numpy as np
from scipy.special import erf, erfc, gamma as Gamma, gammaln, gammainc

SQRT2 = np.sqrt(2.0)
_trapz = getattr(np, "trapezoid", getattr(np, "trapz", None))   # np.trapz renamed in NumPy 2.0


# --------------------------------------------------------------------------- #
#  Generic helpers                                                              #
# --------------------------------------------------------------------------- #

def emp_moments(x):
    x = np.asarray(x, float); m = x.mean(); c2 = x.var()
    return dict(mean=m, var=c2, skew=((x - m) ** 3).mean() / c2 ** 1.5,
                exkurt=((x - m) ** 4).mean() / c2 ** 2 - 3)


def moments_from_pdf(pdf, lo, hi, n=800001):
    """Mean/var/skew/excess-kurtosis by numerical integration of a pdf on a fine
    grid -- an independent cross-check where closed forms exist (and the only route
    for the Hansen skew-t)."""
    x = np.linspace(lo, hi, n)
    px = np.asarray(pdf(x), float)
    Z = _trapz(px, x); px = px / Z
    m = _trapz(x * px, x)
    c2 = _trapz((x - m) ** 2 * px, x)
    c3 = _trapz((x - m) ** 3 * px, x)
    c4 = _trapz((x - m) ** 4 * px, x)
    sd = np.sqrt(c2)
    return dict(mean=m, var=c2, skew=c3 / sd ** 3, exkurt=c4 / c2 ** 2 - 3)


def _grid_cdf(pdf, xs, lo, npts=20001):
    """CDF at each point of xs by cumulative-trapezoid integration of pdf from lo."""
    xs = np.asarray(xs, float)
    out = np.empty_like(xs)
    for i, xi in enumerate(xs):
        g = np.linspace(lo, xi, npts)
        out[i] = _trapz(pdf(g), g)
    return np.clip(out, 0, 1)


def inverse_transform_sample(pdf, lo, hi, n, rng, npts=200001):
    """General inverse-transform sampler for a density with no elementary quantile:
    build the CDF on a fine grid (cumulative trapezoid), then map uniforms through
    the inverse CDF by interpolation."""
    g = np.linspace(lo, hi, npts)
    pg = np.asarray(pdf(g), float)
    cdf = np.concatenate([[0.0], np.cumsum((pg[1:] + pg[:-1]) / 2 * np.diff(g))])
    cdf /= cdf[-1]
    u = rng.random(n)
    return np.interp(u, cdf, g)


# --------------------------------------------------------------------------- #
#  Laplace(loc mu, scale b)  --  double exponential                             #
# --------------------------------------------------------------------------- #

def laplace_pdf(x, mu, b):
    x = np.asarray(x, float)
    return np.exp(-np.abs(x - mu) / b) / (2 * b)

def laplace_cdf(x, mu, b):
    x = np.asarray(x, float); z = (x - mu) / b
    return np.where(x < mu, 0.5 * np.exp(z), 1 - 0.5 * np.exp(-z))

def laplace_rng(mu, b, n, rng):
    """Inverse transform: U-0.5 -> mu - b*sign(U-.5)*ln(1-2|U-.5|)."""
    u = rng.random(n) - 0.5
    return mu - b * np.sign(u) * np.log1p(-2 * np.abs(u))

def laplace_moments(mu, b):
    return dict(mean=mu, var=2 * b ** 2, skew=0.0, exkurt=3.0)


# --------------------------------------------------------------------------- #
#  Logistic(loc mu, scale s)                                                    #
# --------------------------------------------------------------------------- #

def logistic_pdf(x, mu, s):
    z = (np.asarray(x, float) - mu) / s
    e = np.exp(-np.abs(z))                      # stable form
    return e / (s * (1 + e) ** 2)

def logistic_cdf(x, mu, s):
    z = (np.asarray(x, float) - mu) / s
    return 1.0 / (1 + np.exp(-z))

def logistic_rng(mu, s, n, rng):
    """Inverse transform (the logit): mu + s*ln(U/(1-U))."""
    u = rng.random(n)
    return mu + s * np.log(u / (1 - u))

def logistic_moments(mu, s):
    return dict(mean=mu, var=s ** 2 * np.pi ** 2 / 3, skew=0.0, exkurt=1.2)


# --------------------------------------------------------------------------- #
#  Cauchy(loc x0, scale g)  --  the no-moments distribution (t with df=1)        #
# --------------------------------------------------------------------------- #

def cauchy_pdf(x, x0, g):
    z = (np.asarray(x, float) - x0) / g
    return 1.0 / (np.pi * g * (1 + z ** 2))

def cauchy_cdf(x, x0, g):
    z = (np.asarray(x, float) - x0) / g
    return 0.5 + np.arctan(z) / np.pi

def cauchy_rng(x0, g, n, rng):
    """Inverse transform: x0 + g*tan(pi*(U-1/2))."""
    return x0 + g * np.tan(np.pi * (rng.random(n) - 0.5))

def cauchy_moments(x0, g):
    return dict(mean=np.nan, var=np.nan, skew=np.nan, exkurt=np.nan)  # none exist


# --------------------------------------------------------------------------- #
#  Student-t, location-scale (loc mu, scale sig, df nu)  --  tunable tails       #
# --------------------------------------------------------------------------- #

def student_t_ls_pdf(x, mu, sig, nu):
    z = (np.asarray(x, float) - mu) / sig
    logp = gammaln((nu + 1) / 2) - gammaln(nu / 2) - 0.5 * np.log(nu * np.pi) \
           - (nu + 1) / 2 * np.log1p(z ** 2 / nu)
    return np.exp(logp) / sig

def student_t_ls_cdf(x, mu, sig, nu):
    from scipy.special import betainc
    z = (np.asarray(x, float) - mu) / sig
    ib = betainc(nu / 2.0, 0.5, nu / (nu + z ** 2))
    return np.where(z >= 0, 1 - 0.5 * ib, 0.5 * ib)

def student_t_ls_rng(mu, sig, nu, n, rng):
    """mu + sig * Z / sqrt(V/nu),  Z~Normal, V~Chi-square(nu) (built from Gammas)."""
    z = _std_normal(n, rng)
    v = 2.0 * _gamma_shape(nu / 2.0, n, rng)        # Chi-square(nu) = Gamma(nu/2, scale 2)
    return mu + sig * z / np.sqrt(v / nu)

def student_t_ls_moments(mu, sig, nu):
    mean = mu if nu > 1 else np.nan
    var = sig ** 2 * nu / (nu - 2) if nu > 2 else np.nan
    skew = 0.0 if nu > 3 else np.nan
    exkurt = 6.0 / (nu - 4) if nu > 4 else np.nan
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)


# --------------------------------------------------------------------------- #
#  Generalized Error Distribution / exponential power (loc mu, scale a, shape b) #
#     pdf = b / (2 a Gamma(1/b)) * exp( -(|x-mu|/a)^b )                          #
# --------------------------------------------------------------------------- #

def ged_pdf(x, mu, a, b):
    z = np.abs(np.asarray(x, float) - mu) / a
    return b / (2 * a * Gamma(1.0 / b)) * np.exp(-z ** b)

def ged_cdf(x, mu, a, b):
    x = np.asarray(x, float); z = np.abs(x - mu) / a
    half = 0.5 * gammainc(1.0 / b, z ** b)          # regularized lower incomplete gamma
    return np.where(x >= mu, 0.5 + half, 0.5 - half)

def ged_rng(mu, a, b, n, rng):
    """|X-mu| = a * G^(1/b) with G ~ Gamma(1/b, 1); attach a random sign."""
    g = _gamma_shape(1.0 / b, n, rng)
    sign = np.where(rng.random(n) < 0.5, -1.0, 1.0)
    return mu + sign * a * g ** (1.0 / b)

def ged_moments(mu, a, b):
    var = a ** 2 * Gamma(3.0 / b) / Gamma(1.0 / b)
    exkurt = Gamma(5.0 / b) * Gamma(1.0 / b) / Gamma(3.0 / b) ** 2 - 3
    return dict(mean=mu, var=var, skew=0.0, exkurt=exkurt)


# --------------------------------------------------------------------------- #
#  Skew-Normal(loc xi, scale omega, shape alpha)  --  Azzalini                   #
# --------------------------------------------------------------------------- #

def _std_normal_pdf(z):
    return np.exp(-0.5 * z ** 2) / np.sqrt(2 * np.pi)

def _std_normal_cdf(z):
    return 0.5 * (1 + erf(z / SQRT2))

def skewnorm_pdf(x, xi, omega, alpha):
    z = (np.asarray(x, float) - xi) / omega
    return 2.0 / omega * _std_normal_pdf(z) * _std_normal_cdf(alpha * z)

def skewnorm_cdf(x, xi, omega, alpha):
    # no elementary form (needs Owen's T); integrate the density numerically
    return _grid_cdf(lambda t: skewnorm_pdf(t, xi, omega, alpha), x, xi - 40 * omega)

def skewnorm_rng(xi, omega, alpha, n, rng):
    """Azzalini construction: delta=alpha/sqrt(1+alpha^2);
    Z = delta*|U0| + sqrt(1-delta^2)*U1 with U0,U1 iid Normal."""
    delta = alpha / np.sqrt(1 + alpha ** 2)
    u0 = _std_normal(n, rng); u1 = _std_normal(n, rng)
    z = delta * np.abs(u0) + np.sqrt(1 - delta ** 2) * u1
    return xi + omega * z

def skewnorm_moments(xi, omega, alpha):
    delta = alpha / np.sqrt(1 + alpha ** 2)
    b = np.sqrt(2 / np.pi) * delta
    mean = xi + omega * b
    var = omega ** 2 * (1 - b ** 2)
    skew = (4 - np.pi) / 2 * b ** 3 / (1 - b ** 2) ** 1.5
    exkurt = 2 * (np.pi - 3) * b ** 4 / (1 - b ** 2) ** 2
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)


# --------------------------------------------------------------------------- #
#  Pareto Type I (scale xm, shape alpha)  --  the canonical power-law tail       #
# --------------------------------------------------------------------------- #

def pareto_pdf(x, xm, alpha):
    x = np.asarray(x, float)
    return np.where(x >= xm, alpha * xm ** alpha / x ** (alpha + 1), 0.0)

def pareto_cdf(x, xm, alpha):
    x = np.asarray(x, float)
    return np.where(x >= xm, 1 - (xm / x) ** alpha, 0.0)

def pareto_rng(xm, alpha, n, rng):
    """Inverse transform: xm / U^(1/alpha)  (U ~ Uniform(0,1))."""
    return xm / rng.random(n) ** (1.0 / alpha)

def pareto_moments(xm, alpha):
    mean = alpha * xm / (alpha - 1) if alpha > 1 else np.nan
    var = xm ** 2 * alpha / ((alpha - 1) ** 2 * (alpha - 2)) if alpha > 2 else np.nan
    skew = (2 * (1 + alpha) / (alpha - 3) * np.sqrt((alpha - 2) / alpha)) if alpha > 3 else np.nan
    exkurt = (6 * (alpha ** 3 + alpha ** 2 - 6 * alpha - 2) /
              (alpha * (alpha - 3) * (alpha - 4))) if alpha > 4 else np.nan
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)


# --------------------------------------------------------------------------- #
#  Levy(loc mu, scale c)  --  one-sided, the alpha=1/2 stable                    #
# --------------------------------------------------------------------------- #

def levy_pdf(x, mu, c):
    x = np.asarray(x, float); z = x - mu
    zp = np.where(z > 0, z, 1.0)
    return np.where(z > 0, np.sqrt(c / (2 * np.pi)) * np.exp(-c / (2 * zp)) / zp ** 1.5, 0.0)

def levy_cdf(x, mu, c):
    x = np.asarray(x, float); z = x - mu
    zp = np.where(z > 0, z, np.inf)
    return np.where(z > 0, erfc(np.sqrt(c / (2 * zp))), 0.0)

def levy_rng(mu, c, n, rng):
    """X = mu + c / Z^2  with Z ~ Normal(0,1) -- the alpha=1/2, beta=1 stable law."""
    z = _std_normal(n, rng)
    return mu + c / z ** 2

def levy_moments(mu, c):
    return dict(mean=np.nan, var=np.nan, skew=np.nan, exkurt=np.nan)  # even the mean is infinite


# --------------------------------------------------------------------------- #
#  alpha-Stable(alpha, beta)  --  Chambers-Mallows-Stuck sampler                 #
#     Standard S(alpha, beta; loc 0, scale 1).  No elementary pdf in general.    #
# --------------------------------------------------------------------------- #

def stable_rng(alpha, beta, n, rng, loc=0.0, scale=1.0):
    """Chambers, Mallows & Stuck (1976). Draw W~Exp(1), U~Uniform(-pi/2, pi/2)."""
    U = np.pi * (rng.random(n) - 0.5)
    W = -np.log(rng.random(n))
    if abs(alpha - 1.0) > 1e-8:
        zeta = -beta * np.tan(np.pi * alpha / 2)
        xi = np.arctan(-zeta) / alpha
        X = ((1 + zeta ** 2) ** (1 / (2 * alpha))
             * np.sin(alpha * (U + xi)) / np.cos(U) ** (1 / alpha)
             * (np.cos(U - alpha * (U + xi)) / W) ** ((1 - alpha) / alpha))
    else:
        X = (2 / np.pi) * ((np.pi / 2 + beta * U) * np.tan(U)
                           - beta * np.log((np.pi / 2 * W * np.cos(U)) / (np.pi / 2 + beta * U)))
    return loc + scale * X

def stable_cf(t, alpha, beta):
    """Theoretical characteristic function phi(t)=E[e^{itX}] of the standard stable
    law (loc 0, scale 1), used to validate the CMS sampler where no pdf exists."""
    t = np.asarray(t, float)
    if abs(alpha - 1.0) > 1e-8:
        phi = np.tan(np.pi * alpha / 2)
        return np.exp(-np.abs(t) ** alpha * (1 - 1j * beta * np.sign(t) * phi))
    else:
        return np.exp(-np.abs(t) * (1 + 1j * beta * (2 / np.pi) * np.sign(t) * np.log(np.abs(t) + (t == 0))))


# --------------------------------------------------------------------------- #
#  Hansen (1994) skew-t  --  standardized (mean 0, var 1); heavy tails + skew    #
# --------------------------------------------------------------------------- #

def _hansen_abc(eta, lam):
    c = Gamma((eta + 1) / 2) / (np.sqrt(np.pi * (eta - 2)) * Gamma(eta / 2))
    a = 4 * lam * c * (eta - 2) / (eta - 1)
    b = np.sqrt(1 + 3 * lam ** 2 - a ** 2)
    return a, b, c

def hansen_skewt_pdf(z, eta, lam):
    z = np.asarray(z, float)
    a, b, c = _hansen_abc(eta, lam)
    denom = np.where(z < -a / b, 1 - lam, 1 + lam)
    inner = 1 + 1.0 / (eta - 2) * ((b * z + a) / denom) ** 2
    return b * c * inner ** (-(eta + 1) / 2)

def hansen_skewt_cdf(z, eta, lam):
    return _grid_cdf(lambda t: hansen_skewt_pdf(t, eta, lam), z, -60.0)

def hansen_skewt_rng(eta, lam, n, rng):
    """No elementary quantile -> general inverse-transform sampling of the density."""
    return inverse_transform_sample(lambda t: hansen_skewt_pdf(t, eta, lam), -60.0, 60.0, n, rng)

def hansen_skewt_moments(eta, lam):
    # mean 0 and variance 1 by construction; skew/kurtosis by numerical integration
    m = moments_from_pdf(lambda t: hansen_skewt_pdf(t, eta, lam), -80.0, 80.0)
    return dict(mean=0.0, var=1.0, skew=m["skew"], exkurt=m["exkurt"])


# --------------------------------------------------------------------------- #
#  Private RNG primitives (self-contained: Box-Muller normal, Marsaglia-Tsang)  #
# --------------------------------------------------------------------------- #

def _std_normal(n, rng):
    m = (n + 1) // 2
    u1 = rng.random(m); u2 = rng.random(m)
    r = np.sqrt(-2 * np.log(u1))
    z = np.empty(2 * m)
    z[:m] = r * np.cos(2 * np.pi * u2)
    z[m:] = r * np.sin(2 * np.pi * u2)
    return z[:n]

def _gamma_shape(k, n, rng):
    """Gamma(shape k, scale 1) by Marsaglia-Tsang (vectorised batch rejection)."""
    n = int(n)
    if k < 1:
        g = _gamma_shape(k + 1.0, n, rng)
        u = rng.random(n)
        return g * u ** (1.0 / k)
    d = k - 1.0 / 3.0; c = 1.0 / np.sqrt(9.0 * d)
    res = np.empty(n); todo = np.ones(n, bool)
    while todo.any():
        m = int(todo.sum())
        x = _std_normal(m, rng)
        v = (1.0 + c * x) ** 3
        u = rng.random(m)
        ok = (v > 0) & (np.log(u) < 0.5 * x ** 2 + d - d * v + d * np.log(np.where(v > 0, v, 1.0)))
        cur = np.where(todo)[0]; acc = cur[ok]
        res[acc] = d * v[ok]; todo[acc] = False
    return res