"""Bishop (2006), Chapter 2 (Probability Distributions): the algorithms and every numerical figure.

One function per figure (`python3 ch02_probability_distributions.py --list`); the building blocks
(conjugate updates, Gaussian conditionals, Student t, von Mises, density estimators) are small
documented functions that cite the book's equations and are also used by checks_ch02.py.

Choices and what could not be reproduced exactly (the book gives few numbers for this chapter):
* Figure 2.1: N = 10, mu = 0.25 (stated in the text). Figures 2.2 and 2.13: the parameters printed in
  the panels. Figure 2.3: Beta(2, 2) prior and one observation x = 1 (stated in the caption).
* Figure 2.5: Dirichlet with alpha_k = 0.1, 1, 10 (caption); the surface is drawn over the simplex
  (mu_1, mu_2) with mu_3 = 1 - mu_1 - mu_2 and a common height axis cut at 20.
* Figure 2.6: 1000 draws of the mean of N uniform numbers, 20 bins.
* Figure 2.9: a bivariate Gaussian with mean (0.5, 0.5), standard deviations 0.14 and correlation 0.95
  (read off the marginal and conditional peaks 2.8 and 9 in the book's right-hand panel).
* Figure 2.12: data from N(0.8, 0.1) (variance 0.1, caption), prior N(0, 0.1), likelihood variance
  0.1; the first 1, 2, 10 points of one seeded sample.
* Figure 2.14: normal-gamma with mu_0 = 0, beta = 2, a = 5, b = 6 (caption).
* Figure 2.15: Student t with mu = 0, lambda = 1 and nu = 0.1, 1, infinity (caption).
* Figure 2.16: 30 draws from N(0, 1) (a), then three outliers near 9 (b); the t fit is maximum
  likelihood by EM (mu, lambda) with nu from a one-dimensional maximisation of the EM bound.
* Figure 2.19: m = 5, theta_0 = pi/4 and m = 1, theta_0 = 3 pi/4 (legend of the book's figure).
* Figure 2.21: Old Faithful (272 points, common.old_faithful); a single Gaussian by maximum
  likelihood and a two-component mixture by EM started from the two ends of the data.
* Figure 2.22: three one-dimensional Gaussians with weights 0.3, 0.3, 0.4 (shape read off the book;
  the values are ours). Figure 2.23: pi = (0.5, 0.3, 0.2), means (0.2, 0.4), (0.5, 0.5), (0.8, 0.6),
  covariance 0.016 [[1, r], [r, 1]] with r = 0.75, -0.75, 0.75 (read off the contours, as in Chapter 9).
* Figures 2.24 to 2.26: 50 points from 0.35 N(0.3, 0.12^2) + 0.65 N(0.8, 0.1^2) (the book's green curve
  is not tabulated; this mixture has the same two peaks, heights about 1.2 and 2.6).
* Figure 2.27: a small synthetic two-class set (the book's points are a sketch).
* Figure 2.28: the first 200 points of the oil-flow training set (the book's subset is not stated),
  inputs x6 and x7, K = 1, 3, 31, ties broken at random (seeded), which gives the speckled ties of
  the book's K = 3 panel.
* Figures 2.4, 2.7, 2.10, 2.17 and 2.18 are sketches and are not reproduced; Figure 2.11 is replaced by a
  real Robbins-Monro run (an addition, not in the book in this form).
"""
import math

import numpy as np
import matplotlib.pyplot as plt
from scipy import optimize, special

import common
from ch11_sampling import JSRandom  # noqa: F401  (re-exported: the demos' generator)

RED, BLUE, GREEN, MAGENTA = common.RED, common.BLUE, common.GREEN, common.MAGENTA
LOG2PI = math.log(2 * math.pi)


# =========================================================================== Section 2.1: binary variables
def binom_pmf(m, N, mu):
    """Binomial distribution Bin(m | N, mu), eq. (2.9)."""
    m = np.asarray(m)
    return np.exp(special.gammaln(N + 1) - special.gammaln(m + 1) - special.gammaln(N - m + 1)
                  + special.xlogy(m, mu) + special.xlogy(N - m, 1 - mu))


def beta_pdf(mu, a, b):
    """Beta(mu | a, b) = Gamma(a + b) / (Gamma(a) Gamma(b)) mu^(a-1) (1-mu)^(b-1), eq. (2.13)."""
    mu = np.asarray(mu, dtype=float)
    return np.exp(special.gammaln(a + b) - special.gammaln(a) - special.gammaln(b)
                  + special.xlogy(a - 1, mu) + special.xlog1py(b - 1, -mu))


def beta_update(a, b, x):
    """Sequential update of a beta prior by Bernoulli data x (an array of 0 and 1), eq. (2.18):
    a gains the number of ones m, b the number of zeros l = N - m. Returns the list of (a, b) after each point."""
    out = []
    for xi in np.atleast_1d(x):
        a, b = a + xi, b + (1 - xi)
        out.append((a, b))
    return out


def beta_moments(a, b):
    """Mean a/(a+b) (2.15), variance ab/((a+b)^2 (a+b+1)) (2.16), and the mode (a-1)/(a+b-2)."""
    return a / (a + b), a * b / ((a + b) ** 2 * (a + b + 1)), (a - 1) / (a + b - 2) if a + b > 2 else float("nan")


def beta_predictive(a, b):
    """Predictive probability p(x = 1 | D) = a/(a+b), eq. (2.20): the posterior mean of mu."""
    return a / (a + b)


# =========================================================================== Section 2.2: multinomial variables
def dirichlet_logpdf(mu, alpha):
    """ln Dir(mu | alpha), eq. (2.38); mu has shape (..., K) with rows on the simplex."""
    mu, alpha = np.asarray(mu, dtype=float), np.asarray(alpha, dtype=float)
    return (special.gammaln(alpha.sum()) - special.gammaln(alpha).sum() + (special.xlogy(alpha - 1, mu)).sum(-1))


def dirichlet_update(alpha, counts):
    """Posterior Dir(mu | alpha + m), eq. (2.41): add the observed counts m_k."""
    return np.asarray(alpha, dtype=float) + np.asarray(counts)


def dirichlet_moments(alpha):
    """Mean alpha_k / alpha_0, variance and covariance, eqs. (2.273) to (2.276)."""
    a0 = alpha.sum()
    mean = alpha / a0
    cov = (np.diag(alpha) * a0 - np.outer(alpha, alpha)) / (a0 ** 2 * (a0 + 1))   # off-diagonal: -a_j a_k/(a0^2 (a0+1))
    return mean, cov


def multinomial_ml(counts):
    """Maximum likelihood mu_k = m_k / N for the multinomial, eq. (2.33), by the Lagrange multiplier lambda = -N."""
    counts = np.asarray(counts, dtype=float)
    return counts / counts.sum()


# =========================================================================== Section 2.3: the Gaussian
def gauss_pdf(x, mu, var):
    """One-dimensional N(x | mu, var), eq. (1.46)."""
    x = np.asarray(x, dtype=float)
    return np.exp(-0.5 * (x - mu) ** 2 / var) / np.sqrt(2 * np.pi * var)


