"""Bishop (2006), Chapter 10, Sections 10.5 to 10.7: local variational methods,
variational logistic regression and expectation propagation (EP).

One function per figure; run `python3 ch10_local_variational_ep.py --list`.
The figures are written to ../img/ch10/fig10-10.png ... fig10-17.png.

Algorithms, each with the book's equation numbers in its docstring:

  Section 10.5    conjugate_exp_neg, tangent_exp_neg, binary_entropy,
                  sigmoid_upper_bound, jj_lambda, sigmoid_lower_bound
  Section 10.6    vlr_e_step, vlr_m_step, vlr_bound,
                  variational_logistic_regression (EM for xi, fixed prior),
                  variational_logistic_regression_alpha (alpha inferred too),
                  predictive_probability
  Figure 10.14    laplace_toy, kl_qp_gaussian_toy, ep_toy
  Section 10.7.1  ep_site_update, ep_clutter (every update (10.212) to (10.224),
                  evidence included), vb_clutter, laplace_clutter, exact_clutter

Data. The figures and the two browser demos (variational-logistic-demo.html,
ep-clutter-demo.html) draw their data from the same seeded generator: DemoRNG
below is a line-by-line port of DK.rng in ../../demo-kit.js (mulberry32 uniforms
and Box-Muller normals), so a seed gives the same data set in Python and in the
browser.

What could not be reproduced exactly, and the choices made:

* Figure 10.10 and 10.12 are exact (they are plots of formulas). Figure 10.11 is
  a schematic in the book; we draw it for the convex function
  f(x) = 6.03 exp(-1.358 x) - 0.9 and slope eta = -0.43.
* Figure 10.13: the book gives neither the data nor the prior. We draw 25 + 25
  points, class t = 0 (crosses) with x1 ~ N(-1.1, 0.25^2), x2 ~ N(0.8, 1) and
  class t = 1 (circles) with x1 ~ N(1.1, 0.25^2), x2 ~ N(-1, 1), using seed
  LOGIT_SEED, which was picked among seeds giving a nearly vertical boundary
  like the book's. The model has a bias, phi(x) = (1, x1, x2) (the book's five
  sampled boundaries do not share a common point, so it has one too), and the
  prior is N(w | 0, alpha^-1 I) with alpha = LOGIT_ALPHA = 0.03, which gives
  contours like the book's. Inferring alpha as in Section 10.6.3 (a0 = b0 = 0.01)
  gives E[alpha] = 0.11 on these data. EM for xi starts from xi_n = 1 and needs a
  few hundred iterations here: with separable data the weights keep growing until
  the prior stops them. The predictive contours use the probit approximation
  (4.155) of Section 4.5.2, as the book says.
* Figure 10.14 needs no data. The EP approximation of a density with one
  non-Gaussian factor converges after one update to the Gaussian with the mean
  and variance of p (moment matching). The variational (KL(q||p)) and EP curves
  agree with the book's; the book's Laplace curve is narrower (peak about 0.76)
  than the Laplace approximation itself, whose mode is z0 = 0.0775 with
  precision -d2 ln p(z0) = 2.54 (peak 0.64), as in our Figure 10.1.
* EP is not guaranteed to converge to a good solution (Section 10.7). On some
  clutter data sets (for example seeds 6 and 23 with N = 20) the first sweep is
  captured by two far-out clutter points and EP settles at a poor fixed point;
  its evidence estimate is then far below the exact one, which reveals the
  failure. The seeds used here do not do this.
* Figure 10.15: N = 20 points, theta = 3, w = 0.5, a = 10 (seed CLUTTER_SEED).
* Figure 10.16: the book shows two unspecified examples. We take x_n = 2 and the
  cavity distributions N(theta | 0, 0.3^2) (left) and N(theta | 3, 0.3^2)
  (right), which is what the book's panels show, and compute f~_n by the EP
  update (10.216) to (10.222).
* Figure 10.17: N = 200 points, theta = 3 (seed CLUTTER_SEED), D = 1. Minka
  (2001), whose experiment the book reproduces, counted flops with MATLAB. We
  count them explicitly with the convention stated next to FLOP_CONVENTION
  below, so our absolute numbers differ from the book's, but the comparison
  between the methods is like for like. The book plots the absolute error of
  the evidence p(D) (about 10^-200 here); we plot the error of ln p(D), which is
  the relative error of p(D) for small errors and is easier to read.

Errata applied (Bishop's PRML errata list, pages 493 to 513):
* eta replaces lambda in Section 10.5 up to (10.141); (10.141) to (10.143) and
  (10.164) in their corrected forms; (10.177) with a_N, b_N; (10.183) with
  E[w w^T] = Sigma_N + mu_N mu_N^T; the left-hand sides of (10.217), (10.218)
  are m^new and v^new.
* Two further inconsistencies, not in the errata list, are handled here and
  checked numerically in checks_ch10b.py:
  (i) (10.222) is the coefficient of the unnormalised site
      f~_n(theta) = s_n exp(-||theta - m_n||^2 / (2 v_n)) used by Minka (2001),
      not of s_n N(theta | m_n, v_n I) as written in (10.213); in the notation
      of (10.213) the coefficient is s_n = Z_n / N(m_n | m^\\n, (v_n + v^\\n) I).
  (ii) (10.223) and (10.224) must include the prior factor f~_0 = p(theta),
      that is the product and the sum start at n = 0; otherwise the prior's
      normalisation (2 pi b)^(-D/2) is lost. With the sites in Minka's form the
      evidence is
          ln p(D) ~ (D/2) ln(2 pi v) + B/2 - (D/2) ln(2 pi b) + sum_n ln s_n,
      which is exact when w = 0 (a conjugate Gaussian model).
"""
import math

import numpy as np
from scipy import optimize

import common

LN2PI = math.log(2.0 * math.pi)

# Settings shared with the demos (variational-logistic-demo.html, ep-clutter-demo.html).
LOGIT_SEED = 287          # data of Figure 10.13
LOGIT_ALPHA = 0.03        # prior precision alpha of Figure 10.13
LOGIT_SAMPLE_SEED = 5     # the five posterior samples of Figure 10.13 (right)
CLUTTER_SEED = 13         # data of Figures 10.15 and 10.17
CLUTTER_W, CLUTTER_A, CLUTTER_B = 0.5, 10.0, 100.0   # Minka's values, Section 10.7.1


# --------------------------------------------------------------------------- shared random numbers
class DemoRNG:
    """Port of DK.rng (../../demo-kit.js): mulberry32 uniforms, Box-Muller normals.

    DemoRNG(seed).u() and .norm() return exactly the numbers that DK.rng(seed)
    returns in the browser (up to the last bit of sin, cos and log), so the
    figures and the demos share their data sets.
    """

    def __init__(self, seed):
        self.a = self._i32(seed)
        self.spare = None

    @staticmethod
    def _i32(x):
        x &= 0xFFFFFFFF
        return x - (1 << 32) if x & 0x80000000 else x

    def _imul(self, p, q):
        return self._i32((p & 0xFFFFFFFF) * (q & 0xFFFFFFFF))

    def u(self):
        """A uniform number in [0, 1)."""
        self.a = self._i32(self.a + 0x6D2B79F5)
        a = self.a
        t = self._imul(a ^ ((a & 0xFFFFFFFF) >> 15), 1 | a)
        t = self._i32(self._i32(t + self._imul(t ^ ((t & 0xFFFFFFFF) >> 7), 61 | t)) ^ t)
        return ((t ^ ((t & 0xFFFFFFFF) >> 14)) & 0xFFFFFFFF) / 4294967296.0

    def norm(self):
        """A standard normal number (Box-Muller, the second of each pair kept for the next call)."""
        if self.spare is not None:
            s, self.spare = self.spare, None
            return s
        x = 0.0
        while x == 0.0:
            x = self.u()
        r, th = math.sqrt(-2.0 * math.log(x)), 2.0 * math.pi * self.u()
        self.spare = r * math.sin(th)
        return r * math.cos(th)


