"""Bishop (2006), Chapter 12: Continuous latent variables.

One function per figure; run `python3 ch12_continuous_latent.py --list` to see them,
`python3 ch12_continuous_latent.py --fig 12.12` to make one, or no option to make all.

Algorithms (small readable functions, each citing the book's equations):
    pca, pca_small_n, pca_reconstruct, whiten       Section 12.1, (12.1)-(12.30)
    fisher_direction                                 Section 4.1.4, (4.30), for Figure 12.7
    ppca_ml, ppca_loglik, ppca_posterior             Section 12.2.1, (12.36)-(12.48)
    ppca_em, ppca_em_missing                         Section 12.2.2, (12.54)-(12.57)
    em_pca                                           (12.58), (12.59): the sigma^2 -> 0 limit
    bayesian_pca                                     Section 12.2.3, ARD (12.62), (12.63)
    bayesian_pca_gibbs                               Gibbs sampler of Figure 12.15
    factor_analysis_em, fa_loglik                    Section 12.2.4, (12.65)-(12.70)
    gaussian_kernel, kernel_pca, kernel_pca_project  Section 12.3, (12.80)-(12.87)
    gtm_fit                                          Section 12.4.3; Bishop, Svensen and Williams (1998)

What could not be reproduced exactly, and the choices made:
* Figures 12.1, 12.3 to 12.5 use the MNIST test-set threes (1010 images, pixel values 0 to 255).
  The book's eigenvalues 3.4, 2.8, 2.4, 1.6 x 10^5 are reproduced to two figures, and the digit of
  Figures 12.1 and 12.5 is MNIST test image 44 (correlation 0.99999 with the book's picture).
  For Figure 12.1 the 28 x 28 digit is enlarged to 64 x 64 (bilinear) as in the text (p. 559) and
  placed in a 100 x 100 field with a random rotation (+-25 degrees) and shift (+-16 pixels).
* Figure 12.2 and Figure 12.9 are drawn from small made-up configurations, but every line, point
  and contour is computed (the principal axis, the projections, the Gaussian contours).
* Figure 12.6: the right-hand panel whitens the original data (12.24) (whitening the standardised
  data gives the same picture up to a rotation); as in the book's figure, the horizontal axis is the
  whitened coordinate along the minor principal axis and the vertical axis along the major one.
* Figure 12.7: the two-class data are simulated (the book's are not available).
* Figures 12.8, 12.11, 12.21: the oil flow training set (1000 points). Eigenvector signs are chosen
  to match the book's orientation. In 12.11 the PPCA model (M = 2) is fitted to the first 100
  points; for the right panel 30% of the 1200 values are deleted at random (seeded, every point
  loses at least one value) and the exact EM for data missing at random is used.
* Figure 12.12: the ten data points are read off the book's figure (then centred); the initial W
  points along (2, -1) as in panel (b).
* Figure 12.14: 300 points in D = 10, standard deviation 1 along 3 random orthonormal directions
  and 0.5 along the other 7; M = 9. Maximum likelihood PPCA and Bayesian PCA are both run by EM
  from the same random start, so the ML solution comes with an arbitrary rotation R (12.45).
* Figure 12.15: the book gives no settings. We use D = 4, M = 3, N = 50 points from a PPCA model
  with one direction of variance 0.7 plus isotropic noise variance 0.1, vague Gamma(0.001, 0.001)
  priors on the alpha_i and on 1/sigma^2, and mu fixed at the sample mean; 10,000 Gibbs sweeps.
  The barrier between the three modes grows with N: with N = 100 we saw no switch at all.
* Figure 12.17: the toy data of Schoelkopf et al. (1998): three clusters of 30 points (sd 0.1)
  centred at (-0.5, -0.2), (0, 0.6), (0.5, 0), kernel exp(-|x - x'|^2/0.1) (12.86). The titles
  give the eigenvalues of the centred Gram matrix (12.85), which equal N lambda_i by (12.80).
  The points are our own draw, so the eigenvalues differ slightly from the book's.
* Figure 12.21: GTM with a 20 x 20 latent grid on [-1, 1]^2, an 8 x 8 grid of Gaussian basis
  functions plus a bias, width equal to twice the spacing of the basis centres, weight
  regularisation alpha = 0.01, initialised from PCA and run for 200 EM iterations (the book does
  not give its settings; these gave the highest likelihood among the grids we tried).
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.patches import ConnectionPatch, FancyArrowPatch, Rectangle
from scipy import ndimage

import common

BLUE_INK = LinearSegmentedColormap.from_list("white_blue", ["#ffffff", "#1a1aa6"])
# Figure 12.3: blue positive, white zero, yellow negative (book's caption, as corrected in the errata)
YELLOW_BLUE = LinearSegmentedColormap.from_list("yellow_white_blue", ["#c8c800", "#ffffff", "#1a1aa6"])
KPCA_GREEN = LinearSegmentedColormap.from_list("kpca_green", ["#000000", "#00684c", "#00b383", "#3dffc8"])
HINTON_BG = "#8099ff"
OIL_STYLE = {2: dict(marker="x", color=common.RED, label="laminar"),
             0: dict(marker="+", color=common.BLUE, label="homogeneous"),
             1: dict(marker="o", color=common.GREEN, label="annular")}


# =========================================================================== Section 12.1: PCA
def pca(X):
    """Principal component analysis of the rows of X (N x D).

    Returns the sample mean (12.1), the eigenvalues lambda_1 >= ... >= lambda_D of the
    covariance matrix S (12.3) and the orthonormal eigenvectors u_i as the columns of U,
    S u_i = lambda_i u_i (12.5), (12.17).
    """
    xbar = X.mean(0)
    Xc = X - xbar
    S = Xc.T @ Xc / len(X)
    lam, U = np.linalg.eigh(S)
    order = np.argsort(lam)[::-1]
    return xbar, lam[order], U[:, order]


def pca_small_n(X, M):
    """PCA when there are fewer points than dimensions (Section 12.1.4).

    Solves the N x N eigenproblem (1/N) X X^T v_i = lambda_i v_i (12.28) for the centred data
    matrix X and maps back with u_i = X^T v_i / (N lambda_i)^{1/2} (12.30), which has unit length.
    """
    xbar = X.mean(0)
    Xc = X - xbar
    N = len(X)
    lam, V = np.linalg.eigh(Xc @ Xc.T / N)
    order = np.argsort(lam)[::-1][:M]
    lam, V = lam[order], V[:, order]
    U = Xc.T @ V / np.sqrt(N * lam)
    return xbar, lam, U


def pca_reconstruct(X, xbar, U, M):
    """PCA approximation x~_n = xbar + sum_{i<=M} ((x_n - xbar)^T u_i) u_i, eq. (12.20)."""
    Z = (X - xbar) @ U[:, :M]
    return xbar + Z @ U[:, :M].T


def distortion(X, xbar, U, M):
    """Mean squared reconstruction error J (12.11) of the M-dimensional PCA approximation."""
    R = X - pca_reconstruct(X, xbar, U, M)
    return (R ** 2).sum(1).mean()


def whiten(X):
    """Whitening (12.24): y_n = L^{-1/2} U^T (x_n - xbar), so that the y_n have zero mean and
    identity covariance (12.25). Column i of the result is the coordinate along u_i."""
    xbar, lam, U = pca(X)
    return (X - xbar) @ U / np.sqrt(lam)


def fisher_direction(X, t):
    """Fisher's linear discriminant for two classes t in {0, 1}: w proportional to
    S_W^{-1} (m_1 - m_0) (Bishop 2006, eq. 4.30), returned with unit length."""
    m0, m1 = X[t == 0].mean(0), X[t == 1].mean(0)
    SW = (X[t == 0] - m0).T @ (X[t == 0] - m0) + (X[t == 1] - m1).T @ (X[t == 1] - m1)
    w = np.linalg.solve(SW, m1 - m0)
    return w / np.linalg.norm(w)


# =========================================================================== Section 12.2: PPCA
def ppca_ml(X, M):
    """Closed-form maximum likelihood PPCA (Tipping and Bishop 1999).

    mu = xbar, sigma^2_ML = average of the D - M discarded eigenvalues (12.46) and
    W_ML = U_M (L_M - sigma^2 I)^{1/2} R (12.45) with the rotation R = I.
    """
    xbar, lam, U = pca(X)
    sigma2 = lam[M:].mean()
    W = U[:, :M] * np.sqrt(np.maximum(lam[:M] - sigma2, 0.0))
    return xbar, W, sigma2


def ppca_loglik(X, mu, W, sigma2):
    """Log likelihood (12.43) of PPCA with covariance C = W W^T + sigma^2 I (12.36).

    Only M x M matrices are inverted: C^{-1} = sigma^{-2} (I - W M^{-1} W^T) by the Woodbury
    identity (12.40), with M = W^T W + sigma^2 I (12.41), and ln|C| = (D - M) ln sigma^2 + ln|M|.
    """
    N, D = X.shape
    q = W.shape[1]
    Mm = W.T @ W + sigma2 * np.eye(q)
    Xc = X - mu
    A = Xc @ W
    quad = ((Xc ** 2).sum() - np.sum(A * np.linalg.solve(Mm, A.T).T)) / sigma2
    logdetC = (D - q) * np.log(sigma2) + np.linalg.slogdet(Mm)[1]
    return -0.5 * (N * D * np.log(2 * np.pi) + N * logdetC + quad)


def ppca_posterior(X, mu, W, sigma2):
    """Posterior p(z_n | x_n) = N(z | M^{-1} W^T (x_n - mu), sigma^2 M^{-1}) (12.42).

    Returns the posterior means as rows (12.48) and the common covariance sigma^2 M^{-1}.
    """
    Mm = W.T @ W + sigma2 * np.eye(W.shape[1])
    Ez = np.linalg.solve(Mm, W.T @ (X - mu).T).T
    return Ez, sigma2 * np.linalg.inv(Mm)


def ppca_em_step(Xc, W, sigma2):
    """One EM step for PPCA on centred data Xc (mu = xbar already substituted).

    E step (12.54), (12.55): E[z_n] = M^{-1} W^T (x_n - xbar) and
    sum_n E[z_n z_n^T] = N sigma^2 M^{-1} + sum_n E[z_n] E[z_n]^T.
    M step (12.56), (12.57) for W and sigma^2.
    """
    N, D = Xc.shape
    q = W.shape[1]
    Minv = np.linalg.inv(W.T @ W + sigma2 * np.eye(q))
    Ez = Xc @ W @ Minv                                      # rows are E[z_n]^T
    sumEzz = N * sigma2 * Minv + Ez.T @ Ez
    W_new = (Xc.T @ Ez) @ np.linalg.inv(sumEzz)              # (12.56)
    sigma2_new = ((Xc ** 2).sum() - 2 * np.sum(Ez * (Xc @ W_new))
                  + np.trace(sumEzz @ W_new.T @ W_new)) / (N * D)   # (12.57)
    return W_new, sigma2_new


def ppca_em(X, M, n_iter=200, W=None, sigma2=1.0, rng=None, tol=None):
    """EM for probabilistic PCA, iterating ppca_em_step from W (random if None).

    Returns mu, W, sigma^2 and the log likelihood (12.43) before every step plus the final one.
    """
    rng = np.random.default_rng(0) if rng is None else rng
    xbar = X.mean(0)
    Xc = X - xbar
    W = rng.standard_normal((X.shape[1], M)) if W is None else W.copy()
    ll = [ppca_loglik(X, xbar, W, sigma2)]
    for _ in range(n_iter):
        W, sigma2 = ppca_em_step(Xc, W, sigma2)
        ll.append(ppca_loglik(X, xbar, W, sigma2))
        if tol is not None and abs(ll[-1] - ll[-2]) < tol * abs(ll[-1]):
            break
    return xbar, W, sigma2, np.array(ll)


def ppca_loglik_missing(X, observed, mu, W, sigma2):
    """Log likelihood of the observed entries: sum_n ln N(x_{n,o} | mu_o, W_o W_o^T + sigma^2 I)."""
    total = 0.0
    for x, o in zip(X, observed):
        C = W[o] @ W[o].T + sigma2 * np.eye(o.sum())
        total += common.gauss_logpdf(x[o][None, :], mu[o], C)[0]
    return total


def ppca_em_missing(X, observed, M, n_iter=500, rng=None, tol=1e-10):
    """Exact EM for PPCA when values are missing at random (Section 12.2.2, Figure 12.11).

    observed is a boolean N x D mask. Both z_n and the missing values x_{n,m} are latent.
    E step: with o the observed coordinates of x_n, M_n = W_o^T W_o + sigma^2 I and
        E[z_n] = M_n^{-1} W_o^T (x_{n,o} - mu_o),  E[z_n z_n^T] = sigma^2 M_n^{-1} + E[z_n] E[z_n]^T
    (12.54), (12.55) restricted to the observed rows; a missing value has
        E[x_nd | z_n] = w_d^T z_n + mu_d, so E[x_nd z_n^T] = w_d^T E[z_n z_n^T] + mu_d E[z_n]^T and
        E[x_nd^2] = w_d^T E[z_n z_n^T] w_d + 2 mu_d w_d^T E[z_n] + mu_d^2 + sigma^2.
    M step: (12.56) generalised to estimate (W, mu) jointly by regressing x on (z, 1), and
    (12.57) for sigma^2, using these expected sufficient statistics. The observed-data log
    likelihood never decreases.
    """
    rng = np.random.default_rng(0) if rng is None else rng
    N, D = X.shape
    mu = np.array([X[observed[:, d], d].mean() for d in range(D)])
    W = rng.standard_normal((D, M))
    sigma2 = 1.0
    ll = [ppca_loglik_missing(X, observed, mu, W, sigma2)]
    for _ in range(n_iter):
        Sxz = np.zeros((D, M + 1))           # sum_n E[x_n zt_n^T] with zt = (z, 1)
        Szz = np.zeros((M + 1, M + 1))       # sum_n E[zt_n zt_n^T]
        Sxx = np.zeros(D)                    # sum_n E[x_nd^2]
        for x, o in zip(X, observed):
            m = ~o
            Wo = W[o]
            Minv = np.linalg.inv(Wo.T @ Wo + sigma2 * np.eye(M))
            Ez = Minv @ Wo.T @ (x[o] - mu[o])
            Ezz = sigma2 * Minv + np.outer(Ez, Ez)
            Ezt = np.append(Ez, 1.0)
            Eztzt = np.block([[Ezz, Ez[:, None]], [Ez[None, :], np.ones((1, 1))]])
            Ex_zt = np.empty((D, M + 1))
            Ex_zt[o] = np.outer(x[o], Ezt)
            Wm, mum = W[m], mu[m]
            Ex_zt[m] = np.hstack([Wm @ Ezz + np.outer(mum, Ez), (Wm @ Ez + mum)[:, None]])
            Exx = np.empty(D)
            Exx[o] = x[o] ** 2
            Exx[m] = np.einsum("dm,mk,dk->d", Wm, Ezz, Wm) + 2 * mum * (Wm @ Ez) + mum ** 2 + sigma2
            Sxz += Ex_zt
            Szz += Eztzt
            Sxx += Exx
        Wt = Sxz @ np.linalg.inv(Szz)        # rows (w_d^T, mu_d)
        sigma2 = (Sxx.sum() - 2 * np.sum(Wt * Sxz) + np.sum((Wt @ Szz) * Wt)) / (N * D)
        W, mu = Wt[:, :M], Wt[:, M]
        ll.append(ppca_loglik_missing(X, observed, mu, W, sigma2))
        if abs(ll[-1] - ll[-2]) < tol * abs(ll[-1]):
            break
    return mu, W, sigma2, np.array(ll)


def ppca_posterior_missing(X, observed, mu, W, sigma2):
    """Posterior means E[z_n | x_{n,o}] = M_n^{-1} W_o^T (x_{n,o} - mu_o) using observed entries."""
    M = W.shape[1]
    out = np.empty((len(X), M))
    for n, (x, o) in enumerate(zip(X, observed)):
        Wo = W[o]
        out[n] = np.linalg.solve(Wo.T @ Wo + sigma2 * np.eye(M), Wo.T @ (x[o] - mu[o]))
    return out


def em_pca_step(Xc, W):
    """EM for PCA in the limit sigma^2 -> 0 (Roweis 1998), on centred data Xc (N x D).

    E step (12.58): Omega = (W^T W)^{-1} W^T Xc^T, column n = orthogonal projection of x_n.
    M step (12.59): W_new = Xc^T Omega^T (Omega Omega^T)^{-1}, least squares with Omega fixed.
    Returns (Omega^T, W_new) so that the rows of Omega^T are the latent points z_n.
    """
    Omega = np.linalg.solve(W.T @ W, W.T @ Xc.T)
    W_new = Xc.T @ Omega.T @ np.linalg.inv(Omega @ Omega.T)
    return Omega.T, W_new


def em_pca(X, W, n_iter=50):
    """Iterate em_pca_step; the columns of W converge to a basis of the principal subspace."""
    Xc = X - X.mean(0)
    for _ in range(n_iter):
        _, W = em_pca_step(Xc, W)
    return W


def bayesian_pca(X, M, n_iter=500, W=None, sigma2=1.0, rng=None, alpha_max=1e12):
    """Bayesian PCA with automatic relevance determination (Bishop 1999a; Section 12.2.3).

    Prior p(W | alpha) = prod_i N(w_i | 0, alpha_i^{-1} I) (12.60). EM with the E step (12.54),
    (12.55), the M step W_new = [sum (x_n - xbar) E[z_n]^T][sum E[z_n z_n^T] + sigma^2 A]^{-1}
    (12.63) with A = diag(alpha), sigma^2 from (12.57), and alpha_i = D / ||w_i||^2 (12.62)
    after every M step. Columns whose alpha_i runs off to infinity are switched off.
    Returns mu, W, sigma^2, alpha and the history of alpha.
    """
    rng = np.random.default_rng(0) if rng is None else rng
    N, D = X.shape
    xbar = X.mean(0)
    Xc = X - xbar
    W = rng.standard_normal((D, M)) if W is None else W.copy()
    alpha = np.minimum(D / np.maximum((W ** 2).sum(0), 1e-300), alpha_max)
    hist = [alpha.copy()]
    for _ in range(n_iter):
        Minv = np.linalg.inv(W.T @ W + sigma2 * np.eye(M))
        Ez = Xc @ W @ Minv
        sumEzz = N * sigma2 * Minv + Ez.T @ Ez
        W = (Xc.T @ Ez) @ np.linalg.inv(sumEzz + sigma2 * np.diag(alpha))          # (12.63)
        sigma2 = ((Xc ** 2).sum() - 2 * np.sum(Ez * (Xc @ W))
                  + np.trace(sumEzz @ W.T @ W)) / (N * D)                           # (12.57)
        alpha = np.minimum(D / np.maximum((W ** 2).sum(0), 1e-300), alpha_max)      # (12.62)
        hist.append(alpha.copy())
    return xbar, W, sigma2, alpha, np.array(hist)


def bayesian_pca_gibbs(X, M, n_iter, rng, a0=1e-3, b0=1e-3, c0=1e-3, d0=1e-3, W=None):
    """Gibbs sampler for Bayesian PCA (Section 12.2.3 and Section 11.3; Figure 12.15).

    Model: x_n = W z_n + mu + noise, z_n ~ N(0, I) (12.31), noise ~ N(0, sigma^2 I) (12.32),
    w_i ~ N(0, alpha_i^{-1} I) (12.60), alpha_i ~ Gam(a0, b0), tau = 1/sigma^2 ~ Gam(c0, d0), and
    mu fixed at the sample mean. One sweep draws in turn
      z_n | W, tau     ~ N(M^{-1} W^T (x_n - mu), sigma^2 M^{-1})                      (12.42)
      row d of W | Z   ~ N(tau S Z^T x_{:,d}, S),  S = (diag(alpha) + tau Z^T Z)^{-1}
      alpha_i | w_i    ~ Gam(a0 + D/2, b0 + ||w_i||^2 / 2)
      tau | W, Z       ~ Gam(c0 + N D/2, d0 + sum_n ||x_n - mu - W z_n||^2 / 2).
    Returns the N_iter x M trace of ln alpha and the trace of sigma^2.
    """
    N, D = X.shape
    Xc = X - X.mean(0)
    W = rng.standard_normal((D, M)) if W is None else W.copy()
    tau = 1.0 / Xc.var()
    alpha = np.ones(M)
    trace, trace_s2 = np.empty((n_iter, M)), np.empty(n_iter)
    for it in range(n_iter):
        # z_n | W, tau
        Minv = np.linalg.inv(W.T @ W + np.eye(M) / tau)
        Ez = Xc @ W @ Minv
        Lz = np.linalg.cholesky(Minv / tau)
        Z = Ez + rng.standard_normal((N, M)) @ Lz.T
        # rows of W | Z, alpha, tau (all rows share the same covariance S)
        S = np.linalg.inv(np.diag(alpha) + tau * Z.T @ Z)
        Wmean = tau * Xc.T @ Z @ S
        W = Wmean + rng.standard_normal((D, M)) @ np.linalg.cholesky(S).T
        # alpha_i | w_i  (numpy's gamma uses shape and scale = 1 / rate)
        alpha = rng.gamma(a0 + D / 2, 1.0 / (b0 + 0.5 * (W ** 2).sum(0)))
        # tau | W, Z
        R = Xc - Z @ W.T
        tau = rng.gamma(c0 + N * D / 2, 1.0 / (d0 + 0.5 * (R ** 2).sum()))
        trace[it] = np.log(alpha)
        trace_s2[it] = 1.0 / tau
    return trace, trace_s2


def factor_analysis_em(X, M, n_iter=200, W=None, psi=None, rng=None):
    """EM for factor analysis (Rubin and Thayer 1982; Section 12.2.4).

    E step with G = (I + W^T Psi^{-1} W)^{-1} (12.68): E[z_n] = G W^T Psi^{-1} (x_n - xbar) (12.66),
    E[z_n z_n^T] = G + E[z_n] E[z_n]^T (12.67). M step: W_new from (12.69) and
    Psi_new = diag(S - W_new (1/N) sum_n E[z_n] (x_n - xbar)^T) (12.70).
    Returns mu, W, the diagonal of Psi and the log likelihood before every step plus the final one.
    """
    rng = np.random.default_rng(0) if rng is None else rng
    N, D = X.shape
    xbar = X.mean(0)
    Xc = X - xbar
    S = Xc.T @ Xc / N
    W = rng.standard_normal((D, M)) if W is None else W.copy()
    psi = np.diag(S).copy() if psi is None else psi.copy()
    ll = [fa_loglik(X, xbar, W, psi)]
    for _ in range(n_iter):
        G = np.linalg.inv(np.eye(M) + W.T @ (W / psi[:, None]))
        Ez = Xc @ (W / psi[:, None]) @ G                       # (12.66), rows E[z_n]^T
        sumEzz = N * G + Ez.T @ Ez                             # (12.67) summed over n
        W = (Xc.T @ Ez) @ np.linalg.inv(sumEzz)                 # (12.69)
        psi = np.diag(S - W @ (Ez.T @ Xc) / N).copy()          # (12.70)
        ll.append(fa_loglik(X, xbar, W, psi))
    return xbar, W, psi, np.array(ll)


def fa_loglik(X, mu, W, psi):
    """Log likelihood sum_n ln N(x_n | mu, W W^T + Psi), with C from (12.65)."""
    return common.gauss_logpdf(X, mu, W @ W.T + np.diag(psi)).sum()


# =========================================================================== Section 12.3: kernel PCA
def gaussian_kernel(A, B, c=0.1):
    """The 'Gaussian' kernel k(x, x') = exp(-||x - x'||^2 / c) of (12.86), c = 0.1 in the book."""
    d2 = (A ** 2).sum(1)[:, None] + (B ** 2).sum(1)[None, :] - 2 * A @ B.T
    return np.exp(-np.maximum(d2, 0) / c)


def centre_gram(K):
    """Centred Gram matrix K~ = K - 1_N K - K 1_N + 1_N K 1_N (12.85), 1_N = matrix of entries 1/N."""
    N = len(K)
    one = np.full((N, N), 1.0 / N)
    return K - one @ K - K @ one + one @ K @ one


def kernel_pca(X, kernel, n_components, centre=True):
    """Kernel PCA (Schoelkopf, Smola and Mueller 1998; Section 12.3).

    Solves K~ a_i = (N lambda_i) a_i (12.80) and scales a_i so that lambda_i N a_i^T a_i = 1
    (12.81), which makes the feature-space eigenvectors v_i = sum_n a_in phi~(x_n) unit vectors.
    Returns a dict with the eigenvalues of K~ (= N lambda_i), the coefficients A (columns a_i)
    and what kernel_pca_project needs.
    """
    K = kernel(X, X)
    Kt = centre_gram(K) if centre else K
    ev, V = np.linalg.eigh(Kt)
    order = np.argsort(ev)[::-1][:n_components]
    ev, V = ev[order], V[:, order]
    A = V / np.sqrt(ev)                 # unit eigenvectors divided by sqrt(N lambda_i)
    return dict(X=X, K=K, kernel=kernel, eig=ev, A=A, centre=centre)


def kernel_pca_project(model, Xnew):
    """Projections y_i(x) = phi~(x)^T v_i = sum_n a_in k~(x, x_n) (12.82), (12.87), where
    k~(x, x_n) is the kernel between centred feature vectors (as in (12.84), one argument new)."""
    X, K, A = model["X"], model["K"], model["A"]
    k = model["kernel"](Xnew, X)
    if model["centre"]:
        N = len(X)
        k = k - k.mean(1, keepdims=True) - K.mean(0)[None, :] + K.mean()
    return k @ A


# =========================================================================== Section 12.4.3: GTM
def regular_grid(n, lo=-1.0, hi=1.0):
    """n x n grid of points on [lo, hi]^2, rows ordered with x varying fastest."""
    g = np.linspace(lo, hi, n)
    gx, gy = np.meshgrid(g, g)
    return np.column_stack([gx.ravel(), gy.ravel()])


def gtm_fit(X, n_latent=20, n_rbf=8, width_factor=2.0, alpha=0.01, n_iter=200):
    """Generative topographic mapping (Bishop, Svensen and Williams 1998) fitted by EM.

    Latent space: K = n_latent^2 points u_k on a regular grid in [-1, 1]^2 with equal prior mass
    1/K. Mapping y(u) = W^T phi(u) with n_rbf^2 Gaussian basis functions (width = width_factor x
    centre spacing) plus a bias. Data density p(x) = (1/K) sum_k N(x | y(u_k), beta^{-1} I): a
    constrained Gaussian mixture. EM:
      E step  R_kn = N(x_n | y_k, beta^{-1} I) / sum_j N(x_n | y_j, beta^{-1} I)
      M step  (Phi^T G Phi + (alpha/beta) I) W = Phi^T R X,  G = diag(sum_n R_kn),
              1/beta = (1/(N D)) sum_n sum_k R_kn ||x_n - y_k||^2.
    W starts so that the manifold is the plane of the first two principal components (as in the
    original paper and Netlab). Returns a dict with the latent grid, Phi, W, beta, the
    responsibilities and the log likelihood trace.
    """
    N, D = X.shape
    U_lat = regular_grid(n_latent)
    C_rbf = regular_grid(n_rbf)
    s = width_factor * (2.0 / (n_rbf - 1))
    d2 = ((U_lat[:, None, :] - C_rbf[None, :, :]) ** 2).sum(-1)
    Phi = np.hstack([np.exp(-d2 / (2 * s ** 2)), np.ones((len(U_lat), 1))])
    # initialise: map the standardised latent grid onto the principal plane of the data
    xbar, lam, U = pca(X)
    Ustd = (U_lat - U_lat.mean(0)) / U_lat.std(0)
    target = xbar + Ustd @ (U[:, :2] * np.sqrt(lam[:2])).T
    W = np.linalg.lstsq(Phi, target, rcond=None)[0]
    Y = Phi @ W
    dd = ((Y[:, None, :] - Y[None, :, :]) ** 2).sum(-1)
    np.fill_diagonal(dd, np.inf)
    beta = 1.0 / max(lam[2], dd.min(1).mean() / 2)
    ll = []
    for _ in range(n_iter + 1):
        Y = Phi @ W
        D2 = ((Y ** 2).sum(1)[:, None] + (X ** 2).sum(1)[None, :] - 2 * Y @ X.T)  # K x N
        logp = -0.5 * beta * D2 + 0.5 * D * np.log(beta / (2 * np.pi))
        lse = common.logsumexp(logp, axis=0)
        ll.append(np.sum(lse - np.log(len(U_lat))))
        R = np.exp(logp - lse[None, :])
        if len(ll) > n_iter:
            break
        G = np.diag(R.sum(1))
        W = np.linalg.solve(Phi.T @ G @ Phi + (alpha / beta) * np.eye(Phi.shape[1]), Phi.T @ R @ X)
        Y = Phi @ W
        D2 = ((Y ** 2).sum(1)[:, None] + (X ** 2).sum(1)[None, :] - 2 * Y @ X.T)
        beta = N * D / np.sum(R * D2)
    return dict(latent=U_lat, Phi=Phi, W=W, beta=beta, R=R, loglik=np.array(ll), width=s)


# =========================================================================== data helpers
def mnist_threes():
    """The 1010 threes of the MNIST test set as rows of pixel values 0 to 255."""
    X, _ = common.mnist(train=False, digits=(3,))
    return X * 255.0


_THREES_PCA = {}


def threes_pca():
    """PCA of the MNIST threes (computed once and reused by Figures 12.3 to 12.5)."""
    if not _THREES_PCA:
        _THREES_PCA["fit"] = pca(mnist_threes())
    return _THREES_PCA["fit"]


def book_three():
    """MNIST test image 44, the digit three of the book's Figures 12.1 and 12.5 (28 x 28, 0..1)."""
    X, _ = common.mnist(train=False)
    return X[44].reshape(28, 28)


