"""
multivariate.py -- From-scratch MULTIVARIATE distributions.

Backs the notebooks in  Z-Statistical Distributions-Multivariate  (Folder 6 of the
Statistical Distributions catalog).

The multivariate laws: joint distributions over vectors, where the new object is
the DEPENDENCE STRUCTURE (a covariance / precision matrix, a shape vector, a
spatial graph). Everything is implemented from the definitions using only
elementary linear algebra and the log-gamma as a primitive; `scipy.stats`
(Python) and R (R notebook) are used as INDEPENDENT references.

Per distribution, where defined:
    *_pdf / *_pmf (vectorised over rows of points, for contour/surface grids)
    *_rng(..., n, rng)          the sampling ALGORITHM
    (moments / marginals via the helpers, validated in the notebooks)

Highlights (the from-scratch sampling algorithms):
    Multivariate Normal   : mu + L z,  Sigma = L L^T (Cholesky)
    Multivariate t        : mu + sqrt(nu/w) L z,  w ~ chi^2_nu
    Multivariate SkewNorm : Azzalini-Dalla Valle  delta|U0| + V
    Multinomial           : sequential conditional Binomials
    Dirichlet-Multinomial : compound  p ~ Dirichlet, x ~ Multinomial(n, p)
    GMRF / CAR            : precision-based Gaussian, solve L^T x = z with Q = L L^T
"""

import numpy as np
from scipy.special import gammaln

LOG2PI = np.log(2 * np.pi)


# --------------------------------------------------------------------------- #
#  Multivariate Normal(mu, Sigma)                                               #
# --------------------------------------------------------------------------- #

def mvn_pdf(X, mu, Sigma):
    """Density at each row of X (shape (N,k) or (k,)). Vectorised for grids."""
    X = np.atleast_2d(np.asarray(X, float)); mu = np.asarray(mu, float); k = len(mu)
    Si = np.linalg.inv(Sigma); _, logdet = np.linalg.slogdet(Sigma)
    d = X - mu
    quad = np.einsum("ij,jk,ik->i", d, Si, d)
    return np.exp(-0.5 * (k * LOG2PI + logdet + quad))

def mvn_rng(mu, Sigma, n, rng):
    """mu + L z, with Sigma = L L^T (Cholesky) and z ~ N(0, I)."""
    mu = np.asarray(mu, float); k = len(mu)
    L = np.linalg.cholesky(Sigma)
    z = _std_normal(n * k, rng).reshape(n, k)
    return mu + z @ L.T

def mvn_conditional(mu, Sigma, idx_obs, x_obs):
    """Mean/cov of the components NOT in idx_obs, given those in idx_obs = x_obs
    (all Gaussian conditionals are Gaussian -- the property behind Gibbs sampling)."""
    k = len(mu); idx_obs = list(idx_obs)
    idx_un = [i for i in range(k) if i not in idx_obs]
    S = np.asarray(Sigma, float); mu = np.asarray(mu, float)
    Saa = S[np.ix_(idx_un, idx_un)]; Sab = S[np.ix_(idx_un, idx_obs)]
    Sbb = S[np.ix_(idx_obs, idx_obs)]
    Sbb_i = np.linalg.inv(Sbb)
    mean = mu[idx_un] + Sab @ Sbb_i @ (np.asarray(x_obs, float) - mu[idx_obs])
    cov = Saa - Sab @ Sbb_i @ Sab.T
    return idx_un, mean, cov


# --------------------------------------------------------------------------- #
#  Multivariate Student-t(mu, Sigma, nu)                                         #
# --------------------------------------------------------------------------- #

def mvt_pdf(X, mu, Sigma, nu):
    X = np.atleast_2d(np.asarray(X, float)); mu = np.asarray(mu, float); k = len(mu)
    Si = np.linalg.inv(Sigma); _, logdet = np.linalg.slogdet(Sigma)
    d = X - mu
    quad = np.einsum("ij,jk,ik->i", d, Si, d)
    logp = (gammaln((nu + k) / 2) - gammaln(nu / 2) - 0.5 * (k * np.log(nu * np.pi) + logdet)
            - (nu + k) / 2 * np.log1p(quad / nu))
    return np.exp(logp)

