"""
discrete.py -- From-scratch univariate DISCRETE distributions.

Backs the notebooks in  Z-Statistical Distributions-Discrete.

Everything here is implemented from the definitions using only elementary math
and the (log-)gamma function as a mathematical primitive -- no distribution
objects. `scipy.stats` (in the notebook) and R's d/p/q/r functions are then used
as INDEPENDENT references to validate these implementations.

For every distribution we provide:
    *_pmf(k, ...)      probability mass function
    *_cdf(k, ...)      cumulative distribution function
    *_rng(..., n, rng) a random-number generator (the sampling ALGORITHM matters)
    *_moments(...)     mean, variance, skewness, excess kurtosis (closed form)

A generic `moments_from_pmf` computes the four moments by direct summation over
the support -- an exact cross-check on the closed-form formulas.

Conventions match scipy so validation is apples-to-apples:
    geometric  : k = number of TRIALS to first success, k = 1,2,...   (scipy.geom)
    negbinom   : k = number of FAILURES before r successes, k = 0,1,... (scipy.nbinom)
"""

import numpy as np
from scipy.special import gammaln, betaln


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

def _logchoose(n, k):
    return gammaln(n + 1) - gammaln(k + 1) - gammaln(n - k + 1)


def moments_from_pmf(support, probs):
    """Exact mean/variance/skewness/excess-kurtosis by summation over the support."""
    support = np.asarray(support, float); p = np.asarray(probs, float)
    p = p / p.sum()
    m = np.sum(support * p)
    c2 = np.sum((support - m) ** 2 * p)
    c3 = np.sum((support - m) ** 3 * p)
    c4 = np.sum((support - m) ** 4 * p)
    sd = np.sqrt(c2)
    return dict(mean=m, var=c2, skew=c3 / sd ** 3, exkurt=c4 / c2 ** 2 - 3)


def inverse_cdf_sample(support, probs, n, rng):
    """Generic discrete sampler by inverse transform: draw U~Uniform, return the
    smallest support value whose CDF exceeds U."""
    cdf = np.cumsum(np.asarray(probs, float))
    cdf /= cdf[-1]
    return np.asarray(support)[np.searchsorted(cdf, rng.random(n))]


def _cdf_from_pmf(k, pmf_fn, kmax, *args):
    """CDF at integer(s) k by summing the pmf over 0..floor(k)."""
    ks = np.arange(0, kmax + 1)
    csum = np.cumsum(pmf_fn(ks, *args))
    kk = np.floor(np.asarray(k, float)).astype(int)
    out = np.where(kk < 0, 0.0, np.where(kk >= kmax, 1.0, csum[np.clip(kk, 0, kmax)]))
    return out


# --------------------------------------------------------------------------- #
#  Bernoulli(p)                                                                  #
# --------------------------------------------------------------------------- #

def bernoulli_pmf(k, p):
    k = np.asarray(k)
    return np.where(k == 1, p, np.where(k == 0, 1 - p, 0.0))

def bernoulli_cdf(k, p):
    k = np.floor(np.asarray(k, float))
    return np.clip(np.where(k < 0, 0.0, np.where(k < 1, 1 - p, 1.0)), 0, 1)

def bernoulli_rng(p, n, rng):
    return (rng.random(n) < p).astype(int)

def bernoulli_moments(p):
    v = p * (1 - p)
    return dict(mean=p, var=v, skew=(1 - 2 * p) / np.sqrt(v), exkurt=(1 - 6 * v) / v)


# --------------------------------------------------------------------------- #
#  Binomial(n, p)                                                                #
# --------------------------------------------------------------------------- #

def binomial_pmf(k, n, p):
    k = np.asarray(k)
    logp = _logchoose(n, k) + k * np.log(p) + (n - k) * np.log1p(-p)
    return np.where((k >= 0) & (k <= n), np.exp(logp), 0.0)

def binomial_cdf(k, n, p):
    return _cdf_from_pmf(k, binomial_pmf, n, n, p)

def binomial_rng(n, p, size, rng):
    """As a sum of n independent Bernoulli(p) indicators."""
    return (rng.random((size, n)) < p).sum(axis=1)

def binomial_moments(n, p):
    v = n * p * (1 - p)
    return dict(mean=n * p, var=v, skew=(1 - 2 * p) / np.sqrt(v),
                exkurt=(1 - 6 * p * (1 - p)) / v)


# --------------------------------------------------------------------------- #
#  Poisson(lam)                                                                  #
# --------------------------------------------------------------------------- #