def oil_pca_coordinates(X):
    """First two principal-component coordinates, signs chosen to match Figure 12.8."""
    xbar, lam, U = pca(X)
    U = U[:, :2] * np.array([-1.0, 1.0])
    return (X - xbar) @ U


def oil_scatter(ax, Y, t, size=34):
    """Scatter plot in the book's style: laminar red x, homogeneous blue +, annular green o."""
    for k in (1, 0, 2):
        st = OIL_STYLE[k]
        s = t == k
        if st["marker"] == "o":
            ax.scatter(Y[s, 0], Y[s, 1], s=size, facecolors="none", edgecolors=st["color"],
                       linewidths=1.0, label=st["label"])
        else:
            ax.scatter(Y[s, 0], Y[s, 1], s=size, marker=st["marker"], color=st["color"],
                       linewidths=1.1, label=st["label"])


def savefig_compact(fig, name, out_dir=None, max_bytes=250_000):
    """common.savefig, then re-encode with an adaptive 256-colour palette if the PNG is large
    (smooth colour gradients such as Figure 12.17 compress poorly as full-colour PNG)."""
    import os
    from PIL import Image
    path = common.savefig(fig, 12, name, out_dir)
    if os.path.getsize(path) > max_bytes:
        Image.open(path).convert("RGB").quantize(256, method=Image.Quantize.MEDIANCUT).save(path, optimize=True)
    return path