# --------------------------------------------------------------------------- Section 10.5
def log_sigmoid(x):
    """ln sigma(x) = -ln(1 + e^-x), computed without overflow."""
    return -np.logaddexp(0.0, -np.asarray(x, dtype=float))


def tangent_exp_neg(x, xi):
    """Tangent line of f(x) = exp(-x) at x = xi, (10.125) and (10.126): a lower bound on exp(-x)."""
    return np.exp(-xi) - np.exp(-xi) * (np.asarray(x, dtype=float) - xi)


def conjugate_exp_neg(eta):
    """Conjugate (dual) function of exp(-x), (10.131): g(eta) = eta - eta ln(-eta), eta < 0.

    g(eta) = max_x {eta x - exp(-x)} (10.129); the maximiser is x = -ln(-eta).
    """
    eta = np.asarray(eta, dtype=float)
    return eta - eta * np.log(-eta)


def binary_entropy(eta):
    """g(eta) = -eta ln eta - (1 - eta) ln(1 - eta), the conjugate of ln sigma(x), (10.135)."""
    eta = np.asarray(eta, dtype=float)
    return -eta * np.log(eta) - (1.0 - eta) * np.log1p(-eta)


def sigmoid_upper_bound(x, eta):
    """Upper bound sigma(x) <= exp(eta x - g(eta)) from convex duality, (10.136) and (10.137).

    The bound touches sigma at x = ln((1 - eta) / eta), where d ln sigma / dx = eta.
    """
    return np.exp(eta * np.asarray(x, dtype=float) - binary_entropy(eta))


def jj_lambda(xi):
    """lambda(xi) = (sigma(xi) - 1/2) / (2 xi) = tanh(xi / 2) / (4 xi), (10.141) and (10.150).

    Even in xi, decreasing in |xi|, with lambda(0) = 1/8 (the limit as xi -> 0).
    """
    xi = np.abs(np.asarray(xi, dtype=float))
    small = xi < 1e-4
    safe = np.where(small, 1.0, xi)
    return np.where(small, 0.125 - xi ** 2 / 96.0, np.tanh(safe / 2.0) / (4.0 * safe))


def sigmoid_lower_bound(x, xi):
    """Jaakkola-Jordan bound sigma(x) >= sigma(xi) exp{(x - xi)/2 - lambda(xi)(x^2 - xi^2)}, (10.144).

    The bound is exact at x = xi and x = -xi; it is the exponential of a quadratic in x.
    """
    x = np.asarray(x, dtype=float)
    return np.exp(log_sigmoid(xi) + (x - xi) / 2.0 - jj_lambda(xi) * (x ** 2 - xi ** 2))


# --------------------------------------------------------------------------- Section 10.6
def logistic_data(seed=LOGIT_SEED, n0=25, n1=25):
    """Two linearly separable classes in 2-D (Figure 10.13); same generator as the demo.

    Class t = 0: x1 ~ N(-1.1, 0.25^2), x2 ~ N(0.8, 1); class t = 1: x1 ~ N(1.1, 0.25^2),
    x2 ~ N(-1, 1). Returns X (N x 2) and t (N,) with N = n0 + n1.
    """
    r = DemoRNG(seed)
    X, t = [], []
    for k, (c1, c2, n) in enumerate([(-1.1, 0.8, n0), (1.1, -1.0, n1)]):
        for _ in range(n):
            x1 = c1 + 0.25 * r.norm()
            x2 = c2 + 1.0 * r.norm()
            X.append([x1, x2])
            t.append(k)
    return np.array(X), np.array(t, dtype=float)


def design_matrix(X):
    """phi(x) = (1, x1, x2): a bias and the two inputs."""
    X = np.atleast_2d(X)
    return np.column_stack([np.ones(len(X)), X])


def vlr_e_step(Phi, t, xi, m0, S0inv):
    """E step: the Gaussian q(w) = N(w | m_N, S_N) for fixed xi, (10.156) to (10.158).

    S_N^-1 = S_0^-1 + 2 sum_n lambda(xi_n) phi_n phi_n^T            (10.158)
    m_N    = S_N (S_0^-1 m_0 + sum_n (t_n - 1/2) phi_n)            (10.157)
    """
    lam = jj_lambda(xi)
    SNinv = S0inv + 2.0 * (Phi.T * lam) @ Phi
    SN = np.linalg.inv(SNinv)
    SN = 0.5 * (SN + SN.T)
    mN = SN @ (S0inv @ m0 + Phi.T @ (t - 0.5))
    return mN, SN


def vlr_m_step(Phi, mN, SN):
    """M step (10.163): xi_n^2 = phi_n^T E[w w^T] phi_n = phi_n^T (S_N + m_N m_N^T) phi_n."""
    EwwT = SN + np.outer(mN, mN)
    return np.sqrt(np.einsum("ni,ij,nj->n", Phi, EwwT, Phi))


def vlr_bound(xi, mN, SN, m0, S0inv):
    """Lower bound L(xi) <= ln p(t), (10.159) evaluated in closed form, (10.164) as corrected in the errata.

    L(xi) = 1/2 ln(|S_N| / |S_0|) + 1/2 m_N^T S_N^-1 m_N - 1/2 m_0^T S_0^-1 m_0
            + sum_n {ln sigma(xi_n) - xi_n / 2 + lambda(xi_n) xi_n^2},
    where (m_N, S_N) is the E-step posterior for this xi.
    """
    _, ldN = np.linalg.slogdet(SN)
    _, ld0inv = np.linalg.slogdet(S0inv)
    return (0.5 * (ldN + ld0inv) + 0.5 * mN @ np.linalg.solve(SN, mN) - 0.5 * m0 @ S0inv @ m0
            + np.sum(log_sigmoid(xi) - xi / 2.0 + jj_lambda(xi) * xi ** 2))