def gauss_conditional(mu, Sigma, a, b, xb):
    """Conditional p(x_a | x_b) of a joint Gaussian, in covariance form,
    mean mu_a + Sigma_ab Sigma_bb^{-1} (x_b - mu_b) and covariance Sigma_aa - Sigma_ab Sigma_bb^{-1} Sigma_ba,
    eqs. (2.81), (2.82). (The precision form (2.73), (2.75) gives the same numbers.)"""
    mu, Sigma = np.asarray(mu, dtype=float), np.asarray(Sigma, dtype=float)
    Sab, Sbb = Sigma[np.ix_(a, b)], Sigma[np.ix_(b, b)]
    G = Sab @ np.linalg.inv(Sbb)
    return mu[a] + G @ (np.asarray(xb) - mu[b]), Sigma[np.ix_(a, a)] - G @ Sab.T


def gauss_conditional_precision(mu, Sigma, a, b, xb):
    """The same conditional from the precision matrix Lambda = Sigma^{-1}:
    covariance Lambda_aa^{-1} (2.73) and mean mu_a - Lambda_aa^{-1} Lambda_ab (x_b - mu_b) (2.75)."""
    mu = np.asarray(mu, dtype=float)
    L = np.linalg.inv(Sigma)
    Laa_inv = np.linalg.inv(L[np.ix_(a, a)])
    return mu[a] - Laa_inv @ L[np.ix_(a, b)] @ (np.asarray(xb) - mu[b]), Laa_inv


def gauss_marginal(mu, Sigma, a):
    """Marginal p(x_a) = N(x_a | mu_a, Sigma_aa), eqs. (2.92), (2.93)."""
    return np.asarray(mu, dtype=float)[a], np.asarray(Sigma, dtype=float)[np.ix_(a, a)]


def gauss_mean_posterior(x, sigma2, mu0, s02):
    """Posterior over the mean mu of N(x | mu, sigma2) with prior N(mu | mu0, s02):
    mu_N = (sigma2 mu0 + N s02 xbar)/(N s02 + sigma2) (2.141) and 1/s_N^2 = 1/s02 + N/sigma2 (2.142)."""
    x = np.asarray(x, dtype=float)
    N = len(x)
    xbar = x.mean() if N else 0.0
    precN = 1 / s02 + N / sigma2
    muN = (mu0 / s02 + N * xbar / sigma2) / precN
    return muN, 1 / precN


def robbins_monro_mean(x, a_seq=None):
    """Sequential maximum likelihood estimate of a Gaussian mean, eq. (2.126):
    mu_N = mu_{N-1} + (1/N)(x_N - mu_{N-1}); the Robbins-Monro step (2.135) with a_N = sigma^2/N.
    Returns the whole trajectory (mu_1, ..., mu_N). a_seq(N) may replace the 1/N step size."""
    mu, out = 0.0, []
    for n, xn in enumerate(np.asarray(x, dtype=float), start=1):
        step = 1.0 / n if a_seq is None else a_seq(n)
        mu = mu + step * (xn - mu)
        out.append(mu)
    return np.array(out)


def gamma_pdf(lam, a, b):
    """Gam(lambda | a, b) = b^a lambda^(a-1) exp(-b lambda) / Gamma(a), eq. (2.146)."""
    lam = np.asarray(lam, dtype=float)
    return np.exp(a * np.log(b) - special.gammaln(a) + special.xlogy(a - 1, lam) - b * lam)


def gauss_precision_posterior(x, mu, a0, b0):
    """Gamma posterior over the precision lambda of N(x | mu, lambda^-1) with prior Gam(a0, b0):
    a_N = a0 + N/2 (2.150), b_N = b0 + (1/2) sum (x_n - mu)^2 (2.151)."""
    x = np.asarray(x, dtype=float)
    return a0 + len(x) / 2, b0 + 0.5 * ((x - mu) ** 2).sum()


def normal_gamma_posterior(x, mu0, beta0, a0, b0):
    """Posterior of (mu, lambda) under the normal-gamma prior N(mu | mu0, (beta0 lambda)^-1) Gam(lambda | a0, b0), eq. (2.154):
    beta_N = beta0 + N, mu_N = (beta0 mu0 + N xbar)/beta_N, a_N = a0 + N/2,
    b_N = b0 + (1/2) sum (x_n - xbar)^2 + beta0 N (xbar - mu0)^2 / (2 (beta0 + N)) (Exercise 2.44)."""
    x = np.asarray(x, dtype=float)
    N, xbar = len(x), x.mean()
    betaN = beta0 + N
    muN = (beta0 * mu0 + N * xbar) / betaN
    aN = a0 + N / 2
    bN = b0 + 0.5 * ((x - xbar) ** 2).sum() + beta0 * N * (xbar - mu0) ** 2 / (2 * betaN)
    return muN, betaN, aN, bN


def normal_gamma_pdf(mu, lam, mu0, beta, a, b):
    """N(mu | mu0, (beta lambda)^-1) Gam(lambda | a, b), eq. (2.154)."""
    return gauss_pdf(mu, mu0, 1.0 / (beta * lam)) * gamma_pdf(lam, a, b)


def student_pdf(x, mu=0.0, lam=1.0, nu=1.0):
    """St(x | mu, lambda, nu) = Gamma(nu/2+1/2)/Gamma(nu/2) (lambda/(pi nu))^(1/2) [1 + lambda (x-mu)^2/nu]^(-nu/2-1/2), eq. (2.159)."""
    x = np.asarray(x, dtype=float)
    return np.exp(special.gammaln(nu / 2 + 0.5) - special.gammaln(nu / 2) + 0.5 * np.log(lam / (np.pi * nu))
                  - (nu / 2 + 0.5) * np.log1p(lam * (x - mu) ** 2 / nu))


def student_mixture_pdf(x, mu, lam, nu, n=4000):
    """St(x | mu, lambda, nu) as the infinite mixture of Gaussians int N(x | mu, (eta lambda)^-1) Gam(eta | nu/2, nu/2) d eta,
    eq. (2.158), by the midpoint rule on a grid of quantiles of the gamma distribution."""
    from scipy.special import gammaincinv
    u = (np.arange(n) + 0.5) / n
    eta = gammaincinv(nu / 2, u) / (nu / 2)             # quantiles of Gam(nu/2, rate nu/2)
    x = np.asarray(x, dtype=float)
    return np.array([np.mean(gauss_pdf(xi, mu, 1.0 / (eta * lam))) for xi in np.atleast_1d(x)])