def image_panel(ax, img, cmap=BLUE_INK, vmin=0, vmax=None, title=None):
    ax.imshow(img, cmap=cmap, vmin=vmin, vmax=vmax, interpolation="nearest")
    ax.set_xticks([])
    ax.set_yticks([])
    if title:
        ax.set_title(title, fontsize=11)


def axes_arrows(ax, xlab, ylab, xmax, ymax, lw=1.6):
    """Draw x and y axes as arrows from the origin (schematic figures 12.2 and 12.9)."""
    for (x1, y1) in ((xmax, 0), (0, ymax)):
        ax.annotate("", xy=(x1, y1), xytext=(0, 0),
                    arrowprops=dict(arrowstyle="-|>", lw=lw, color="k", mutation_scale=14))
    ax.text(xmax * 0.97, -0.07 * ymax, xlab, ha="right", va="top", fontsize=13)
    ax.text(-0.04 * xmax, ymax * 0.93, ylab, ha="right", va="top", fontsize=13)
    ax.set_axis_off()


def hinton(ax, W, max_abs):
    """Hinton diagram: one square per element, area proportional to |w|, white for positive
    and black for negative values, on the book's light blue background (Figure 12.14)."""
    D, M = W.shape
    ax.add_patch(Rectangle((-0.5, -0.5), M, D, color=HINTON_BG, zorder=0))
    for i in range(D):
        for j in range(M):
            side = 0.92 * np.sqrt(abs(W[i, j]) / max_abs)
            if side < 0.01:
                continue
            ax.add_patch(Rectangle((j - side / 2, D - 1 - i - side / 2), side, side,
                                   color="white" if W[i, j] > 0 else "black", zorder=1))
    ax.set_xlim(-0.5, M - 0.5)
    ax.set_ylim(-0.5, D - 0.5)
    ax.set_aspect("equal")
    ax.set_xticks([])
    ax.set_yticks([])