def poisson_pmf(k, lam):
    k = np.asarray(k)
    return np.where(k >= 0, np.exp(k * np.log(lam) - lam - gammaln(k + 1)), 0.0)

def poisson_cdf(k, lam):
    kmax = int(lam + 10 * np.sqrt(lam) + 20)
    return _cdf_from_pmf(k, poisson_pmf, kmax, lam)

def poisson_rng(lam, n, rng):
    """Knuth's algorithm: count Uniforms multiplied until the product falls
    below e^{-lam} (equivalently, exponential inter-arrivals in unit time)."""
    L = np.exp(-lam); out = np.zeros(n, int)
    for i in range(n):
        k = 0; prod = rng.random()
        while prod > L:
            k += 1; prod *= rng.random()
        out[i] = k
    return out

def poisson_moments(lam):
    return dict(mean=lam, var=lam, skew=1 / np.sqrt(lam), exkurt=1 / lam)


# --------------------------------------------------------------------------- #
#  Geometric(p)  --  number of TRIALS to first success, k = 1,2,...              #
# --------------------------------------------------------------------------- #

def geometric_pmf(k, p):
    k = np.asarray(k)
    return np.where(k >= 1, (1 - p) ** (k - 1) * p, 0.0)

def geometric_cdf(k, p):
    k = np.floor(np.asarray(k, float))
    return np.where(k < 1, 0.0, 1 - (1 - p) ** k)

def geometric_rng(p, n, rng):
    """Inverse transform in closed form: k = ceil( ln(U) / ln(1-p) )."""
    return np.ceil(np.log(rng.random(n)) / np.log(1 - p)).astype(int)

def geometric_moments(p):
    return dict(mean=1 / p, var=(1 - p) / p ** 2, skew=(2 - p) / np.sqrt(1 - p),
                exkurt=6 + p ** 2 / (1 - p))


# --------------------------------------------------------------------------- #
#  Negative Binomial(r, p)  --  number of FAILURES before r successes, k=0,1,... #
# --------------------------------------------------------------------------- #

def nbinom_pmf(k, r, p):
    k = np.asarray(k)
    logp = _logchoose(k + r - 1, k) + r * np.log(p) + k * np.log1p(-p)
    return np.where(k >= 0, np.exp(logp), 0.0)

def nbinom_cdf(k, r, p):
    mean = r * (1 - p) / p; sd = np.sqrt(r * (1 - p)) / p
    kmax = int(mean + 12 * sd + 30)
    return _cdf_from_pmf(k, nbinom_pmf, kmax, r, p)

def nbinom_rng(r, p, n, rng):
    """Poisson-Gamma mixture: draw lam ~ Gamma(r, (1-p)/p), then k ~ Poisson(lam).
    Makes the over-dispersion mechanism explicit."""
    lam = rng.gamma(shape=r, scale=(1 - p) / p, size=n)
    return rng.poisson(lam)

def nbinom_moments(r, p):
    m = r * (1 - p) / p; v = r * (1 - p) / p ** 2
    return dict(mean=m, var=v, skew=(2 - p) / np.sqrt(r * (1 - p)),
                exkurt=6 / r + p ** 2 / (r * (1 - p)))


# --------------------------------------------------------------------------- #
#  Discrete Uniform {a, ..., b}                                                  #
# --------------------------------------------------------------------------- #

def duniform_pmf(k, a, b):
    k = np.asarray(k); N = b - a + 1
    return np.where((k >= a) & (k <= b), 1.0 / N, 0.0)

def duniform_cdf(k, a, b):
    k = np.floor(np.asarray(k, float)); N = b - a + 1
    return np.clip((k - a + 1) / N, 0, 1)

def duniform_rng(a, b, n, rng):
    return a + np.floor(rng.random(n) * (b - a + 1)).astype(int)

def duniform_moments(a, b):
    N = b - a + 1
    return dict(mean=(a + b) / 2, var=(N ** 2 - 1) / 12, skew=0.0,
                exkurt=-6 * (N ** 2 + 1) / (5 * (N ** 2 - 1)))


# --------------------------------------------------------------------------- #
#  Categorical(probs)                                                            #
# --------------------------------------------------------------------------- #

def categorical_pmf(k, probs):
    probs = np.asarray(probs, float); k = np.asarray(k)
    return np.where((k >= 0) & (k < len(probs)), probs[np.clip(k, 0, len(probs) - 1)], 0.0)

def categorical_cdf(k, probs):
    c = np.cumsum(np.asarray(probs, float)); k = np.floor(np.asarray(k, float)).astype(int)
    return np.where(k < 0, 0.0, np.where(k >= len(probs), 1.0, c[np.clip(k, 0, len(probs) - 1)]))