def student_t_mle(x, iters=300, nu0=5.0):
    """Maximum likelihood fit of mu, lambda, nu of a one-dimensional Student t by EM (scale-mixture view, (2.158)).
    E step: E[eta_n] = (nu+1)/(nu + lambda (x_n - mu)^2). M step: mu = sum E[eta_n] x_n / sum E[eta_n],
    1/lambda = (1/N) sum E[eta_n] (x_n - mu)^2, and nu maximises N[(nu/2) ln(nu/2) - ln Gamma(nu/2)]
    + (nu/2 - 1) sum E[ln eta_n] - (nu/2) sum E[eta_n]. Returns (mu, lambda, nu, log-likelihood)."""
    x = np.asarray(x, dtype=float)
    N = len(x)
    mu, lam, nu = x.mean(), 1 / x.var(), nu0
    ll_old = -np.inf
    for _ in range(iters):
        d2 = lam * (x - mu) ** 2
        Eeta = (nu + 1) / (nu + d2)
        Elog = special.digamma((nu + 1) / 2) - np.log((nu + d2) / 2)
        mu = (Eeta * x).sum() / Eeta.sum()
        lam = N / (Eeta * (x - mu) ** 2).sum()
        negQ = lambda v: -(N * ((v / 2) * np.log(v / 2) - special.gammaln(v / 2)) + (v / 2 - 1) * Elog.sum() - (v / 2) * Eeta.sum())
        nu = optimize.minimize_scalar(negQ, bounds=(0.05, 500.0), method="bounded").x
        ll = np.log(student_pdf(x, mu, lam, nu)).sum()
        if abs(ll - ll_old) < 1e-12:
            break
        ll_old = ll
    return mu, lam, nu, ll


def von_mises_pdf(theta, theta0, m):
    """p(theta | theta0, m) = exp(m cos(theta - theta0)) / (2 pi I0(m)), eq. (2.179)."""
    return np.exp(m * np.cos(np.asarray(theta) - theta0)) / (2 * np.pi * special.i0(m))


def bessel_ratio_A(m):
    """A(m) = I1(m) / I0(m), eq. (2.186)."""
    return special.ive(1, m) / special.ive(0, m)   # exponentially scaled: stable for large m


def von_mises_mle(theta):
    """Maximum likelihood theta0 = atan2(sum sin, sum cos) (2.185) and m from A(m) = rbar (2.186),
    where rbar = |mean of the unit vectors| = (1/N) sum cos(theta_n - theta0)."""
    theta = np.asarray(theta, dtype=float)
    C, S = np.cos(theta).mean(), np.sin(theta).mean()
    theta0, rbar = math.atan2(S, C), math.hypot(C, S)
    m = optimize.brentq(lambda mm: bessel_ratio_A(mm) - rbar, 1e-9, 1e4) if rbar < 1 else float("inf")
    return theta0, m


# --------------------------------------------------------------------------- Gaussian mixtures (Section 2.3.9)
def gmm_logpdf(X, pis, mus, Sigmas):
    """ln of sum_k pi_k N(x | mu_k, Sigma_k), eq. (2.188), for the rows of X."""
    comp = np.stack([np.log(p) + common.gauss_logpdf(X, m, S) for p, m, S in zip(pis, mus, Sigmas)], 1)
    return common.logsumexp(comp, axis=1)


def gmm_em(X, mus, Sigmas=None, iters=200, tol=1e-9):
    """Maximum likelihood for a Gaussian mixture by EM (the algorithm of Chapter 9), started from the means mus.
    Used for the right-hand panel of Figure 2.21. Returns (pis, mus, Sigmas, log-likelihood)."""
    X = np.asarray(X, dtype=float)
    N, D = X.shape
    K = len(mus)
    mus = np.array(mus, dtype=float)
    Sigmas = np.array([np.cov(X.T, bias=True) for _ in range(K)]) if Sigmas is None else np.array(Sigmas, dtype=float)
    pis = np.ones(K) / K
    ll_old = -np.inf
    for _ in range(iters):
        comp = np.stack([np.log(pis[k]) + common.gauss_logpdf(X, mus[k], Sigmas[k]) for k in range(K)], 1)
        lse = common.logsumexp(comp, axis=1, keepdims=True)
        R = np.exp(comp - lse)
        Nk = R.sum(0)
        pis, mus = Nk / N, (R.T @ X) / Nk[:, None]
        for k in range(K):
            d = X - mus[k]
            Sigmas[k] = (R[:, k, None] * d).T @ d / Nk[k] + 1e-8 * np.eye(D)
        ll = lse.sum()
        if abs(ll - ll_old) < tol:
            break
        ll_old = ll
    return pis, mus, Sigmas, ll


# =========================================================================== Section 2.4: exponential family
def exp_family_gaussian_natural(mu, sigma2):
    """Natural parameters eta = (mu/sigma^2, -1/(2 sigma^2)) of the Gaussian in the form (2.194),
    with u(x) = (x, x^2), g(eta) = (-2 eta2)^(1/2) exp(eta1^2/(4 eta2)) (2.221) and h(x) = (2 pi)^(-1/2)."""
    return np.array([mu / sigma2, -1 / (2 * sigma2)])


def exp_family_log_g(eta):
    """ln g(eta) for the Gaussian: 0.5 ln(-2 eta2) + eta1^2 / (4 eta2); then -grad ln g = E[u(x)] (2.226)."""
    return 0.5 * np.log(-2 * eta[1]) + eta[0] ** 2 / (4 * eta[1])


# =========================================================================== Section 2.5: nonparametric methods
def histogram_density(x, grid, delta, lo=0.0):
    """Histogram estimate p_i = n_i / (N Delta), eq. (2.241), on bins [lo + i Delta, lo + (i+1) Delta); evaluated at grid."""
    x = np.asarray(x, dtype=float)
    idx = np.floor((x - lo) / delta).astype(int)
    nb = int(np.ceil((grid.max() - lo) / delta - 1e-12))
    counts = np.bincount(np.clip(idx, 0, max(nb - 1, 0)), minlength=nb)
    gi = np.clip(np.floor((grid - lo) / delta).astype(int), 0, nb - 1)
    return counts[gi] / (len(x) * delta)


def kde_gaussian(x, grid, h):
    """Gaussian kernel density estimate p(x) = (1/N) sum_n (2 pi h^2)^(-D/2) exp(-|x - x_n|^2 / (2 h^2)), eq. (2.250)."""
    x, grid = np.asarray(x, dtype=float), np.asarray(grid, dtype=float)
    if x.ndim == 1:
        return gauss_pdf(grid[:, None], x[None, :], h * h).mean(1)
    D = x.shape[1]
    d2 = ((grid[:, None, :] - x[None, :, :]) ** 2).sum(-1)
    return (np.exp(-d2 / (2 * h * h)) / (2 * np.pi * h * h) ** (D / 2)).mean(1)


def knn_density(x, grid, K):
    """K-nearest-neighbour density p(x) = K / (N V) with V the length of the smallest interval around x that
    contains K data points, V = 2 r_K, eq. (2.246) in one dimension."""
    x = np.sort(np.asarray(x, dtype=float))
    d = np.sort(np.abs(grid[:, None] - x[None, :]), axis=1)[:, K - 1]
    return K / (len(x) * 2 * d)