# =========================================================================== figures
def fig12_1(out_dir=None):
    """Figure 12.1: one digit three, randomly rotated and shifted within a 100 x 100 field."""
    big = ndimage.zoom(book_three(), 64 / 28, order=1)          # 64 x 64, as in the text
    rng = np.random.default_rng(1201)
    fig, axes = plt.subplots(1, 5, figsize=(11, 2.5))
    for ax in axes:
        rot = np.clip(ndimage.rotate(big, rng.uniform(-25, 25), reshape=True, order=1), 0, 1)
        field = np.zeros((100, 100))
        h, w = rot.shape
        r0 = int(round(50 + rng.uniform(-16, 16) - h / 2))
        c0 = int(round(50 + rng.uniform(-16, 16) - w / 2))
        # paste with clipping at the borders of the field
        rs, cs = max(r0, 0), max(c0, 0)
        re, ce = min(r0 + h, 100), min(c0 + w, 100)
        field[rs:re, cs:ce] = np.maximum(field[rs:re, cs:ce], rot[rs - r0:re - r0, cs - c0:ce - c0])
        image_panel(ax, field, vmax=1)
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-1", out_dir)


def fig12_2(out_dir=None):
    """Figure 12.2: the principal subspace (magenta), orthogonal projections (green) and the
    projection errors (blue) for a small 2-D data set (red)."""
    X = np.array([[1.2, 4.6], [4.5, 8.9], [10.4, 8.9], [5.2, 3.0], [8.3, 9.9]])
    xbar, lam, U = pca(X)
    u1 = U[:, 0] * np.sign(U[0, 0])
    Xt = pca_reconstruct(X, xbar, U, 1)                      # projections (12.20) with M = 1
    fig, ax = plt.subplots(figsize=(4.6, 4.4))
    t = np.array([-6.3, 6.0])
    line = xbar + t[:, None] * u1
    ax.plot(line[:, 0], line[:, 1], color=common.MAGENTA, lw=4, solid_capstyle="round", zorder=1)
    for x, xt in zip(X, Xt):
        ax.plot([x[0], xt[0]], [x[1], xt[1]], color=common.BLUE, lw=2.2, zorder=2)
    ax.scatter(X[:, 0], X[:, 1], s=150, color=common.RED, zorder=3)
    ax.scatter(Xt[:, 0], Xt[:, 1], s=150, color=common.GREEN, zorder=3)
    k = 1
    ax.text(X[k, 0] + 0.15, X[k, 1] + 0.35, r"$\mathbf{x}_n$", fontsize=15)
    ax.text(Xt[k, 0] + 0.35, Xt[k, 1] - 0.85, r"$\widetilde{\mathbf{x}}_n$", fontsize=15)
    a0 = xbar + 3.3 * u1 + np.array([-1.0, 1.3])
    ax.annotate("", xy=a0 + 1.5 * u1, xytext=a0,
                arrowprops=dict(arrowstyle="-|>", lw=2.4, color="k", mutation_scale=18))
    ax.text(*(a0 + 1.65 * u1 + np.array([0.15, 0.0])), r"$\mathbf{u}_1$", fontsize=15)
    axes_arrows(ax, r"$x_1$", r"$x_2$", 12.0, 12.2)
    ax.set_xlim(-0.6, 12.3)
    ax.set_ylim(-0.6, 12.6)
    ax.set_aspect("equal")
    common.savefig(fig, 12, "fig12-2", out_dir)


