"""
bayespriors.py -- From-scratch BAYESIAN PRIOR & DATA-AUGMENTATION distributions.

Backs the notebooks in  Z-Statistical Distributions-BayesianPriors-Augmentation
(Folder 5 of the Statistical Distributions catalog).

These are the distributions that power modern Bayesian *computation* rather than
data modelling: constrained and weakly-informative priors, the positive mixing
laws behind scale-mixture and shrinkage models, the latent-variable distribution
that makes logistic likelihoods conditionally Gaussian, and the prior on the
probability simplex.

Everything is implemented from the definitions using only elementary math and a
few special functions as primitives (erf/erfinv, the modified Bessel K, the
log-gamma). No distribution objects. `scipy.stats` (Python) and R (R notebook) are
used as INDEPENDENT references; where a law has no closed-form density
(Polya-Gamma) we validate through its moments and Laplace transform.

Per distribution, where defined:
    *_pdf, *_cdf, *_rng(..., n, rng), *_moments

Highlights (the from-scratch RNG algorithms):
    Truncated Normal   : inverse-CDF between the truncation points
    Half-*             : fold a symmetric law at 0
    Inverse Gaussian   : Michael-Schucany-Haas (exact)
    GIG                : numerical inverse-transform (Bessel-K density)
    Polya-Gamma        : infinite Gamma-sum representation (Polson-Scott-Windle)
    Dirichlet          : normalised independent Gammas

Parameterisations match scipy where a match exists:
    truncnorm(mu, sigma, a, b)   scipy truncnorm((a-mu)/sigma,(b-mu)/sigma,mu,sigma)
    halfnorm(sigma)              scipy halfnorm(scale=sigma)
    halfcauchy(scale)            scipy halfcauchy(scale=scale)
    invgauss(mu, lam)            scipy invgauss(mu/lam, scale=lam)
    gig(p, a, b)                 scipy geninvgauss(p, sqrt(a*b), scale=sqrt(b/a))
    dirichlet(alpha)             scipy dirichlet(alpha)
    polyagamma(b, c)             no scipy; validated via moments and Laplace transform
"""

import numpy as np
from scipy.special import erf, erfinv, gammaln, kv          # kv = modified Bessel K

SQRT2 = np.sqrt(2.0)
_trapz = getattr(np, "trapezoid", getattr(np, "trapz", None))


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

def _phi(z):  return np.exp(-0.5 * z ** 2) / np.sqrt(2 * np.pi)
def _Phi(z):  return 0.5 * (1 + erf(np.asarray(z, float) / SQRT2))
def _Phinv(p): return SQRT2 * erfinv(2 * np.asarray(p, float) - 1)

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):
    x = np.linspace(lo, hi, n); px = np.asarray(pdf(x), float)
    px = px / _trapz(px, x)
    m1 = _trapz(x * px, x); c2 = _trapz((x - m1) ** 2 * px, x)
    c3 = _trapz((x - m1) ** 3 * px, x); c4 = _trapz((x - m1) ** 4 * px, x)
    sd = np.sqrt(c2)
    return dict(mean=m1, var=c2, skew=c3 / sd ** 3, exkurt=c4 / c2 ** 2 - 3)

def inverse_transform_sample(pdf, lo, hi, n, rng, npts=200001):
    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]
    return np.interp(rng.random(n), cdf, g)


# ===========================================================================  #
#  PART A -- TRUNCATED & HALF DISTRIBUTIONS (constrained / scale priors)        #
# ===========================================================================  #

# --- Truncated Normal(mu, sigma; a, b) --------------------------------------- #
def truncnorm_pdf(x, mu, sigma, a, b):
    x = np.asarray(x, float)
    al, be = (a - mu) / sigma, (b - mu) / sigma
    Z = _Phi(be) - _Phi(al)
    return np.where((x >= a) & (x <= b), _phi((x - mu) / sigma) / (sigma * Z), 0.0)