def knn_classify(Xtr, ytr, Xte, K, rng=None):
    """K-nearest-neighbour classifier: the class with most votes among the K nearest training points,
    p(C_k | x) = K_k / K (2.256); ties are broken at random when rng is given, otherwise at the smallest label."""
    Xtr, Xte = np.asarray(Xtr, dtype=float), np.asarray(Xte, dtype=float)
    ncls = int(ytr.max()) + 1
    d2 = ((Xte[:, None, :] - Xtr[None, :, :]) ** 2).sum(-1)
    nn = np.argpartition(d2, K - 1, axis=1)[:, :K] if K < len(Xtr) else np.tile(np.arange(len(Xtr)), (len(Xte), 1))
    votes = np.stack([(ytr[nn] == c).sum(1) for c in range(ncls)], 1).astype(float)
    if rng is not None:
        votes = votes + 1e-3 * rng.random(votes.shape)       # random tie-breaking
    return votes.argmax(1), votes


# =========================================================================== data generators
def fig224_sample(n, rng):
    """50 points of Figures 2.24 to 2.26: 0.35 N(0.3, 0.12^2) + 0.65 N(0.8, 0.1^2), clipped to [0, 1] by rejection."""
    out = []
    while len(out) < n:
        k = rng.random() < 0.35
        v = rng.normal(0.3, 0.12) if k else rng.normal(0.8, 0.1)
        if 0 <= v <= 1:
            out.append(v)
    return np.array(out)


def fig224_density(x):
    """The generating density (green curve) of Figures 2.24 to 2.26 (up to the truncation to [0, 1])."""
    return 0.35 * gauss_pdf(x, 0.3, 0.12 ** 2) + 0.65 * gauss_pdf(x, 0.8, 0.1 ** 2)


FIG223 = {"pi": np.array([0.5, 0.3, 0.2]),
          "mu": np.array([[0.2, 0.4], [0.5, 0.5], [0.8, 0.6]]),
          "Sigma": 0.016 * np.array([[[1, .75], [.75, 1]], [[1, -.75], [-.75, 1]], [[1, .75], [.75, 1]]])}


def _box(ax, ylim=None, xlim=None):
    if ylim is not None:
        ax.set_ylim(*ylim)
    if xlim is not None:
        ax.set_xlim(*xlim)


# =========================================================================== figures
def fig2_1(out_dir=None):
    """Figure 2.1: histogram of the binomial distribution, N = 10, mu = 0.25."""
    fig, ax = common.panel_grid(1, 1, size=(3.6, 2.6))
    ax = ax[0, 0]
    m = np.arange(11)
    ax.bar(m, binom_pmf(m, 10, 0.25), width=0.8, color=BLUE, edgecolor="k", lw=0.6)
    ax.set_xlabel("$m$"); ax.set_xticks(m); ax.set_xlim(-0.6, 10.6)
    common.savefig(fig, 2, "fig2-1", out_dir)


def fig2_2(out_dir=None):
    """Figure 2.2: beta distributions Beta(mu | a, b) for (a, b) = (0.1, 0.1), (1, 1), (2, 3), (8, 4)."""
    mu = np.linspace(0.0005, 0.9995, 2000)
    fig, axes = common.panel_grid(1, 4, size=(2.5, 2.3))
    for ax, (a, b), lab in zip(axes[0], [(0.1, 0.1), (1, 1), (2, 3), (8, 4)], "abcd"):
        ax.plot(mu, beta_pdf(mu, a, b), color=RED)
        ax.set_ylim(0, 3); ax.set_xlim(0, 1); ax.set_xticks([0, 0.5, 1]); ax.set_xlabel(r"$\mu$")
        ax.text(0.35, 0.86, "$a = %g$\n$b = %g$" % (a, b), transform=ax.transAxes, va="top")
        common.panel_label(ax, lab, "lower right" if (a, b) == (0.1, 0.1) else "upper right")
    fig.subplots_adjust(wspace=0.25)
    common.savefig(fig, 2, "fig2-2", out_dir)


def fig2_3(out_dir=None):
    """Figure 2.3: one step of sequential Bayesian inference: Beta(2, 2) prior, one observation x = 1, Beta(3, 2) posterior."""
    mu = np.linspace(0, 1, 500)
    (a, b), = beta_update(2, 2, [1])
    fig, axes = common.panel_grid(1, 3, size=(2.7, 2.3))
    data = [("prior: Beta(2, 2)", beta_pdf(mu, 2, 2), RED, (0, 2)),
            ("likelihood function $\\mu^1$", mu, BLUE, (0, 2)),
            ("posterior: Beta(%d, %d)" % (a, b), beta_pdf(mu, a, b), RED, (0, 2))]
    for ax, (t, y, c, yl), lab in zip(axes[0], data, "abc"):
        ax.plot(mu, y, color=c, lw=2); ax.set_xlim(0, 1); ax.set_ylim(*yl); ax.set_xlabel(r"$\mu$")
        ax.set_title(t, fontsize=9.5); common.panel_label(ax, lab, "upper left")
    fig.subplots_adjust(wspace=0.3)
    common.savefig(fig, 2, "fig2-3", out_dir)


def _simplex_grid(n=140):
    """Points (mu1, mu2) of the simplex mu1 + mu2 <= 1 on a regular grid, NaN outside."""
    g = np.linspace(0, 1, n)
    M1, M2 = np.meshgrid(g, g)
    M3 = 1 - M1 - M2
    inside = M3 >= 0
    return M1, M2, np.where(inside, M3, np.nan), inside


def fig2_5(out_dir=None):
    """Figure 2.5: the Dirichlet density on the simplex for alpha_k = 0.1, 1, 10 (surface over the plane of the simplex)."""
    M1, M2, M3, inside = _simplex_grid(160)
    fig = plt.figure(figsize=(8.0, 2.8))
    top = 0
    zs = []
    for al in (0.1, 1.0, 10.0):
        mm = np.clip(np.stack([M1, M2, np.nan_to_num(M3)], -1), 1e-3, 1)
        z = np.exp(dirichlet_logpdf(mm, np.full(3, al)))
        z[~inside] = np.nan
        zs.append(z)
    top = 20.0
    for i, (al, z) in enumerate(zip((0.1, 1.0, 10.0), zs)):
        ax = fig.add_subplot(1, 3, i + 1, projection="3d")
        ax.plot_surface(M1, M2, np.minimum(z, top), cmap="cool", rstride=2, cstride=2, linewidth=0, antialiased=True, rasterized=True)
        ax.set_zlim(0, top); ax.view_init(elev=32, azim=-62)
        ax.set_xlabel(r"$\mu_1$", labelpad=-6); ax.set_ylabel(r"$\mu_2$", labelpad=-6)
        ax.set_xticks([0, 0.5, 1]); ax.set_yticks([0, 0.5, 1]); ax.set_zticks([0, 10, 20]); ax.tick_params(labelsize=7, pad=-3)
        ax.set_title(r"$\alpha_k = %g$" % al, fontsize=10, y=0.97)
        ax.text2D(0.04, 0.9, "(%s)" % "abc"[i], transform=ax.transAxes)
    fig.subplots_adjust(left=0.0, right=1.0, wspace=0.0, top=1.0, bottom=0.0)
    common.savefig(fig, 2, "fig2-5", out_dir)


