"""
gaussexp.py -- From-scratch Gaussian & exponential-family CONTINUOUS distributions.

Backs the notebooks in  Z-Statistical Distributions-Gaussian-ExpFamily  (Folder 2 of
the Statistical Distributions catalog).

Everything is implemented from the definitions using only elementary math and a
handful of special functions as mathematical primitives -- erf, the regularized
incomplete gamma (gammainc / gammaincc) and the regularized incomplete beta
(betainc). No distribution objects. `scipy.stats` (Python notebook) and R's
d/p/q/r (R notebook) are used as INDEPENDENT references.

For every distribution we provide:
    *_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)

The RNGs form a dependency chain that is itself the story:
    Uniform --(Box-Muller)--> Normal
    Uniform --(inverse-transform)--> Exponential
    Normal + Uniform --(Marsaglia-Tsang rejection)--> Gamma
    Gamma --> Chi-square, Inverse-Gamma, Beta (ratio of two Gammas)
    Normal / Chi-square --> Student-t;  Chi-square ratio --> F;  exp(Normal) --> Log-normal

Parameterisations match scipy so validation is apples-to-apples:
    normal      : (mu, sigma)          -- sigma is the STANDARD DEVIATION
    exponential : rate = 1/scale       -- scipy uses scale, so compare to scale=1/rate
    gamma       : (shape k, scale th)  -- mean = k*th          (scipy gamma(a, scale=th))
    invgamma    : (shape a, scale b)   -- pdf prop x^{-a-1} e^{-b/x}  (scipy invgamma(a, scale=b))
    beta        : (a, b) on (0,1)
    chi2        : df
    student_t   : df                   -- standard t (loc 0, scale 1)
    f           : (d1, d2)
    lognormal   : (mu, sigma) of the underlying Normal  (scipy lognorm(s=sigma, scale=exp(mu)))
"""

import numpy as np
from scipy.special import erf, gammainc, gammaincc, betainc, gammaln, betaln

SQRT2 = np.sqrt(2.0)


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

def moments_from_pdf(pdf, lo, hi, n=400001):
    """Mean/variance/skewness/excess-kurtosis by numerical integration of a pdf
    on a fine grid -- an independent cross-check on the closed-form formulas.
    (Grid-based, so best for light/moderate tails.)"""
    x = np.linspace(lo, hi, n)
    px = np.asarray(pdf(x), float)
    Z = np.trapz(px, x)
    px = px / Z
    m = np.trapz(x * px, x)
    c2 = np.trapz((x - m) ** 2 * px, x)
    c3 = np.trapz((x - m) ** 3 * px, x)
    c4 = np.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 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)


# --------------------------------------------------------------------------- #
#  Continuous Uniform(a, b)                                                      #
# --------------------------------------------------------------------------- #

def uniform_pdf(x, a, b):
    x = np.asarray(x, float)
    return np.where((x >= a) & (x <= b), 1.0 / (b - a), 0.0)

def uniform_cdf(x, a, b):
    return np.clip((np.asarray(x, float) - a) / (b - a), 0.0, 1.0)

def uniform_rng(a, b, n, rng):
    return a + (b - a) * rng.random(n)

def uniform_moments(a, b):
    return dict(mean=(a + b) / 2, var=(b - a) ** 2 / 12, skew=0.0, exkurt=-6.0 / 5)


# --------------------------------------------------------------------------- #
#  Normal(mu, sigma)  --  sampled by Box-Muller                                  #
# --------------------------------------------------------------------------- #

def normal_pdf(x, mu, sigma):
    x = np.asarray(x, float)
    return np.exp(-0.5 * ((x - mu) / sigma) ** 2) / (sigma * np.sqrt(2 * np.pi))

def normal_cdf(x, mu, sigma):
    return 0.5 * (1 + erf((np.asarray(x, float) - mu) / (sigma * SQRT2)))

