"""
compound_np.py -- From-scratch COMPOUND & NONPARAMETRIC distributions.

Backs the notebooks in  Z-Statistical Distributions-Compound-Nonparametric
(Folder 8 of the Statistical Distributions catalog -- the final folder).

Here the "parameter" grows up: from a fixed vector to an entire random
DISTRIBUTION or an entire random FUNCTION.
  * COMPOUND -- the Tweedie (compound Poisson-Gamma): a random sum of a random
    number of Gamma jumps, giving a distribution with a point mass at zero and a
    continuous positive part.
  * NONPARAMETRIC -- the Dirichlet process / stick-breaking (GEM), a prior over
    probability distributions, and the Gaussian process, a prior over functions.

Everything is implemented from the definitions using elementary math and the
Gamma / Normal samplers built in the earlier folders. Where a law has no ordinary
density (Tweedie, the DP) we validate through its CONSTRUCTION and its moments.

Highlights (the from-scratch algorithms):
    Tweedie           : N ~ Poisson(lambda), Y = sum of N Gamma jumps
    GEM / stick-break : w_k = beta_k prod_{j<k}(1-beta_j),  beta_k ~ Beta(1, alpha)
    Chinese Restaurant: sequential seating with prob proportional to table size / alpha
    Gaussian process  : f ~ N(0, K) via Cholesky of a kernel matrix; closed-form posterior
"""

import numpy as np


# ===========================================================================  #
#  COMPOUND -- Tweedie (compound Poisson-Gamma), 1 < p < 2                       #
# ===========================================================================  #

def tweedie_params(mu, phi, p):
    """Map the Tweedie (mean mu, dispersion phi, power p) to the compound
    Poisson-Gamma pieces: Poisson rate lambda, Gamma shape a, Gamma scale theta."""
    lam = mu ** (2 - p) / (phi * (2 - p))
    a = (2 - p) / (p - 1)
    theta = phi * (p - 1) * mu ** (p - 1)
    return lam, a, theta

def tweedie_rng(mu, phi, p, n, rng):
    """N ~ Poisson(lambda); Y = sum of N i.i.d. Gamma(a, theta) jumps (Y=0 when N=0).
    (The Poisson count uses numpy's sampler; the Poisson RNG itself was built from
    scratch in Folder 1, the Gamma jumps from the Marsaglia-Tsang sampler below.)"""
    lam, a, theta = tweedie_params(mu, phi, p)
    N = rng.poisson(lam, n)
    Ntot = int(N.sum())
    jumps = theta * _gamma_shape(a, Ntot, rng) if Ntot > 0 else np.zeros(0)
    Y = np.zeros(n)
    idx = np.repeat(np.arange(n), N)
    np.add.at(Y, idx, jumps)
    return Y

def tweedie_mean_var(mu, phi, p):
    return dict(mean=mu, var=phi * mu ** p)

def tweedie_p0(mu, phi, p):
    """Probability of an exact zero: P(N=0) = exp(-lambda)."""
    lam, _, _ = tweedie_params(mu, phi, p)
    return np.exp(-lam)


# ===========================================================================  #
#  NONPARAMETRIC (i) -- Dirichlet process: stick-breaking (GEM) & the CRP        #
# ===========================================================================  #

def gem_rng(alpha, K, n, rng):
    """Stick-breaking weights (GEM(alpha)), truncated at K sticks, for n draws.
    beta_k ~ Beta(1, alpha) = 1 - U^{1/alpha};  w_k = beta_k * prod_{j<k}(1-beta_j).
    Returns (n, K); each row is a probability vector (up to the truncation remainder)."""
    b = 1.0 - rng.random((n, K)) ** (1.0 / alpha)            # Beta(1, alpha), closed form
    one = np.ones((n, 1))
    rem = np.cumprod(1 - b, axis=1)                          # prod_{j<=k}(1-beta_j)
    rem_shift = np.hstack([one, rem[:, :-1]])                # prod_{j<k}(1-beta_j)
    return b * rem_shift

def crp_rng(alpha, n_customers, rng):
    """Chinese Restaurant Process: customer i joins an occupied table with probability
    proportional to its size, or starts a new one with probability proportional to alpha.
    Returns the table assignment of each customer and the final table sizes."""
    sizes = []
    assign = np.zeros(n_customers, int)
    for i in range(n_customers):
        w = np.array(sizes + [alpha], float); w /= w.sum()
        j = np.searchsorted(np.cumsum(w), rng.random())
        if j == len(sizes):
            sizes.append(1)
        else:
            sizes[j] += 1
        assign[i] = j
    return assign, np.array(sizes)

