"""Bishop (2006), Chapter 1: Introduction.

Reproduces the numerical figures of Chapter 1 of C. M. Bishop, Pattern Recognition and Machine
Learning (Springer, 2006), pp. 1-58, written from scratch with numpy, scipy and matplotlib. Every
computation is a small function meant to be read next to the lecture notes
(ch01-introduction.html):

* Section 1.1, polynomial curve fitting: design_matrix, fit_poly (the normal equations (1.122)
  solved stably by least squares, and the regularised error (1.4) by stacking sqrt(lambda) I under
  the design matrix), sse, rms_error (1.3).
* Section 1.2.5-1.2.6, probabilistic and Bayesian curve fitting: map_poly (MAP = regularised least
  squares with lambda = alpha/beta, (1.67)), bayes_predictive (the predictive mean and variance
  (1.69)-(1.72), with S^{-1} = alpha I + beta sum phi phi^T).
* Section 1.2.4, the Gaussian: gauss_pdf, gauss_ml (1.55)-(1.56), expected_ml_variance (1.58).
* Section 1.4: shell_fraction (1.76 and Fig. 1.22), radial_density (Fig. 1.23).
* Section 1.5, decision theory: bayes_decision (minimum expected loss (1.81)), expected_loss,
  reject_decision and decide_with_reject (Section 1.5.3, threshold theta, reject cost lambda = 1 - theta), error_probability (1.78 for two classes).
* Section 1.6, information theory: entropy (1.93), kl_discrete (1.113), kl_gauss (closed form for
  two univariate Gaussians), discretised_gaussian_hist (Fig. 1.30).
* js_noisy_sine: the data generator. It uses the seeded generator JSRandom (the generator DK.rng of
  the browser demos, imported from ch11_sampling.py), so curve-fitting-demo.html draws exactly the
  same data sets as the figures below.

Figures: 1.1-1.8, 1.11-1.17, 1.19-1.24, 1.26-1.31, plus "T1.1" (Table 1.1: the coefficients).
Run `python3 ch01_introduction.py --list`, `python3 ch01_introduction.py` (all, about a minute)
or `python3 ch01_introduction.py --fig 1.4 1.17`. Files go to ../img/ch01/.

Settings chosen where the book gives none (the book's data sets are not available):
* Data (Figures 1.2-1.8, 1.16, 1.17, Table 1.1): t = sin(2 pi x) + Gaussian noise with standard
  deviation 0.3 (so beta = 1/0.3^2 = 11.1, as in the caption of Fig. 1.17), x_n = (n-1)/(N-1)
  equally spaced on [0, 1]. Training set N = 10 (DATA_SEED = %(seed)d, JSRandom); N = 15 and N = 100
  sets of Fig. 1.6 use the same generator and seed; the test set of Figs 1.5 and 1.8 has 100 points
  (TEST_SEED), generated the same way. DATA_SEED is the first seed whose M = 9 fit over-fits as
  strongly as in the book (test RMS error of M = 9 above 0.8).
* Figure 1.8 plots ln lambda from -35 to -5 (the book shows -35 to -20 and the book's test error is smallest near the right-hand end of that
  range; with our data it is smallest near ln lambda = -10, where E_RMS(test) = 0.37).
* Fig 1.11: the 60 points are not recoverable, so we draw 60 points from the 9 x 2 table of counts
  read off the book's figure (COUNTS_1_11), with uniform jitter inside each cell.
* Fig 1.15: three samples of size N = 2 from N(0, 1) (the book's are not given), plus a fourth panel
  (d) with the average of sigma^2_ML over 10^5 data sets against N (not in the book).
* Fig 1.19-1.20: the three-phase oil flow data of Appendix A (1000 training points, x6 against x7), the
  test point x = (0.4, 1.3), 4 x 4 cells of width 0.25 and 0.5; 1.21 draws 4^D cells (D = 1, 2, 3).
* Fig 1.23: Gaussian with sigma = 1/2 (the peak heights 1.60 for D = 1 and 1.2 for D = 2 in the book's axes imply it).
* Fig 1.24: schematic in the book; we use p(x, C1) = 0.4 N(3, 1.3^2) + 0.2 N(7.2, 1^2) and
  p(x, C2) = 0.4 N(7.3, 1^2) on [0, 11] and x_hat = 7.3 (the peak of p(x, C2)); the shaded areas are
  computed. Fig 1.26 uses p(C2|x) = 0.72 s(6x) + 0.28 s(4(x - 3)) (a made-up posterior with a plateau).
* Fig 1.27: p(x|C1) = 0.25 N(0.2, 0.075^2) + 0.75 N(0.5, 0.14^2), p(x|C2) = N(0.7, 0.1^2), equal priors: the book's
  mixture parameters are not given, so these resemble its curves (the book's p(x|C1) is not exactly
  normalised).
* Fig 1.28: y(x) = sin(2 pi x) with Gaussian noise of standard deviation 0.3 around it.
* Fig 1.30: discretised Gaussians over 30 bins with standard deviations (in bins) 1.42 and 5.49,
  which give H = 1.77 and H = 3.09 as in the book (the uniform distribution has ln 30 = 3.40).
* Schematic figures not reproduced (they carry no numbers): 1.9, 1.10, 1.18 (cross-validation).
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from scipy import stats
from scipy.special import gammaln

import common
from common import RED, BLUE, GREEN, GREY
from ch11_sampling import JSRandom

DATA_SEED = 6        # chosen by choose_seed(); see the module docstring
TEST_SEED = 1001
SIGMA = 0.3          # noise standard deviation of the curve-fitting data
ALPHA, BETA = 5e-3, 11.1   # Fig. 1.17

__doc__ = __doc__ % {"seed": DATA_SEED}


# =========================================================================== data
def js_noisy_sine(N, seed, sigma=SIGMA):
    """x_n = (n-1)/(N-1) and t_n = sin(2 pi x_n) + N(0, sigma^2), noise from the demos' JSRandom(seed)."""
    rng = JSRandom(seed)
    x = np.linspace(0.0, 1.0, N)
    t = np.sin(2 * np.pi * x) + sigma * np.array([rng.norm() for _ in range(N)])
    return x, t


