"""Bishop (2006), Chapter 10, Sections 10.1 to 10.4: variational inference.

Reproduces the numerical figures 10.1, 10.2, 10.3, 10.4, 10.6, 10.7 and 10.9 of
C. M. Bishop, Pattern Recognition and Machine Learning (Springer, 2006),
pages 461 to 491. One function per figure; run

    python3 ch10_variational_inference.py --list
    python3 ch10_variational_inference.py --fig 10.6 10.7

The algorithms are small functions whose docstrings cite the book's equations,
so that they can be read next to the lecture notes (ch10-approximate-inference.html):

  skewed_log_density, laplace_1d, gauss_kl_qp_1d, moments_1d   Section 10.1, Figure 10.1
  mean_field_gaussian, factorised_kl_pq                        Section 10.1.2, Figure 10.2
  mixture_log_density, gauss_kl_qp_2d, moment_match_mixture    Section 10.1.2, Figure 10.3
  normal_gamma_posterior, vb_gaussian, vb_gaussian_bound       Section 10.1.3, Figure 10.4
  vbgmm_init, vbgmm_e_step, vbgmm_m_step, vbgmm_bound, vb_gmm,
  vbgmm_predictive                                             Section 10.2, Figures 10.6, 10.7
  vb_linear_regression, vblr_bound                             Section 10.3, Figure 10.9

Corrections to the first printing used here (Svensen and Bishop's errata, and our own derivation):
  * (10.63) is nu_k = nu_0 + N_k. The first printing has nu_0 + N_k + 1, which is not the
    maximiser of the bound (the solution of Exercise 10.13 gives nu_0 + N_k); checks_ch10a.py
    confirms by finite differences that nu_0 + N_k is the stationary point.
  * (10.69): E[pi_k] = (alpha_0 + N_k) / (K alpha_0 + N) (alpha_0 in the numerator).
  * (10.110): the last term is ln Gamma(a_0), not ln Gamma(a_N).
  * (10.29) has (N + 1)/2, and (10.31) and (10.33) read
    1/E[tau] = N/(N + 1) (mean(x^2) - 2 xbar E[mu] + E[mu^2]) and 1/E[tau] = mean(x^2) - xbar^2.

What the book does not specify, and the choices made here:
  * Fig 10.1: our Laplace approximation uses the exact second derivative at the mode
    z0 = 0.0775 (precision A = 2.546, sd 0.627, peak 0.636). The red curve of the book (here
    and in Fig 4.14) is narrower (sd about 0.52, peak 0.76), apparently a less accurate Hessian.
    The variational Gaussian (mean 0.707, sd 0.494) agrees with the book's green curve.
  * Fig 10.2: p has mean (0.5, 0.5), standard deviations 0.15 and correlation 0.95 (read off
    the book's 3-sd contours).
  * Fig 10.3: the bimodal p is our own equal mixture of two tilted Gaussians placed as in the
    book; contour levels are fractions of each density's maximum.
  * Fig 10.4: the book's data and prior are not given. We use N = 15 values drawn from N(0, 1)
    (numpy seed below, rounded to 3 decimals; the demo uses the same numbers), the prior
    mu_0 = 0, lambda_0 = 1, a_0 = b_0 = 1, and an initial guess placed as in panel (a).
    Contours are drawn at the same fractions of each density's maximum.
  * Figs 10.6 and 10.7: the book gives alpha_0 = 1e-3 and (typically) m_0 = 0. We take
    beta_0 = 1, W_0 = I and nu_0 = D = 2 on the standardised data. Starting point: unit
    Gaussian bumps on K data points chosen at random (a seeded K-means-like soft start),
    then one M step ('iteration 0'). The six starting rows of Fig 10.6 are fixed below and
    are also used by the demo vb-gmm-demo.html.
  * Fig 10.7: 100 random starts for each K = 1, ..., 6, with the ln K! term added as in the
    book, which does not say which alpha_0 it used. With alpha_0 = 1e-3 (or even 1) surplus
    components switch off, every K >= 2 reaches the same two-component fit and ln K! would then
    reward large K for nothing (it counts the K! relabellings of K distinct components in use).
    We use alpha_0 = 10, for which every component stays in use (as the book notes for
    Fig 10.6), so the correction applies; the bound then peaks at K = 2 as in the book.
    In our runs every start reached the same bound for each K (the book shows a few
    poorer local maxima), and K = 1 lies much further below than in the book's picture.
  * Fig 10.9: the cubic, the inputs (uniform on (-5, 5)) and the noise draw are ours;
    beta = 1/0.09 is known; with a_0 = b_0 = 0 the improper prior's normalising constant
    (the same for every degree) is omitted, so only differences between degrees matter.
"""
import os

for _v in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ.setdefault(_v, "1")   # the matrices here are tiny: BLAS threads only add overhead

import numpy as np  # noqa: E402
import matplotlib.pyplot as plt  # noqa: E402
from matplotlib.patches import Polygon  # noqa: E402
from scipy import optimize  # noqa: E402
from scipy.special import digamma, gammaln, log_expit, logsumexp  # noqa: E402

import common  # noqa: E402

TINY = 1e-300
YELLOW = "#e6c75a"   # the book's filled density colour in Figs 4.14 and 10.1


# =========================================================================== quadrature
def gauss_grid(n=61, half_width=6.0):
    """Nodes u and weights w for E[f(u)], u ~ N(0, 1): a uniform grid on [-h, h] with
    weights phi(u) du normalised to sum to 1 (spectrally accurate for smooth f).
    The demos use exactly the same rule, so Python and JavaScript agree to rounding."""
    u = np.linspace(-half_width, half_width, n)
    w = np.exp(-0.5 * u ** 2)
    return u, w / w.sum()


# =========================================================================== Section 10.1, Fig 10.1
SKEW_SLOPE, SKEW_SHIFT = 20.0, 4.0


def skewed_log_density(z, a=SKEW_SLOPE, b=SKEW_SHIFT):
    """ln p~(z) = -z^2/2 + ln sigma(a z + b): the unnormalised density of Figs 4.14 and 10.1."""
    return -0.5 * z ** 2 + log_expit(a * z + b)


def laplace_1d(a=SKEW_SLOPE, b=SKEW_SHIFT):
    """Laplace approximation (Section 4.4, eqs 4.127 to 4.134) to p(z) proportional to
    exp(-z^2/2) sigma(a z + b). The mode solves d/dz ln p~ = -z + a (1 - sigma(a z + b)) = 0
    (Newton's method); the precision is A = -d^2/dz^2 ln p~ = 1 + a^2 sigma (1 - sigma).
    Returns (z0, sd)."""
    z = 0.0
    for _ in range(100):
        s = common.sigmoid(a * z + b)
        g, h = -z + a * (1 - s), -1 - a * a * s * (1 - s)
        step = g / h
        z -= step
        if abs(step) < 1e-14:
            break
    s = common.sigmoid(a * z + b)
    return z, 1.0 / np.sqrt(1 + a * a * s * (1 - s))