def mvt_rng(mu, Sigma, nu, n, rng):
    """mu + sqrt(nu/w) * (L z),  z ~ N(0,I), w ~ chi^2_nu -- a Gaussian scaled by an
    independent chi-square, so the marginals are heavy-tailed but still elliptical."""
    mu = np.asarray(mu, float); k = len(mu)
    L = np.linalg.cholesky(Sigma)
    z = _std_normal(n * k, rng).reshape(n, k)
    w = 2.0 * _gamma_shape(nu / 2.0, n, rng)          # chi^2_nu
    return mu + np.sqrt(nu / w)[:, None] * (z @ L.T)


# --------------------------------------------------------------------------- #
#  Multivariate Skew-Normal(xi, Omega, alpha)  --  Azzalini                      #
# --------------------------------------------------------------------------- #

def mvskewnorm_pdf(X, xi, Omega, alpha):
    """2 * phi_k(x; xi, Omega) * Phi( alpha^T omega^{-1} (x-xi) )."""
    from scipy.special import erf
    X = np.atleast_2d(np.asarray(X, float)); xi = np.asarray(xi, float)
    Omega = np.asarray(Omega, float); alpha = np.asarray(alpha, float)
    omega = np.sqrt(np.diag(Omega))
    base = mvn_pdf(X, xi, Omega)
    z = (X - xi) / omega
    lin = z @ alpha
    Phi = 0.5 * (1 + erf(lin / np.sqrt(2)))
    return 2.0 * base * Phi

def mvskewnorm_rng(xi, Omega, alpha, n, rng):
    """Azzalini-Dalla Valle: with correlation Omega_bar and delta = Omega_bar alpha /
    sqrt(1+alpha^T Omega_bar alpha),  Y = delta|U0| + V,  V ~ N(0, Omega_bar - delta delta^T);
    then X = xi + omega Y."""
    xi = np.asarray(xi, float); Omega = np.asarray(Omega, float); alpha = np.asarray(alpha, float)
    omega = np.sqrt(np.diag(Omega)); k = len(xi)
    Ob = Omega / np.outer(omega, omega)                # correlation matrix
    denom = np.sqrt(1 + alpha @ Ob @ alpha)
    delta = (Ob @ alpha) / denom
    Cv = Ob - np.outer(delta, delta)
    Lv = np.linalg.cholesky(Cv + 1e-12 * np.eye(k))
    u0 = np.abs(_std_normal(n, rng))
    zz = _std_normal(n * k, rng).reshape(n, k)
    Y = np.outer(u0, delta) + zz @ Lv.T
    return xi + omega * Y


# --------------------------------------------------------------------------- #
#  Multinomial(n, p)  --  discrete multivariate counts                           #
# --------------------------------------------------------------------------- #

def multinomial_pmf(x, n, p):
    x = np.atleast_2d(np.asarray(x, float)); p = np.asarray(p, float)
    logc = gammaln(n + 1) - np.sum(gammaln(x + 1), axis=1)
    logp = logc + np.sum(x * np.log(p), axis=1)
    ok = (x.sum(axis=1) == n)
    return np.where(ok, np.exp(logp), 0.0)

def multinomial_rng(n, p, size, rng):
    """Sequential conditional Binomials: x_1 ~ Bin(n, p_1);
    x_2 ~ Bin(n-x_1, p_2/(1-p_1)); ... (the stick-breaking of a count)."""
    p = np.asarray(p, float); K = len(p)
    out = np.zeros((size, K), int)
    remaining = np.full(size, n, int); prem = 1.0
    for j in range(K - 1):
        pj = np.clip(p[j] / prem, 0, 1)
        out[:, j] = _binomial_rng(remaining, pj, rng)
        remaining = remaining - out[:, j]; prem = prem - p[j]
    out[:, K - 1] = remaining
    return out