def normal_rng(mu, sigma, n, rng):
    """Box-Muller: two uniforms U1,U2 -> two standard normals
    R*cos(2 pi U2), R*sin(2 pi U2) with R = sqrt(-2 ln U1)."""
    m = (n + 1) // 2
    u1 = rng.random(m); u2 = rng.random(m)
    r = np.sqrt(-2.0 * 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 mu + sigma * z[:n]

def normal_moments(mu, sigma):
    return dict(mean=mu, var=sigma ** 2, skew=0.0, exkurt=0.0)


# --------------------------------------------------------------------------- #
#  Exponential(rate)  --  inverse transform                                      #
# --------------------------------------------------------------------------- #

def exponential_pdf(x, rate):
    x = np.asarray(x, float)
    return np.where(x >= 0, rate * np.exp(-rate * x), 0.0)

def exponential_cdf(x, rate):
    x = np.asarray(x, float)
    return np.where(x >= 0, 1 - np.exp(-rate * x), 0.0)

def exponential_rng(rate, n, rng):
    """Inverse transform: X = -ln(U)/rate."""
    return -np.log(rng.random(n)) / rate

def exponential_moments(rate):
    return dict(mean=1 / rate, var=1 / rate ** 2, skew=2.0, exkurt=6.0)


# --------------------------------------------------------------------------- #
#  Gamma(shape k, scale theta)  --  Marsaglia-Tsang rejection sampler            #
# --------------------------------------------------------------------------- #

def gamma_pdf(x, k, theta):
    x = np.asarray(x, float)
    logp = (k - 1) * np.log(np.where(x > 0, x, 1.0)) - x / theta - k * np.log(theta) - gammaln(k)
    return np.where(x > 0, np.exp(logp), 0.0)

def gamma_cdf(x, k, theta):
    x = np.asarray(x, float)
    return np.where(x > 0, gammainc(k, x / theta), 0.0)   # regularized lower incomplete gamma P(k, x/theta)

def gamma_rng(k, theta, n, rng):
    """Marsaglia & Tsang (2000). For k>=1: d=k-1/3, c=1/sqrt(9d); propose a
    standard normal x, set v=(1+c x)^3, accept d*v when ln U < x^2/2 + d - d v + d ln v.
    For k<1 use the boosting trick  G_k = G_{k+1} * U^{1/k}. Vectorised batch rejection."""
    n = int(n)
    if k < 1:
        g = gamma_rng(k + 1.0, 1.0, n, rng)
        u = rng.random(n)
        return theta * 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 = normal_rng(0.0, 1.0, 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 theta * res

def gamma_moments(k, theta):
    return dict(mean=k * theta, var=k * theta ** 2, skew=2.0 / np.sqrt(k), exkurt=6.0 / k)


# --------------------------------------------------------------------------- #
#  Inverse-Gamma(shape a, scale b)  --  reciprocal of a Gamma                    #
# --------------------------------------------------------------------------- #

def invgamma_pdf(x, a, b):
    x = np.asarray(x, float)
    logp = a * np.log(b) - gammaln(a) - (a + 1) * np.log(np.where(x > 0, x, 1.0)) - b / np.where(x > 0, x, 1.0)
    return np.where(x > 0, np.exp(logp), 0.0)

def invgamma_cdf(x, a, b):
    x = np.asarray(x, float)
    return np.where(x > 0, gammaincc(a, b / np.where(x > 0, x, 1.0)), 0.0)  # upper incomplete gamma Q(a, b/x)

def invgamma_rng(a, b, n, rng):
    """If Y ~ Gamma(shape a, rate b) then 1/Y ~ Inverse-Gamma(a, b)."""
    y = gamma_rng(a, 1.0 / b, n, rng)     # Gamma(shape a, scale 1/b) == rate b
    return 1.0 / y

def invgamma_moments(a, b):
    mean = b / (a - 1) if a > 1 else np.nan
    var = b ** 2 / ((a - 1) ** 2 * (a - 2)) if a > 2 else np.nan
    skew = 4 * np.sqrt(a - 2) / (a - 3) if a > 3 else np.nan
    exkurt = 6 * (5 * a - 11) / ((a - 3) * (a - 4)) if a > 4 else np.nan
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)


# --------------------------------------------------------------------------- #
#  Chi-square(df)  --  Gamma(df/2, scale 2)                                      #
# --------------------------------------------------------------------------- #

def chi2_pdf(x, df):
    return gamma_pdf(x, df / 2.0, 2.0)

def chi2_cdf(x, df):
    return gamma_cdf(x, df / 2.0, 2.0)

def chi2_rng(df, n, rng):
    return gamma_rng(df / 2.0, 2.0, n, rng)

def chi2_moments(df):
    return dict(mean=df, var=2.0 * df, skew=np.sqrt(8.0 / df), exkurt=12.0 / df)


# --------------------------------------------------------------------------- #
#  Beta(a, b) on (0, 1)  --  ratio of two Gammas                                 #
# --------------------------------------------------------------------------- #

def beta_pdf(x, a, b):
    x = np.asarray(x, float)
    logp = (a - 1) * np.log(np.where((x > 0) & (x < 1), x, 0.5)) + \
           (b - 1) * np.log(np.where((x > 0) & (x < 1), 1 - x, 0.5)) - betaln(a, b)
    return np.where((x > 0) & (x < 1), np.exp(logp), 0.0)

def beta_cdf(x, a, b):
    x = np.asarray(x, float)
    return np.clip(betainc(a, b, np.clip(x, 0, 1)), 0, 1)   # regularized incomplete beta I_x(a,b)

def beta_rng(a, b, n, rng):
    """X ~ Gamma(a,1), Y ~ Gamma(b,1) => X/(X+Y) ~ Beta(a,b)."""
    x = gamma_rng(a, 1.0, n, rng); y = gamma_rng(b, 1.0, n, rng)
    return x / (x + y)

def beta_moments(a, b):
    s = a + b
    mean = a / s
    var = a * b / (s ** 2 * (s + 1))
    skew = 2 * (b - a) * np.sqrt(s + 1) / ((s + 2) * np.sqrt(a * b))
    exkurt = 6 * ((a - b) ** 2 * (s + 1) - a * b * (s + 2)) / (a * b * (s + 2) * (s + 3))
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)