def gauss_kl_qp_1d(log_ptilde, m_init, s_init, n_grid=2001):
    """The Gaussian q(z) = N(z | m, s^2) minimising KL(q || p), Section 10.1 and Fig 10.1.

    By (10.2), minimising KL(q || p) is maximising the lower bound (10.3),
        L(m, s) = E_q[ln p~(z)] + H[q],   H[q] = (1/2) ln(2 pi e s^2),
    a free-form optimisation restricted to the Gaussian family (the book's 'parametric
    distribution q(Z | omega)'). The expectation uses gauss_grid on [-8, 8] (the sigmoid
    is steep); the optimum over (m, ln s) is found with BFGS. Returns (m, s)."""
    u, w = gauss_grid(n_grid, 8.0)

    def neg_bound(th):
        m, s = th[0], np.exp(th[1])
        return -(np.sum(w * log_ptilde(m + s * u)) + 0.5 * np.log(2 * np.pi * np.e * s * s))

    res = optimize.minimize(neg_bound, [m_init, np.log(s_init)], method="BFGS", options={"gtol": 1e-10})
    return res.x[0], float(np.exp(res.x[1]))


def moments_1d(log_ptilde, lo=-12.0, hi=12.0, n=240001):
    """Mean, sd and ln Z of p(z) = p~(z)/Z by the trapezoid rule on a fine grid. The Gaussian
    with these moments minimises KL(p || q) (moment matching, Section 10.7, eq. 10.17)."""
    z = np.linspace(lo, hi, n)
    lp = log_ptilde(z)
    mx = lp.max()
    f = np.exp(lp - mx)
    Z = np.trapz(f, z)
    mean = np.trapz(z * f, z) / Z
    var = np.trapz((z - mean) ** 2 * f, z) / Z
    return mean, np.sqrt(var), np.log(Z) + mx


def gauss_logpdf_1d(z, m, s):
    return -0.5 * ((z - m) / s) ** 2 - np.log(s * np.sqrt(2 * np.pi))


# =========================================================================== Section 10.1.2, Fig 10.2
def mean_field_gaussian(mu, Lam, m_init, n_sweeps=50):
    """Factorised approximation q1(z1) q2(z2) to p(z) = N(z | mu, Lam^{-1}) minimising KL(q || p).

    The optimal factors are Gaussian with precisions Lam_11 and Lam_22 (10.12), (10.14); their
    means are updated in turn by (10.13) and (10.15),
        m1 = mu1 - Lam_11^{-1} Lam_12 (E[z2] - mu2),   m2 = mu2 - Lam_22^{-1} Lam_21 (E[z1] - mu1),
    which converge to m = mu (Exercise 10.2). Returns the path of the means, one row per
    single-factor update, and the factor variances."""
    m = np.array(m_init, dtype=float)
    path = [m.copy()]
    for _ in range(n_sweeps):
        m[0] = mu[0] - Lam[0, 1] / Lam[0, 0] * (m[1] - mu[1])   # (10.13)
        path.append(m.copy())
        m[1] = mu[1] - Lam[1, 0] / Lam[1, 1] * (m[0] - mu[0])   # (10.15)
        path.append(m.copy())
    return np.array(path), np.array([1 / Lam[0, 0], 1 / Lam[1, 1]])


def factorised_kl_pq(mu, Sigma):
    """Factorised q minimising KL(p || q): each factor is the marginal of p (10.17), so for a
    Gaussian p the factors are N(z_i | mu_i, Sigma_ii) (using 2.98). Returns (means, variances)."""
    return np.array(mu, dtype=float), np.diag(Sigma).copy()


def kl_gauss(m0, S0, m1, S1):
    """KL(N(m0, S0) || N(m1, S1)) in closed form."""
    D = len(m0)
    S1i = np.linalg.inv(S1)
    d = np.asarray(m1) - np.asarray(m0)
    return 0.5 * (np.trace(S1i @ S0) + d @ S1i @ d - D + np.linalg.slogdet(S1)[1] - np.linalg.slogdet(S0)[1])


# =========================================================================== Fig 10.3
def rot_cov(theta_deg, s_major, s_minor):
    t = np.radians(theta_deg)
    R = np.array([[np.cos(t), -np.sin(t)], [np.sin(t), np.cos(t)]])
    return R @ np.diag([s_major ** 2, s_minor ** 2]) @ R.T


# The bimodal p(Z) of Fig 10.3: an equal mixture of two tilted Gaussians (our choice).
MIX_W = np.array([0.5, 0.5])
MIX_MU = np.array([[-0.4, -0.42], [0.42, 0.4]])
MIX_COV = np.array([rot_cov(62, 0.23, 0.16), rot_cov(-40, 0.22, 0.17)])


def mixture_log_density(Z, w=MIX_W, mu=MIX_MU, cov=MIX_COV):
    """ln p(z) for a Gaussian mixture, z given as rows of Z (any leading shape)."""
    Z = np.asarray(Z, dtype=float)
    flat = Z.reshape(-1, 2)
    comp = np.stack([np.log(w[k]) + common.gauss_logpdf(flat, mu[k], cov[k]) for k in range(len(w))], 1)
    return logsumexp(comp, axis=1).reshape(Z.shape[:-1])


def chol_from_params(th):
    """Lower-triangular L = [[exp(l1), 0], [c, exp(l2)]] from th = (m1, m2, l1, c, l2)."""
    return np.array([[np.exp(th[2]), 0.0], [th[3], np.exp(th[4])]])


def gauss_kl_qp_2d(logp, m_init, cov_init, n_grid=61):
    """A 2-D Gaussian q = N(m, L L^T) minimising KL(q || p) by numerical optimisation (Fig 10.3b, c).

    KL(q || p) = -H[q] - E_q[ln p(z)] with H[q] = ln(2 pi e) + ln|L|; the expectation
    E_u[ln p(m + L u)], u ~ N(0, I), uses the product rule of gauss_grid (61 x 61 nodes).
    Different starting points reach different local minima, one per mode. Returns (m, Sigma, KL)."""
    u, w = gauss_grid(n_grid, 6.0)
    U1, U2 = np.meshgrid(u, u, indexing="ij")
    W2 = np.outer(w, w)

    def kl(th):
        L = chol_from_params(th)
        z1 = th[0] + L[0, 0] * U1
        z2 = th[1] + L[1, 0] * U1 + L[1, 1] * U2
        Elogp = np.sum(W2 * logp(np.stack([z1, z2], -1)))
        return -(np.log(2 * np.pi * np.e) + th[2] + th[4]) - Elogp

    L0 = np.linalg.cholesky(cov_init)
    th0 = np.array([m_init[0], m_init[1], np.log(L0[0, 0]), L0[1, 0], np.log(L0[1, 1])])
    res = optimize.minimize(kl, th0, method="BFGS", options={"gtol": 1e-9})
    L = chol_from_params(res.x)
    return res.x[:2], L @ L.T, res.fun