def fig2_6(out_dir=None):
    """Figure 2.6: histograms of the mean of N uniform numbers, N = 1, 2, 10 (central limit theorem)."""
    rng = np.random.default_rng(6)
    fig, axes = common.panel_grid(1, 3, size=(2.6, 2.3))
    for ax, N, lab in zip(axes[0], (1, 2, 10), "abc"):
        means = rng.random((1000, N)).mean(1)
        ax.hist(means, bins=np.linspace(0, 1, 21), density=True, color=BLUE, edgecolor="k", lw=0.5)
        ax.set_xlim(0, 1); ax.set_ylim(0, 3.2); ax.set_xticks([0, 0.5, 1]); ax.text(0.05, 0.85, "$N = %d$" % N, transform=ax.transAxes, fontsize=11)
        common.panel_label(ax, lab, "upper right")
    fig.subplots_adjust(wspace=0.25)
    common.savefig(fig, 2, "fig2-6", out_dir)


def fig2_8(out_dir=None):
    """Figure 2.8: contours of constant density of a two-dimensional Gaussian with (a) general, (b) diagonal, (c) isotropic covariance."""
    covs = [np.array([[1.0, 0.6], [0.6, 1.0]]) * 0.8, np.array([[1.6, 0.0], [0.0, 0.55]]), np.eye(2) * 0.9]
    fig, axes = common.panel_grid(1, 3, size=(2.5, 2.4))
    t = np.linspace(0, 2 * np.pi, 200)
    for ax, S, lab in zip(axes[0], covs, "abc"):
        for r in (0.6, 1.2, 1.8, 2.4):
            x, y = common.ellipse_xy(np.zeros(2), S, r)
            ax.plot(x, y, color=RED, lw=1.6)
        ax.set_xlim(-3, 3); ax.set_ylim(-3, 3); ax.set_aspect("equal"); ax.set_xticks([]); ax.set_yticks([])
        ax.set_xlabel("$x_1$"); ax.set_ylabel("$x_2$"); common.panel_label(ax, lab, "upper left")
    common.savefig(fig, 2, "fig2-8", out_dir)


def fig2_9(out_dir=None):
    """Figure 2.9: contours of a bivariate Gaussian with the line x_b = 0.7 (left); marginal p(x_a) and conditional p(x_a | x_b = 0.7) (right)."""
    mu = np.array([0.5, 0.5])
    s, rho = 0.14, 0.95
    Sigma = s * s * np.array([[1, rho], [rho, 1]])
    cm, cS = gauss_conditional(mu, Sigma, [0], [1], [0.7])
    mm, mS = gauss_marginal(mu, Sigma, [0])
    fig, axes = common.panel_grid(1, 2, size=(3.3, 3.0))
    ax = axes[0, 0]
    g = np.linspace(0, 1, 200)
    XA, XB = np.meshgrid(g, g)
    P = np.exp(common.gauss_logpdf(np.stack([XA.ravel(), XB.ravel()], 1), mu, Sigma)).reshape(XA.shape)
    ax.contour(XA, XB, P, levels=np.array([0.2, 0.5, 0.85]) * P.max(), colors=GREEN, linewidths=1.2)
    ax.axhline(0.7, color=RED, lw=1.6)
    ax.set_xlabel("$x_a$"); ax.set_ylabel("$x_b$"); ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.set_xticks([0, 0.5, 1]); ax.set_yticks([0, 0.5, 1])
    ax.text(0.2, 0.74, "$x_b = 0.7$", fontsize=9); ax.text(0.55, 0.22, "$p(x_a, x_b)$", fontsize=9, color="k")
    common.panel_label(ax, "a", "upper left")
    ax = axes[0, 1]
    ax.plot(g, gauss_pdf(g, mm[0], mS[0, 0]), color=BLUE); ax.plot(g, gauss_pdf(g, cm[0], cS[0, 0]), color=RED)
    ax.text(0.12, 3.3, "$p(x_a)$", color=BLUE, fontsize=9); ax.text(0.5, 9.4, "$p(x_a \\mid x_b = 0.7)$", color=RED, fontsize=9)
    ax.set_xlim(0, 1); ax.set_ylim(0, 10.5); ax.set_xlabel("$x_a$"); ax.set_xticks([0, 0.5, 1]); common.panel_label(ax, "b", "upper left")
    fig.subplots_adjust(wspace=0.3)
    print("   conditional mean %.4f, sd %.4f (peak %.2f); marginal sd %.3f (peak %.2f)"
          % (cm[0], math.sqrt(cS[0, 0]), 1 / math.sqrt(2 * np.pi * cS[0, 0]), math.sqrt(mS[0, 0]), 1 / math.sqrt(2 * np.pi * mS[0, 0])))
    common.savefig(fig, 2, "fig2-9", out_dir)


def fig2_11(out_dir=None):
    """Figure 2.11 (computed variant): the Robbins-Monro procedure (2.129), (2.135) finds the root of a regression function
    from noisy values: (a) the regression function f(theta) = E[z | theta] with samples of z; (b) the iterates theta_N for
    f(theta) = theta_star - theta... here the sequential Gaussian mean (2.126): theta_N approaches the sample mean."""
    rng = np.random.default_rng(11)
    mu_true = 0.8
    x = rng.normal(mu_true, math.sqrt(0.1), 400)
    traj = robbins_monro_mean(x)
    fig, axes = common.panel_grid(1, 2, size=(3.4, 2.5))
    ax = axes[0, 0]
    th = np.linspace(-0.2, 1.8, 200)
    ax.plot(th, mu_true - th, color=RED, lw=1.8, label=r"$f(\theta)=\mu-\theta$")
    zs = mu_true - th[::6] + rng.normal(0, math.sqrt(0.1), len(th[::6]))
    ax.plot(th[::6], zs, "o", color=BLUE, ms=3, label=r"$z$ sampled at $\theta$")
    ax.axhline(0, color="k", lw=0.6); ax.axvline(mu_true, color=GREEN, ls="--", lw=1)
    ax.set_xlabel(r"$\theta$"); ax.set_ylabel("$z$"); ax.legend(frameon=False, loc="upper right", fontsize=8); common.panel_label(ax, "a", "lower left")
    ax = axes[0, 1]
    ax.plot(np.arange(1, len(traj) + 1), traj, color=RED)
    ax.axhline(mu_true, color=GREEN, ls="--", lw=1); ax.axhline(x.mean(), color="k", lw=0.6)
    ax.set_xscale("log"); ax.set_xlabel("$N$"); ax.set_ylabel(r"$\mu^{(N)}_{\mathrm{ML}}$"); common.panel_label(ax, "b", "lower right")
    fig.subplots_adjust(wspace=0.35)
    common.savefig(fig, 2, "fig2-11", out_dir)