def fig12_3(out_dir=None):
    """Figure 12.3: mean and first four eigenvectors of the MNIST threes, with eigenvalues."""
    xbar, lam, U = threes_pca()
    signs = [1, 1, 1, -1]                            # eigenvector signs are arbitrary: as in the book
    fig, axes = plt.subplots(1, 5, figsize=(11, 2.75))
    image_panel(axes[0], xbar.reshape(28, 28), vmax=255, title="Mean")
    for i in range(4):
        u = signs[i] * U[:, i].reshape(28, 28)
        m = np.abs(u).max()
        mant, ex = ("%.1e" % lam[i]).split("e")
        image_panel(axes[i + 1], u, cmap=YELLOW_BLUE, vmin=-m, vmax=m,
                    title=r"$\lambda_%d = %s \cdot 10^{%d}$" % (i + 1, mant, int(ex)))
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-3", out_dir)


def fig12_4(out_dir=None):
    """Figure 12.4: (a) eigenvalue spectrum of the threes; (b) distortion J against M (12.18)."""
    _, lam, _ = threes_pca()
    D = len(lam)
    J = np.array([lam[M:].sum() for M in range(1, D + 1)])       # J = sum_{i > M} lambda_i
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.6))
    axes[0].plot(np.arange(1, D + 1), lam / 1e5, color=common.BLUE, lw=2)
    axes[0].set_xlabel("$i$")
    axes[0].set_ylabel(r"$\lambda_i$  ($\times 10^5$)")
    axes[1].plot(np.arange(1, D + 1), J / 1e6, color=common.BLUE, lw=2)
    axes[1].set_xlabel("$M$")
    axes[1].set_ylabel(r"$J$  ($\times 10^6$)")
    axes[0].set_ylim(0, 3.6)
    axes[1].set_ylim(0, 3.0)
    for ax, lab in zip(axes, "ab"):
        ax.set_xlim(0, D)
        common.panel_label(ax, lab, "upper right")
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-4", out_dir)