def moment_match_mixture(w=MIX_W, mu=MIX_MU, cov=MIX_COV):
    """The Gaussian minimising KL(p || q) for a mixture p (Fig 10.3a): the mean and covariance
    of p, mean = sum_k w_k mu_k, cov = sum_k w_k (Sigma_k + mu_k mu_k^T) - mean mean^T."""
    m = w @ mu
    S = sum(w[k] * (cov[k] + np.outer(mu[k], mu[k])) for k in range(len(w))) - np.outer(m, m)
    return m, S


# =========================================================================== Section 10.1.3, Fig 10.4
def normal_gamma_posterior(x, mu0, lam0, a0, b0):
    """Exact posterior for the Gaussian mean and precision under the conjugate prior (10.22),
    (10.23): p(mu, tau | D) = N(mu | m_N, (lam_N tau)^{-1}) Gam(tau | a_N, b_N) (Exercise 2.44) with
    lam_N = lam0 + N, m_N = (lam0 mu0 + N xbar)/lam_N, a_N = a0 + N/2,
    b_N = b0 + (1/2) sum (x_n - xbar)^2 + lam0 N (xbar - mu0)^2 / (2 lam_N)."""
    x = np.asarray(x, dtype=float)
    N, xbar = len(x), x.mean()
    lamN = lam0 + N
    return dict(m=(lam0 * mu0 + N * xbar) / lamN, lam=lamN, a=a0 + N / 2,
                b=b0 + 0.5 * np.sum((x - xbar) ** 2) + lam0 * N * (xbar - mu0) ** 2 / (2 * lamN))


def normal_gamma_logpdf(mu, tau, post):
    """ln N(mu | m, (lam tau)^{-1}) + ln Gam(tau | a, b) for the dictionary of normal_gamma_posterior."""
    m, lam, a, b = post["m"], post["lam"], post["a"], post["b"]
    tau = np.maximum(tau, TINY)
    return (0.5 * np.log(lam * tau / (2 * np.pi)) - 0.5 * lam * tau * (mu - m) ** 2
            + a * np.log(b) - gammaln(a) + (a - 1) * np.log(tau) - b * tau)


def gaussian_log_evidence(x, mu0, lam0, a0, b0):
    """ln p(D) for the Normal-Gamma model (the exact quantity that the bound L(q) approximates):
    ln Gamma(a_N) - ln Gamma(a0) + a0 ln b0 - a_N ln b_N + (1/2) ln(lam0/lam_N) - (N/2) ln(2 pi)."""
    post = normal_gamma_posterior(x, mu0, lam0, a0, b0)
    N = len(x)
    return (gammaln(post["a"]) - gammaln(a0) + a0 * np.log(b0) - post["a"] * np.log(post["b"])
            + 0.5 * np.log(lam0 / post["lam"]) - 0.5 * N * np.log(2 * np.pi))


def vb_update_mu(x, prior, E_tau):
    """q(mu) = N(mu | mu_N, lambda_N^{-1}) with (10.26) mu_N = (lam0 mu0 + N xbar)/(lam0 + N)
    and (10.27) lambda_N = (lam0 + N) E[tau]."""
    N = len(x)
    return dict(m=(prior["lam0"] * prior["mu0"] + np.sum(x)) / (prior["lam0"] + N), lam=(prior["lam0"] + N) * E_tau)


def vb_update_tau(x, prior, qmu):
    """q(tau) = Gam(tau | a_N, b_N) with (10.29) a_N = a0 + (N + 1)/2 (corrected) and (10.30)
    b_N = b0 + (1/2) E_mu[ sum_n (x_n - mu)^2 + lam0 (mu - mu0)^2 ],
    where E[mu] = mu_N and E[mu^2] = mu_N^2 + 1/lambda_N."""
    N = len(x)
    Emu2 = qmu["m"] ** 2 + 1 / qmu["lam"]
    Esq = np.sum(x ** 2) - 2 * qmu["m"] * np.sum(x) + N * Emu2
    Eprior = (qmu["m"] - prior["mu0"]) ** 2 + 1 / qmu["lam"]
    return dict(a=prior["a0"] + (N + 1) / 2, b=prior["b0"] + 0.5 * (Esq + prior["lam0"] * Eprior))


def vb_gaussian_bound(x, prior, qmu, qtau):
    """Lower bound L(q) = E[ln p(D, mu, tau)] - E[ln q(mu)] - E[ln q(tau)] for Section 10.1.3:
      E[ln p(D | mu, tau)] = (N/2)(E[ln tau] - ln 2 pi) - (E[tau]/2) E[sum (x_n - mu)^2]
      E[ln p(mu | tau)]    = (1/2)(ln(lam0/2 pi) + E[ln tau]) - (lam0 E[tau]/2) E[(mu - mu0)^2]
      E[ln p(tau)]         = a0 ln b0 - ln Gamma(a0) + (a0 - 1) E[ln tau] - b0 E[tau]
      -E[ln q(mu)]         = (1/2)(1 + ln 2 pi - ln lambda_N)
      -E[ln q(tau)]        = ln Gamma(a_N) - (a_N - 1) psi(a_N) - ln b_N + a_N   (B.31)
    with E[tau] = a_N/b_N and E[ln tau] = psi(a_N) - ln b_N (B.27), (B.30)."""
    N = len(x)
    mu0, lam0, a0, b0 = prior["mu0"], prior["lam0"], prior["a0"], prior["b0"]
    Et, Elt = qtau["a"] / qtau["b"], digamma(qtau["a"]) - np.log(qtau["b"])
    Emu2 = qmu["m"] ** 2 + 1 / qmu["lam"]
    Esq = np.sum(x ** 2) - 2 * qmu["m"] * np.sum(x) + N * Emu2
    lik = 0.5 * N * (Elt - np.log(2 * np.pi)) - 0.5 * Et * Esq
    pmu = 0.5 * (np.log(lam0 / (2 * np.pi)) + Elt) - 0.5 * lam0 * Et * ((qmu["m"] - mu0) ** 2 + 1 / qmu["lam"])
    ptau = a0 * np.log(b0) - gammaln(a0) + (a0 - 1) * Elt - b0 * Et
    hmu = 0.5 * (1 + np.log(2 * np.pi) - np.log(qmu["lam"]))
    htau = gammaln(qtau["a"]) - (qtau["a"] - 1) * digamma(qtau["a"]) - np.log(qtau["b"]) + qtau["a"]
    return lik + pmu + ptau + hmu + htau