def choose_seed(lo=1, hi=60):
    """First seed in [lo, hi] whose M = 9 fit has test RMS error above 0.8 (over-fits as in Fig. 1.5)."""
    xt, tt = js_noisy_sine(100, TEST_SEED)
    for s in range(lo, hi + 1):
        x, t = js_noisy_sine(10, s)
        w = fit_poly(x, t, 9)
        if rms_error(xt, tt, w) > 0.8 and rms_error(xt, tt, fit_poly(x, t, 3)) < 0.45:
            return s
    return None


# =========================================================================== Section 1.1
def design_matrix(x, M):
    """Phi_{ni} = x_n^i, i = 0..M (the polynomial y(x, w) = sum_j w_j x^j of (1.1))."""
    return np.vander(np.asarray(x, dtype=float), M + 1, increasing=True)


def fit_poly(x, t, M, lam=0.0):
    """Minimise E(w) = 1/2 sum (y(x_n, w) - t_n)^2 + lam/2 ||w||^2, (1.2) and (1.4).

    The minimiser solves (lam I + Phi^T Phi) w = Phi^T t (normal equations, (1.122)). Solving them
    directly squares the condition number (about 1e12 for M = 9), so we solve the equivalent
    least-squares problem [Phi; sqrt(lam) I] w ~ [t; 0] with an SVD (numpy lstsq).
    """
    Phi = design_matrix(x, M)
    if lam > 0:
        Phi = np.vstack([Phi, np.sqrt(lam) * np.eye(M + 1)])
        t = np.concatenate([t, np.zeros(M + 1)])
    return np.linalg.lstsq(Phi, t, rcond=None)[0]


def sse(x, t, w):
    """E(w) = 1/2 sum_n (y(x_n, w) - t_n)^2, (1.2)."""
    r = design_matrix(x, len(w) - 1) @ w - t
    return 0.5 * float(r @ r)


def rms_error(x, t, w):
    """E_RMS = sqrt(2 E(w*) / N), (1.3)."""
    return float(np.sqrt(2.0 * sse(x, t, w) / len(x)))


# =========================================================================== Sections 1.2.5, 1.2.6
def map_poly(x, t, M, alpha, beta):
    """MAP estimate (1.67): minimise beta/2 sum (y - t)^2 + alpha/2 w^T w, i.e. (1.4) with lambda = alpha/beta."""
    return fit_poly(x, t, M, lam=alpha / beta)


def bayes_predictive(x_train, t_train, x_new, M, alpha, beta):
    """Predictive distribution of Bayesian curve fitting, (1.68)-(1.72).

    S^{-1} = alpha I + beta sum_n phi(x_n) phi(x_n)^T, m(x) = beta phi(x)^T S sum_n phi(x_n) t_n,
    s^2(x) = 1/beta + phi(x)^T S phi(x). Returns (mean, variance) at x_new.
    """
    Phi = design_matrix(x_train, M)
    S = np.linalg.inv(alpha * np.eye(M + 1) + beta * Phi.T @ Phi)
    phi = design_matrix(x_new, M)
    mean = beta * phi @ S @ Phi.T @ t_train
    var = 1.0 / beta + np.einsum("ni,ij,nj->n", phi, S, phi)
    return mean, var


def _save(fig, name, out_dir=None):
    """Tidy the panel spacing of multi-panel figures, then save PNG, SVG and PDF through common.savefig."""
    if len(fig.axes) > 1 and not any(a.name == "3d" for a in fig.axes):
        fig.tight_layout()
    return common.savefig(fig, 1, name, out_dir)


def style(ax, xlim=(-0.05, 1.05), ylim=(-1.6, 1.6)):
    ax.set_xlim(*xlim)
    ax.set_ylim(*ylim)
    ax.set_xlabel("$x$")
    ax.set_ylabel("$t$", rotation=0, labelpad=8)


def data_axes(ax, x, t, ylim=(-1.6, 1.6), curve=True):
    """Blue circles for the data and the green sin(2 pi x), the book's colours."""
    xs = np.linspace(0, 1, 400)
    if curve:
        ax.plot(xs, np.sin(2 * np.pi * xs), color=GREEN, lw=2, zorder=1)
    ax.plot(x, t, "o", mfc="none", mec=BLUE, mew=1.8, ms=7, zorder=3)
    style(ax, ylim=ylim)


def fig1_1(out_dir=None):
    """Figure 1.1: examples of handwritten digits (one MNIST example of each digit 0-9)."""
    X, y = common.mnist(train=False)
    fig, axes = plt.subplots(2, 5, figsize=(6.2, 2.9))
    for d, ax in enumerate(axes.ravel()):
        ax.imshow(X[np.flatnonzero(y == d)[0]].reshape(28, 28), cmap="gray_r")
        ax.axis("off")
    _save(fig, "fig1-1", out_dir)


def fig1_2(out_dir=None):
    """Figure 1.2: the training set, N = 10 points, and the function sin(2 pi x) that generated it."""
    x, t = js_noisy_sine(10, DATA_SEED)
    fig, ax = plt.subplots(figsize=(4.2, 3.2))
    data_axes(ax, x, t, ylim=(-1.5, 1.5))
    _save(fig, "fig1-2", out_dir)


def fig1_3(out_dir=None):
    """Figure 1.3: the sum-of-squares error (1.2) is half the sum of squared vertical displacements."""
    x, t = js_noisy_sine(10, DATA_SEED)
    w = fit_poly(x, t, 3)
    xs = np.linspace(0, 1, 300)
    fig, ax = plt.subplots(figsize=(4.4, 3.2))
    ax.plot(xs, design_matrix(xs, 3) @ w, color=RED, lw=2)
    resid = np.abs(design_matrix(x, 3) @ w - t)
    idx = sorted(np.argsort(resid)[-3:])
    for i in idx:
        yi = design_matrix(x[i:i + 1], 3) @ w
        ax.plot([x[i], x[i]], [yi[0], t[i]], color=GREEN, lw=2.2)
        ax.plot(x[i], yi[0], "o", color=GREEN, ms=6)
        ax.plot(x[i], t[i], "o", mfc="none", mec=BLUE, mew=1.8, ms=7)
    j = idx[-1]
    ax.text(x[j] + 0.025, t[j], "$t_n$", va="center")
    ax.text(x[j] + 0.025, (design_matrix(x[j:j + 1], 3) @ w)[0] + 0.1, "$y(x_n,\\mathbf{w})$", va="bottom", ha="left", fontsize=9)
    style(ax, ylim=(-1.5, 1.5))
    ax.set_title("$E(\\mathbf{w})=\\frac{1}{2}\\sum_n (y(x_n,\\mathbf{w})-t_n)^2$ (green lengths squared)", fontsize=8.5)
    _save(fig, "fig1-3", out_dir)