def fig12_5(out_dir=None):
    """Figure 12.5: a three and its PCA reconstructions (12.20) with M = 1, 10, 50, 250."""
    xbar, lam, U = threes_pca()
    x = book_three().ravel()[None, :] * 255.0
    fig, axes = plt.subplots(1, 5, figsize=(11, 2.75))
    image_panel(axes[0], x.reshape(28, 28), vmax=255, title="Original")
    for ax, M in zip(axes[1:], (1, 10, 50, 250)):
        image_panel(ax, pca_reconstruct(x, xbar, U, M).reshape(28, 28), vmax=255, title="$M = %d$" % M)
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-5", out_dir)


def fig12_6(out_dir=None):
    """Figure 12.6: Old Faithful: original, standardised with principal axes, and whitened."""
    X = common.old_faithful()
    Z = common.old_faithful(standardize=True)
    _, lam, U = pca(Z)
    Y = whiten(X)                                 # (12.24) applied to the original data, as in the book
    Y = np.column_stack([Y[:, 1], Y[:, 0]])      # book's order: minor axis across, major axis up
    fig, axes = plt.subplots(1, 3, figsize=(11, 3.7))
    style = dict(s=26, facecolors="none", edgecolors=common.GREEN, linewidths=1.1)
    axes[0].scatter(X[:, 0], X[:, 1], **style)
    axes[0].set_xlim(0.7, 6.3)
    axes[0].set_ylim(37, 105)
    axes[0].set_xlabel("duration (min)")
    axes[0].set_ylabel("waiting time (min)")
    axes[1].scatter(Z[:, 0], Z[:, 1], **style)
    for i in range(2):                            # principal axes over +- lambda_i^{1/2}
        e = np.sqrt(lam[i]) * U[:, i]
        axes[1].plot([-e[0], e[0]], [-e[1], e[1]], color=common.RED, lw=3.2, solid_capstyle="butt")
    axes[2].scatter(Y[:, 0], Y[:, 1], **style)
    axes[2].set_xlabel(r"$y_2$ (minor axis)")
    axes[2].set_ylabel(r"$y_1$ (major axis)")
    for ax in axes[1:]:
        ax.set_xlim(-2.5, 2.5)
        ax.set_ylim(-2.5, 2.5)
        ax.set_aspect("equal")
    axes[1].set_xlabel("standardised duration")
    axes[1].set_ylabel("standardised waiting time")
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-6", out_dir)


def fig12_7(out_dir=None):
    """Figure 12.7: PCA (magenta) versus Fisher's linear discriminant (green) for two classes."""
    rng = np.random.default_rng(127)
    n = 100
    cov = np.array([[3.2, 0.05], [0.05, 0.085]])
    Xb = rng.multivariate_normal([-0.3, 1.0], cov, n)
    Xr = rng.multivariate_normal([0.4, -1.05], cov, n)
    X = np.vstack([Xb, Xr])
    t = np.r_[np.zeros(n, int), np.ones(n, int)]
    _, lam, U = pca(X)
    u = U[:, 0] * np.sign(U[0, 0])
    w = fisher_direction(X, t)
    w = w * np.sign(w[1])
    fig, ax = plt.subplots(figsize=(5.4, 3.9))
    ax.scatter(Xb[:, 0], Xb[:, 1], s=12, color=common.BLUE, zorder=3)
    ax.scatter(Xr[:, 0], Xr[:, 1], s=12, color=common.RED, zorder=3)
    s = np.array([-10, 10])
    ax.plot(s * u[0], s * u[1], color=common.MAGENTA, lw=2.6, label="PCA direction")
    ax.plot(s * w[0], s * w[1], color=common.GREEN, lw=2.6, label="Fisher direction")
    ax.set_xlim(-5.5, 5.5)
    ax.set_ylim(-2, 2)
    ax.legend(loc="lower left", fontsize=8.5, framealpha=0.9)
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-7", out_dir)


def fig12_8(out_dir=None):
    """Figure 12.8: the oil flow data projected onto the first two principal components."""
    X, t = common.oil_flow("trn")
    Y = oil_pca_coordinates(X)
    fig, ax = plt.subplots(figsize=(4.6, 4.6))
    oil_scatter(ax, Y, t)
    ax.set_xlim(-3.2, 3.2)
    ax.set_ylim(-3.2, 3.2)
    ax.set_aspect("equal")
    ax.set_xlabel("first principal component")
    ax.set_ylabel("second principal component")
    ax.legend(loc="upper left", fontsize=8, framealpha=0.9, handletextpad=0.2)
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-8", out_dir)