def truncnorm_cdf(x, mu, sigma, a, b):
    x = np.asarray(x, float)
    al, be = (a - mu) / sigma, (b - mu) / sigma
    Z = _Phi(be) - _Phi(al)
    out = (_Phi((np.clip(x, a, b) - mu) / sigma) - _Phi(al)) / Z
    return np.where(x < a, 0.0, np.where(x > b, 1.0, out))

def truncnorm_rng(mu, sigma, a, b, n, rng):
    """Inverse-CDF between the truncation points: draw U, map through
    Phi^{-1}(Phi(alpha) + U*(Phi(beta)-Phi(alpha)))."""
    al, be = (a - mu) / sigma, (b - mu) / sigma
    Pa, Pb = _Phi(al), _Phi(be)
    return mu + sigma * _Phinv(Pa + rng.random(n) * (Pb - Pa))

def truncnorm_moments(mu, sigma, a, b):
    al, be = (a - mu) / sigma, (b - mu) / sigma
    Z = _Phi(be) - _Phi(al)
    pa, pb = _phi(al), _phi(be)
    mean = mu + sigma * (pa - pb) / Z
    # skew/kurt are messy in closed form; integrate the (bounded-ish) density
    lo = a if np.isfinite(a) else mu - 40 * sigma
    hi = b if np.isfinite(b) else mu + 40 * sigma
    m = moments_from_pdf(lambda t: truncnorm_pdf(t, mu, sigma, a, b), lo, hi)
    m["mean"] = mean
    return m


# --- Half-Normal(sigma) ------------------------------------------------------ #
def halfnorm_pdf(x, sigma):
    x = np.asarray(x, float)
    return np.where(x >= 0, np.sqrt(2 / np.pi) / sigma * np.exp(-x ** 2 / (2 * sigma ** 2)), 0.0)

def halfnorm_cdf(x, sigma):
    x = np.asarray(x, float)
    return np.where(x >= 0, 2 * _Phi(x / sigma) - 1, 0.0)

def halfnorm_rng(sigma, n, rng):
    return np.abs(sigma * _std_normal(n, rng))

def halfnorm_moments(sigma):
    mean = sigma * np.sqrt(2 / np.pi)
    var = sigma ** 2 * (1 - 2 / np.pi)
    skew = np.sqrt(2) * (4 - np.pi) / (np.pi - 2) ** 1.5
    exkurt = 8 * (np.pi - 3) / (np.pi - 2) ** 2
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)


# --- Half-Cauchy(scale) : Gelman's weakly-informative scale prior ------------ #
def halfcauchy_pdf(x, s):
    x = np.asarray(x, float)
    return np.where(x >= 0, 2.0 / (np.pi * s * (1 + (x / s) ** 2)), 0.0)

def halfcauchy_cdf(x, s):
    x = np.asarray(x, float)
    return np.where(x >= 0, 2.0 / np.pi * np.arctan(x / s), 0.0)

def halfcauchy_rng(s, n, rng):
    return np.abs(s * np.tan(np.pi * (rng.random(n) - 0.5)))

def halfcauchy_moments(s):
    return dict(mean=np.nan, var=np.nan, skew=np.nan, exkurt=np.nan)   # no moments


# --- Half-Student-t(sigma, nu) : general scale prior ------------------------- #
def _t_pdf(x, nu):
    return np.exp(gammaln((nu + 1) / 2) - gammaln(nu / 2) - 0.5 * np.log(nu * np.pi)
                  - (nu + 1) / 2 * np.log1p(np.asarray(x, float) ** 2 / nu))

def halfstudent_pdf(x, sigma, nu):
    x = np.asarray(x, float)
    return np.where(x >= 0, 2.0 / sigma * _t_pdf(x / sigma, nu), 0.0)

def halfstudent_cdf(x, sigma, nu):
    from scipy.special import betainc
    x = np.asarray(x, float); z = np.maximum(x, 0) / sigma
    tcdf = 1 - 0.5 * betainc(nu / 2, 0.5, nu / (nu + z ** 2))
    return np.where(x >= 0, 2 * tcdf - 1, 0.0)

