"""
dpmix.py -- DIRICHLET-PROCESS & PITMAN-YOR MIXTURES, the random-measure view (from scratch).

Backs the notebooks in  "Dirichlet-Process & Pitman-Yor Mixtures".

This project is the MULTIVARIATE / random-measure companion to the univariate DP-mixture
notebook in the variable-selection arc (`Dirichlet-Process Mixtures -- How Many Components?`,
which fits the 1-D galaxy data and asks "how many components?"). Here we instead:

  1. show what a Dirichlet process IS -- a prior over DISCRETE random probability measures,
     built by STICK-BREAKING:  G = sum_k w_k delta_{theta_k},  theta_k ~ G0,
     w_k = v_k prod_{j<k}(1-v_j),  v_k ~ Beta(1-d, alpha + k d);
  2. fit a BIVARIATE DP mixture (Normal-InverseWishart base) for 2-D density estimation and
     clustering (Old Faithful), with a multivariate Student-t predictive;
  3. generalise the DP to the two-parameter PITMAN-YOR process (discount d), whose cluster
     sizes follow a POWER law -- more, smaller clusters than the DP.

Stick-breaking weights, DP vs Pitman-Yor:
    d = 0                -> Dirichlet process,  E[#clusters] ~ alpha * log n
    0 < d < 1            -> Pitman-Yor,         E[#clusters] ~ (alpha/d?) n^d   (power law)

The mixture sampler is the collapsed CRP Gibbs of Neal (2000, Algorithm 3): the component
means/covariances (mu_k, Sigma_k) are integrated out under the conjugate Normal-InverseWishart
base, so a point's predictive under a component is a multivariate Student-t and we resample only
the labels. Data are standardised per dimension so the NIW hyperparameters are portable; the
density is mapped back to the original scale.
"""

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


# --------------------------------------------------------------------------- #
#  Stick-breaking: DRAW a random probability measure G ~ DP / Pitman-Yor        #
# --------------------------------------------------------------------------- #

def stick_breaking(alpha, K, rng, d=0.0):
    """Return the first K stick-breaking weights of a DP(alpha) (d=0) or Pitman-Yor(alpha,d).
    v_k ~ Beta(1-d, alpha + k*d);  w_k = v_k * prod_{j<k}(1-v_j)."""
    k = np.arange(K)
    v = rng.beta(1.0 - d, alpha + (k + 1) * d)          # (index k+1 so first is Beta(1-d, alpha+d))
    v[-1] = 1.0                                          # truncate: last stick takes the remainder
    w = v * np.concatenate([[1.0], np.cumprod(1.0 - v)[:-1]])
    return w

def expected_clusters(alpha, n, d=0.0):
    """E[number of occupied clusters] after n draws from a DP / Pitman-Yor CRP (exact recursion)."""
    ek = 0.0
    for i in range(n):
        ek += (alpha + d * ek) / (alpha + i)            # prob draw i+1 starts a new cluster
    return ek


# --------------------------------------------------------------------------- #
#  Multivariate Student-t (the Normal-InverseWishart marginal)                  #
# --------------------------------------------------------------------------- #

def _mvt_logpdf(X, df, loc, scale):
    """log pdf of a multivariate Student-t at rows of X. scale is the D x D scale matrix."""
    X = np.atleast_2d(X); D = X.shape[1]
    L = np.linalg.cholesky(scale); logdet = 2 * np.log(np.diag(L)).sum()
    z = np.linalg.solve(L, (X - loc).T)                 # (D, n)
    maha = (z * z).sum(0)
    return (gammaln((df + D) / 2) - gammaln(df / 2) - 0.5 * D * np.log(df * np.pi)
            - 0.5 * logdet - (df + D) / 2 * np.log1p(maha / df))

def _niw_pred(n, s, S, m0, k0, nu0, Psi0, D):
    """multivariate-t predictive (df, loc, scale) for a cluster with count n, sum vector s,
    sum-of-outer-products S; n=0 gives the prior predictive under the NIW base."""
    if n == 0:
        kn, nun, mn, Psin = k0, nu0, m0, Psi0
    else:
        xbar = s / n
        kn = k0 + n; nun = nu0 + n
        mn = (k0 * m0 + s) / kn
        Cov = S - np.outer(s, xbar) - np.outer(xbar, s) + n * np.outer(xbar, xbar)   # scatter
        Psin = Psi0 + Cov + (k0 * n / kn) * np.outer(xbar - m0, xbar - m0)
    df = nun - D + 1
    scale = Psin * (kn + 1) / (kn * df)
    return df, mn, scale


# --------------------------------------------------------------------------- #
#  Collapsed CRP Gibbs for a multivariate DP / Pitman-Yor mixture               #
# --------------------------------------------------------------------------- #