def fig1_4(out_dir=None):
    """Figure 1.4: polynomials of order M = 0, 1, 3, 9 fitted to the N = 10 data of Fig 1.2."""
    x, t = js_noisy_sine(10, DATA_SEED)
    xs = np.linspace(0, 1, 400)
    fig, axes = common.panel_grid(2, 2, size=(3.3, 2.6))
    for ax, M in zip(axes.ravel(), (0, 1, 3, 9)):
        w = fit_poly(x, t, M)
        data_axes(ax, x, t)
        ax.plot(xs, design_matrix(xs, M) @ w, color=RED, lw=2, zorder=2)
        ax.text(0.97, 0.95, "$M=%d$" % M, transform=ax.transAxes, ha="right", va="top", fontsize=11)
    _save(fig, "fig1-4", out_dir)


def rms_curves(M_values=range(10)):
    x, t = js_noisy_sine(10, DATA_SEED)
    xt, tt = js_noisy_sine(100, TEST_SEED)
    tr, te = [], []
    for M in M_values:
        w = fit_poly(x, t, M)
        tr.append(rms_error(x, t, w))
        te.append(rms_error(xt, tt, w))
    return np.array(tr), np.array(te)


def fig1_5(out_dir=None):
    """Figure 1.5: training and test RMS error (1.3) against the order M."""
    tr, te = rms_curves()
    fig, ax = plt.subplots(figsize=(4.6, 3.2))
    ax.plot(range(10), tr, "o-", color=BLUE, mfc="none", mew=1.5, label="Training")
    ax.plot(range(10), te, "o-", color=RED, mfc="none", mew=1.5, label="Test")
    ax.set_xlabel("$M$")
    ax.set_ylabel("$E_{\\mathrm{RMS}}$")
    ax.set_xticks([0, 3, 6, 9])
    ax.set_ylim(0, 1.3)
    ax.annotate("test, M = 9: %.2f" % te[9], (9, te[9]), xytext=(4.3, 1.1), fontsize=8,
                arrowprops=dict(arrowstyle="->", lw=0.7))
    ax.legend(loc="upper center", frameon=False, bbox_to_anchor=(0.3, 1.0))
    _save(fig, "fig1-5", out_dir)
    return dict(train=tr, test=te)


def fig1_6(out_dir=None):
    """Figure 1.6: the M = 9 solution for N = 15 and N = 100 data points (more data cures over-fitting)."""
    xs = np.linspace(0, 1, 400)
    fig, axes = common.panel_grid(1, 2, size=(3.5, 2.8))
    out = {}
    for ax, N in zip(axes.ravel(), (15, 100)):
        x, t = js_noisy_sine(N, DATA_SEED)
        w = fit_poly(x, t, 9)
        data_axes(ax, x, t)
        ax.plot(xs, design_matrix(xs, 9) @ w, color=RED, lw=2, zorder=2)
        ax.text(0.97, 0.95, "$N=%d$" % N, transform=ax.transAxes, ha="right", va="top", fontsize=11)
        out[N] = w
    _save(fig, "fig1-6", out_dir)
    return out


def fig1_7(out_dir=None):
    """Figure 1.7: M = 9 polynomials fitted with the regularised error (1.4), ln lambda = -18 and 0."""
    x, t = js_noisy_sine(10, DATA_SEED)
    xs = np.linspace(0, 1, 400)
    fig, axes = common.panel_grid(1, 2, size=(3.5, 2.8))
    for ax, ll in zip(axes.ravel(), (-18, 0)):
        w = fit_poly(x, t, 9, np.exp(ll))
        data_axes(ax, x, t)
        ax.plot(xs, design_matrix(xs, 9) @ w, color=RED, lw=2, zorder=2)
        ax.text(0.97, 0.95, "$\\ln\\lambda=%d$" % ll, transform=ax.transAxes, ha="right", va="top", fontsize=11)
    _save(fig, "fig1-7", out_dir)


def fig1_8(out_dir=None):
    """Figure 1.8: E_RMS against ln lambda for the M = 9 polynomial (training and test)."""
    x, t = js_noisy_sine(10, DATA_SEED)
    xt, tt = js_noisy_sine(100, TEST_SEED)
    lls = np.linspace(-35, -5, 61)
    tr, te = [], []
    for ll in lls:
        w = fit_poly(x, t, 9, np.exp(ll))
        tr.append(rms_error(x, t, w))
        te.append(rms_error(xt, tt, w))
    fig, ax = plt.subplots(figsize=(4.6, 3.2))
    ax.plot(lls, tr, "-", color=BLUE, label="Training")
    ax.plot(lls, te, "-", color=RED, label="Test")
    ax.set_xlabel("$\\ln\\lambda$")
    ax.set_ylabel("$E_{\\mathrm{RMS}}$")
    ax.set_ylim(0, 1.05)
    ax.legend(loc="upper left", frameon=False)
    _save(fig, "fig1-8", out_dir)
    return dict(lnlam=lls, train=np.array(tr), test=np.array(te), best=lls[int(np.argmin(te))])


def table1_1(out_dir=None):
    """Table 1.1: coefficients w* of the fitted polynomials (M = 0, 1, 3, 9, and M = 9 with ln lambda = -18, 0)."""
    x, t = js_noisy_sine(10, DATA_SEED)
    rows = {"M=0": fit_poly(x, t, 0), "M=1": fit_poly(x, t, 1), "M=3": fit_poly(x, t, 3), "M=9": fit_poly(x, t, 9),
            "M=9, ln lambda=-18": fit_poly(x, t, 9, np.exp(-18)), "M=9, ln lambda=0": fit_poly(x, t, 9, 1.0)}
    for k, w in rows.items():
        print("%-20s" % k, " ".join("%10.2f" % v for v in w))
    return rows