def variational_logistic_regression(Phi, t, alpha=LOGIT_ALPHA, m0=None, S0=None, xi0=1.0,
                                    max_iter=5000, tol=1e-9):
    """Bayesian logistic regression with the local variational bound, EM for xi (Section 10.6.2).

    Prior N(w | m0, S0), by default N(w | 0, alpha^-1 I). Starting from xi_n = xi0 the
    algorithm alternates the E step (10.156) to (10.158) and the M step (10.163); the
    bound (10.164) is recorded after every E step and never decreases. Stops when the
    bound changes by less than tol. Returns a dict with m_N, S_N, xi and the bound trace.
    """
    N, M = Phi.shape
    m0 = np.zeros(M) if m0 is None else np.asarray(m0, dtype=float)
    S0 = np.eye(M) / alpha if S0 is None else np.asarray(S0, dtype=float)
    S0inv = np.linalg.inv(S0)
    xi = np.full(N, float(xi0)) if np.isscalar(xi0) else np.asarray(xi0, dtype=float)
    bounds = []
    for it in range(max_iter):
        mN, SN = vlr_e_step(Phi, t, xi, m0, S0inv)
        bounds.append(vlr_bound(xi, mN, SN, m0, S0inv))
        if it > 0 and abs(bounds[-1] - bounds[-2]) < tol:
            break
        xi = vlr_m_step(Phi, mN, SN)
    return {"mN": mN, "SN": SN, "xi": xi, "bound": np.array(bounds), "iterations": len(bounds)}


def variational_logistic_regression_alpha(Phi, t, a0=1e-2, b0=1e-2, xi0=1.0, max_iter=5000, tol=1e-9):
    """Logistic regression with alpha inferred as well: Section 10.6.3, (10.165) to (10.183).

    Prior p(w | alpha) = N(w | 0, alpha^-1 I) (10.165), hyperprior p(alpha) = Gam(alpha | a0, b0)
    (10.166), factorised q(w) q(alpha) (10.173) and local bounds with parameters xi. One
    cycle updates, in turn,
      q(w) = N(w | mu_N, Sigma_N):  Sigma_N^-1 = E[alpha] I + 2 sum lambda(xi_n) phi_n phi_n^T (10.176),
                                    mu_N = Sigma_N sum (t_n - 1/2) phi_n                     (10.175)
      q(alpha) = Gam(a_N, b_N):     a_N = a0 + M/2 (10.178),  b_N = b0 + E[w^T w]/2 (10.179)
      xi:                           xi_n^2 = phi_n^T (Sigma_N + mu_N mu_N^T) phi_n          (10.181)
    with E[alpha] = a_N / b_N (10.182). Each update maximises the bound L~(q, xi) of (10.172)
    over one block, so the recorded bound never decreases.
    """
    N, M = Phi.shape
    xi = np.full(N, float(xi0))
    aN, bN = a0 + 0.5 * M, b0          # (10.178); b_N starts at b0, i.e. E[alpha] = aN / b0
    bounds, Ealphas = [], []
    for it in range(max_iter):
        Ealpha = aN / bN                                           # (10.182)
        lam = jj_lambda(xi)
        Sinv = Ealpha * np.eye(M) + 2.0 * (Phi.T * lam) @ Phi     # (10.176)
        S = np.linalg.inv(Sinv)
        S = 0.5 * (S + S.T)
        mu = S @ (Phi.T @ (t - 0.5))                               # (10.175)
        EwTw = mu @ mu + np.trace(S)
        bN = b0 + 0.5 * EwTw                                       # (10.179)
        xi = np.sqrt(np.einsum("ni,ij,nj->n", Phi, S + np.outer(mu, mu), Phi))   # (10.181)
        bounds.append(alpha_bound(Phi, t, xi, mu, S, aN, bN, a0, b0))
        Ealphas.append(aN / bN)
        if it > 0 and abs(bounds[-1] - bounds[-2]) < tol:
            break
    return {"mN": mu, "SN": S, "xi": xi, "aN": aN, "bN": bN, "Ealpha": np.array(Ealphas),
            "bound": np.array(bounds), "iterations": len(bounds)}


def alpha_bound(Phi, t, xi, mu, S, aN, bN, a0, b0):
    """The bound L~(q, xi) of (10.172) for q(w) = N(mu, S) and q(alpha) = Gam(aN, bN).

    L~ = E[ln h(w, xi)] + E[ln p(w | alpha)] + E[ln p(alpha)] - E[ln q(w)] - E[ln q(alpha)],
    all expectations in closed form (E[ln alpha] = psi(aN) - ln bN).
    """
    from scipy.special import digamma, gammaln
    M = len(mu)
    Ea, Elna = aN / bN, digamma(aN) - np.log(bN)
    a = Phi @ mu
    Ea2 = np.einsum("ni,ij,nj->n", Phi, S + np.outer(mu, mu), Phi)
    E_ln_h = np.sum(log_sigmoid(xi) + a * t - (a + xi) / 2.0 - jj_lambda(xi) * (Ea2 - xi ** 2))
    EwTw = mu @ mu + np.trace(S)
    E_ln_pw = 0.5 * M * (Elna - LN2PI) - 0.5 * Ea * EwTw
    E_ln_pa = a0 * np.log(b0) - gammaln(a0) + (a0 - 1.0) * Elna - b0 * Ea
    H_qw = 0.5 * np.linalg.slogdet(S)[1] + 0.5 * M * (1.0 + LN2PI)
    H_qa = gammaln(aN) - (aN - 1.0) * digamma(aN) - np.log(bN) + aN
    return E_ln_h + E_ln_pw + E_ln_pa + H_qw + H_qa


def predictive_probability(Phi, mN, SN):
    """p(t = 1 | phi) ~ sigma(kappa(sigma_a^2) mu_a), the probit approximation (4.153) to (4.155).

    mu_a = m_N^T phi, sigma_a^2 = phi^T S_N phi, kappa(s2) = (1 + pi s2 / 8)^(-1/2).
    """
    mu_a = Phi @ mN
    s2 = np.einsum("ni,ij,nj->n", Phi, SN, Phi)
    return common.sigmoid(mu_a / np.sqrt(1.0 + np.pi * s2 / 8.0))


# --------------------------------------------------------------------------- Figure 10.14: a toy density
def toy_log_density(z):
    """Unnormalised ln p(z) = -z^2/2 + ln sigma(20 z + 4) (Figures 4.14, 10.1 and 10.14)."""
    z = np.asarray(z, dtype=float)
    return -0.5 * z ** 2 + log_sigmoid(20.0 * z + 4.0)


def _toy_grid(n=40001):
    z = np.linspace(-10.0, 10.0, n)
    lp = toy_log_density(z)
    dz = z[1] - z[0]
    p = np.exp(lp)
    Z = p.sum() * dz
    return z, p / Z, dz


def laplace_toy():
    """Laplace approximation (Section 4.4): Newton's method for the mode z0, then A = -d2 ln p(z0).

    d ln p / dz = -z + 20 sigma(-(20 z + 4)); d2 ln p / dz2 = -1 - 400 sigma(1 - sigma).
    Returns (mean, variance) = (z0, 1 / A).
    """
    z = 0.0
    for _ in range(100):
        s = common.sigmoid(20.0 * z + 4.0)
        g = -z + 20.0 * (1.0 - s)
        h = -1.0 - 400.0 * s * (1.0 - s)
        z -= g / h
        if abs(g) < 1e-14:
            break
    s = common.sigmoid(20.0 * z + 4.0)
    return z, 1.0 / (1.0 + 400.0 * s * (1.0 - s))