def halfstudent_rng(sigma, nu, n, rng):
    z = _std_normal(n, rng); v = 2.0 * _gamma_shape(nu / 2.0, n, rng)  # chi2(nu)=Gamma(nu/2,2)
    return np.abs(sigma * z / np.sqrt(v / nu))

def halfstudent_moments(sigma, nu):
    return moments_from_pdf(lambda t: halfstudent_pdf(t, sigma, nu), 0.0, 400.0 * sigma)


# ===========================================================================  #
#  PART B -- POSITIVE MIXING DISTRIBUTIONS (scale mixtures / shrinkage)         #
# ===========================================================================  #

# --- Inverse Gaussian / Wald(mu, lam) ---------------------------------------- #
def invgauss_pdf(x, mu, lam):
    x = np.asarray(x, float); xp = np.where(x > 0, x, 1.0)
    return np.where(x > 0, np.sqrt(lam / (2 * np.pi * xp ** 3)) *
                    np.exp(-lam * (xp - mu) ** 2 / (2 * mu ** 2 * xp)), 0.0)

def invgauss_cdf(x, mu, lam):
    x = np.asarray(x, float); xp = np.where(x > 0, x, 1e-300)
    r = np.sqrt(lam / xp)
    out = _Phi(r * (xp / mu - 1)) + np.exp(2 * lam / mu) * _Phi(-r * (xp / mu + 1))
    return np.where(x > 0, out, 0.0)

def invgauss_rng(mu, lam, n, rng):
    """Michael-Schucany-Haas: y=Z^2; x=mu+mu^2 y/(2lam)-(mu/2lam)sqrt(4 mu lam y+mu^2 y^2);
    return x with prob mu/(mu+x), else mu^2/x."""
    y = _std_normal(n, rng) ** 2
    x = mu + mu ** 2 * y / (2 * lam) - (mu / (2 * lam)) * np.sqrt(4 * mu * lam * y + mu ** 2 * y ** 2)
    u = rng.random(n)
    return np.where(u <= mu / (mu + x), x, mu ** 2 / x)

def invgauss_moments(mu, lam):
    return dict(mean=mu, var=mu ** 3 / lam, skew=3 * np.sqrt(mu / lam), exkurt=15 * mu / lam)


# --- Generalized Inverse Gaussian GIG(p, a, b) ------------------------------- #
def gig_pdf(x, p, a, b):
    x = np.asarray(x, float); xp = np.where(x > 0, x, 1.0)
    w = np.sqrt(a * b)
    const = (a / b) ** (p / 2) / (2 * kv(p, w))
    return np.where(x > 0, const * xp ** (p - 1) * np.exp(-(a * xp + b / xp) / 2), 0.0)

def gig_cdf(x, p, a, b):
    return _grid_cdf(lambda t: gig_pdf(t, p, a, b), x, 1e-9)

def gig_rng(p, a, b, n, rng):
    """No elementary quantile; sample by numerical inverse-transform of the Bessel-K
    density. (Specialised exact algorithms, e.g. Devroye 2014, also exist.)"""
    hi = gig_moments(p, a, b)["mean"] * 30 + 20
    return inverse_transform_sample(lambda t: gig_pdf(t, p, a, b), 1e-6, hi, n, rng)

def gig_moments(p, a, b):
    w = np.sqrt(a * b); r = np.sqrt(b / a)
    def raw(k): return r ** k * kv(p + k, w) / kv(p, w)
    m1, m2, m3, m4 = raw(1), raw(2), raw(3), raw(4)
    var = m2 - m1 ** 2
    c3 = m3 - 3 * m1 * m2 + 2 * m1 ** 3
    c4 = m4 - 4 * m1 * m3 + 6 * m1 ** 2 * m2 - 3 * m1 ** 4
    sd = np.sqrt(var)
    return dict(mean=m1, var=var, skew=c3 / sd ** 3, exkurt=c4 / var ** 2 - 3)