def vb_gaussian(x, prior, qmu, qtau, n_iter=100, tol=1e-12):
    """Coordinate ascent of Section 10.1.3: alternately re-estimate q(mu) (10.26, 10.27) and
    q(tau) (10.29, 10.30), starting from the factors (qmu, qtau). Returns the list of states
    (qmu, qtau) after every single-factor update, the first entry being the start."""
    states = [(dict(qmu), dict(qtau))]
    L_old = vb_gaussian_bound(x, prior, qmu, qtau)
    for _ in range(n_iter):
        qmu = vb_update_mu(x, prior, qtau["a"] / qtau["b"])
        states.append((qmu, dict(qtau)))
        qtau = vb_update_tau(x, prior, qmu)
        states.append((dict(qmu), qtau))
        L = vb_gaussian_bound(x, prior, qmu, qtau)
        if abs(L - L_old) < tol:
            break
        L_old = L
    return states


def gamma_logpdf(tau, a, b):
    tau = np.maximum(tau, TINY)
    return a * np.log(b) - gammaln(a) + (a - 1) * np.log(tau) - b * tau


# Fig 10.4 and vb-gaussian-demo.html: data, prior and initial factors (see the module docstring).
VBG_SEED = 287
VBG_DATA = np.round(np.random.default_rng(VBG_SEED).standard_normal(100), 3)
VBG_N = 15
VBG_PRIOR = dict(mu0=0.0, lam0=1.0, a0=1.0, b0=1.0)
VBG_INIT = (dict(m=0.55, lam=1 / 0.17 ** 2), dict(a=25.0, b=25.0 / 1.3))


# =========================================================================== Section 10.2, VB-GMM
def default_gmm_prior(D, alpha0=1e-3):
    """Prior of (10.39), (10.40): alpha0, m0 = 0, beta0 = 1, W0 = I, nu0 = D (see the docstring)."""
    return dict(alpha0=alpha0, beta0=1.0, m0=np.zeros(D), W0=np.eye(D), nu0=float(D))


def wishart_log_B(W, nu):
    """ln B(W, nu) of (B.79): -(nu/2) ln|W| - (nu D/2) ln 2 - (D(D-1)/4) ln pi - sum_i ln Gamma((nu + 1 - i)/2).
    W may be a stack (K, D, D) with nu of shape (K,)."""
    W = np.asarray(W)
    D = W.shape[-1]
    nu = np.asarray(nu, dtype=float)
    i = np.arange(1, D + 1)
    return (-0.5 * nu * np.linalg.slogdet(W)[1] - 0.5 * nu * D * np.log(2) - 0.25 * D * (D - 1) * np.log(np.pi)
            - gammaln(0.5 * (nu[..., None] + 1 - i)).sum(-1))


def vbgmm_expectations(q):
    """ln pi~_k = E[ln pi_k] = psi(alpha_k) - psi(sum_j alpha_j)                    (10.66)
       ln Lambda~_k = E[ln|Lambda_k|] = sum_i psi((nu_k + 1 - i)/2) + D ln 2 + ln|W_k|  (10.65)"""
    D = q["m"].shape[1]
    i = np.arange(1, D + 1)
    ln_pi = digamma(q["alpha"]) - digamma(q["alpha"].sum())
    ln_lam = digamma(0.5 * (q["nu"][:, None] + 1 - i)).sum(1) + D * np.log(2) + np.linalg.slogdet(q["W"])[1]
    return ln_pi, ln_lam


def vbgmm_e_step(X, q):
    """Variational E step: responsibilities r_nk = rho_nk / sum_j rho_nj (10.49) with
    ln rho_nk = E[ln pi_k] + (1/2) E[ln|Lambda_k|] - (D/2) ln 2 pi - (1/2) E[(x_n - mu_k)^T Lambda_k (x_n - mu_k)]
    (10.46), the last expectation being D/beta_k + nu_k (x_n - m_k)^T W_k (x_n - m_k) (10.64)."""
    N, D = X.shape
    ln_pi, ln_lam = vbgmm_expectations(q)
    quad = np.empty((N, len(ln_pi)))
    for k in range(len(ln_pi)):
        d = X - q["m"][k]
        quad[:, k] = D / q["beta"][k] + q["nu"][k] * np.sum((d @ q["W"][k]) * d, 1)   # (10.64)
    ln_rho = ln_pi + 0.5 * ln_lam - 0.5 * D * np.log(2 * np.pi) - 0.5 * quad          # (10.46)
    return np.exp(ln_rho - logsumexp(ln_rho, axis=1, keepdims=True))                # (10.49)


def vbgmm_stats(X, r):
    """N_k (10.51), xbar_k (10.52) and S_k (10.53): the responsibility-weighted statistics."""
    Nk = r.sum(0)
    xbar = (r.T @ X) / np.maximum(Nk, TINY)[:, None]
    S = np.empty((len(Nk), X.shape[1], X.shape[1]))
    for k in range(len(Nk)):
        d = X - xbar[k]
        S[k] = (r[:, k, None] * d).T @ d / max(Nk[k], TINY)
    return Nk, xbar, S


def vbgmm_m_step(X, r, prior):
    """Variational M step: q(pi) = Dir(pi | alpha) and q(mu_k, Lambda_k) = N(mu_k | m_k, (beta_k Lambda_k)^{-1})
    W(Lambda_k | W_k, nu_k) (10.57), (10.59) with
      alpha_k = alpha0 + N_k (10.58),  beta_k = beta0 + N_k (10.60),  m_k = (beta0 m0 + N_k xbar_k)/beta_k (10.61),
      W_k^{-1} = W0^{-1} + N_k S_k + beta0 N_k/(beta0 + N_k) (xbar_k - m0)(xbar_k - m0)^T (10.62),
      nu_k = nu0 + N_k (10.63, corrected: the first printing has an extra + 1)."""
    Nk, xbar, S = vbgmm_stats(X, r)
    b0, m0 = prior["beta0"], prior["m0"]
    beta = b0 + Nk
    dm = xbar - m0
    Winv = (np.linalg.inv(prior["W0"])[None] + Nk[:, None, None] * S
            + (b0 * Nk / (b0 + Nk))[:, None, None] * np.einsum("ki,kj->kij", dm, dm))
    return dict(alpha=prior["alpha0"] + Nk, beta=beta, m=(b0 * m0 + Nk[:, None] * xbar) / beta[:, None],
                W=np.linalg.inv(Winv), nu=prior["nu0"] + Nk)