def kl_qp_gaussian_toy():
    """The Gaussian q = N(mu, s^2) minimising KL(q || p) (the global variational approximation).

    KL(q || p) = -H[q] - E_q[ln p~(z)] + ln Z, with E_q computed by 80-point Gauss-Hermite
    quadrature; minimised over (mu, ln s) with BFGS. Returns (mean, variance).
    """
    nodes, weights = np.polynomial.hermite_e.hermegauss(80)
    weights = weights / weights.sum()

    def kl(par):
        mu, ls = par
        z = mu + np.exp(ls) * nodes
        return -ls - weights @ toy_log_density(z)

    res = optimize.minimize(kl, x0=[0.5, np.log(0.4)], method="BFGS", options={"gtol": 1e-10})
    return res.x[0], np.exp(2 * res.x[1])


def ep_toy(n_iter=3):
    """EP for p(z) proportional to f0(z) f1(z), f0 = exp(-z^2/2), f1 = sigma(20 z + 4) (Section 10.7).

    f0 is Gaussian and is kept exactly; the site f~1 starts at 1. One update: cavity
    q^\\1 = q / f~1 (10.205), moments of the tilted distribution f1 q^\\1 / Z1 by quadrature
    (10.206), new site f~1 = Z1 q^new / q^\\1 (10.207). The second update changes nothing,
    so EP returns the Gaussian with the mean and variance of p (moment matching, (10.187)).
    Returns (mean, variance).
    """
    z = np.linspace(-10.0, 10.0, 40001)
    dz = z[1] - z[0]
    tau1, nu1 = 0.0, 0.0                      # site f~1 in natural parameters (precision, precision x mean)
    for _ in range(n_iter):
        tau, nu = 1.0 + tau1, nu1             # q = f0 f~1, normalised
        tau_c, nu_c = tau - tau1, nu - nu1    # cavity (10.205)
        m_c, v_c = nu_c / tau_c, 1.0 / tau_c
        tilt = common.sigmoid(20.0 * z + 4.0) * np.exp(-0.5 * (z - m_c) ** 2 / v_c)
        tilt /= tilt.sum() * dz
        m_new = (tilt * z).sum() * dz
        v_new = (tilt * (z - m_new) ** 2).sum() * dz
        tau1, nu1 = 1.0 / v_new - tau_c, m_new / v_new - nu_c     # (10.207)
    return m_new, v_new


# --------------------------------------------------------------------------- Section 10.7.1: the clutter problem
def clutter_data(N, theta=3.0, w=CLUTTER_W, a=CLUTTER_A, seed=CLUTTER_SEED):
    """N draws from (1 - w) N(x | theta, 1) + w N(x | 0, a), (10.209) with D = 1.

    For each point: u = U(0, 1), z = N(0, 1); x = sqrt(a) z if u < w (clutter), else theta + z.
    Same generator as ep-clutter-demo.html.
    """
    r = DemoRNG(seed)
    x = []
    for _ in range(N):
        u = r.u()
        z = r.norm()
        x.append(math.sqrt(a) * z if u < w else theta + z)
    return np.array(x)


def ln_gauss_spherical(d2, var, D):
    """ln N(x | m, var I) in D dimensions as a function of d2 = ||x - m||^2."""
    return -0.5 * (D * (LN2PI + np.log(var)) + d2 / var)


def ln_clutter_terms(x, w, a):
    """ln {w N(x_n | 0, a I)} for every n: the clutter part of f_n, fixed during the iterations."""
    x = _as_points(x)
    D = x.shape[1]
    return np.log(w) + ln_gauss_spherical((x ** 2).sum(1), a, D)


def _as_points(x):
    x = np.asarray(x, dtype=float)
    return x[:, None] if x.ndim == 1 else x


# FLOP_CONVENTION: each +, -, *, /, comparison, exp, log, log1p and sqrt counts as one
# floating-point operation. Constants that never change (ln w, ln(1 - w), 1/b, ln 2 pi and
# the clutter terms ln{w N(x_n | 0, aI)}) are computed once and charged to the set-up.
# Only the work of each method is counted; evaluating the errors and the evidence at every
# recorded point (for the plots) is not, because a user would do it once at the end.
def flops_setup(N, D=1):
    """Set-up shared by EP, VB and Laplace: the N clutter terms ln{w N(x_n | 0, aI)}."""
    return N * (2 * D + 3) + 3


def flops_ep_site(D=1):
    """One EP site update, (10.214) to (10.222): 14 D + 45 flops (itemised in ep_site_update)."""
    return 14 * D + 45


def flops_vb_iteration(N, D=1):
    """One VB iteration (q(theta) then all q(z_n)): N (5 D + 7) + 2 flops (itemised in vb_clutter)."""
    return N * (5 * D + 7) + 2


def flops_laplace_eval(N):
    """ln p(theta, D) with its first and second derivatives, D = 1: 19 N + 8 flops."""
    return 19 * N + 8


def ep_site_update(x_n, ln_clutter_n, m, v, tau_n, m_n, w):
    """Refine the factor f~_n of the clutter problem: (10.214) to (10.222).

    The site is f~_n(theta) = s_n exp{-||theta - m_n||^2 / (2 v_n)} with precision
    tau_n = 1 / v_n (tau_n = 0 is the initial flat site, v_n = infinity; tau_n < 0 is
    allowed). q(theta) = N(theta | m, v I) is the current approximation (10.212).
    Returns (m_new, v_new, tau_n, m_n, ln s_n, ln Z_n), or None when the cavity variance
    is not positive, in which case the update is skipped (Minka 2001).
    Flop count per line in the comments; 14 D + 45 in total.
    """
    D = m.shape[0]
    tau_cav = 1.0 / v - tau_n                        # (10.215): 2 flops, + 1 for the test
    if tau_cav <= 0.0:
        return None
    v_cav = 1.0 / tau_cav                            # 1
    m_cav = m + (v_cav * tau_n) * (m - m_n)          # (10.214) m + v^\n v_n^-1 (m - m_n): 3D + 1
    d = x_n - m_cav                                  # D
    d2 = d @ d                                       # 2D - 1
    s = v_cav + 1.0                                  # 1
    ln_signal = math.log(1.0 - w) + ln_gauss_spherical(d2, s, D)    # 7 (ln(1 - w) is a constant)
    ln_Z = np.logaddexp(ln_signal, ln_clutter_n)     # (10.216) Z_n: 5
    rho = math.exp(ln_signal - ln_Z)                 # (10.219) rho_n = 1 - w N(x_n|0,aI)/Z_n: 2
    k = v_cav / s                                    # 1
    m_new = m_cav + (rho * k) * d                    # (10.217) m^new: 2D + 1
    v_new = (v_cav - rho * v_cav * k                 # (10.218) v^new: 10
             + rho * (1.0 - rho) * k * k * d2 / D)
    tau_n = 1.0 / v_new - tau_cav                    # (10.220) v_n^-1: 2, + 1 for the test below
    if tau_n == 0.0:                                 # the site stays flat (x_n far out in the clutter)
        return m_new, v_new, 0.0, np.zeros(D), ln_Z, ln_Z
    v_n = 1.0 / tau_n                                # 1
    m_n = m_cav + ((v_n + v_cav) * tau_cav) * (m_new - m_cav)       # (10.221): 3D + 2
    e = m_n - m_cav                                  # D
    # (10.222) s_n = Z_n / ((2 pi v_n)^(D/2) N(m_n | m^\n, (v_n + v^\n) I)), written with
    # (v_n + v^\n) / v_n = v^\n / v^new so that it stays real when v_n < 0: 2D + 6
    ln_s = ln_Z + 0.5 * D * math.log(v_cav / v_new) + 0.5 * (e @ e) / (v_n + v_cav)
    return m_new, v_new, tau_n, m_n, ln_s, ln_Z