# --------------------------------------------------------------------------- #
#  Dirichlet-Multinomial(n, alpha)  --  overdispersed multinomial                #
# --------------------------------------------------------------------------- #

def dirmult_pmf(x, n, alpha):
    x = np.atleast_2d(np.asarray(x, float)); alpha = np.asarray(alpha, float); A = alpha.sum()
    logc = gammaln(n + 1) - np.sum(gammaln(x + 1), axis=1)
    logk = gammaln(A) - gammaln(n + A)
    logt = np.sum(gammaln(x + alpha) - gammaln(alpha), axis=1)
    ok = (x.sum(axis=1) == n)
    return np.where(ok, np.exp(logc + logk + logt), 0.0)

def dirmult_rng(n, alpha, size, rng):
    """Compound: p ~ Dirichlet(alpha), then x ~ Multinomial(n, p) -- the mixing over p
    adds over-dispersion relative to a plain Multinomial."""
    alpha = np.asarray(alpha, float); K = len(alpha)
    G = np.empty((size, K))
    for i in range(K):
        G[:, i] = _gamma_shape(alpha[i], size, rng)
    P = G / G.sum(axis=1, keepdims=True)
    out = np.zeros((size, K), int)
    for r in range(size):
        out[r] = multinomial_rng(n, P[r], 1, rng)[0]
    return out


# --------------------------------------------------------------------------- #
#  Gaussian Markov Random Field / CAR  --  precision-based spatial Gaussian       #
# --------------------------------------------------------------------------- #

def grid_adjacency(nr, nc):
    """4-neighbour adjacency matrix W of an nr x nc lattice."""
    N = nr * nc; W = np.zeros((N, N))
    for i in range(nr):
        for j in range(nc):
            a = i * nc + j
            for di, dj in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                ii, jj = i + di, j + dj
                if 0 <= ii < nr and 0 <= jj < nc:
                    W[a, ii * nc + jj] = 1
    return W

def car_precision(W, tau, rho):
    """Proper CAR precision  Q = tau (D - rho W),  D = diag(row sums of W).
    rho in [0,1) controls spatial dependence (rho=0 -> independent)."""
    D = np.diag(W.sum(axis=1))
    return tau * (D - rho * W)

def gmrf_rng(Q, n, rng):
    """Sample x ~ N(0, Q^{-1}) from the PRECISION: with Q = L L^T (Cholesky),
    solve L^T x = z for z ~ N(0, I). Never forms the dense covariance."""
    N = Q.shape[0]
    L = np.linalg.cholesky(Q)
    Z = _std_normal(n * N, rng).reshape(n, N)
    # solve L^T x = z  for each sample (upper-triangular back-substitution)
    from scipy.linalg import solve_triangular
    X = solve_triangular(L.T, Z.T, lower=False).T
    return X


# --------------------------------------------------------------------------- #
#  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 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

def _binomial_rng(n_arr, p, rng):
    """Vectorised Binomial with per-element n (counts), constant p, by inverse-CDF
    search on the Binomial CDF built from its PMF up to max(n)."""
    n_arr = np.asarray(n_arr, int)
    if p <= 0: return np.zeros_like(n_arr)
    if p >= 1: return n_arr.copy()
    nmax = int(n_arr.max()) if n_arr.size else 0
    out = np.empty_like(n_arr)
    U = rng.random(len(n_arr))
    # group by n for efficiency
    for nn in np.unique(n_arr):
        mask = n_arr == nn
        ks = np.arange(nn + 1)
        logpmf = (gammaln(nn + 1) - gammaln(ks + 1) - gammaln(nn - ks + 1)
                  + ks * np.log(p) + (nn - ks) * np.log1p(-p))
        cdf = np.cumsum(np.exp(logpmf))
        out[mask] = np.searchsorted(cdf, U[mask])
        np.clip(out, 0, nn, out=out)
    return out