# =========================================================================== Section 1.2: probability
COUNTS_1_11 = np.array([[3, 6, 7, 8, 5, 3, 1, 0, 0],     # Y = 1 (bottom row in the book), 33 points
                        [0, 0, 0, 1, 4, 5, 8, 6, 3]])    # Y = 2, 27 points


def fig1_11(out_dir=None):
    """Figure 1.11: joint distribution of X (9 values) and Y (2 values) from 60 points; marginals; p(X | Y = 1)."""
    rng = np.random.default_rng(1)
    C = COUNTS_1_11
    fig, axes = common.panel_grid(2, 2, size=(3.3, 2.5))
    a, b, c, d = axes.ravel()
    for y in (0, 1):
        for i in range(9):
            for _ in range(C[y, i]):
                a.plot(i + rng.uniform(0.15, 0.85), y + rng.uniform(0.15, 0.85), "o", color=BLUE, ms=5)
    for i in range(10):
        a.axvline(i, color=RED, lw=1.5)
    for y in range(3):
        a.axhline(y, color=RED, lw=1.5)
    a.set_xlim(0, 9)
    a.set_ylim(0, 2)
    a.set_xticks([])
    a.set_yticks([0.5, 1.5])
    a.set_yticklabels(["$Y=1$", "$Y=2$"])
    a.set_title("$p(X,Y)$", fontsize=10)
    pY = C.sum(1) / C.sum()
    pX = C.sum(0) / C.sum()
    pXY1 = C[0] / C[0].sum()
    b.bar([1, 2], pY, color=BLUE, width=0.8)
    b.set_title("$p(Y)$", fontsize=10)
    b.set_xticks([1, 2])
    b.set_xticklabels(["$Y=1$", "$Y=2$"])
    c.bar(np.arange(1, 10), pX, color=BLUE, width=0.8)
    c.set_title("$p(X)$", fontsize=10)
    d.bar(np.arange(1, 10), pXY1, color=BLUE, width=0.8)
    d.set_title("$p(X\\,|\\,Y=1)$", fontsize=10)
    for ax in (c, d):
        ax.set_xticks([])
    for ax in (a, b, c, d):
        ax.tick_params(length=0)
    _save(fig, "fig1-11", out_dir)
    return dict(pY=pY, pX=pX, pXY1=pXY1)


def fig1_12(out_dir=None):
    """Figure 1.12: a probability density p(x) and its cumulative distribution P(x), (1.25)-(1.26)."""
    x = np.linspace(-4, 4, 400)
    fig, ax = plt.subplots(figsize=(4.4, 3))
    ax.plot(x, stats.norm.pdf(x, 0, 1), color=RED, label="$p(x)$")
    ax.plot(x, stats.norm.cdf(x, 0, 1), color=BLUE, label="$P(x)$")
    ax.fill_between(x, stats.norm.pdf(x), where=(x > -0.5) & (x < 0.8), color=GREEN, alpha=0.4)
    ax.text(0.15, 0.1, "$\\delta x$", ha="center", fontsize=9)
    ax.set_xlabel("$x$")
    ax.legend(frameon=False, loc="upper left")
    _save(fig, "fig1-12", out_dir)


def gauss_pdf(x, mu, var):
    """N(x | mu, sigma^2), (1.46)."""
    return np.exp(-0.5 * (x - mu) ** 2 / var) / np.sqrt(2 * np.pi * var)


def fig1_13(out_dir=None):
    """Figure 1.13: the univariate Gaussian N(x | mu, sigma^2) with its mean and its width 2 sigma."""
    x = np.linspace(-4, 4, 400)
    fig, ax = plt.subplots(figsize=(4.4, 3))
    ax.plot(x, gauss_pdf(x, 0, 1), color=RED)
    h = gauss_pdf(1, 0, 1)
    ax.annotate("", (1, h), (-1, h), arrowprops=dict(arrowstyle="<->", lw=0.9))
    ax.text(0, h + 0.015, "$2\\sigma$", ha="center")
    ax.axvline(0, color=GREY, lw=0.8, ls=":")
    ax.set_xticks([0])
    ax.set_xticklabels(["$\\mu$"])
    ax.set_xlabel("$x$")
    ax.set_ylabel("$\\mathcal{N}(x|\\mu,\\sigma^2)$")
    _save(fig, "fig1-13", out_dir)


def gauss_ml(x):
    """Maximum-likelihood mean and variance of a Gaussian sample: (1.55) mu_ML = mean, (1.56) sigma2_ML = mean (x - mu_ML)^2."""
    x = np.asarray(x, dtype=float)
    mu = x.mean()
    return mu, float(np.mean((x - mu) ** 2))


def expected_ml_variance(sigma2, N):
    """E[sigma2_ML] = (N-1)/N sigma^2, (1.58): the ML variance is biased low by the factor (N-1)/N."""
    return (N - 1) / N * sigma2


def fig1_14(out_dir=None):
    """Figure 1.14: the likelihood of a Gaussian is the product of the density values at the data (blue)."""
    pts = np.array([-1.6, -0.3, 0.7, 1.4])
    mu, var = gauss_ml(pts)
    x = np.linspace(-4, 4, 400)
    fig, ax = plt.subplots(figsize=(4.8, 3))
    ax.plot(x, gauss_pdf(x, mu, var), color=RED)
    for p in pts:
        ax.plot([p, p], [0, gauss_pdf(p, mu, var)], color=BLUE, lw=1.6, ls="--")
        ax.plot(p, gauss_pdf(p, mu, var), "o", color=BLUE, ms=5)
        ax.plot(p, 0, "o", color="black", ms=5)
    ax.set_ylim(-0.03, 0.62)
    ax.set_xlabel("$x$")
    ax.set_ylabel("$p(x)$")
    ax.text(0.97, 0.95, "$\\mu_{ML}=%.2f,\\ \\sigma^2_{ML}=%.2f$" % (mu, var), transform=ax.transAxes, ha="right", va="top", fontsize=8.5)
    _save(fig, "fig1-14", out_dir)