def vbgmm_bound_terms(X, r, q, prior):
    """The seven terms of the lower bound (10.70), evaluated by (10.71) to (10.77), for any
    responsibilities r and any q(pi), q(mu, Lambda) (not only at an update):
      (10.71) E[ln p(X|Z,mu,Lambda)] = (1/2) sum_k N_k { ln Lam~_k - D/beta_k - nu_k Tr(S_k W_k)
                                        - nu_k (xbar_k - m_k)^T W_k (xbar_k - m_k) - D ln 2 pi }
      (10.72) E[ln p(Z|pi)] = sum_nk r_nk ln pi~_k
      (10.73) E[ln p(pi)] = ln C(alpha0) + (alpha0 - 1) sum_k ln pi~_k
      (10.74) E[ln p(mu,Lambda)] = (1/2) sum_k { D ln(beta0/2 pi) + ln Lam~_k - D beta0/beta_k
              - beta0 nu_k (m_k - m0)^T W_k (m_k - m0) } + K ln B(W0, nu0)
              + (nu0 - D - 1)/2 sum_k ln Lam~_k - (1/2) sum_k nu_k Tr(W0^{-1} W_k)
      (10.75) E[ln q(Z)] = sum_nk r_nk ln r_nk
      (10.76) E[ln q(pi)] = sum_k (alpha_k - 1) ln pi~_k + ln C(alpha)
      (10.77) E[ln q(mu,Lambda)] = sum_k { (1/2) ln Lam~_k + (D/2) ln(beta_k/2 pi) - D/2 - H[q(Lambda_k)] }
    with C(alpha) of (B.23), B(W, nu) of (B.79) and the Wishart entropy (B.82)
      H[Lambda] = -ln B(W, nu) - (nu - D - 1)/2 E[ln|Lambda|] + nu D/2."""
    N, D = X.shape
    K = r.shape[1]
    Nk, xbar, S = vbgmm_stats(X, r)
    alpha, beta, m, W, nu = q["alpha"], q["beta"], q["m"], q["W"], q["nu"]
    a0, b0, m0, W0, nu0 = prior["alpha0"], prior["beta0"], prior["m0"], prior["W0"], prior["nu0"]
    ln_pi, ln_lam = vbgmm_expectations(q)
    dx = xbar - m
    trSW = np.einsum("kij,kji->k", S, W)
    lik = 0.5 * np.sum(Nk * (ln_lam - D / beta - nu * trSW - nu * np.einsum("ki,kij,kj->k", dx, W, dx)
                             - D * np.log(2 * np.pi)))
    pz = np.sum(r * ln_pi)
    ppi = gammaln(K * a0) - K * gammaln(a0) + (a0 - 1) * ln_pi.sum()
    dm = m - m0
    pml = (0.5 * np.sum(D * np.log(b0 / (2 * np.pi)) + ln_lam - D * b0 / beta
                        - b0 * nu * np.einsum("ki,kij,kj->k", dm, W, dm))
           + K * wishart_log_B(W0, nu0) + 0.5 * (nu0 - D - 1) * ln_lam.sum()
           - 0.5 * np.sum(nu * np.einsum("ij,kji->k", np.linalg.inv(W0), W)))
    qz = np.sum(r * np.log(np.maximum(r, TINY)))
    qpi = np.sum((alpha - 1) * ln_pi) + gammaln(alpha.sum()) - gammaln(alpha).sum()
    H = -wishart_log_B(W, nu) - 0.5 * (nu - D - 1) * ln_lam + 0.5 * nu * D
    qml = np.sum(0.5 * ln_lam + 0.5 * D * np.log(beta / (2 * np.pi)) - 0.5 * D - H)
    return dict(lik=lik, pz=pz, ppi=ppi, pml=pml, qz=qz, qpi=qpi, qml=qml)


def vbgmm_bound(X, r, q, prior):
    """The lower bound L of (10.70): the sum of the expectations of the log joint factors minus
    the expectations of the log q factors (vbgmm_bound_terms)."""
    t = vbgmm_bound_terms(X, r, q, prior)
    return t["lik"] + t["pz"] + t["ppi"] + t["pml"] - t["qz"] - t["qpi"] - t["qml"]


def vbgmm_init(X, idx, prior):
    """Starting point ('iteration 0'): q(mu_k, Lambda_k) centred on the K data points X[idx] with
    unit expected covariance, E[Lambda_k]^{-1} = (nu_k W_k)^{-1} = I, and every component credited
    with N/K points: alpha_k = alpha0 + N/K, beta_k = beta0 + N/K, nu_k = nu0 + N/K, m_k = X[idx[k]].
    The matching responsibilities are those of the variational E step for this q."""
    N, D = X.shape
    K = len(idx)
    n0 = np.full(K, N / K)
    nu = prior["nu0"] + n0
    q = dict(alpha=prior["alpha0"] + n0, beta=prior["beta0"] + n0, m=X[np.asarray(idx)].astype(float),
             W=np.eye(D)[None] / nu[:, None, None], nu=nu)
    return vbgmm_e_step(X, q), q


def vb_gmm(X, idx, prior, max_iter=1000, tol=1e-9, keep=()):
    """Variational Bayes for the Gaussian mixture (Section 10.2). The start (vbgmm_init) already
    contains an E step, so each iteration is a variational M step (10.57) to (10.63) followed by a
    variational E step (10.46) to (10.49); the bound (10.70) is evaluated after each iteration and
    the loop stops when it changes by less than tol (relative), but not before the last iteration
    in keep. Returns a dict with the final r and q, the bound after every iteration (entry 0 = start),
    copies of q at the iterations in keep and 'converged', the first iteration that met the test."""
    r, q = vbgmm_init(X, idx, prior)
    bounds = [vbgmm_bound(X, r, q, prior)]
    snaps = {0: q} if 0 in keep else {}
    converged = None
    for it in range(1, max_iter + 1):
        q = vbgmm_m_step(X, r, prior)
        r = vbgmm_e_step(X, q)
        bounds.append(vbgmm_bound(X, r, q, prior))
        if it in keep:
            snaps[it] = q
        if converged is None and abs(bounds[-1] - bounds[-2]) < tol * abs(bounds[-1]):
            converged = it
        if converged is not None and it >= max(keep, default=0):
            break
    return dict(r=r, q=q, bounds=np.array(bounds), snaps=snaps, iters=len(bounds) - 1, converged=converged)