def dpmix_gibbs(Y, rng, alpha=1.0, d=0.0, m0=None, k0=0.1, nu0=None, Psi0=None,
                draws=2000, burn=1000, grid=None):
    """Collapsed Gibbs for a DP(alpha) / Pitman-Yor(alpha,d) multivariate Normal mixture.
    Y:(N,D) standardised internally. Returns the posterior of the number of clusters K, a
    representative label vector, the cluster occupancy, and -- if a 2-D `grid` (G,2) is given --
    the posterior-predictive density on that grid (original units)."""
    Y = np.asarray(Y, float); N, D = Y.shape
    mu_, sd_ = Y.mean(0), Y.std(0)
    Ys = (Y - mu_) / sd_
    if m0 is None: m0 = np.zeros(D)
    if nu0 is None: nu0 = D + 2.0
    if Psi0 is None: Psi0 = np.eye(D)
    if grid is not None:
        grid = np.asarray(grid, float); gs = (grid - mu_) / sd_; dens = np.zeros(len(grid))
    z = np.zeros(N, int)
    cnt = [float(N)]; sm = [Ys.sum(0)]; SS = [Ys.T @ Ys]
    Ks = []; nsave = 0
    for it in range(draws + burn):
        for i in range(N):
            yi = Ys[i]; c = z[i]
            cnt[c] -= 1; sm[c] -= yi; SS[c] -= np.outer(yi, yi)
            if cnt[c] == 0:
                del cnt[c]; del sm[c]; del SS[c]; z[z > c] -= 1
            K = len(cnt); logp = np.empty(K + 1)
            for k in range(K):
                df, loc, sc = _niw_pred(cnt[k], sm[k], SS[k], m0, k0, nu0, Psi0, D)
                logp[k] = np.log(cnt[k] - d) + _mvt_logpdf(yi, df, loc, sc)[0]
            df, loc, sc = _niw_pred(0, None, None, m0, k0, nu0, Psi0, D)
            logp[K] = np.log(alpha + d * K) + _mvt_logpdf(yi, df, loc, sc)[0]
            logp -= logp.max(); p = np.exp(logp); p /= p.sum()
            new = rng.choice(K + 1, p=p)
            if new == K:
                cnt.append(0.0); sm.append(np.zeros(D)); SS.append(np.zeros((D, D)))
            z[i] = new; cnt[new] += 1; sm[new] += yi; SS[new] += np.outer(yi, yi)
        if it >= burn:
            Ks.append(len(cnt)); nsave += 1
            if grid is not None:
                K = len(cnt); denom = N + alpha; f = np.zeros(len(grid))
                for k in range(K):
                    df, loc, sc = _niw_pred(cnt[k], sm[k], SS[k], m0, k0, nu0, Psi0, D)
                    f += (cnt[k] - d) * np.exp(_mvt_logpdf(gs, df, loc, sc))
                df, loc, sc = _niw_pred(0, None, None, m0, k0, nu0, Psi0, D)
                f += (alpha + d * K) * np.exp(_mvt_logpdf(gs, df, loc, sc))
                dens += (f / denom) / np.prod(sd_)
    out = dict(K=np.array(Ks), z=z.copy(), counts=np.array(cnt))
    if grid is not None:
        out["grid"] = grid; out["density"] = dens / nsave
    return out


# --------------------------------------------------------------------------- #
#  simulation helpers                                                           #
# --------------------------------------------------------------------------- #

def simulate_mv_mixture(N, weights, means, covs, rng):
    """draw N points from a finite multivariate Normal mixture (the 'truth')."""
    weights = np.asarray(weights, float); weights /= weights.sum()
    k = rng.choice(len(weights), size=N, p=weights)
    means = np.asarray(means); covs = np.asarray(covs)
    return np.array([rng.multivariate_normal(means[kk], covs[kk]) for kk in k]), k

def crp_simulate_K(alpha, n, rng, d=0.0, reps=200):
    """simulate the number of clusters formed by n CRP draws (DP if d=0, else Pitman-Yor),
    returning the mean over `reps` -- used to contrast DP (log n) with PY (power-law) growth."""
    Ks = np.empty(reps)
    for r in range(reps):
        counts = []
        for i in range(n):
            probs = np.array([c - d for c in counts] + [alpha + d * len(counts)])
            probs /= probs.sum()
            j = rng.choice(len(counts) + 1, p=probs)
            if j == len(counts): counts.append(1.0)
            else: counts[j] += 1.0
        Ks[r] = len(counts)
    return Ks.mean()

def adjusted_rand(a, b):
    """adjusted Rand index between two label vectors (clustering-recovery check)."""
    a = np.asarray(a); b = np.asarray(b); n = len(a)
    ca = {v: i for i, v in enumerate(np.unique(a))}; cb = {v: i for i, v in enumerate(np.unique(b))}
    M = np.zeros((len(ca), len(cb)))
    for i in range(n): M[ca[a[i]], cb[b[i]]] += 1
    su = M.sum(1); sv = M.sum(0)
    idx = (M * (M - 1) / 2).sum(); ei = (su * (su - 1) / 2).sum(); ej = (sv * (sv - 1) / 2).sum()
    exp = ei * ej / (n * (n - 1) / 2); mx = 0.5 * (ei + ej)
    return (idx - exp) / (mx - exp) if mx != exp else 1.0