def fig1_15(out_dir=None):
    """Figure 1.15: bias of the ML variance. (a)-(c) three data sets of N = 2 from N(0, 1); (d) E[sigma2_ML] against N."""
    rng = np.random.default_rng(5)
    x = np.linspace(-4, 4, 400)
    fig, axes = common.panel_grid(2, 2, size=(3.4, 2.4))
    for lab, ax in zip("abc", axes.ravel()[:3]):
        s = rng.normal(0, 1, 2)
        mu, var = gauss_ml(s)
        ax.plot(x, gauss_pdf(x, 0, 1), color=GREEN, lw=2)
        ax.plot(x, gauss_pdf(x, mu, var), color=RED, lw=2)
        ax.plot(s, [0, 0], "o", color=BLUE, ms=6, zorder=5)
        ax.axvline(0, ymax=0.05, color="black")
        ax.set_yticks([])
        ax.set_ylim(-0.02, 1.7)
        panel_title = "$\\hat\\sigma^2=%.2f$" % var
        ax.text(0.97, 0.95, panel_title, transform=ax.transAxes, ha="right", va="top", fontsize=9)
        common.panel_label(ax, lab)
    d = axes.ravel()[3]
    Ns = np.arange(2, 21)
    sim = [np.mean([gauss_ml(rng.normal(0, 1, N))[1] for _ in range(20000)]) for N in Ns]
    d.plot(Ns, [expected_ml_variance(1.0, N) for N in Ns], "-", color=GREEN, lw=2, label="$(N-1)/N$")
    d.plot(Ns, sim, "o", color=RED, ms=4, label="simulation")
    d.axhline(1, color=GREY, lw=0.8, ls=":")
    d.set_xlabel("$N$")
    d.set_ylabel("$E[\\sigma^2_{ML}]$")
    d.legend(frameon=False, fontsize=8, loc="lower right")
    common.panel_label(d, "d")
    _save(fig, "fig1-15", out_dir)


def fig1_16(out_dir=None):
    """Figure 1.16: the probabilistic model p(t | x0, w, beta) = N(t | y(x0, w), beta^-1), (1.60)."""
    xs = np.linspace(0, 1, 300)
    x0 = 0.6
    y0 = np.sin(2 * np.pi * x0)
    fig, ax = plt.subplots(figsize=(4.6, 3.2))
    ax.plot(xs, np.sin(2 * np.pi * xs), color=RED, lw=2, label="$y(x,\\mathbf{w})$")
    tt = np.linspace(-1.6, 1.6, 300)
    dens = gauss_pdf(tt, y0, SIGMA ** 2)
    ax.plot(x0 + 0.1 * dens / dens.max(), tt, color=BLUE, lw=1.8)
    ax.plot([x0, x0], [y0 - SIGMA, y0 + SIGMA], color=GREEN, lw=3)
    ax.plot(x0, y0, "o", color=RED, ms=5)
    ax.text(x0 - 0.03, y0 + SIGMA + 0.1, "$2\\beta^{-1/2}$", ha="right", fontsize=9)
    ax.annotate("$p(t|x_0,\\mathbf{w},\\beta)$", (x0 + 0.05, y0 + 0.4), xytext=(0.74, 0.95), fontsize=9,
                arrowprops=dict(arrowstyle="->", lw=0.7))
    ax.set_xticks([0, x0, 1])
    ax.set_xticklabels(["0", "$x_0$", "1"])
    style(ax, xlim=(-0.05, 1.1), ylim=(-1.6, 1.6))
    _save(fig, "fig1-16", out_dir)


def fig1_17(out_dir=None):
    """Figure 1.17: Bayesian predictive distribution, M = 9, alpha = 5e-3, beta = 11.1: mean (red) and +- 1 sd (pink)."""
    x, t = js_noisy_sine(10, DATA_SEED)
    xs = np.linspace(0, 1, 300)
    m, v = bayes_predictive(x, t, xs, 9, ALPHA, BETA)
    s = np.sqrt(v)
    fig, ax = plt.subplots(figsize=(4.6, 3.3))
    ax.fill_between(xs, m - s, m + s, color="#f4b6c2", zorder=0)
    data_axes(ax, x, t, ylim=(-1.5, 1.5))
    ax.plot(xs, m, color=RED, lw=2, zorder=2)
    _save(fig, "fig1-17", out_dir)
    return dict(mean=m, sd=s)


# =========================================================================== Section 1.4: dimensionality
def oil_xy():
    X, y = common.oil_flow("trn")
    return X[:, 5:7], y     # x6 and x7 (1-based)


OIL_COLOURS = ("#ff0000", "#00dd00", "#0000ff")   # homogeneous red, annular green, laminar blue, as in the book


def scatter_oil(ax, cross=True):
    P, y = oil_xy()
    for k in (0, 1, 2):
        ax.scatter(P[y == k, 0], P[y == k, 1], s=22, facecolors="none", edgecolors=OIL_COLOURS[k], linewidths=1.2, rasterized=True)
    if cross:
        ax.plot(0.4, 1.3, "x", color="black", ms=8, mew=2.5, zorder=5)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 2)
    ax.set_xlabel("$x_6$")
    ax.set_ylabel("$x_7$", rotation=0, labelpad=8)


def fig1_19(out_dir=None):
    """Figure 1.19: the oil-flow data, x6 against x7, with the new test point marked by a cross."""
    fig, ax = plt.subplots(figsize=(4.6, 4.4))
    scatter_oil(ax)
    _save(fig, "fig1-19", out_dir)


def fig1_20(out_dir=None):
    """Figure 1.20: the cell-based classifier. Each 4 x 4 cell takes the majority class of its training points."""
    P, y = oil_xy()
    fig, ax = plt.subplots(figsize=(4.6, 4.4))
    pale = ("#ffb0b0", "#b0ffb0", "#b0b0ff")
    for i in range(4):
        for j in range(4):
            x0, y0 = 0.25 * i, 0.5 * j
            inside = (P[:, 0] >= x0) & (P[:, 0] < x0 + 0.25 + (i == 3) * 1e-9) & (P[:, 1] >= y0) & (P[:, 1] < y0 + 0.5 + (j == 3) * 1e-9)
            if inside.any():
                k = np.bincount(y[inside], minlength=3).argmax()
                ax.add_patch(Rectangle((x0, y0), 0.25, 0.5, fc=pale[k], ec="black", lw=0.8, zorder=0))
    scatter_oil(ax)
    ax.set_xticks([0, 0.25, 0.5, 0.75, 1])
    _save(fig, "fig1-20", out_dir)