def fig12_9(out_dir=None):
    """Figure 12.9: the generative view of PPCA (D = 2, M = 1), with computed densities:
    p(z), p(x | z_hat) = N(x | w z_hat + mu, sigma^2 I) and p(x) = N(x | mu, w w^T + sigma^2 I)."""
    mu = np.array([3.0, 3.9])
    w = np.array([1.75, 1.3])
    sigma = 0.6
    z_hat = 1.0
    C = np.outer(w, w) + sigma ** 2 * np.eye(2)                    # (12.36)
    fig, axes = plt.subplots(1, 3, figsize=(12, 3.9), gridspec_kw=dict(width_ratios=[1, 1, 1]))
    # left: the latent prior p(z) = N(z | 0, 1)
    ax = axes[0]
    z = np.linspace(-4, 4, 400)
    pz = np.exp(-0.5 * z ** 2) / np.sqrt(2 * np.pi)
    ax.plot(z, pz * 9, color=common.MAGENTA, lw=2.5)
    ax.annotate("", xy=(4.3, 0), xytext=(-4.2, 0),
                arrowprops=dict(arrowstyle="-|>", lw=1.6, color="k", mutation_scale=14))
    ax.plot([0, 0], [-0.12, 0.12], color="k", lw=1.4)
    ax.scatter([z_hat], [0], s=45, color="k", zorder=5)
    ax.text(z_hat, -0.35, r"$\widehat{z}$", ha="center", va="top", fontsize=14)
    ax.text(4.0, -0.35, "$z$", ha="center", va="top", fontsize=14)
    ax.text(-2.0, 2.4, "$p(z)$", fontsize=14, ha="right")
    ax.set_xlim(-4.4, 4.5)
    ax.set_ylim(-1.0, 8.6)
    ax.set_axis_off()
    # middle: the line mu + w z, and p(x | z_hat) as circles of radius sigma/2 and sigma
    ax = axes[1]
    g = np.linspace(-0.5, 8.5, 400)
    GX, GY = np.meshgrid(g, g)
    P = np.column_stack([GX.ravel(), GY.ravel()])
    m_hat = mu + w * z_hat
    dens = np.exp(common.gauss_logpdf(P, m_hat, sigma ** 2 * np.eye(2))).reshape(GX.shape)
    peak = 1 / (2 * np.pi * sigma ** 2)
    ax.contour(GX, GY, dens, levels=sorted(peak * np.exp(-0.5 * np.array([0.5, 1.0]) ** 2)),
               colors=common.RED, linewidths=2.2)
    s = np.array([-mu[0] / w[0], 2.7])
    ax.plot(mu[0] + s * w[0], mu[1] + s * w[1], color=common.BLUE, lw=3)
    ax.scatter(*mu, s=50, color="k", zorder=5)
    ax.scatter(*m_hat, s=50, color="k", zorder=5)
    ax.text(mu[0] - 0.25, mu[1] + 0.25, r"$\boldsymbol{\mu}$", fontsize=15, ha="right")
    ax.text(m_hat[0] - 1.4, m_hat[1] + 0.85, r"$p(\mathbf{x}|\widehat{z})$", fontsize=14, ha="right")
    nrm = np.array([w[1], -w[0]]) / np.linalg.norm(w)         # unit normal below the line
    b0, b1 = mu + 0.45 * nrm, m_hat + 0.45 * nrm
    ax.annotate("", xy=b1, xytext=b0, arrowprops=dict(arrowstyle="<|-|>", lw=1.3, color="k",
                                                      mutation_scale=10, shrinkA=0, shrinkB=0))
    for e in (mu, m_hat):
        ax.plot([e[0] + 0.2 * nrm[0], e[0] + 0.6 * nrm[0]], [e[1] + 0.2 * nrm[1], e[1] + 0.6 * nrm[1]],
                color="k", lw=1.0)
    bm = 0.5 * (b0 + b1) + 0.3 * nrm
    ax.text(bm[0], bm[1], r"$\widehat{z}\,|\mathbf{w}|$", fontsize=14, ha="left", va="top")
    a0 = mu + 2.35 * w + np.array([-0.55, 0.45])
    ax.annotate("", xy=a0 + 0.55 * w, xytext=a0,
                arrowprops=dict(arrowstyle="-|>", lw=2.2, color="k", mutation_scale=16))
    ax.text(*(a0 + 0.35 * w + np.array([-0.2, 0.55])), r"$\mathbf{w}$", fontsize=15)
    axes_arrows(ax, r"$x_1$", r"$x_2$", 8.4, 8.4)
    ax.set_xlim(-0.6, 8.8)
    ax.set_ylim(-0.6, 8.8)
    ax.set_aspect("equal")
    # right: the marginal p(x) = N(x | mu, C), contours at Mahalanobis distances 0.5 and 1
    ax = axes[2]
    dens = np.exp(common.gauss_logpdf(P, mu, C)).reshape(GX.shape)
    peak = 1 / (2 * np.pi * np.sqrt(np.linalg.det(C)))
    ax.contour(GX, GY, dens, levels=sorted(peak * np.exp(-0.5 * np.array([0.5, 1.0]) ** 2)),
               colors=common.GREEN, linewidths=2.2)
    ax.scatter(*mu, s=50, color="k", zorder=5)
    ax.text(mu[0] - 0.25, mu[1] + 0.25, r"$\boldsymbol{\mu}$", fontsize=15, ha="right")
    ax.text(mu[0] + 1.3, mu[1] - 1.6, r"$p(\mathbf{x})$", fontsize=14)
    axes_arrows(ax, r"$x_1$", r"$x_2$", 8.4, 8.4)
    ax.set_xlim(-0.6, 8.8)
    ax.set_ylim(-0.6, 8.8)
    ax.set_aspect("equal")
    # dashed arrow from z_hat in latent space to w z_hat + mu in data space
    con = ConnectionPatch(xyA=(z_hat + 0.15, 0.25), coordsA=axes[0].transData,
                          xyB=(m_hat[0] - 0.22, m_hat[1]), coordsB=axes[1].transData,
                          arrowstyle="-|>", ls=(0, (5, 4)), lw=1.8, color="k",
                          connectionstyle="arc3,rad=-0.3", mutation_scale=14)
    fig.add_artist(con)
    fig.subplots_adjust(wspace=0.05, left=0.01, right=0.99, top=0.98, bottom=0.04)
    common.savefig(fig, 12, "fig12-9", out_dir)


def missing_mask(N, D, frac, rng):
    """Random mask with a fraction frac of values missing; every row loses at least one value
    and keeps at least one (as stated in the caption of Figure 12.11)."""
    observed = rng.random((N, D)) >= frac
    bad = observed.all(1) | ~observed.any(1)
    while bad.any():                     # redraw only the offending rows
        observed[bad] = rng.random((bad.sum(), D)) >= frac
        bad = observed.all(1) | ~observed.any(1)
    return observed


def align_latent(Z, Zref):
    """Flip the signs of the columns of Z to agree with a reference (PPCA is defined up to R)."""
    return Z * np.sign(np.sum(Z * Zref, 0))


def fig12_11(out_dir=None):
    """Figure 12.11: PPCA posterior means for the first 100 oil flow points; then with 30% of
    the values removed at random and the EM algorithm for missing data."""
    X, t = common.oil_flow("trn")
    X, t = X[:100], t[:100]
    # complete data: closed-form ML (12.45), (12.46) and posterior means (12.48)
    mu, W, s2 = ppca_ml(X, 2)
    Zc = ppca_posterior(X, mu, W, s2)[0]
    Zc = align_latent(Zc, oil_pca_coordinates(common.oil_flow("trn")[0])[:100])
    # 30% missing at random, exact EM, then rotate W to orthogonal columns (choice of R)
    rng = np.random.default_rng(1211)
    observed = missing_mask(*X.shape, 0.3, rng)
    mu_m, W_m, s2_m, ll = ppca_em_missing(X, observed, 2, n_iter=3000, rng=rng)
    Uw, Sw, Vt = np.linalg.svd(W_m, full_matrices=False)
    W_m = Uw * Sw                                   # W R with R = V: orthogonal columns
    Zm = align_latent(ppca_posterior_missing(X, observed, mu_m, W_m, s2_m), Zc)
    print("  Figure 12.11: %.1f%% of values missing; EM ran %d iterations, log likelihood %.2f -> %.2f"
          % (100 * (1 - observed.mean()), len(ll) - 1, ll[0], ll[-1]))
    fig, axes = plt.subplots(1, 2, figsize=(9, 4.6))
    for ax, Z, title in zip(axes, (Zc, Zm), ("complete data", "30% of the values missing")):
        oil_scatter(ax, Z, t, size=46)
        ax.set_xlim(-3, 3)
        ax.set_ylim(-3, 3)
        ax.set_aspect("equal")
        ax.set_title(title)
        ax.set_xlabel(r"$\mathbb{E}[z_{n1}]$")
        ax.set_ylabel(r"$\mathbb{E}[z_{n2}]$")
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-11", out_dir)


# the ten points of Figure 12.12, read off the book's figure (then centred)
EM_PCA_POINTS = np.array([[1.84, 2.06], [0.37, 1.30], [-0.58, 0.34], [-0.47, -0.19], [-0.16, -0.19],
                          [-0.25, -0.58], [0.07, -0.81], [1.14, -0.15], [-0.87, -0.77], [-1.14, -0.99]])