# --------------------------------------------------------------------------- #
#  Student-t(df)  --  Normal / sqrt(Chi-square/df)                               #
# --------------------------------------------------------------------------- #

def student_t_pdf(x, df):
    x = np.asarray(x, float)
    logp = gammaln((df + 1) / 2) - gammaln(df / 2) - 0.5 * np.log(df * np.pi) \
           - (df + 1) / 2 * np.log1p(x ** 2 / df)
    return np.exp(logp)

def student_t_cdf(x, df):
    x = np.asarray(x, float)
    ib = betainc(df / 2.0, 0.5, df / (df + x ** 2))       # I_{df/(df+x^2)}(df/2, 1/2)
    return np.where(x >= 0, 1 - 0.5 * ib, 0.5 * ib)

def student_t_rng(df, n, rng):
    """Z ~ Normal(0,1), V ~ Chi-square(df)  =>  Z / sqrt(V/df) ~ t_df."""
    z = normal_rng(0.0, 1.0, n, rng); v = chi2_rng(df, n, rng)
    return z / np.sqrt(v / df)

def student_t_moments(df):
    mean = 0.0 if df > 1 else np.nan
    var = df / (df - 2) if df > 2 else np.nan
    skew = 0.0 if df > 3 else np.nan
    exkurt = 6.0 / (df - 4) if df > 4 else np.nan
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)


# --------------------------------------------------------------------------- #
#  F(d1, d2)  --  ratio of two scaled Chi-squares                               #
# --------------------------------------------------------------------------- #

def f_pdf(x, d1, d2):
    x = np.asarray(x, float)
    xp = np.where(x > 0, x, 1.0)
    logp = 0.5 * (d1 * np.log(d1) + d2 * np.log(d2)) + (d1 / 2 - 1) * np.log(xp) \
           - 0.5 * (d1 + d2) * np.log(d1 * xp + d2) - betaln(d1 / 2, d2 / 2)
    return np.where(x > 0, np.exp(logp), 0.0)

def f_cdf(x, d1, d2):
    x = np.asarray(x, float)
    return np.where(x > 0, betainc(d1 / 2, d2 / 2, d1 * x / (d1 * x + d2)), 0.0)

def f_rng(d1, d2, n, rng):
    """(U1/d1)/(U2/d2) with U1 ~ Chi-square(d1), U2 ~ Chi-square(d2)."""
    u1 = chi2_rng(d1, n, rng); u2 = chi2_rng(d2, n, rng)
    return (u1 / d1) / (u2 / d2)

def f_moments(d1, d2):
    mean = d2 / (d2 - 2) if d2 > 2 else np.nan
    var = (2 * d2 ** 2 * (d1 + d2 - 2) / (d1 * (d2 - 2) ** 2 * (d2 - 4))) if d2 > 4 else np.nan
    skew = ((2 * d1 + d2 - 2) * np.sqrt(8 * (d2 - 4)) /
            ((d2 - 6) * np.sqrt(d1 * (d1 + d2 - 2)))) if d2 > 6 else np.nan
    exkurt = (12 * (d1 * (5 * d2 - 22) * (d1 + d2 - 2) + (d2 - 4) * (d2 - 2) ** 2) /
              (d1 * (d2 - 6) * (d2 - 8) * (d1 + d2 - 2))) if d2 > 8 else np.nan
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)


# --------------------------------------------------------------------------- #
#  Log-normal(mu, sigma)  --  exp of a Normal                                    #
# --------------------------------------------------------------------------- #

def lognormal_pdf(x, mu, sigma):
    x = np.asarray(x, float); xp = np.where(x > 0, x, 1.0)
    logp = -np.log(xp * sigma * np.sqrt(2 * np.pi)) - (np.log(xp) - mu) ** 2 / (2 * sigma ** 2)
    return np.where(x > 0, np.exp(logp), 0.0)

def lognormal_cdf(x, mu, sigma):
    x = np.asarray(x, float); xp = np.where(x > 0, x, 1.0)
    return np.where(x > 0, 0.5 * (1 + erf((np.log(xp) - mu) / (sigma * SQRT2))), 0.0)

def lognormal_rng(mu, sigma, n, rng):
    """exp(Normal(mu, sigma)) -- the multiplicative analogue of the Gaussian."""
    return np.exp(normal_rng(mu, sigma, n, rng))

def lognormal_moments(mu, sigma):
    s2 = sigma ** 2
    mean = np.exp(mu + s2 / 2)
    var = (np.exp(s2) - 1) * np.exp(2 * mu + s2)
    skew = (np.exp(s2) + 2) * np.sqrt(np.exp(s2) - 1)
    exkurt = np.exp(4 * s2) + 2 * np.exp(3 * s2) + 3 * np.exp(2 * s2) - 6
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)