def fig1_21(out_dir=None):
    """Figure 1.21: dividing each axis into 4 intervals gives 4^D cells: 4, 16, 64 for D = 1, 2, 3."""
    fig = plt.figure(figsize=(8.4, 3.0))
    a = fig.add_subplot(1, 3, 1)
    for k in range(5):
        a.plot([k, k], [-0.1, 0.1], color=RED, lw=2)
    a.axhline(0, color="black", lw=1.5)
    a.set_xlim(-0.3, 4.5)
    a.set_ylim(-1, 1)
    a.axis("off")
    a.set_title("$D=1$: 4 cells", fontsize=10)
    b = fig.add_subplot(1, 3, 2)
    for k in range(5):
        b.plot([k, k], [0, 4], color=RED, lw=1.6)
        b.plot([0, 4], [k, k], color=RED, lw=1.6)
    b.set_aspect("equal")
    b.axis("off")
    b.set_title("$D=2$: 16 cells", fontsize=10)
    c = fig.add_subplot(1, 3, 3, projection="3d")
    filled = np.ones((4, 4, 4), dtype=bool)
    c.voxels(filled, facecolors=(1, 0, 0, 0.06), edgecolor=(1, 0, 0, 0.9), linewidth=0.8)
    c.set_axis_off()
    c.set_title("$D=3$: 64 cells", fontsize=10)
    _save(fig, "fig1-21", out_dir)


def shell_fraction(eps, D):
    """Fraction of the volume of a D-sphere of radius 1 lying between r = 1 - eps and r = 1: 1 - (1 - eps)^D, (1.76)."""
    return 1.0 - (1.0 - eps) ** D


def fig1_22(out_dir=None):
    """Figure 1.22: the volume fraction of the outer shell 1 - eps < r < 1 for D = 1, 2, 5, 20."""
    eps = np.linspace(0, 1, 400)
    fig, ax = plt.subplots(figsize=(4.4, 3.4))
    for D in (1, 2, 5, 20):
        ax.plot(eps, shell_fraction(eps, D), color=RED, lw=1.8)
        xx = {1: 0.5, 2: 0.55, 5: 0.35, 20: 0.1}[D]
        ax.text(xx, shell_fraction(xx, D) + (0.04 if D == 1 else 0.03), "$D=%d$" % D, fontsize=9, color=RED)
    ax.set_xlabel("$\\epsilon$")
    ax.set_ylabel("volume fraction")
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    _save(fig, "fig1-22", out_dir)


def radial_density(r, D, sigma=0.5):
    """Density of the radius r = ||x|| for x ~ N(0, sigma^2 I_D): p(r) = S_D r^{D-1} exp(-r^2/(2 sigma^2)) / (2 pi sigma^2)^{D/2}, (1.79)."""
    r = np.asarray(r, dtype=float)
    logS = np.log(2) + 0.5 * D * np.log(np.pi) - gammaln(0.5 * D)     # log of the surface area S_D of the unit sphere
    return np.exp(logS + (D - 1) * np.log(np.maximum(r, 1e-300)) - r ** 2 / (2 * sigma ** 2) - 0.5 * D * np.log(2 * np.pi * sigma ** 2))


def fig1_23(out_dir=None):
    """Figure 1.23: radial density p(r) of a Gaussian (sigma = 1/2, as the book's axes imply) for D = 1, 2, 20: the mass sits in a shell at r ~ sigma sqrt(D - 1)."""
    r = np.linspace(0, 4, 500)
    fig, ax = plt.subplots(figsize=(4.4, 3.4))
    for D, col, dx in ((1, RED, 0.15), (2, GREEN, 0.45), (20, BLUE, 0.35)):
        ax.plot(r, radial_density(r, D), color=col, lw=2)
        k = np.argmax(radial_density(r, D))
        ax.text(r[k] + dx, radial_density(r, D)[k] * (0.97 if D < 20 else 0.8), "$D=%d$" % D, fontsize=10, color=col)
    ax.set_xlabel("$r$")
    ax.set_ylabel("$p(r)$")
    ax.set_xlim(0, 4)
    ax.set_ylim(0, 2)
    _save(fig, "fig1-23", out_dir)


# =========================================================================== Section 1.5: decision theory
def bayes_decision(post, loss):
    """Minimum-expected-loss decision (1.81): for each row of post (N x K posteriors p(C_k|x)) pick j minimising sum_k L_kj p(C_k|x).

    loss[k, j] is the loss for choosing class j when the truth is k.
    """
    return np.argmin(np.asarray(post) @ np.asarray(loss), axis=1)


def expected_loss(post, loss, decision):
    """Expected loss of each decision, sum_k L_{k, d} p(C_k | x)."""
    return np.einsum("nk,kn->n", post, np.asarray(loss)[:, decision])


def reject_decision(post, theta):
    """Reject option (Section 1.5.3): -1 (reject) where max_k p(C_k|x) <= theta, else the most probable class."""
    post = np.asarray(post)
    return np.where(post.max(1) > theta, post.argmax(1), -1)


def decide_with_reject(post, loss, lam):
    """Bayes rule with a reject option of cost lam: -1 (reject) where the smallest expected loss exceeds lam, else argmin_j (1.81).

    With 0-1 loss and lam = 1 - theta this is reject_decision (reject where max_k p(C_k|x) < theta).
    """
    el = np.asarray(post) @ np.asarray(loss)
    return np.where(el.min(1) > lam, -1, el.argmin(1))


def error_probability(x, joint1, joint2, xhat):
    """p(mistake) = int_{R1} p(x, C2) dx + int_{R2} p(x, C1) dx with R1 = (-inf, xhat), R2 = [xhat, inf), (1.78)."""
    dx = x[1] - x[0]
    return float(joint2[x < xhat].sum() * dx + joint1[x >= xhat].sum() * dx)