def fig12_12(out_dir=None):
    """Figure 12.12: EM for PCA, (12.58) and (12.59), step by step on ten points in 2-D."""
    X = EM_PCA_POINTS - EM_PCA_POINTS.mean(0)
    _, lam, U = pca(X)
    W0 = np.array([[2.0], [-1.0]])
    Z1, W1 = em_pca_step(X, W0)              # (b) E step with W0, then (c) the M step gives W1
    Z2, W2 = em_pca_step(X, W1)              # (d) E step with W1, (e) M step gives W2
    W = W2
    for _ in range(200):                      # (f) converged
        Zc, W = em_pca_step(X, W)
    Zf = np.linalg.solve(W.T @ W, W.T @ X.T).T
    panels = [(None, None), (W0, Z1), (W1, Z1), (W1, Z2), (W2, Z2), (W, Zf)]
    fig, axes = common.panel_grid(2, 3, size=(3.0, 3.0))
    for ax, (Wk, Zk), lab in zip(axes.ravel(), panels, "abcdef"):
        if Wk is None:
            for i in range(2):
                e = np.sqrt(lam[i]) * U[:, i]
                ax.plot([-e[0], e[0]], [-e[1], e[1]], "k--", lw=2, dashes=(4, 3))
        else:
            d = Wk[:, 0] / np.linalg.norm(Wk[:, 0])
            ax.plot([-6 * d[0], 6 * d[0]], [-6 * d[1], 6 * d[1]], color=common.RED, lw=4, zorder=1)
            P = Zk @ Wk.T                    # latent points mapped to data space, Z W^T
            for x, p in zip(X, P):
                ax.plot([x[0], p[0]], [x[1], p[1]], color=common.BLUE, lw=2, zorder=2)
            ax.scatter(P[:, 0], P[:, 1], s=80, color=common.CYAN, zorder=3)
        ax.scatter(X[:, 0], X[:, 1], s=80, color=common.GREEN, zorder=4)
        ax.set_xlim(-3, 3)
        ax.set_ylim(-3, 3)
        ax.set_aspect("equal")
        ax.set_xticks([-2, 0, 2])
        ax.set_yticks([-2, 0, 2])
        common.panel_label(ax, lab)
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-12", out_dir)


def bpca_data(N=300, D=10, rng=None):
    """Synthetic data of Figure 12.14: sd 1 along 3 random orthonormal directions, 0.5 along 7."""
    rng = np.random.default_rng(1214) if rng is None else rng
    Q, _ = np.linalg.qr(rng.standard_normal((D, D)))
    sd = np.r_[np.ones(3), 0.5 * np.ones(D - 3)]
    return (rng.standard_normal((N, D)) * sd) @ Q.T


def fig12_14(out_dir=None):
    """Figure 12.14: Hinton diagrams of W for ML PPCA and for Bayesian PCA (M = 9, D = 10)."""
    rng = np.random.default_rng(1214)
    X = bpca_data(rng=rng)
    W0 = rng.standard_normal((10, 9))
    _, W_ml, s2_ml, ll = ppca_em(X, 9, n_iter=2000, W=W0)
    _, W_b, s2_b, alpha, _ = bayesian_pca(X, 9, n_iter=2000, W=W0)
    _, Wcf, s2cf = ppca_ml(X, 9)
    print("  Figure 12.14: ML sigma^2 by EM %.4f, closed form %.4f; |WW^T - W_ML W_ML^T| = %.1e"
          % (s2_ml, s2cf, np.abs(W_ml @ W_ml.T - Wcf @ Wcf.T).max()))
    print("  Bayesian PCA: alpha =", np.array2string(alpha, precision=3, max_line_width=200),
          " sigma^2 = %.4f" % s2_b)
    m = max(np.abs(W_ml).max(), np.abs(W_b).max())
    fig, axes = plt.subplots(1, 2, figsize=(7.6, 4.4))
    hinton(axes[0], W_ml, m)
    hinton(axes[1], W_b, m)
    axes[0].set_title("maximum likelihood PPCA")
    axes[1].set_title("Bayesian PCA")
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-14", out_dir)


def fig12_15(out_dir=None):
    """Figure 12.15: Gibbs sampling for Bayesian PCA: ln alpha_i against iteration (D = 4, M = 3)."""
    rng = np.random.default_rng(1215)
    N, D = 50, 4
    w_true = np.sqrt(0.7) * np.array([1.0, 0.0, 0.0, 0.0])     # one direction of high variance
    X = np.outer(rng.standard_normal(N), w_true) + np.sqrt(0.1) * rng.standard_normal((N, D))
    trace, s2 = bayesian_pca_gibbs(X, 3, 10000, rng)
    fig, axes = plt.subplots(3, 1, figsize=(8.5, 6.2), sharex=True)
    for ax, i, col in zip(axes, range(3), (common.RED, common.GREEN, common.BLUE)):
        ax.plot(np.arange(1, len(trace) + 1), trace[:, i], color=col, lw=0.6)
        ax.set_ylim(0, 10)
        ax.set_ylabel(r"$\ln\alpha_%d$" % (i + 1))
    axes[-1].set_xlabel("Gibbs iteration")
    axes[-1].set_xlim(0, len(trace))
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-15", out_dir)


def kpca_toy_data(rng=None):
    """Three Gaussian clusters of 30 points with sd 0.1 (Schoelkopf et al. 1998), as in Fig 12.17."""
    rng = np.random.default_rng(1217) if rng is None else rng
    centres = np.array([[-0.5, -0.2], [0.0, 0.6], [0.5, 0.0]])
    return np.vstack([c + 0.1 * rng.standard_normal((30, 2)) for c in centres])


def fig12_17(out_dir=None):
    """Figure 12.17: kernel PCA with a Gaussian kernel: contours of the first eight components."""
    X = kpca_toy_data()
    model = kernel_pca(X, gaussian_kernel, 8)
    gx = np.linspace(-1, 1, 121)
    gy = np.linspace(-0.5, 1.5, 121)
    GX, GY = np.meshgrid(gx, gy)
    Y = kernel_pca_project(model, np.column_stack([GX.ravel(), GY.ravel()]))
    fig, axes = plt.subplots(2, 4, figsize=(11, 5.9))
    for i, ax in enumerate(axes.ravel()):
        F = Y[:, i].reshape(GX.shape)
        ax.pcolormesh(GX, GY, F, cmap=KPCA_GREEN, shading="gouraud", rasterized=True)
        ax.contour(GX, GY, F, 9, colors="#0018c8", linewidths=0.9, linestyles="solid")
        ax.scatter(X[:, 0], X[:, 1], s=22, facecolors="none", edgecolors=common.RED, linewidths=1.3)
        ax.set_xlim(-1, 1)
        ax.set_ylim(-0.5, 1.5)
        ax.set_aspect("equal")
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_title("Eigenvalue = %.2f" % model["eig"][i], loc="left", fontsize=10.5)
    fig.tight_layout()
    savefig_compact(fig, "fig12-17", out_dir)


def fig12_21(out_dir=None):
    """Figure 12.21: the oil flow data visualised with PCA (left) and GTM posterior means (right)."""
    X, t = common.oil_flow("trn")
    g = gtm_fit(X)
    Zmean = g["R"].T @ g["latent"]                    # posterior mean in latent space
    print("  Figure 12.21: GTM log likelihood %.1f -> %.1f in %d EM iterations, 1/beta = %.4f"
          % (g["loglik"][0], g["loglik"][-1], len(g["loglik"]) - 1, 1 / g["beta"]))
    fig, axes = plt.subplots(1, 2, figsize=(9, 4.6))
    oil_scatter(axes[0], oil_pca_coordinates(X), t, size=30)
    axes[0].set_xlim(-3.2, 3.2)
    axes[0].set_ylim(-3.2, 3.2)
    axes[0].set_title("PCA")
    oil_scatter(axes[1], Zmean, t, size=30)
    axes[1].set_xlim(-1.1, 1.1)
    axes[1].set_ylim(-1.1, 1.1)
    axes[1].set_title("GTM (posterior means)")
    for ax in axes:
        ax.set_aspect("equal")
    fig.tight_layout()
    common.savefig(fig, 12, "fig12-21", out_dir)


FIGURES = {
    "12.1": fig12_1, "12.2": fig12_2, "12.3": fig12_3, "12.4": fig12_4, "12.5": fig12_5,
    "12.6": fig12_6, "12.7": fig12_7, "12.8": fig12_8, "12.9": fig12_9, "12.11": fig12_11,
    "12.12": fig12_12, "12.14": fig12_14, "12.15": fig12_15, "12.17": fig12_17, "12.21": fig12_21,
}

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