def fig2_12(out_dir=None):
    """Figure 2.12: Bayesian inference for the mean of a Gaussian with known variance: prior (N = 0) and posteriors for N = 1, 2, 10."""
    rng = np.random.default_rng(12)
    x = rng.normal(0.8, math.sqrt(0.1), 10)
    mu = np.linspace(-1, 1, 800)
    fig, ax = common.panel_grid(1, 1, size=(4.2, 3.0))
    ax = ax[0, 0]
    for N, c in ((0, "k"), (1, GREEN), (2, BLUE), (10, RED)):
        mN, vN = gauss_mean_posterior(x[:N], 0.1, 0.0, 0.1)
        ax.plot(mu, gauss_pdf(mu, mN, vN), color=c)
        peak = gauss_pdf(mN, mN, vN)
        ax.text(mN + (0.02 if N else -0.35), peak + 0.12 if N else peak + 0.1, "$N = %d$" % N, color=c, fontsize=9)
    ax.set_xlim(-1, 1); ax.set_ylim(0, 5); ax.set_xlabel(r"$\mu$")
    common.savefig(fig, 2, "fig2-12", out_dir)
    return x


def fig2_13(out_dir=None):
    """Figure 2.13: gamma distributions Gam(lambda | a, b) for (a, b) = (0.1, 0.1), (1, 1), (4, 6)."""
    lam = np.linspace(0.001, 2, 1000)
    fig, axes = common.panel_grid(1, 3, size=(2.7, 2.3))
    for ax, (a, b), lab in zip(axes[0], [(0.1, 0.1), (1, 1), (4, 6)], "abc"):
        ax.plot(lam, gamma_pdf(lam, a, b), color=RED, lw=2)
        ax.set_xlim(0, 2); ax.set_ylim(0, 2); ax.set_xticks([0, 1, 2]); ax.set_xlabel(r"$\lambda$")
        ax.text(0.55, 0.85, "$a = %g$\n$b = %g$" % (a, b), transform=ax.transAxes, va="top"); common.panel_label(ax, lab, "upper left")
    fig.subplots_adjust(wspace=0.25)
    common.savefig(fig, 2, "fig2-13", out_dir)


def fig2_14(out_dir=None):
    """Figure 2.14: contours of the normal-gamma distribution (2.154) with mu_0 = 0, beta = 2, a = 5, b = 6."""
    mu, lam = np.meshgrid(np.linspace(-2, 2, 400), np.linspace(0.005, 2, 400))
    P = normal_gamma_pdf(mu, lam, 0.0, 2.0, 5.0, 6.0)
    fig, ax = common.panel_grid(1, 1, size=(4.0, 2.6))
    ax = ax[0, 0]
    ax.contour(mu, lam, P, levels=P.max() * np.array([0.05, 0.15, 0.3, 0.5, 0.7, 0.85, 0.97]), colors=RED, linewidths=1.0)
    ax.set_xlabel(r"$\mu$"); ax.set_ylabel(r"$\lambda$"); ax.set_xlim(-2, 2); ax.set_ylim(0, 2); ax.set_yticks([0, 1, 2])
    common.savefig(fig, 2, "fig2-14", out_dir)


def fig2_15(out_dir=None):
    """Figure 2.15: Student's t-distribution (2.159) with mu = 0, lambda = 1 for nu = 0.1, 1, and the Gaussian limit nu -> infinity."""
    x = np.linspace(-5, 5, 1000)
    fig, ax = common.panel_grid(1, 1, size=(4.0, 2.7))
    ax = ax[0, 0]
    ax.plot(x, gauss_pdf(x, 0, 1), color=GREEN, label=r"$\nu \to \infty$")
    ax.plot(x, student_pdf(x, 0, 1, 1.0), color=BLUE, label=r"$\nu = 1.0$")
    ax.plot(x, student_pdf(x, 0, 1, 0.1), color=RED, label=r"$\nu = 0.1$")
    ax.set_xlim(-5, 5); ax.set_ylim(0, 0.5); ax.legend(frameon=False, loc="upper right")
    common.savefig(fig, 2, "fig2-15", out_dir)


def fig2_16(out_dir=None):
    """Figure 2.16: robustness of the t-distribution: (a) 30 Gaussian points, (b) the same plus three outliers; ML fits of a t (red) and a Gaussian (green)."""
    rng = np.random.default_rng(16)
    x = rng.normal(0, 1, 30)
    xo = np.concatenate([x, [8.6, 9.2, 9.8]])
    g = np.linspace(-5, 10, 1000)
    fig, axes = common.panel_grid(1, 2, size=(3.5, 2.6))
    out = []
    for ax, data, lab in zip(axes[0], (x, xo), "ab"):
        mu, lam, nu, _ = student_t_mle(data)
        ax.hist(data, bins=np.arange(-2.5, 10.6, 0.5), density=True, color="#9ea7e0", edgecolor="#4a4a8a", lw=0.5)
        ax.plot(g, gauss_pdf(g, data.mean(), data.var()), color=GREEN, lw=1.4)
        ax.plot(g, student_pdf(g, mu, lam, nu), color=RED, lw=1.4)
        ax.set_xlim(-5, 10); ax.set_ylim(0, 0.5); common.panel_label(ax, lab, "upper right")
        out.append((mu, lam, nu, data.mean(), data.std()))
    fig.subplots_adjust(wspace=0.2)
    for tag, (mu, lam, nu, gm, gs) in zip("ab", out):
        print("   (%s) t fit: mu %.3f, sd-scale %.3f, nu %.2f; Gaussian fit: mean %.3f, sd %.3f" % (tag, mu, 1 / math.sqrt(lam), nu, gm, gs))
    common.savefig(fig, 2, "fig2-16", out_dir)


def fig2_19(out_dir=None):
    """Figure 2.19: the von Mises distribution for (m, theta_0) = (5, pi/4) and (1, 3 pi/4): Cartesian and polar plots."""
    th = np.linspace(0, 2 * np.pi, 600)
    pars = [(5, np.pi / 4, RED, r"$m = 5,\ \theta_0 = \pi/4$"), (1, 3 * np.pi / 4, BLUE, r"$m = 1,\ \theta_0 = 3\pi/4$")]
    fig = plt.figure(figsize=(6.6, 2.8))
    ax = fig.add_subplot(1, 2, 1)
    for m, t0, c, lab in pars:
        ax.plot(th, von_mises_pdf(th, t0, m), color=c, label=lab)
    ax.set_xlim(0, 2 * np.pi); ax.set_xticks([0, np.pi / 2, np.pi, 3 * np.pi / 2, 2 * np.pi]); ax.set_xticklabels(["0", r"$\pi/2$", r"$\pi$", r"$3\pi/2$", r"$2\pi$"])
    ax.set_xlabel(r"$\theta$"); ax.legend(frameon=False, fontsize=8, loc="upper right"); common.panel_label(ax, "a", "upper left")
    ax = fig.add_subplot(1, 2, 2, projection="polar")
    for m, t0, c, lab in pars:
        ax.plot(th, von_mises_pdf(th, t0, m), color=c)
    ax.set_yticks([0.5, 1.0]); ax.tick_params(labelsize=7)
    ax.text(0.0, 1.0, "(b)", transform=ax.transAxes, va="top")
    fig.subplots_adjust(wspace=0.3)
    common.savefig(fig, 2, "fig2-19", out_dir)