def expected_pi(q):
    """E[pi_k] = alpha_k / sum_j alpha_j = (alpha0 + N_k)/(K alpha0 + N) (10.69, corrected)."""
    return q["alpha"] / q["alpha"].sum()


def student_t_logpdf(X, mu, L, nu):
    """ln St(x | mu, L, nu) of (B.68) for rows of X: ln Gamma(D/2 + nu/2) - ln Gamma(nu/2)
    + (1/2) ln|L| - (D/2) ln(pi nu) - (D + nu)/2 ln(1 + (x - mu)^T L (x - mu)/nu)."""
    D = X.shape[1]
    d = X - mu
    q = np.einsum("ni,ij,nj->n", d, L, d)
    return (gammaln(0.5 * (D + nu)) - gammaln(0.5 * nu) + 0.5 * np.linalg.slogdet(L)[1]
            - 0.5 * D * np.log(np.pi * nu) - 0.5 * (D + nu) * np.log1p(q / nu))


def vbgmm_predictive(Xnew, q):
    """Approximate predictive density (10.81): sum_k (alpha_k/alpha_hat) St(x | m_k, L_k, nu_k + 1 - D)
    with L_k = (nu_k + 1 - D) beta_k / (1 + beta_k) W_k (10.82)."""
    D = q["m"].shape[1]
    w = expected_pi(q)
    out = np.zeros(len(Xnew))
    for k in range(len(w)):
        dof = q["nu"][k] + 1 - D
        L = dof * q["beta"][k] / (1 + q["beta"][k]) * q["W"][k]
        out += w[k] * np.exp(student_t_logpdf(Xnew, q["m"][k], L, dof))
    return out


# Fig 10.6 and vb-gmm-demo.html (preset 1): the six starting rows (0-based) of the data file.
FIG10_6_START = [140, 210, 217, 54, 79, 172]


# =========================================================================== Section 10.3, Fig 10.9
def vblr_bound(Phi, t, beta, a0, b0, mN, SN, aN, bN):
    """Lower bound (10.107) for variational linear regression, the sum of
      (10.108) E[ln p(t|w)] = (N/2) ln(beta/2 pi) - (beta/2) t^T t + beta m_N^T Phi^T t
                              - (beta/2) Tr(Phi^T Phi (m_N m_N^T + S_N))
      (10.109) E[ln p(w|alpha)] = -(M/2) ln 2 pi + (M/2)(psi(a_N) - ln b_N) - a_N/(2 b_N)(m_N^T m_N + Tr S_N)
      (10.110) E[ln p(alpha)] = a0 ln b0 + (a0 - 1)(psi(a_N) - ln b_N) - b0 a_N/b_N - ln Gamma(a0)  (corrected)
      (10.111) -E[ln q(w)] = (1/2) ln|S_N| + (M/2)(1 + ln 2 pi)
      (10.112) -E[ln q(alpha)] = ln Gamma(a_N) - (a_N - 1) psi(a_N) - ln b_N + a_N.
    For the improper prior a0 = 0 the constant a0 ln b0 - ln Gamma(a0) is dropped (it is the same
    for every model, so comparisons between models are unaffected)."""
    N, M = Phi.shape
    PtP, Ptt = Phi.T @ Phi, Phi.T @ t
    Elna = digamma(aN) - np.log(bN)
    Eww = mN @ mN + np.trace(SN)
    l1 = 0.5 * N * np.log(beta / (2 * np.pi)) - 0.5 * beta * t @ t + beta * mN @ Ptt - 0.5 * beta * np.trace(PtP @ (np.outer(mN, mN) + SN))
    l2 = -0.5 * M * np.log(2 * np.pi) + 0.5 * M * Elna - aN / (2 * bN) * Eww
    l3 = (a0 - 1) * Elna - b0 * aN / bN + (a0 * np.log(b0) - gammaln(a0) if a0 > 0 else 0.0)
    l4 = 0.5 * np.linalg.slogdet(SN)[1] + 0.5 * M * (1 + np.log(2 * np.pi))
    l5 = gammaln(aN) - (aN - 1) * digamma(aN) - np.log(bN) + aN
    return l1 + l2 + l3 + l4 + l5


def vb_linear_regression(Phi, t, beta, a0=0.0, b0=0.0, E_alpha=1.0, max_iter=5000, tol=1e-10):
    """Variational Bayes for linear regression with known noise precision beta (Section 10.3):
    alternate q(w) = N(w | m_N, S_N) with S_N = (E[alpha] I + beta Phi^T Phi)^{-1}, m_N = beta S_N Phi^T t
    (10.99) to (10.101), and q(alpha) = Gam(alpha | a_N, b_N) with a_N = a0 + M/2,
    b_N = b0 + E[w^T w]/2 (10.93) to (10.95), E[alpha] = a_N/b_N (10.102), E[w w^T] = m_N m_N^T + S_N (10.103).
    Returns (m_N, S_N, a_N, b_N, bounds after every cycle)."""
    N, M = Phi.shape
    PtP, Ptt = Phi.T @ Phi, Phi.T @ t
    aN = a0 + 0.5 * M
    bounds = []
    for _ in range(max_iter):
        SN = np.linalg.inv(E_alpha * np.eye(M) + beta * PtP)
        mN = beta * SN @ Ptt
        bN = b0 + 0.5 * (mN @ mN + np.trace(SN))
        E_alpha = aN / bN
        bounds.append(vblr_bound(Phi, t, beta, a0, b0, mN, SN, aN, bN))
        if len(bounds) > 1 and abs(bounds[-1] - bounds[-2]) < tol:
            break
    return mN, SN, aN, bN, np.array(bounds)


# Fig 10.9: the cubic, the noise level and the data (our choice, see the module docstring).
FIG10_9_COEF = (1.0, 0.5, -0.04, -0.03)    # t = w0 + w1 x + w2 x^2 + w3 x^3 + noise
FIG10_9_SEED = 18
FIG10_9_NOISE_VAR = 0.09


def fig10_9_data():
    rng = np.random.default_rng(FIG10_9_SEED)
    x = rng.uniform(-5, 5, 10)
    t = np.polyval(FIG10_9_COEF[::-1], x) + rng.normal(0, np.sqrt(FIG10_9_NOISE_VAR), 10)
    return x, t