# ===========================================================================  #
#  PART C -- POLYA-GAMMA (the logistic / negative-binomial augmentation)        #
# ===========================================================================  #

def polyagamma_rng(b, c, n, rng, K=200):
    """Polson-Scott-Windle infinite-sum representation:
        omega = (1/(2 pi^2)) * sum_{k>=1} G_k / ((k-1/2)^2 + c^2/(4 pi^2)),  G_k ~ Gamma(b,1).
    Truncating the sum at K terms gives an accurate draw (tail terms ~ 1/k^2)."""
    k = np.arange(1, K + 1)
    denom = (k - 0.5) ** 2 + c ** 2 / (4 * np.pi ** 2)
    G = np.empty((n, K))
    for j in range(K):
        G[:, j] = _gamma_shape(b, n, rng)          # Gamma(b, 1)
    return (G / denom).sum(axis=1) / (2 * np.pi ** 2)

def polyagamma1_pdf(x, c, nterms=200):
    """Density of PG(1, c) via the alternating (Jacobi-theta) series, exponentially
    tilted from PG(1,0):  f(x|1,c) = cosh(c/2) exp(-c^2 x/2) f(x|1,0)."""
    x = np.asarray(x, float); xp = np.where(x > 0, x, 1.0)
    ncoef = np.arange(nterms)
    terms = ((-1) ** ncoef) * (2 * ncoef + 1)
    S = np.zeros_like(xp)
    for t, nc in zip(terms, 2 * ncoef + 1):
        S = S + t / np.sqrt(2 * np.pi * xp ** 3) * np.exp(-nc ** 2 / (8 * xp))
    base = np.where(x > 0, S, 0.0)
    return np.cosh(c / 2) * np.exp(-c ** 2 * x / 2) * base

def polyagamma_laplace(t, b, c):
    """Theoretical Laplace transform E[exp(-t*omega)] of PG(b,c) -- the rigorous check
    for a law with an awkward density:  cosh^b(c/2) / cosh^b(sqrt((t+c^2/2)/2))."""
    t = np.asarray(t, float)
    return np.cosh(c / 2) ** b / np.cosh(np.sqrt((t + c ** 2 / 2) / 2)) ** b

def polyagamma_moments(b, c):
    if abs(c) < 1e-9:
        mean = b / 4.0
    else:
        mean = b / (2 * c) * np.tanh(c / 2)
    return dict(mean=mean, var=np.nan, skew=np.nan, exkurt=np.nan)   # mean is the key summary


# ===========================================================================  #
#  PART D -- DIRICHLET (the prior on the probability simplex)                   #
# ===========================================================================  #

def dirichlet_pdf(x, alpha):
    """Density on the simplex; x is (..., K) with rows summing to 1."""
    x = np.asarray(x, float); alpha = np.asarray(alpha, float)
    logB = np.sum(gammaln(alpha)) - gammaln(np.sum(alpha))
    return np.exp(np.sum((alpha - 1) * np.log(x), axis=-1) - logB)

def dirichlet_rng(alpha, n, rng):
    """Normalised independent Gammas: g_i ~ Gamma(alpha_i, 1), x = g / sum(g)."""
    alpha = np.asarray(alpha, float); K = len(alpha)
    G = np.empty((n, K))
    for i in range(K):
        G[:, i] = _gamma_shape(alpha[i], n, rng)
    return G / G.sum(axis=1, keepdims=True)

def dirichlet_marginal_beta(alpha, i):
    """Component i of a Dirichlet is Beta(alpha_i, alpha_0 - alpha_i)."""
    a0 = np.sum(alpha)
    return alpha[i], a0 - alpha[i]


# --------------------------------------------------------------------------- #
#  Private helpers used above                                                   #
# --------------------------------------------------------------------------- #

def _grid_cdf(pdf, xs, lo, npts=20001):
    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 _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):
    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