def fig2_20(out_dir=None):
    """Figure 2.20: the Bessel function I_0(m) (2.180) and A(m) = I_1(m)/I_0(m) (2.186)."""
    m = np.linspace(0, 10, 400)
    fig, axes = common.panel_grid(1, 2, size=(3.3, 2.5))
    axes[0, 0].plot(m, special.i0(m), color=RED); axes[0, 0].set_xlabel("$m$"); axes[0, 0].set_ylabel("$I_0(m)$"); axes[0, 0].set_xlim(0, 10); axes[0, 0].set_ylim(0, 3000)
    axes[0, 1].plot(m, bessel_ratio_A(m), color=RED); axes[0, 1].set_xlabel("$m$"); axes[0, 1].set_ylabel("$A(m)$"); axes[0, 1].set_xlim(0, 10); axes[0, 1].set_ylim(0, 1)
    for ax, lab in zip(axes[0], "ab"):
        common.panel_label(ax, lab, "upper left" if lab == "a" else "lower right")
    fig.subplots_adjust(wspace=0.4)
    print("   I0(10) = %.1f, A(10) = %.4f" % (special.i0(10), bessel_ratio_A(10.0)))
    common.savefig(fig, 2, "fig2-20", out_dir)


def fig2_21(out_dir=None):
    """Figure 2.21: Old Faithful with (a) a single Gaussian fitted by maximum likelihood and (b) a mixture of two Gaussians fitted by EM."""
    X = common.old_faithful()
    mu1, S1 = X.mean(0), np.cov(X.T, bias=True)
    pis, mus, Ss, ll2 = gmm_em(X, [X[X[:, 0].argmin()], X[X[:, 0].argmax()]])
    ll1 = common.gauss_logpdf(X, mu1, S1).sum()
    gx, gy = np.meshgrid(np.linspace(0.8, 6, 200), np.linspace(36, 100, 200))
    G = np.stack([gx.ravel(), gy.ravel()], 1)
    fig, axes = common.panel_grid(1, 2, size=(3.4, 3.0))
    dens = [np.exp(common.gauss_logpdf(G, mu1, S1)), np.exp(gmm_logpdf(G, pis, mus, Ss))]
    for ax, P, lab in zip(axes[0], dens, "ab"):
        ax.plot(X[:, 0], X[:, 1], "o", mfc="none", mec=GREEN, ms=3.5, mew=1.0)
        ax.contour(gx, gy, P.reshape(gx.shape), levels=P.max() * np.array([0.1, 0.3, 0.6, 0.9]), colors=BLUE, linewidths=1.1)
        ax.set_xlim(1, 6); ax.set_ylim(40, 100); ax.set_xlabel("eruption duration (min)"); common.panel_label(ax, lab, "lower right")
    axes[0, 0].set_ylabel("waiting time (min)")
    fig.subplots_adjust(wspace=0.25)
    print("   ln L: single Gaussian %.2f, two-component mixture %.2f; mixture pi = %s, means %s"
          % (ll1, ll2, np.round(pis, 3), np.round(mus, 2).tolist()))
    common.savefig(fig, 2, "fig2-21", out_dir)


def fig2_22(out_dir=None):
    """Figure 2.22: a one-dimensional mixture of three Gaussians (components scaled by their weights in blue, the sum in red)."""
    x = np.linspace(0, 1, 800)
    pis, mus, sds = [0.35, 0.25, 0.4], [0.2, 0.55, 0.82], [0.045, 0.1, 0.045]
    comps = [p * gauss_pdf(x, m, s * s) for p, m, s in zip(pis, mus, sds)]
    fig, ax = common.panel_grid(1, 1, size=(4.2, 2.5))
    ax = ax[0, 0]
    for c in comps:
        ax.plot(x, c, color=BLUE)
    ax.plot(x, sum(comps), color=RED, lw=2.2)
    ax.set_xlim(0, 1); ax.set_ylim(0, None); ax.set_xticks([]); ax.set_yticks([]); ax.set_xlabel("$x$"); ax.set_ylabel("$p(x)$")
    common.savefig(fig, 2, "fig2-22", out_dir)


def fig2_23(out_dir=None):
    """Figure 2.23: a mixture of three Gaussians in two dimensions: (a) components, (b) contours of p(x), (c) surface plot."""
    p = FIG223
    g = np.linspace(-0.2, 1.2, 160)
    GX, GY = np.meshgrid(g, g)
    G = np.stack([GX.ravel(), GY.ravel()], 1)
    P = np.exp(gmm_logpdf(G, p["pi"], p["mu"], p["Sigma"])).reshape(GX.shape)
    fig = plt.figure(figsize=(8.4, 2.9))
    ax = fig.add_subplot(1, 3, 1)
    for k, c in enumerate((RED, GREEN, BLUE)):
        for r in (1.0, 2.0):
            x, y = common.ellipse_xy(p["mu"][k], p["Sigma"][k], r)
            ax.plot(x, y, color=c, lw=1.8)
        ax.text(p["mu"][k][0] + (0.08 if k == 2 else 0.0), p["mu"][k][1] - 0.27 if k != 2 else p["mu"][k][1] - 0.22, "%.1f" % p["pi"][k], color=c)
    ax.set_xlim(-0.2, 1.2); ax.set_ylim(-0.2, 1.2); ax.set_aspect("equal"); ax.set_xticks([0, 0.5, 1]); ax.set_yticks([0, 0.5, 1]); common.panel_label(ax, "a", "upper left")
    ax = fig.add_subplot(1, 3, 2)
    ax.contour(GX, GY, P, levels=P.max() * np.array([0.08, 0.2, 0.35, 0.5, 0.7, 0.9]), colors=MAGENTA, linewidths=1.3)
    ax.set_xlim(-0.2, 1.2); ax.set_ylim(-0.2, 1.2); ax.set_aspect("equal"); ax.set_xticks([0, 0.5, 1]); ax.set_yticks([0, 0.5, 1]); common.panel_label(ax, "b", "upper left")
    ax = fig.add_subplot(1, 3, 3, projection="3d")
    ax.plot_surface(GX, GY, P, cmap="copper", rstride=2, cstride=2, linewidth=0, antialiased=True, rasterized=True)
    ax.view_init(elev=38, azim=-125); ax.set_axis_off(); ax.text2D(0.05, 0.9, "(c)", transform=ax.transAxes)
    fig.subplots_adjust(left=0.04, right=0.99, wspace=0.12, top=0.97, bottom=0.04)
    common.savefig(fig, 2, "fig2-23", out_dir)


def fig2_24(out_dir=None):
    """Figure 2.24: histogram density estimates of 50 points for bin widths 0.04, 0.08, 0.25 (generating density in green)."""
    x = fig224_sample(50, np.random.default_rng(24))
    g = np.linspace(0, 1, 1000)
    fig, axes = common.panel_grid(3, 1, size=(3.8, 1.45), sharex=True)
    for ax, d in zip(axes[:, 0], (0.04, 0.08, 0.25)):
        h = histogram_density(x, g, d)
        ax.fill_between(g, h, step=None, color="#8f8fe8", alpha=0.0)
        edges = np.arange(0, 1 + d / 2, d)
        ax.hist(x, bins=edges, density=True, color="#8f8fe8", edgecolor="k", lw=0.5)
        ax.plot(g, fig224_density(g), color=GREEN, lw=1.5)
        ax.set_xlim(0, 1); ax.set_ylim(0, 5); ax.set_yticks([0, 5]); ax.text(0.04, 0.78, r"$\Delta = %g$" % d, transform=ax.transAxes)
    axes[-1, 0].set_xticks([0, 0.5, 1])
    fig.subplots_adjust(hspace=0.1)
    common.savefig(fig, 2, "fig2-24", out_dir)