def dp_expected_clusters(alpha, n):
    """E[#clusters] = sum_{i=1}^n alpha/(alpha+i-1)  (~ alpha*log(1+n/alpha))."""
    i = np.arange(1, n + 1)
    return np.sum(alpha / (alpha + i - 1))

def dp_random_measure(alpha, atoms_rng, K, rng):
    """One draw of a Dirichlet-process random distribution truncated at K: stick-breaking
    weights on atoms drawn from the base measure atoms_rng(K). Returns (locations, weights)."""
    w = gem_rng(alpha, K, 1, rng)[0]
    locs = atoms_rng(K)
    return locs, w


# ===========================================================================  #
#  NONPARAMETRIC (ii) -- the Gaussian process (a prior over functions)          #
# ===========================================================================  #

def rbf_kernel(X1, X2, lengthscale=1.0, variance=1.0):
    X1 = np.atleast_2d(X1); X2 = np.atleast_2d(X2)
    d2 = np.sum(X1 ** 2, 1)[:, None] + np.sum(X2 ** 2, 1)[None, :] - 2 * X1 @ X2.T
    return variance * np.exp(-0.5 * np.maximum(d2, 0) / lengthscale ** 2)

def matern32_kernel(X1, X2, lengthscale=1.0, variance=1.0):
    X1 = np.atleast_2d(X1); X2 = np.atleast_2d(X2)
    d2 = np.sum(X1 ** 2, 1)[:, None] + np.sum(X2 ** 2, 1)[None, :] - 2 * X1 @ X2.T
    r = np.sqrt(np.maximum(d2, 0)); s = np.sqrt(3) * r / lengthscale
    return variance * (1 + s) * np.exp(-s)

def matern12_kernel(X1, X2, lengthscale=1.0, variance=1.0):
    X1 = np.atleast_2d(X1); X2 = np.atleast_2d(X2)
    d2 = np.sum(X1 ** 2, 1)[:, None] + np.sum(X2 ** 2, 1)[None, :] - 2 * X1 @ X2.T
    return variance * np.exp(-np.sqrt(np.maximum(d2, 0)) / lengthscale)

def gp_prior_sample(X, kernel, n, rng, jitter=1e-8):
    """Draw n functions from the GP prior at inputs X: f ~ N(0, K), K = kernel(X,X),
    sampled by Cholesky. Returns (n, len(X))."""
    X = np.atleast_2d(X); m = X.shape[0]
    K = kernel(X, X) + jitter * np.eye(m)
    L = np.linalg.cholesky(K)
    Z = _std_normal(n * m, rng).reshape(n, m)
    return Z @ L.T

def gp_posterior(Xtr, ytr, Xte, kernel, noise=1e-6):
    """Closed-form GP-regression posterior at test inputs Xte given training (Xtr, ytr):
    mean = Ks^T (K+sigma^2 I)^{-1} y,  cov = Kss - Ks^T (K+sigma^2 I)^{-1} Ks.

    Solved through the Cholesky factor rather than by forming K^{-1}
    (Rasmussen & Williams, Alg. 2.1). Same answer, but stable when K is
    ill-conditioned -- which a smooth kernel on closely spaced inputs always
    is -- and it returns an exactly symmetric covariance."""
    Xtr = np.atleast_2d(Xtr); Xte = np.atleast_2d(Xte)
    K = kernel(Xtr, Xtr) + noise * np.eye(Xtr.shape[0])
    Ks = kernel(Xtr, Xte); Kss = kernel(Xte, Xte)
    L = np.linalg.cholesky(K)
    alpha = np.linalg.solve(L.T, np.linalg.solve(L, ytr))   # K^{-1} y, without K^{-1}
    V = np.linalg.solve(L, Ks)                              # L^{-1} Ks
    mean = Ks.T @ alpha
    cov = Kss - V.T @ V
    return mean, cov

def gp_posterior_sample(Xtr, ytr, Xte, kernel, n, rng, noise=1e-6, jitter=1e-8):
    mean, cov = gp_posterior(Xtr, ytr, Xte, kernel, noise)
    L = np.linalg.cholesky(cov + jitter * np.eye(cov.shape[0]))
    Z = _std_normal(n * len(mean), rng).reshape(n, len(mean))
    return mean + Z @ L.T


# --------------------------------------------------------------------------- #
#  Private RNG primitives                                                        #
# --------------------------------------------------------------------------- #

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 n == 0:
        return np.zeros(0)
    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