# =========================================================================== figures
def fig10_1(out_dir=None):
    """Figure 10.1: p(z) proportional to exp(-z^2/2) sigma(20z + 4) with its Laplace and variational Gaussians."""
    z0, s_lap = laplace_1d()
    m_var, s_var = gauss_kl_qp_1d(skewed_log_density, z0, s_lap)
    _, _, lnZ = moments_1d(skewed_log_density)
    z = np.linspace(-2, 4, 1201)
    lp = skewed_log_density(z) - lnZ
    fig, axes = plt.subplots(1, 2, figsize=(9.0, 3.3))
    ax = axes[0]
    ax.fill_between(z, np.exp(lp), color=YELLOW, lw=0, label="$p(z)$")
    ax.plot(z, np.exp(gauss_logpdf_1d(z, z0, s_lap)), color=common.RED, lw=2, label="Laplace")
    ax.plot(z, np.exp(gauss_logpdf_1d(z, m_var, s_var)), color=common.GREEN, lw=2, label=r"variational, min KL$(q\Vert p)$")
    ax.set_xlim(-2, 4)
    ax.set_ylim(0, 1)
    ax.set_xlabel("$z$")
    ax.legend(loc="upper right", frameon=False, fontsize=8.5)
    ax = axes[1]
    ax.plot(z, -lp, color=YELLOW, lw=2.2)
    ax.plot(z, -gauss_logpdf_1d(z, z0, s_lap), color=common.RED, lw=2)
    ax.plot(z, -gauss_logpdf_1d(z, m_var, s_var), color=common.GREEN, lw=2)
    ax.set_xlim(-2, 4)
    ax.set_ylim(0, 40)
    ax.set_xlabel("$z$")
    ax.set_ylabel(r"$-\ln$ density")
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-1", out_dir)
    print("  Laplace: z0 = %.4f, sd = %.4f; variational: m = %.4f, sd = %.4f" % (z0, s_lap, m_var, s_var))


FIG10_2_MU = np.array([0.5, 0.5])
FIG10_2_SIGMA = 0.0225 * np.array([[1.0, 0.95], [0.95, 1.0]])


def _ellipses(ax, mu, cov, color, nsds=(1, 2, 3), lw=2.0):
    for c in nsds:
        ex, ey = common.ellipse_xy(mu, cov, c)
        ax.plot(ex, ey, color=color, lw=lw)


def fig10_2(out_dir=None):
    """Figure 10.2: a correlated Gaussian (green) and the factorised q minimising KL(q||p) (a) or KL(p||q) (b)."""
    mu, Sigma = FIG10_2_MU, FIG10_2_SIGMA
    Lam = np.linalg.inv(Sigma)
    path, var_qp = mean_field_gaussian(mu, Lam, [0.2, 0.8])
    m_pq, var_pq = factorised_kl_pq(mu, Sigma)
    fig, axes = plt.subplots(1, 2, figsize=(6.6, 3.5))
    for ax, (m, v), lab, title in zip(axes, [(path[-1], var_qp), (m_pq, var_pq)], "ab",
                                      [r"KL$(q\Vert p)$", r"KL$(p\Vert q)$"]):
        _ellipses(ax, mu, Sigma, common.GREEN)
        _ellipses(ax, m, np.diag(v), common.RED)
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        ax.set_aspect("equal")
        ax.set_xticks([0, 0.5, 1])
        ax.set_yticks([0, 0.5, 1])
        ax.set_xlabel("$z_1$")
        ax.set_ylabel("$z_2$", rotation=0, labelpad=8)
        ax.set_title("(%s) minimise %s" % (lab, title), fontsize=10)
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-2", out_dir)
    print("  q sd under KL(q||p): %.4f, under KL(p||q): %.4f" % (np.sqrt(var_qp[0]), np.sqrt(var_pq[0])))


def fig10_3(out_dir=None):
    """Figure 10.3: a two-mode p (blue) and the single Gaussian minimising KL(p||q) (a) and two local minima of KL(q||p) (b, c)."""
    g = np.linspace(-1, 1, 301)
    G1, G2 = np.meshgrid(g, g)
    P = np.exp(mixture_log_density(np.stack([G1, G2], -1)))
    m_pq, S_pq = moment_match_mixture()
    fits = [(m_pq, S_pq)]
    for k in range(2):
        m, S, kl = gauss_kl_qp_2d(mixture_log_density, MIX_MU[k], 0.5 * MIX_COV[k])
        fits.append((m, S))
        print("  KL(q||p) local minimum near mode %d: mean (%.3f, %.3f), KL = %.4f" % (k + 1, m[0], m[1], kl))
    fig, axes = plt.subplots(1, 3, figsize=(9.6, 3.5))
    p_levels = P.max() * np.array([0.02, 0.08, 0.25, 0.6])
    for ax, (m, S), lab in zip(axes, fits, "abc"):
        ax.contour(G1, G2, P, levels=p_levels, colors=common.BLUE, linewidths=2.0)
        Q = np.exp(common.gauss_logpdf(np.stack([G1.ravel(), G2.ravel()], 1), m, S)).reshape(G1.shape)
        ax.contour(G1, G2, Q, levels=Q.max() * np.linspace(0.08, 0.92, 8), colors=common.RED, linewidths=1.8)
        ax.set_xlim(-1, 1)
        ax.set_ylim(-1, 1)
        ax.set_aspect("equal")
        ax.set_xticks([0])
        ax.set_yticks([0])
        ax.set_xticklabels([])
        ax.set_yticklabels([])
        ax.set_xlabel("(%s)" % lab)
    axes[0].set_title(r"minimise KL$(p\Vert q)$", fontsize=10)
    axes[1].set_title(r"minimise KL$(q\Vert p)$: one local minimum", fontsize=10)
    axes[2].set_title(r"minimise KL$(q\Vert p)$: another", fontsize=10)
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-3", out_dir)


