"""Unknown number of mixture components -- from scratch (collapsed sampler).

The mixture analogue of the change-point problem (Part 5): how many normal components
generated the data?  Congdon (2005) Example 3.8 fits K = 2,3,4,5 separately and compares;
the trans-dimensional answer lets the number of components K be a parameter.

Two standard routes exist.  Green's reversible-jump MCMC (Richardson & Green 1997) uses
split/merge moves with explicit Jacobians.  Here we take the *collapsed* route -- the same
trick that made the change-point sampler clean: with a conjugate Normal-Inverse-Gamma prior
on each component's (mean, variance) the component parameters integrate out analytically,
and the sampler moves over the discrete *allocation* of points to components.  Under a
Dirichlet-process (Chinese-restaurant) prior, the number of occupied components changes as
points are reassigned -- a point can open a brand-new component (a "birth") or vacate the
last seat of one (a "death").  This is Neal's (2000) Algorithm 3; no Jacobian is needed.

Component model: y_i | z_i=k ~ N(mu_k, s2_k), with (mu_k, s2_k) ~ NIG(m0, kappa0, a0, b0).
The collapsed predictive density of a point given a component's current members is Student-t.
Concentration alpha (which governs the prior on K) is given a Gamma hyperprior and sampled
(Escobar & West 1995), so the posterior on K is not pinned by a fixed alpha.
"""
import numpy as np
from scipy.special import gammaln


def _t_logpred(y, n, s1, s2, m0, k0, a0, b0):
    """log Student-t predictive density of y given a component's sufficient stats
    (count n, sum s1, sum-of-squares s2) under the NIG prior."""
    kn = k0 + n
    mn = (k0 * m0 + s1) / kn
    an = a0 + n / 2.0
    bn = b0 + 0.5 * (s2 + k0 * m0 ** 2 - kn * mn ** 2)
    nu = 2 * an
    scale2 = bn * (kn + 1) / (an * kn)
    return (gammaln((nu + 1) / 2) - gammaln(nu / 2) - 0.5 * np.log(np.pi * nu * scale2)
            - (nu + 1) / 2 * np.log1p((y - mn) ** 2 / (nu * scale2)))


def dp_mixture(y, m0=None, k0=0.05, a0=2.0, b0=None, a_alpha=2.0, b_alpha=2.0,
               ndraw=8000, burn=4000, seed=0):
    """Collapsed Gibbs sampler (Neal Algorithm 3) for a DP mixture of normals.

    Returns the posterior over the number of occupied components K, the concentration
    alpha, and a posterior-mean predictive density on a grid.
    """
    rng = np.random.default_rng(seed)
    y = np.asarray(y, float)
    n = len(y)
    if m0 is None: m0 = y.mean()
    if b0 is None: b0 = 0.5 * y.var()                       # weakly informative scale

    z = np.zeros(n, int)                                     # start: everyone in one component
    # per-component sufficient statistics stored in dicts keyed by label
    N = {0: n}; S1 = {0: y.sum()}; S2 = {0: (y ** 2).sum()}
    alpha = 1.0
    grid = np.linspace(y.min() - 2, y.max() + 2, 200)
    Ktr = np.zeros(ndraw, int); Atr = np.zeros(ndraw); dens = np.zeros(len(grid)); nkeep = 0

    def nextlabel():
        k = 0
        while k in N: k += 1
        return k

    for it in range(ndraw + burn):
        for i in range(n):
            k = z[i]                                          # remove point i from its component
            N[k] -= 1; S1[k] -= y[i]; S2[k] -= y[i] ** 2
            if N[k] == 0:
                del N[k], S1[k], S2[k]
            labels = list(N.keys())
            logp = np.array([np.log(N[k]) + _t_logpred(y[i], N[k], S1[k], S2[k], m0, k0, a0, b0)
                             for k in labels])
            lognew = np.log(alpha) + _t_logpred(y[i], 0, 0.0, 0.0, m0, k0, a0, b0)
            allp = np.append(logp, lognew)
            allp -= allp.max(); p = np.exp(allp); p /= p.sum()
            j = rng.choice(len(allp), p=p)
            if j == len(labels):                              # opened a NEW component (birth)
                k = nextlabel(); N[k] = 0; S1[k] = 0.0; S2[k] = 0.0
            else:
                k = labels[j]
            z[i] = k; N[k] += 1; S1[k] += y[i]; S2[k] += y[i] ** 2

        # ---- sample concentration alpha (Escobar-West 1995 augmentation) ----
        K = len(N)
        eta = rng.beta(alpha + 1, n)
        odds = (a_alpha + K - 1) / (n * (b_alpha - np.log(eta)))
        pi_eta = odds / (1 + odds)
        alpha = (rng.gamma(a_alpha + K, 1 / (b_alpha - np.log(eta))) if rng.random() < pi_eta
                 else rng.gamma(a_alpha + K - 1, 1 / (b_alpha - np.log(eta))))

        if it >= ndraw * 0 + burn:
            k_ = it - burn
            Ktr[k_] = len(N); Atr[k_] = alpha
            for lab in N:                                     # accumulate posterior predictive density
                dens += (N[lab] / n) * np.exp([_t_logpred(g, N[lab], S1[lab], S2[lab], m0, k0, a0, b0) for g in grid])
            nkeep += 1

    kpost = np.bincount(Ktr, minlength=15) / nkeep
    return dict(K=Ktr, kpost=kpost, alpha=Atr, grid=grid, dens=dens / nkeep, n=n)