def fig1_24(out_dir=None):
    """Figure 1.24: joint densities p(x, C1), p(x, C2), a decision boundary x_hat and the error areas."""
    x = np.linspace(0, 11, 2200)
    j1 = 0.40 * gauss_pdf(x, 3.0, 1.3 ** 2) + 0.20 * gauss_pdf(x, 7.2, 1.0)
    j2 = 0.40 * gauss_pdf(x, 7.3, 1.0)
    idx = np.flatnonzero(np.diff(np.sign(j1 - j2)) != 0)
    x0 = x[idx[idx > 1000][0]]
    xhat = 7.3
    fig, ax = plt.subplots(figsize=(6.0, 3.3))
    ax.plot(x, j1, "k", lw=2)
    ax.plot(x, j2, "k", lw=2)
    low = np.where(x < x0, j2, j1)
    m = (x < xhat)
    ax.fill_between(x, 0, low, where=m, color="#80ff80")
    ax.fill_between(x, j1, j2, where=(x >= x0) & m, color="#ff8080")
    ax.fill_between(x, 0, j1, where=x >= xhat, color="#8080ff")
    ax.axvline(x0, color="black", ls="--", lw=1.8, ymax=0.8)
    ax.axvline(xhat, color="black", lw=2, ymax=0.8)
    ax.text(x0, 0.175, "$x_0$", ha="center", fontsize=11)
    ax.text(xhat, 0.175, "$\\widehat{x}$", ha="center", fontsize=11)
    ax.text(1.3, 0.125, "$p(x,\\mathcal{C}_1)$", fontsize=10)
    ax.text(8.6, 0.09, "$p(x,\\mathcal{C}_2)$", fontsize=10)
    ax.annotate("", (0, -0.012), (xhat, -0.012), arrowprops=dict(arrowstyle="<->"))
    ax.annotate("", (xhat, -0.012), (11, -0.012), arrowprops=dict(arrowstyle="<->"))
    ax.text(xhat / 2, -0.03, "$\\mathcal{R}_1$", ha="center", fontsize=11)
    ax.text((xhat + 11) / 2, -0.03, "$\\mathcal{R}_2$", ha="center", fontsize=11)
    ax.set_xlim(0, 11)
    ax.set_ylim(-0.045, 0.19)
    ax.set_yticks([])
    ax.set_xlabel("$x$")
    _save(fig, "fig1-24", out_dir)
    return dict(x0=x0, p_err_xhat=error_probability(x, j1, j2, xhat), p_err_x0=error_probability(x, j1, j2, x0))


def fig1_26(out_dir=None):
    """Figure 1.26: the reject option. Inputs where max_k p(C_k|x) <= theta are rejected."""
    x = np.linspace(-3, 6, 900)
    p2 = 0.72 * common.sigmoid(6 * x) + 0.28 * common.sigmoid(4 * (x - 3))
    p1 = 1 - p2
    theta = 0.85
    rej = np.max([p1, p2], axis=0) <= theta
    xl, xr = x[rej][0], x[rej][-1]
    fig, ax = plt.subplots(figsize=(5, 3.3))
    ax.plot(x, p1, color=BLUE, lw=2)
    ax.plot(x, p2, color=RED, lw=2)
    ax.axhline(theta, color=GREEN, ls="--", lw=1.8)
    ax.axvline(xl, color=GREEN, lw=1.8, ymax=theta / 1.08)
    ax.axvline(xr, color=GREEN, lw=1.8, ymax=theta / 1.08)
    ax.annotate("", (xl, -0.05), (xr, -0.05), arrowprops=dict(arrowstyle="<->"))
    ax.text((xl + xr) / 2, -0.15, "reject region", ha="center")
    ax.text(-2.9, 1.02, "$p(\\mathcal{C}_1|x)$", color=BLUE)
    ax.text(3.7, 1.02, "$p(\\mathcal{C}_2|x)$", color=RED)
    ax.set_yticks([0, theta, 1])
    ax.set_yticklabels(["0.0", "$\\theta$", "1.0"])
    ax.set_xticks([])
    ax.set_ylim(-0.2, 1.1)
    ax.set_xlabel("$x$")
    _save(fig, "fig1-26", out_dir)
    return dict(reject=(xl, xr))


def class_densities_1_27(x):
    p1 = 0.25 * gauss_pdf(x, 0.2, 0.075 ** 2) + 0.75 * gauss_pdf(x, 0.5, 0.14 ** 2)
    p2 = gauss_pdf(x, 0.7, 0.1 ** 2)
    return p1, p2


def fig1_27(out_dir=None):
    """Figure 1.27: class-conditional densities (a) and posterior probabilities (b), equal priors, with the minimum-error boundary."""
    x = np.linspace(0, 1, 800)
    p1, p2 = class_densities_1_27(x)
    post1 = p1 / (p1 + p2)
    cross = x[np.flatnonzero(np.diff(np.sign(post1 - 0.5)) != 0)]
    fig, axes = common.panel_grid(1, 2, size=(3.6, 3.3))
    a, b = axes.ravel()
    a.plot(x, p1, color=BLUE, lw=2)
    a.plot(x, p2, color=RED, lw=2)
    a.text(0.1, 1.45, "$p(x|\\mathcal{C}_1)$", color=BLUE)
    a.text(0.55, 4.25, "$p(x|\\mathcal{C}_2)$", color=RED)
    a.set_ylim(0, 5)
    a.set_xlabel("$x$")
    a.set_ylabel("class densities")
    common.panel_label(a, "a", "upper left")
    b.plot(x, post1, color=BLUE, lw=2)
    b.plot(x, 1 - post1, color=RED, lw=2)
    for c in cross[-1:]:
        b.axvline(c, color=GREEN, lw=1.8)
    b.text(0.05, 0.9, "$p(\\mathcal{C}_1|x)$", color=BLUE)
    b.text(0.62, 0.9, "$p(\\mathcal{C}_2|x)$", color=RED)
    b.set_ylim(0, 1.05)
    b.set_xlabel("$x$")
    b.set_ylabel("posterior probabilities")
    common.panel_label(b, "b", "lower left")
    _save(fig, "fig1-27", out_dir)
    return dict(boundary=float(cross[-1]))