def fig2_25(out_dir=None):
    """Figure 2.25: Gaussian kernel density estimates (2.250) of the same 50 points for h = 0.005, 0.07, 0.2."""
    x = fig224_sample(50, np.random.default_rng(24))
    g = np.linspace(0, 1, 1000)
    fig, axes = common.panel_grid(3, 1, size=(3.8, 1.45), sharex=True)
    for ax, h in zip(axes[:, 0], (0.005, 0.07, 0.2)):
        ax.plot(g, fig224_density(g), color=GREEN, lw=1.5)
        ax.plot(g, kde_gaussian(x, g, h), color=BLUE, lw=1.4)
        ax.set_xlim(0, 1); ax.set_ylim(0, 5); ax.set_yticks([0, 5]); ax.text(0.04, 0.78, "$h = %g$" % h, transform=ax.transAxes)
    axes[-1, 0].set_xticks([0, 0.5, 1])
    fig.subplots_adjust(hspace=0.1)
    common.savefig(fig, 2, "fig2-25", out_dir)


def fig2_26(out_dir=None):
    """Figure 2.26: K-nearest-neighbour density estimates (2.246) of the same 50 points for K = 1, 5, 30."""
    x = fig224_sample(50, np.random.default_rng(24))
    g = np.linspace(0.0005, 0.9995, 2000)
    fig, axes = common.panel_grid(3, 1, size=(3.8, 1.45), sharex=True)
    for ax, K in zip(axes[:, 0], (1, 5, 30)):
        ax.plot(g, fig224_density(g), color=GREEN, lw=1.5)
        ax.plot(g, knn_density(x, g, K), color=BLUE, lw=1.4)
        ax.set_xlim(0, 1); ax.set_ylim(0, 5); ax.set_yticks([0, 5]); ax.text(0.04, 0.78, "$K = %d$" % K, transform=ax.transAxes)
    axes[-1, 0].set_xticks([0, 0.5, 1])
    fig.subplots_adjust(hspace=0.1)
    common.savefig(fig, 2, "fig2-26", out_dir)


def fig227_data(seed=27):
    """The small two-class data set drawn in Figure 2.27 (synthetic): 12 red and 12 blue points in the unit square."""
    rng = np.random.default_rng(seed)
    red = np.column_stack([rng.normal(0.32, 0.14, 12), rng.normal(0.65, 0.16, 12)])
    blue = np.column_stack([rng.normal(0.68, 0.14, 12), rng.normal(0.4, 0.16, 12)])
    X = np.vstack([red, blue])
    return X, np.r_[np.zeros(12, int), np.ones(12, int)]


def fig2_27(out_dir=None):
    """Figure 2.27: (a) the K = 3 neighbours of a test point (black diamond), (b) the K = 1 decision boundary."""
    X, y = fig227_data()
    q = np.array([0.5, 0.52])
    d = np.linalg.norm(X - q, axis=1)
    nn = np.argsort(d)[:3]
    cols = np.array([RED, BLUE])
    fig, axes = common.panel_grid(1, 2, size=(3.3, 3.2))
    for ax in axes[0]:
        ax.scatter(X[:, 0], X[:, 1], c=cols[y], s=14, zorder=3)
        ax.set_xlim(-0.05, 1.05); ax.set_ylim(-0.05, 1.05); ax.set_xticks([]); ax.set_yticks([]); ax.set_xlabel("$x_1$"); ax.set_ylabel("$x_2$")
    ax = axes[0, 0]
    for i in nn:
        ax.plot([q[0], X[i, 0]], [q[1], X[i, 1]], color=GREEN, lw=1.3, zorder=2)
    ax.plot(*q, "D", color="k", mfc="white", ms=6, zorder=4); common.panel_label(ax, "a", "lower right")
    ax = axes[0, 1]
    g = np.linspace(-0.05, 1.05, 400)
    GX, GY = np.meshgrid(g, g)
    lab, _ = knn_classify(X, y, np.stack([GX.ravel(), GY.ravel()], 1), 1)
    ax.contour(GX, GY, lab.reshape(GX.shape).astype(float), levels=[0.5], colors=GREEN, linewidths=1.5)
    common.panel_label(ax, "b", "lower right")
    fig.subplots_adjust(wspace=0.15)
    common.savefig(fig, 2, "fig2-27", out_dir)


def fig2_28(out_dir=None):
    """Figure 2.28: K-nearest-neighbour classification of the oil-flow data (x6 against x7, 200 points), K = 1, 3, 31."""
    X, lab = common.oil_flow("trn")
    X, lab = X[:200, 5:7], lab[:200]
    colour = {0: RED, 1: GREEN, 2: BLUE}    # colours chosen to match the book: the top cluster (label 2 in our file order) is blue
    shade = {0: "#ffc9c9", 1: "#c9ffc9", 2: "#c9c9ff"}
    g = np.linspace(0, 2, 160)
    GX, GY = np.meshgrid(g, g)
    G = np.stack([GX.ravel(), GY.ravel()], 1)
    fig, axes = common.panel_grid(1, 3, size=(2.7, 2.7))
    for ax, K in zip(axes[0], (1, 3, 31)):
        pred, _ = knn_classify(X, lab, G, K, rng=np.random.default_rng(28))
        img = np.array([[int(shade[c][i:i + 2], 16) for i in (1, 3, 5)] for c in range(3)], dtype=np.uint8)[pred].reshape(GX.shape + (3,))
        ax.imshow(img, extent=(0, 2, 0, 2), origin="lower", aspect="auto", interpolation="nearest", rasterized=True)
        for c in range(3):
            ax.plot(X[lab == c, 0], X[lab == c, 1], "o", color=colour[c], ms=3)
        ax.set_xlim(0, 2); ax.set_ylim(0, 2); ax.set_title("$K = %d$" % K); ax.set_xlabel("$x_6$"); ax.set_xticks([0, 1, 2]); ax.set_yticks([0, 1, 2])
    axes[0, 0].set_ylabel("$x_7$")
    fig.subplots_adjust(wspace=0.25)
    common.savefig(fig, 2, "fig2-28", out_dir)


FIGURES = {"2.1": fig2_1, "2.2": fig2_2, "2.3": fig2_3, "2.5": fig2_5, "2.6": fig2_6, "2.8": fig2_8, "2.9": fig2_9,
           "2.11": fig2_11, "2.12": fig2_12, "2.13": fig2_13, "2.14": fig2_14, "2.15": fig2_15, "2.16": fig2_16,
           "2.19": fig2_19, "2.20": fig2_20, "2.21": fig2_21, "2.22": fig2_22, "2.23": fig2_23, "2.24": fig2_24,
           "2.25": fig2_25, "2.26": fig2_26, "2.27": fig2_27, "2.28": fig2_28}

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