def ep_evidence(m, v, b, tau, msite, ln_s):
    """EP approximation to ln p(D): (10.223) and (10.224) with the prior factor included.

    ln p(D) ~ (D/2) ln(2 pi v) + B/2 - (D/2) ln(2 pi b) + sum_n ln s_n,
    B = ||m||^2 / v - sum_n ||m_n||^2 / v_n   (sites in Minka's form, see the module notes).
    """
    D = m.shape[0]
    B = m @ m / v - np.sum(tau * (msite ** 2).sum(1))
    return 0.5 * D * math.log(2 * math.pi * v) + 0.5 * B - 0.5 * D * math.log(2 * math.pi * b) + np.sum(ln_s)


def ep_clutter(x, w=CLUTTER_W, a=CLUTTER_A, b=CLUTTER_B, max_sweeps=100, tol=1e-12, record_every=None):
    """Expectation propagation for the clutter problem, Section 10.7.1 (Minka 2001).

    x: N data points (N,) or (N, D). Prior p(theta) = N(theta | 0, b I) (10.210) is kept
    exactly as f~_0; the sites f~_n, n = 1..N, start flat (tau_n = 0, ln s_n = 0), so the
    initial q(theta) is the prior. Each sweep refines f~_1, ..., f~_N in turn with
    ep_site_update ((10.214) to (10.222)); sweeps stop when m and v change by less than
    tol over a sweep. The first sweep alone is assumed density filtering (ADF).
    record_every = k stores (flops, m, v, ln p(D) estimate) at the end of the first sweep and
    then after every k site updates (default: once per sweep). Before the first sweep is
    complete q does not yet contain all the data, so nothing is recorded. Returns a dict.
    """
    x = _as_points(x)
    N, D = x.shape
    lnc = ln_clutter_terms(x, w, a)
    m, v = np.zeros(D), float(b)
    tau, msite, ln_s, ln_Z = np.zeros(N), np.zeros((N, D)), np.zeros(N), np.zeros(N)
    flops = flops_setup(N, D)
    every = record_every or N
    history, updates, skipped, sweeps = [], 0, 0, 0
    for sweep in range(max_sweeps):
        m_old, v_old = m.copy(), v
        for n in range(N):
            out = ep_site_update(x[n], lnc[n], m, v, tau[n], msite[n], w)
            flops += flops_ep_site(D) if out is not None else 3
            if out is None:
                skipped += 1
            else:
                m, v, tau[n], msite[n], ln_s[n], ln_Z[n] = out
            updates += 1
            if updates >= N and (updates - N) % every == 0:
                history.append((flops, m[0] if D == 1 else m.copy(), v, ep_evidence(m, v, b, tau, msite, ln_s)))
        sweeps = sweep + 1
        if np.max(np.abs(m - m_old)) < tol and abs(v - v_old) < tol * max(1.0, v):
            break
    return {"m": m[0] if D == 1 else m, "v": v, "tau": tau, "msite": msite, "ln_s": ln_s, "ln_Z": ln_Z,
            "ln_evidence": ep_evidence(m, v, b, tau, msite, ln_s), "sweeps": sweeps, "skipped": skipped,
            "flops": flops, "history": history}


def vb_clutter(x, w=CLUTTER_W, a=CLUTTER_A, b=CLUTTER_B, max_iter=200, tol=1e-12):
    """Mean-field variational Bayes for the clutter problem (Sections 10.1 and 10.2 applied to (10.209)).

    Latent z_n = 1 if x_n comes from N(theta, I), 0 if it is clutter; q(theta) prod_n q(z_n)
    with q(theta) = N(theta | m, v I) and r_n = q(z_n = 1), from the general result (10.9):
      q(theta): 1 / v = 1 / b + sum r_n,  m = v sum r_n x_n
      q(z_n):   r_n = sigma(l1_n - l0_n), l1_n = ln(1-w) - (D/2) ln 2pi - (||x_n - m||^2 + D v) / 2,
                l0_n = ln{w N(x_n | 0, a I)}.
    Starts from r_n = 1 - w (the prior probability of not being clutter). The bound
    L(q) <= ln p(D) is recorded after every iteration and never decreases.
    Returns a dict with m, v, r, the bound and a history of (flops, m, v, bound).
    """
    x = _as_points(x)
    N, D = x.shape
    l0 = ln_clutter_terms(x, w, a)
    c1 = math.log(1.0 - w) - 0.5 * D * LN2PI
    r = np.full(N, 1.0 - w)
    flops = flops_setup(N, D)
    history, L_old = [], -np.inf
    for it in range(max_iter):
        # q(theta): N - 1 + 2 + (2N - 1) D + D = N + 1 + 2 N D flops
        v = 1.0 / (1.0 / b + r.sum())
        m = v * (r @ x)
        # q(z_n): per point d (D), ||d||^2 (2D - 1), l1 (3), r (4): N (3D + 6) + 1 flops (D v once)
        l1 = c1 - 0.5 * (((x - m) ** 2).sum(1) + D * v)
        r = 1.0 / (1.0 + np.exp(l0 - l1))
        flops += flops_vb_iteration(N, D)
        L = vb_bound(m, v, r, l1, l0, b)
        history.append((flops, m[0] if D == 1 else m.copy(), v, L))
        if abs(L - L_old) < tol:
            break
        L_old = L
    return {"m": m[0] if D == 1 else m, "v": v, "r": r, "bound": L, "iterations": it + 1,
            "flops": flops, "history": history}


def vb_bound(m, v, r, l1, l0, b):
    """L(q) = E[ln p(theta)] + sum_n {r_n l1_n + (1 - r_n) l0_n + H[q(z_n)]} + H[q(theta)]."""
    D = m.shape[0]
    with np.errstate(divide="ignore", invalid="ignore"):
        Hz = -np.where(r > 0, r * np.log(r), 0.0) - np.where(r < 1, (1 - r) * np.log1p(-r), 0.0)
    E_ln_prior = -0.5 * D * (LN2PI + math.log(b)) - (m @ m + D * v) / (2.0 * b)
    H_theta = 0.5 * D * (1.0 + LN2PI + math.log(v))
    return E_ln_prior + np.sum(r * l1 + (1.0 - r) * l0 + Hz) + H_theta