def categorical_rng(probs, n, rng):
    return inverse_cdf_sample(np.arange(len(probs)), probs, n, rng)


# --------------------------------------------------------------------------- #
#  Hypergeometric(M, K, n)  --  M population, K successes, n draws (no replace)   #
# --------------------------------------------------------------------------- #

def hypergeom_pmf(k, M, K, n):
    k = np.asarray(k)
    logp = _logchoose(K, k) + _logchoose(M - K, n - k) - _logchoose(M, n)
    ok = (k >= max(0, n - (M - K))) & (k <= min(n, K))
    return np.where(ok, np.exp(logp), 0.0)

def hypergeom_cdf(k, M, K, n):
    ks = np.arange(0, min(n, K) + 1)
    return _cdf_from_pmf(k, hypergeom_pmf, min(n, K), M, K, n)

def hypergeom_rng(M, K, n, size, rng):
    """Draw n items without replacement from an urn of K successes and M-K
    failures; count the successes."""
    urn = np.array([1] * K + [0] * (M - K))
    return np.array([rng.permutation(urn)[:n].sum() for _ in range(size)])

def hypergeom_moments(M, K, n):
    p = K / M; m = n * p; v = n * p * (1 - p) * (M - n) / (M - 1)
    ks = np.arange(0, min(n, K) + 1)
    return moments_from_pmf(ks, hypergeom_pmf(ks, M, K, n))


# --------------------------------------------------------------------------- #
#  Beta-Binomial(n, a, b)  --  Binomial with a Beta(a,b)-distributed p           #
# --------------------------------------------------------------------------- #

def betabinom_pmf(k, n, a, b):
    k = np.asarray(k)
    logp = _logchoose(n, k) + betaln(k + a, n - k + b) - betaln(a, b)
    return np.where((k >= 0) & (k <= n), np.exp(logp), 0.0)

def betabinom_cdf(k, n, a, b):
    return _cdf_from_pmf(k, betabinom_pmf, n, n, a, b)

def betabinom_rng(n, a, b, size, rng):
    """Compound: draw p ~ Beta(a,b), then k ~ Binomial(n, p)."""
    p = rng.beta(a, b, size); return rng.binomial(n, p)

def betabinom_moments(n, a, b):
    ks = np.arange(0, n + 1)
    return moments_from_pmf(ks, betabinom_pmf(ks, n, a, b))


# --------------------------------------------------------------------------- #
#  Zero-Inflated Poisson(pi, lam)                                                #
# --------------------------------------------------------------------------- #

def zip_pmf(k, pi, lam):
    k = np.asarray(k)
    base = poisson_pmf(k, lam)
    return np.where(k == 0, pi + (1 - pi) * np.exp(-lam), (1 - pi) * base)

def zip_cdf(k, pi, lam):
    kmax = int(lam + 10 * np.sqrt(lam) + 20)
    return _cdf_from_pmf(k, zip_pmf, kmax, pi, lam)

def zip_rng(pi, lam, n, rng):
    """Mixture: with prob pi a structural zero, else a Poisson(lam) count."""
    struct = rng.random(n) < pi
    return np.where(struct, 0, rng.poisson(lam, n))

def zip_moments(pi, lam):
    m = (1 - pi) * lam
    v = m * (1 + pi * lam)
    kmax = int(lam + 12 * np.sqrt(lam) + 30); ks = np.arange(0, kmax + 1)
    return moments_from_pmf(ks, zip_pmf(ks, pi, lam))


# --------------------------------------------------------------------------- #
#  Skellam(mu1, mu2)  --  difference of two independent Poissons                 #
# --------------------------------------------------------------------------- #

def skellam_pmf(k, mu1, mu2):
    from scipy.special import ive          # modified Bessel I, exponentially scaled
    k = np.asarray(k, float)
    return np.exp(-(mu1 + mu2)) * (mu1 / mu2) ** (k / 2) * \
        ive(np.abs(k), 2 * np.sqrt(mu1 * mu2)) * np.exp(2 * np.sqrt(mu1 * mu2))

def skellam_cdf(k, mu1, mu2):
    sd = np.sqrt(mu1 + mu2); lo = int((mu1 - mu2) - 12 * sd - 30)
    ks = np.arange(lo, int((mu1 - mu2) + 12 * sd + 30))
    csum = np.cumsum(skellam_pmf(ks, mu1, mu2))
    return np.interp(np.floor(np.asarray(k, float)), ks, csum)