def fig10_4(out_dir=None):
    """Figure 10.4: VB for the mean and precision of a Gaussian: exact posterior (green) and q(mu)q(tau) at four stages."""
    x = VBG_DATA[:VBG_N]
    post = normal_gamma_posterior(x, **VBG_PRIOR)
    states = vb_gaussian(x, VBG_PRIOR, *VBG_INIT)
    mus = np.linspace(-1, 1, 241)
    taus = np.linspace(0.0005, 2, 241)
    M, T = np.meshgrid(mus, taus)
    P = np.exp(normal_gamma_logpdf(M, T, post))
    frac = np.array([0.1, 0.3, 0.5, 0.7, 0.9])
    fig, axes = plt.subplots(2, 2, figsize=(8.2, 5.6))
    stages = [(states[0], common.BLUE), (states[1], common.BLUE), (states[2], common.BLUE), (states[-1], common.RED)]
    for ax, ((qmu, qtau), col), lab in zip(axes.ravel(), stages, "abcd"):
        ax.contour(M, T, P, levels=P.max() * frac, colors=common.GREEN, linewidths=2.0)
        Q = np.exp(gauss_logpdf_1d(M, qmu["m"], 1 / np.sqrt(qmu["lam"])) + gamma_logpdf(T, qtau["a"], qtau["b"]))
        ax.contour(M, T, Q, levels=Q.max() * frac, colors=col, linewidths=2.0)
        ax.set_xlim(-1, 1)
        ax.set_ylim(0, 2)
        ax.set_xticks([-1, 0, 1])
        ax.set_yticks([0, 1, 2])
        ax.set_xlabel(r"$\mu$")
        ax.set_ylabel(r"$\tau$", rotation=0, labelpad=8)
        common.panel_label(ax, lab)
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-4", out_dir)
    qmu, qtau = states[-1]
    print("  data: N = %d, mean %.3f, sd (divisor N) %.3f; %d single-factor updates to converge"
          % (len(x), x.mean(), x.std(), len(states) - 1))
    print("  exact E[mu] = %.4f, E[tau] = %.4f; VB E[mu] = %.4f, E[tau] = %.4f"
          % (post["m"], post["a"] / post["b"], qmu["m"], qtau["a"] / qtau["b"]))


def _draw_vbgmm(ax, X, q, label, threshold=1e-3):
    ax.plot(X[:, 0], X[:, 1], "o", mfc="none", mec="#18c018", ms=4.2, mew=0.9)
    w = expected_pi(q)
    Sig = np.linalg.inv(q["nu"][:, None, None] * q["W"])     # E[Lambda_k]^{-1}
    for k in np.argsort(-w):
        if w[k] < threshold:
            continue
        ex, ey = common.ellipse_xy(q["m"][k], Sig[k], 1.0)
        ax.add_patch(Polygon(np.c_[ex, ey], closed=True, fc=(1, 0, 0, min(1.0, 0.9 * w[k] + 0.04)), ec="none"))
        ax.plot(ex, ey, color="#ff1a1a", lw=3.2)
        ax.plot(q["m"][k, 0], q["m"][k, 1], "X", color="#ff1a1a", ms=10, mec="#ff1a1a")
    ax.set_xlim(-2.5, 2.5)
    ax.set_ylim(-2.5, 2.5)
    ax.set_aspect("equal")
    ax.set_xticks([])
    ax.set_yticks([])
    ax.text(0.05, 0.95, label, transform=ax.transAxes, ha="left", va="top", fontsize=17)


def fig10_6(out_dir=None):
    """Figure 10.6: VB mixture of K = 6 Gaussians on the re-scaled Old Faithful data, alpha0 = 1e-3."""
    X = common.old_faithful(standardize=True)
    prior = default_gmm_prior(2, 1e-3)
    its = [0, 15, 60, 120]
    res = vb_gmm(X, FIG10_6_START, prior, keep=its)
    fig, axes = plt.subplots(2, 2, figsize=(6.6, 6.6))
    for ax, it in zip(axes.ravel(), its):
        _draw_vbgmm(ax, X, res["snaps"][it], str(it))
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-6", out_dir)
    w = expected_pi(res["q"])
    print("  bound converged (relative change < 1e-9) at iteration %s, bound %.4f; E[pi_k] = %s"
          % (res["converged"], res["bounds"][-1], np.array2string(np.sort(w)[::-1], precision=4)))
    for it in its:
        print("  iteration %3d: bound %.4f, components with E[pi_k] > 1e-3: %d"
              % (it, res["bounds"][it], int(np.sum(expected_pi(res["snaps"][it]) > 1e-3))))


FIG10_7_ALPHA0 = 10.0
FIG10_7_STARTS = 100


def fig10_7(out_dir=None):
    """Figure 10.7: the lower bound (plus ln K!) against K for 100 random starts each, Old Faithful."""
    X = common.old_faithful(standardize=True)
    prior = default_gmm_prior(2, FIG10_7_ALPHA0)
    rng = np.random.default_rng(107)
    fig, ax = plt.subplots(figsize=(4.6, 4.4))
    best = {}
    for K in range(1, 7):
        vals = []
        for _ in range(FIG10_7_STARTS):
            idx = rng.choice(len(X), K, replace=False)
            vals.append(vb_gmm(X, idx, prior, tol=1e-7)["bounds"][-1] + gammaln(K + 1))
        vals = np.array(vals)
        best[K] = vals.max()
        ax.plot(K + rng.uniform(-0.2, 0.2, len(vals)), vals, "+", color=common.RED, ms=9, mew=0.9)
        print("  K = %d: best bound + ln K! = %.3f; %d of %d starts within 0.01 of the best"
              % (K, vals.max(), np.sum(vals > vals.max() - 0.01), len(vals)))
    ax.set_xlim(0, 7)
    ax.set_xticks(range(1, 7))
    ax.set_xlabel("$K$")
    ax.set_ylabel(r"$\mathcal{L} + \ln K!$  (approximates $\ln p(\mathcal{D}\,|\,K)$)")
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-7", out_dir)
    print("  highest bound at K = %d" % max(best, key=best.get))


def fig10_9(out_dir=None):
    """Figure 10.9: the variational lower bound against the polynomial degree, data from a cubic, a0 = b0 = 0."""
    x, t = fig10_9_data()
    beta = 1 / FIG10_9_NOISE_VAR
    degs = np.arange(1, 10)
    L = []
    for M in degs:
        Phi = np.vander(x, M + 1, increasing=True)
        L.append(vb_linear_regression(Phi, t, beta, 0.0, 0.0)[-1][-1])
    L = np.array(L)
    fig, ax = plt.subplots(figsize=(5.4, 3.6))
    ax.plot(degs, L, "-", color=common.BLUE, lw=1.4)
    ax.plot(degs, L, "o", mfc="white", mec=common.BLUE, mew=2.2, ms=9)
    ax.set_xticks([1, 3, 5, 7, 9])
    ax.set_xlim(0.5, 9.5)
    ax.set_xlabel("degree $M$ of the polynomial")
    ax.set_ylabel(r"lower bound $\mathcal{L}(q)$")
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-9", out_dir)
    print("  bound by degree: %s; best degree %d" % (np.array2string(L - L.max(), precision=2), degs[np.argmax(L)]))


FIGURES = {"10.1": fig10_1, "10.2": fig10_2, "10.3": fig10_3, "10.4": fig10_4,
           "10.6": fig10_6, "10.7": fig10_7, "10.9": fig10_9}

if __name__ == "__main__":
    common.run(FIGURES, 10)