def laplace_clutter(x, w=CLUTTER_W, a=CLUTTER_A, b=CLUTTER_B, theta0=None, max_iter=100, tol=1e-12):
    """Laplace approximation (Section 4.4) for the clutter posterior, D = 1.

    Damped Newton's method for the mode of ln p(theta, D): from theta0 (default: the sample
    mean) take the step -g/H (or a unit step uphill if H >= 0) and halve it until the log
    joint increases (up to rounding, 1e-12 relative). At theta, with r_n = (1 - w) N(x_n | theta, 1) / p(x_n | theta) and
    d_n = x_n - theta,
        g = -theta / b + sum r_n d_n,   H = -1 / b + sum {r_n (d_n^2 - 1) - (r_n d_n)^2}.
    q(theta) = N(theta | theta*, -1/H) and ln p(D) ~ ln p(theta*, D) + ln(2 pi)/2 - ln(-H)/2 (4.135).
    Returns a dict with the mode, variance, evidence and a history of (flops, mean, var, evidence).
    """
    x = np.asarray(x, dtype=float).ravel()
    N = len(x)
    lnc = ln_clutter_terms(x, w, a)
    c1 = math.log(1.0 - w) - 0.5 * LN2PI
    c0 = -0.5 * (LN2PI + math.log(b))

    def evaluate(th):
        d = x - th
        ls = c1 - 0.5 * d * d
        lp = np.logaddexp(ls, lnc)
        r = np.exp(ls - lp)
        f = c0 - 0.5 * th * th / b + lp.sum()
        g = -th / b + np.sum(r * d)
        H = -1.0 / b + np.sum(r * (d * d - 1.0) - (r * d) ** 2)
        return f, g, H

    th = float(np.mean(x)) if theta0 is None else float(theta0)
    flops = flops_setup(N, 1) + (N if theta0 is None else 0)
    f, g, H = evaluate(th)
    flops += flops_laplace_eval(N)

    def summary():
        if H < 0:
            return th, -1.0 / H, f + 0.5 * LN2PI - 0.5 * math.log(-H)
        return th, np.nan, np.nan

    history = [(flops,) + summary()]
    for it in range(max_iter):
        step = -g / H if H < 0 else math.copysign(1.0, g)
        flops += 4
        while True:
            f_new, g_new, H_new = evaluate(th + step)
            flops += flops_laplace_eval(N) + 2
            if f_new >= f - 1e-12 * abs(f) or abs(step) < 1e-15:     # allow for rounding near the mode
                break
            step *= 0.5
        th, f, g, H = th + step, f_new, g_new, H_new
        history.append((flops,) + summary())
        if abs(step) < tol * max(1.0, abs(th)):
            break
    mode, var, lnZ = summary()
    return {"m": mode, "v": var, "ln_evidence": lnZ, "iterations": it + 1, "flops": flops, "history": history}


def clutter_log_joint(theta, x, w=CLUTTER_W, a=CLUTTER_A, b=CLUTTER_B):
    """ln p(theta) + sum_n ln p(x_n | theta) on a grid of theta values (D = 1), (10.209) to (10.211)."""
    theta = np.asarray(theta, dtype=float)
    lp = -0.5 * (LN2PI + math.log(b)) - 0.5 * theta ** 2 / b
    lnc = ln_clutter_terms(x, w, a)
    c1 = math.log(1.0 - w) - 0.5 * LN2PI
    for xn, cn in zip(np.asarray(x, dtype=float).ravel(), lnc):
        lp = lp + np.logaddexp(c1 - 0.5 * (xn - theta) ** 2, cn)
    return lp


def exact_clutter(x, w=CLUTTER_W, a=CLUTTER_A, b=CLUTTER_B, n=20001):
    """Exact posterior mean, variance and ln p(D) by quadrature (D = 1).

    The posterior is a mixture of 2^N Gaussians, but in one dimension it is easy to
    integrate numerically: a first grid over [-10 sqrt(b) - max|x|, +] locates the region
    where ln p(theta, D) is within 46 of its maximum (relative density 1e-20), and Simpson's
    rule on a second grid of n points over that region gives the moments and the evidence.
    Returns a dict with mean, var, ln_evidence and the grid and density.
    """
    x = np.asarray(x, dtype=float).ravel()
    lim = 10.0 * math.sqrt(b) + np.abs(x).max()
    lo, hi = -lim, lim
    for _ in range(2):
        th = np.linspace(lo, hi, n)
        lp = clutter_log_joint(th, x, w, a, b)
        keep = np.nonzero(lp > lp.max() - 46.0)[0]
        lo, hi = th[max(keep[0] - 1, 0)], th[min(keep[-1] + 1, n - 1)]
    h = th[1] - th[0]
    wts = np.ones(n)
    wts[1:-1:2], wts[2:-1:2] = 4.0, 2.0
    wts *= h / 3.0
    lnZ = common.logsumexp(lp + np.log(wts))
    p = np.exp(lp - lnZ)
    mean = np.sum(wts * p * th)
    var = np.sum(wts * p * (th - mean) ** 2)
    return {"mean": mean, "var": var, "ln_evidence": lnZ, "theta": th, "density": p}


# --------------------------------------------------------------------------- figures
def _xlabel_at(ax, x, text, y=-0.035):
    """Write an axis label at data position x just below the axis, as in the book's figures."""
    ax.text(x, y, text, transform=ax.get_xaxis_transform(), ha="center", va="top", fontsize=11)


def fig10_10(out_dir=None):
    """Figure 10.10: exp(-x) with tangent lines (left) and eta xi - g(eta) against eta, xi = 1 (right)."""
    fig, axes = common.panel_grid(1, 2, size=(3.3, 3.1))
    ax = axes[0, 0]
    x = np.linspace(0, 3, 400)
    for xi, col in [(0.3, common.GREEN), (2.0, common.GREEN), (1.0, common.BLUE)]:
        ax.plot(x, tangent_exp_neg(x, xi), color=col, lw=2)
    ax.plot(x, np.exp(-x), color=common.RED, lw=2.4)
    ax.plot([1, 1], [0, math.exp(-1)], color="k", lw=2)
    ax.set_xlim(0, 3)
    ax.set_ylim(0, 1)
    ax.set_xticks([0, 1, 1.5, 3])
    ax.set_xticklabels(["0", "ξ", "1.5", "3"])
    ax.set_yticks([0, 0.5, 1])
    ax.set_yticklabels(["0", "0.5", "1"])
    _xlabel_at(ax, 2.4, "$x$")
    ax = axes[0, 1]
    eta = np.linspace(-1, -1e-6, 600)
    ax.plot(eta, eta * 1.0 - conjugate_exp_neg(eta), color=common.RED, lw=2.4)
    ax.plot([-1 / math.e] * 2, [0, 1 / math.e], color="k", lw=2)
    ax.set_xlim(-1, 0)
    ax.set_ylim(0, 0.4)
    ax.set_xticks([-1, -0.5, 0])
    ax.set_xticklabels(["−1", "−0.5", "0"])
    ax.set_yticks([0, 0.2, 0.4])
    ax.set_yticklabels(["0", "0.2", "0.4"])
    _xlabel_at(ax, -0.14, r"$\eta$")
    ax.text(-0.95, 0.375, r"$\eta\xi - g(\eta)$", fontsize=12, va="top")
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-10", out_dir)


def _arrow_axes(ax, xmax, ymin, ymax):
    """Axes through the origin with arrow heads, as in the book's schematic Figure 10.11."""
    for s in ax.spines.values():
        s.set_visible(False)
    ax.set_xticks([])
    ax.set_yticks([])
    kw = dict(arrowstyle="-|>,head_width=0.35,head_length=0.7", lw=2.2, color="k")
    ax.annotate("", xy=(xmax, 0), xytext=(0, 0), arrowprops=kw)
    ax.annotate("", xy=(0, ymax), xytext=(0, ymin), arrowprops=kw)
    ax.text(xmax - 0.15, -0.18, "$x$", fontsize=16, ha="center", va="top")
    ax.text(-0.12, ymax - 0.25, "$y$", fontsize=16, ha="right", va="center")