def skellam_rng(mu1, mu2, n, rng):
    return rng.poisson(mu1, n) - rng.poisson(mu2, n)

def skellam_moments(mu1, mu2):
    s = mu1 + mu2
    return dict(mean=mu1 - mu2, var=s, skew=(mu1 - mu2) / s ** 1.5, exkurt=1 / s)


# --------------------------------------------------------------------------- #
#  Logarithmic(p)  --  log-series, k = 1,2,...                                   #
# --------------------------------------------------------------------------- #

def logarithmic_pmf(k, p):
    k = np.asarray(k)
    return np.where(k >= 1, -(p ** k) / (k * np.log1p(-p)), 0.0)

def logarithmic_cdf(k, p):
    kmax = 2000; ks = np.arange(1, kmax + 1)
    csum = np.cumsum(logarithmic_pmf(ks, p))
    kk = np.floor(np.asarray(k, float)).astype(int)
    return np.where(kk < 1, 0.0, csum[np.clip(kk - 1, 0, kmax - 1)])

def logarithmic_rng(p, n, rng):
    ks = np.arange(1, 5000)
    return inverse_cdf_sample(ks, logarithmic_pmf(ks, p), n, rng)

def logarithmic_moments(p):
    ks = np.arange(1, 20000)
    return moments_from_pmf(ks, logarithmic_pmf(ks, p))


# --------------------------------------------------------------------------- #
#  Conway-Maxwell-Poisson(lam, nu)  --  flexible over/under-dispersion           #
# --------------------------------------------------------------------------- #

def _com_Z(lam, nu, kmax=2000):
    ks = np.arange(0, kmax)
    return np.sum(np.exp(ks * np.log(lam) - nu * gammaln(ks + 1)))

def compoisson_pmf(k, lam, nu, kmax=2000):
    k = np.asarray(k); Z = _com_Z(lam, nu, kmax)
    return np.where(k >= 0, np.exp(k * np.log(lam) - nu * gammaln(k + 1)) / Z, 0.0)

def compoisson_cdf(k, lam, nu):
    kmax = 2000; ks = np.arange(0, kmax)
    csum = np.cumsum(compoisson_pmf(ks, lam, nu, kmax))
    kk = np.floor(np.asarray(k, float)).astype(int)
    return np.where(kk < 0, 0.0, csum[np.clip(kk, 0, kmax - 1)])

def compoisson_rng(lam, nu, n, rng):
    ks = np.arange(0, 2000)
    return inverse_cdf_sample(ks, compoisson_pmf(ks, lam, nu), n, rng)

def compoisson_moments(lam, nu):
    ks = np.arange(0, 2000)
    return moments_from_pmf(ks, compoisson_pmf(ks, lam, nu))


# --------------------------------------------------------------------------- #
#  Yule-Simon(rho)  --  power-law, k = 1,2,...                                   #
# --------------------------------------------------------------------------- #

def yule_pmf(k, rho):
    from scipy.special import beta as Bfn
    k = np.asarray(k)
    return np.where(k >= 1, rho * Bfn(k, rho + 1), 0.0)

def yule_cdf(k, rho):
    kmax = 200000; ks = np.arange(1, kmax + 1)
    csum = np.cumsum(yule_pmf(ks, rho))
    kk = np.floor(np.asarray(k, float)).astype(int)
    return np.where(kk < 1, 0.0, csum[np.clip(kk - 1, 0, kmax - 1)])

def yule_rng(rho, n, rng):
    """Mixture: W ~ Exponential(rho), then X | W ~ Geometric(p = e^{-W}) on {1,2,...}
    -- that mixture is exactly the Yule law (note p = e^{-W}, NOT 1-e^{-W})."""
    w = rng.exponential(1.0 / rho, n)
    return np.ceil(np.log(rng.random(n)) / np.log1p(-np.exp(-w))).astype(int)

def yule_moments(rho):
    """Closed-form mean/variance; skewness and excess kurtosis by summation.
    The m-th moment of Yule-Simon exists only for rho > m, so the mean needs
    rho > 1, the variance rho > 2, the skewness rho > 3 and the excess
    kurtosis rho > 4 -- outside those ranges the moment is infinite."""
    mean = rho / (rho - 1) if rho > 1 else np.inf
    var = rho ** 2 / ((rho - 1) ** 2 * (rho - 2)) if rho > 2 else np.inf
    skew = exkurt = np.inf
    if rho > 3:
        ks = np.arange(1, 3000000)
        m = moments_from_pmf(ks, yule_pmf(ks, rho))
        skew = m["skew"]
        if rho > 4:
            exkurt = m["exkurt"]
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)