def fig1_28(out_dir=None):
    """Figure 1.28: the regression function y(x) minimising expected squared loss is the conditional mean E[t|x]."""
    xs = np.linspace(0, 1, 300)
    x0 = 0.3
    fig, ax = plt.subplots(figsize=(4.8, 3.2))
    ax.plot(xs, np.sin(2 * np.pi * xs), color=RED, lw=2)
    tt = np.linspace(-1.5, 1.5, 300)
    ax.plot(x0 + 0.1 * gauss_pdf(tt, np.sin(2 * np.pi * x0), SIGMA ** 2) / gauss_pdf(0, 0, SIGMA ** 2), tt, color=BLUE, lw=1.8)
    ax.plot(x0, np.sin(2 * np.pi * x0), "o", color=RED, ms=5)
    ax.text(0.9, -0.6, "$y(x)$", color=RED)
    ax.text(x0 + 0.06, -1.1, "$p(t|x_0)$", color=BLUE, fontsize=9)
    ax.set_xticks([0, x0, 1])
    ax.set_xticklabels(["0", "$x_0$", "1"])
    style(ax, xlim=(-0.02, 1.02), ylim=(-1.5, 1.5))
    _save(fig, "fig1-28", out_dir)


def fig1_29(out_dir=None):
    """Figure 1.29: the loss L_q = |y - t|^q for q = 0.3, 1, 2, 10."""
    d = np.linspace(-2, 2, 800)
    fig, axes = common.panel_grid(2, 2, size=(3.0, 2.5))
    for ax, q in zip(axes.ravel(), (0.3, 1, 2, 10)):
        ax.plot(d, np.abs(d) ** q, color=RED, lw=2)
        ax.set_xlim(-2, 2)
        ax.set_ylim(0, 2)
        ax.set_xticks([-2, -1, 0, 1, 2])
        ax.set_xlabel("$y-t$")
        ax.set_ylabel("$|y-t|^q$")
        ax.text(0.05, 0.9, "$q=%g$" % q, transform=ax.transAxes, fontsize=10, va="top")
    fig.tight_layout()
    _save(fig, "fig1-29", out_dir)


# =========================================================================== Section 1.6: information theory
def entropy(p):
    """H[p] = -sum_i p_i ln p_i (nats), (1.93); 0 ln 0 = 0."""
    p = np.asarray(p, dtype=float)
    p = p[p > 0]
    return float(-(p * np.log(p)).sum())


def kl_discrete(p, q):
    """KL(p || q) = sum_i p_i ln(p_i / q_i) (>= 0, = 0 iff p = q), (1.113)."""
    p, q = np.asarray(p, dtype=float), np.asarray(q, dtype=float)
    m = p > 0
    return float((p[m] * np.log(p[m] / q[m])).sum())


def kl_gauss(m1, v1, m2, v2):
    """KL(N(m1, v1) || N(m2, v2)) = 1/2 [ln(v2/v1) + (v1 + (m1 - m2)^2)/v2 - 1] for univariate Gaussians."""
    return 0.5 * (np.log(v2 / v1) + (v1 + (m1 - m2) ** 2) / v2 - 1.0)


def discretised_gaussian_hist(s, n=30):
    """p_i proportional to exp(-(i - (n-1)/2)^2 / (2 s^2)), i = 0..n-1, normalised: histograms of Fig. 1.30."""
    k = np.arange(n)
    p = np.exp(-0.5 * ((k - (n - 1) / 2) / s) ** 2)
    return p / p.sum()


def fig1_30(out_dir=None):
    """Figure 1.30: two distributions over 30 bins: the broader one has the larger entropy (the maximum is ln 30 = 3.40)."""
    fig, axes = common.panel_grid(1, 2, size=(3.2, 3.0))
    for ax, s in zip(axes.ravel(), (1.42, 5.49)):
        p = discretised_gaussian_hist(s)
        ax.bar(np.arange(30), p, width=1.0, color="#8080ff", edgecolor="black", lw=0.6)
        ax.set_ylim(0, 0.5)
        ax.set_xticks([])
        ax.set_ylabel("probabilities")
        ax.text(0.5, 0.65, "H = %.2f" % entropy(p), transform=ax.transAxes, ha="center", fontsize=11)
    _save(fig, "fig1-30", out_dir)
    return dict(H_narrow=entropy(discretised_gaussian_hist(1.42)), H_broad=entropy(discretised_gaussian_hist(5.49)), H_uniform=np.log(30))


def fig1_31(out_dir=None):
    """Figure 1.31: a convex function (x ln x): every chord lies above the function, f(lam a + (1-lam) b) <= lam f(a) + (1-lam) f(b)."""
    f = lambda x: x * np.log(x)
    x = np.linspace(0.05, 3, 300)
    a, b, lam = 0.4, 2.4, 0.45
    xl = lam * a + (1 - lam) * b
    fig, ax = plt.subplots(figsize=(4.4, 3.2))
    ax.plot(x, f(x), color=RED, lw=2)
    ax.plot([a, b], [f(a), f(b)], color=BLUE, lw=2)
    ax.plot([xl, xl], [f(xl), lam * f(a) + (1 - lam) * f(b)], color=GREEN, lw=2)
    ax.set_xticks([a, xl, b])
    ax.set_xticklabels(["$a$", "$x_\\lambda$", "$b$"])
    ax.set_yticks([])
    ax.text(2.55, f(2.55) - 0.55, "$f(x)$", color=RED, fontsize=11)
    ax.text(0.95, 1.0, "chord", color=BLUE, fontsize=10, ha="right")
    ax.set_xlim(0, 3)
    ax.set_ylim(-0.6, 3.5)
    _save(fig, "fig1-31", out_dir)


FIGURES = {
    "1.1": fig1_1, "1.2": fig1_2, "1.3": fig1_3, "1.4": fig1_4, "1.5": fig1_5, "1.6": fig1_6, "1.7": fig1_7, "1.8": fig1_8,
    "T1.1": table1_1, "1.11": fig1_11, "1.12": fig1_12, "1.13": fig1_13, "1.14": fig1_14, "1.15": fig1_15, "1.16": fig1_16,
    "1.17": fig1_17, "1.19": fig1_19, "1.20": fig1_20, "1.21": fig1_21, "1.22": fig1_22, "1.23": fig1_23, "1.24": fig1_24,
    "1.26": fig1_26, "1.27": fig1_27, "1.28": fig1_28, "1.29": fig1_29, "1.30": fig1_30, "1.31": fig1_31,
}

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