def fig10_11(out_dir=None):
    """Figure 10.11: a convex f(x), a linear lower bound eta x and the dual g(eta) (tangent intercept)."""
    A, k, c, eta = 6.03, 1.358, -0.9, -0.43
    f = lambda x: A * np.exp(-k * x) + c
    xstar = -math.log(-eta / (A * k)) / k          # f'(x*) = eta
    g = eta * xstar - f(xstar)                     # g(eta) = max_x {eta x - f(x)} (10.129)
    x = np.linspace(0.33, 3.8, 300)
    fig, axes = common.panel_grid(1, 2, size=(3.6, 2.9))
    for j, ax in enumerate(axes[0]):
        _arrow_axes(ax, 4.25, -1.85, 3.3)
        ax.plot(x, f(x), color=common.RED, lw=2.6)
        ax.text(1.0, 2.0, "$f(x)$", fontsize=15)
        ax.set_xlim(-0.75, 4.35)
        ax.set_ylim(-1.9, 3.35)
    ax = axes[0, 0]
    xl = np.linspace(0, 3.8, 10)
    ax.plot(xl, eta * xl, color=common.BLUE, lw=2.6)
    ax.text(0.9, -1.0, r"$\eta x$", fontsize=15)
    for xv in (0.5, 1.3, xstar, 3.5):
        ax.plot([xv, xv], [eta * xv, f(xv)], ls=(0, (4, 3)), color=common.GREEN, lw=2.2)
    ax = axes[0, 1]
    ax.plot(xl, eta * xl - g, color=common.BLUE, lw=2.6)
    ax.text(-0.08, -g, r"$-g(\eta)$", fontsize=14, ha="right", va="center")
    ax.text(2.0, -1.45, r"$\eta x - g(\eta)$", fontsize=14)
    ax.plot([xstar], [f(xstar)], "o", color="k", ms=4)
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-11", out_dir)


def fig10_12(out_dir=None):
    """Figure 10.12: sigma(x) with exponential upper bounds (10.137) and the Gaussian-form lower bound (10.144)."""
    fig, axes = common.panel_grid(1, 2, size=(3.6, 2.7))
    x = np.linspace(-6, 6, 600)
    ax = axes[0, 0]
    ax.plot(x, common.sigmoid(x), color=common.RED, lw=2.2)
    for eta, pos in [(0.2, (1.4, 0.73)), (0.7, (-0.85, 0.24))]:
        ax.plot(x, sigmoid_upper_bound(x, eta), color=common.BLUE, lw=2.2)
        ax.text(pos[0], pos[1], r"$\eta = %g$" % eta, fontsize=12)
    ax = axes[0, 1]
    xi = 2.5
    ax.plot(x, common.sigmoid(x), color=common.RED, lw=2.2)
    ax.plot(x, sigmoid_lower_bound(x, xi), color=common.BLUE, lw=2.2)
    for s in (-xi, xi):
        ax.plot([s, s], [0, common.sigmoid(s)], ls=(0, (5, 4)), color=common.GREEN, lw=2)
    ax.text(-2.5, 0.78, r"$\xi = 2.5$", fontsize=12)
    for j, ax in enumerate(axes[0]):
        ax.set_xlim(-6, 6)
        ax.set_ylim(0, 1)
        ax.set_yticks([0, 0.5, 1])
        ax.set_yticklabels(["0", "0.5", "1"])
        if j == 0:
            ax.set_xticks([-6, 0, 6])
            ax.set_xticklabels(["−6", "0", "6"])
        else:
            ax.set_xticks([-6, -xi, 0, xi, 6])
            ax.set_xticklabels(["−6", "−ξ", "0", "ξ", "6"])
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-12", out_dir)


def fig10_13(out_dir=None):
    """Figure 10.13: variational logistic regression on separable 2-D data, predictive contours and sampled boundaries."""
    X, t = logistic_data()
    Phi = design_matrix(X)
    res = variational_logistic_regression(Phi, t, alpha=LOGIT_ALPHA)
    mN, SN = res["mN"], res["SN"]
    g1, g2 = np.meshgrid(np.linspace(-4, 4, 321), np.linspace(-6, 6, 481))
    P = predictive_probability(design_matrix(np.c_[g1.ravel(), g2.ravel()]), mN, SN).reshape(g1.shape)
    fig, axes = common.panel_grid(1, 2, size=(3.7, 2.75))
    levels = [0.01, 0.1, 0.25, 0.5, 0.75, 0.9, 0.99]
    jet = common.plt.get_cmap("jet")
    colour = {lv: jet((lv - 0.01) / 0.98) for lv in levels}
    ax = axes[0, 0]
    ax.contour(g1, g2, P, levels=levels, colors=[colour[lv] for lv in levels], linewidths=1.3)
    labelled = [0.01, 0.25, 0.75, 0.99]
    cs = ax.contour(g1, g2, P, levels=labelled, colors=[colour[lv] for lv in labelled], linewidths=1.3)
    ax.clabel(cs, fmt="%g", fontsize=9, colors="k", inline_spacing=3,
              manual=[(-0.45, 3.6), (-0.05, 3.6), (0.4, 3.6), (1.15, 3.6)])
    ax = axes[0, 1]
    r = DemoRNG(LOGIT_SAMPLE_SEED)
    L = np.linalg.cholesky(SN)
    yy = np.array([-6.0, 6.0])
    for col in ("#b000b0", common.BLUE, "#00c000", "#00b0b0", "#b8b800"):
        wk = mN + L @ np.array([r.norm() for _ in range(3)])
        ax.plot(-(wk[0] + wk[2] * yy) / wk[1], yy, color=col, lw=1.8)
    for ax in axes[0]:
        ax.plot(X[t == 0, 0], X[t == 0, 1], "x", color=common.RED, ms=6, mew=1.8)
        ax.plot(X[t == 1, 0], X[t == 1, 1], "o", mfc="none", mec=common.RED, ms=6.5, mew=1.6)
        ax.set_xlim(-4, 4)
        ax.set_ylim(-6, 6)
        ax.set_xticks([-4, -2, 0, 2, 4])
        ax.set_yticks([-6, -4, -2, 0, 2, 4, 6])
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-13", out_dir)
    return res


def fig10_14(out_dir=None):
    """Figure 10.14: p(z) proportional to exp(-z^2/2) sigma(20z + 4) with Laplace, variational and EP Gaussians."""
    z, p, dz = _toy_grid()
    approx = [("Laplace", laplace_toy(), common.RED), ("variational", kl_qp_gaussian_toy(), common.GREEN),
              ("EP", ep_toy(), common.BLUE)]
    zs = np.linspace(-2, 4, 1201)
    pz = np.exp(toy_log_density(zs)) / (np.exp(toy_log_density(z)).sum() * dz)
    fig, axes = common.panel_grid(1, 2, size=(3.6, 2.7))
    ax = axes[0, 0]
    ax.fill_between(zs, 0, pz, color="#e6cc5c", lw=0)
    for name, (mu, var), col in approx:
        ax.plot(zs, np.exp(-0.5 * (zs - mu) ** 2 / var) / math.sqrt(2 * math.pi * var), color=col, lw=2)
    ax.set_ylim(0, 1)
    ax.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1])
    ax.set_yticklabels(["0", "0.2", "0.4", "0.6", "0.8", "1"])
    ax = axes[0, 1]
    for name, (mu, var), col in approx:
        ax.plot(zs, 0.5 * (zs - mu) ** 2 / var + 0.5 * math.log(2 * math.pi * var), color=col, lw=2)
    ax.plot(zs, -np.log(pz), color="#e6cc5c", lw=2.4)
    ax.set_ylim(0, 40)
    ax.set_yticks([0, 10, 20, 30, 40])
    for ax in axes[0]:
        ax.set_xlim(-2, 4)
        ax.set_xticks([-2, -1, 0, 1, 2, 3, 4])
        ax.set_xticklabels(["−2", "−1", "0", "1", "2", "3", "4"])
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-14", out_dir)
    return approx


def fig10_15(out_dir=None):
    """Figure 10.15: the clutter problem in D = 1, data (crosses) and the two mixture components."""
    theta, x = 3.0, clutter_data(20, 3.0)
    xs = np.linspace(-5, 10, 600)
    fig, axes = common.panel_grid(1, 1, size=(4.2, 3.0))
    ax = axes[0, 0]
    ax.plot(xs, np.exp(-0.5 * (xs - theta) ** 2) / math.sqrt(2 * math.pi), color=common.GREEN, lw=2.2)
    ax.plot(xs, np.exp(-0.5 * xs ** 2 / CLUTTER_A) / math.sqrt(2 * math.pi * CLUTTER_A), color=common.RED, lw=2.2)
    ax.plot(x, np.full(len(x), 0.04), "x", color="k", ms=5, mew=1.4)
    ax.set_xlim(-5, 10)
    ax.set_ylim(0, 0.405)
    ax.set_yticks([])
    ax.set_xticks([-5, 0, theta, 5, 10])
    ax.set_xticklabels(["−5", "0", "θ", "5", "10"])
    _xlabel_at(ax, 8.5, "$x$")
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-15", out_dir)
    return x


def site_example(x_n, m_cav, v_cav, w=CLUTTER_W, a=CLUTTER_A):
    """The EP site f~_n produced from a given cavity N(theta | m_cav, v_cav) and data point x_n (D = 1).

    Equivalent to ep_site_update with q = cavity times the old site; returns the new
    site parameters (tau_n, m_n, ln s_n) and (m_new, v_new, ln Z_n).
    """
    lnc = ln_clutter_terms(np.array([x_n]), w, a)[0]
    # with tau_n = 0 the cavity equals q, so ep_site_update starts from exactly this cavity
    m_new, v_new, tau_n, m_n, ln_s, ln_Z = ep_site_update(np.array([x_n]), lnc, np.array([m_cav]), v_cav,
                                                         0.0, np.zeros(1), w)
    return tau_n, m_n[0], ln_s, m_new[0], v_new, ln_Z


def fig10_16(out_dir=None):
    """Figure 10.16: f_n (blue), the site approximation f~_n (red) and the cavity q^\\n (green), two examples."""
    th = np.linspace(-5, 10, 1201)
    x_n = 2.0
    f_n = (1 - CLUTTER_W) * np.exp(-0.5 * (x_n - th) ** 2) / math.sqrt(2 * math.pi) \
        + CLUTTER_W * math.exp(-0.5 * x_n ** 2 / CLUTTER_A) / math.sqrt(2 * math.pi * CLUTTER_A)
    fig, axes = common.panel_grid(1, 2, size=(3.6, 2.6))
    for ax, m_cav in zip(axes[0], (0.0, 3.0)):
        v_cav = 0.09
        tau_n, m_n, ln_s, m_new, v_new, ln_Z = site_example(x_n, m_cav, v_cav)
        site = np.exp(ln_s - 0.5 * tau_n * (th - m_n) ** 2)
        cav = np.exp(-0.5 * (th - m_cav) ** 2 / v_cav) / math.sqrt(2 * math.pi * v_cav)
        ax.plot(th, cav, color=common.GREEN, lw=1.8)
        ax.plot(th, f_n, color="#2a2a9f", lw=1.8)
        ax.plot(th, site, color=common.RED, lw=1.8)
        ax.set_xlim(-5, 10)
        ax.set_ylim(0, 0.3)
        ax.set_yticks([])
        ax.set_xticks([-5, 0, 5, 10])
        ax.set_xticklabels(["−5", "0", "5", "10"])
        _xlabel_at(ax, 7.5, r"$\theta$")
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-16", out_dir)


def clutter_comparison(N=200, theta=3.0, seed=CLUTTER_SEED):
    """Run EP, VB and Laplace on one clutter data set and the exact answer by quadrature (Figure 10.17).

    Returns the exact values and, per method, arrays of (flops, |mean error|, |ln p(D) error|).
    """
    x = clutter_data(N, theta, seed=seed)
    ex = exact_clutter(x)
    ep = ep_clutter(x, max_sweeps=8, tol=0.0, record_every=max(1, N // 5))
    vb = vb_clutter(x, max_iter=60, tol=0.0)
    lap = laplace_clutter(x, max_iter=10, tol=0.0)
    curves = {}
    for name, res in (("ep", ep), ("vb", vb), ("laplace", lap)):
        h = np.array([(f, abs(m - ex["mean"]), abs(lz - ex["ln_evidence"])) for f, m, v, lz in res["history"]])
        curves[name] = h
    return x, ex, {"ep": ep, "vb": vb, "laplace": lap}, curves


def fig10_17(out_dir=None):
    """Figure 10.17: EP, VB and Laplace on the clutter problem: error against flops (explicit counts)."""
    x, ex, res, curves = clutter_comparison()
    cols = {"laplace": common.BLUE, "vb": common.RED, "ep": common.GREEN}
    fig, axes = common.panel_grid(1, 2, size=(3.7, 3.0))
    for j, (ax, title) in enumerate(zip(axes[0], ("Posterior mean", "Evidence"))):
        for name in ("laplace", "vb", "ep"):
            h = curves[name]
            ok = np.isfinite(h[:, 1 + j]) & (h[:, 1 + j] > 0)
            ax.loglog(h[ok, 0], h[ok, 1 + j], color=cols[name], lw=2)
            k = np.nonzero(ok)[0][-1]
            ax.text(h[k, 0] * 1.15, h[k, 1 + j], name, color=cols[name], fontsize=11, va="center")
        ax.set_title(title)
        ax.set_xlabel("flops")
        ax.set_ylabel("error in the mean" if j == 0 else r"error in $\ln p(\mathcal{D})$")
        ax.set_xlim(2e3, 4e5)
    fig.tight_layout()
    common.savefig(fig, 10, "fig10-17", out_dir)
    return ex, res


FIGURES = {
    "10.10": fig10_10,
    "10.11": fig10_11,
    "10.12": fig10_12,
    "10.13": fig10_13,
    "10.14": fig10_14,
    "10.15": fig10_15,
    "10.16": fig10_16,
    "10.17": fig10_17,
}

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