"""Bishop (2006), Chapter 4: Linear Models for Classification.

Reproduces the numerical figures of Chapter 4 of C. M. Bishop, Pattern Recognition and Machine
Learning (Springer, 2006), pp. 179-224, written from scratch with numpy, scipy and matplotlib.
Every algorithm is a small function meant to be read next to the lecture notes
(ch04-linear-classification.html):

* Section 4.1, discriminant functions:
  least_squares_classifier (the pseudo-inverse solution (4.16)), fisher_two_class ((4.28)-(4.30)),
  fisher_criterion ((4.25)-(4.27)), fisher_multiclass ((4.40)-(4.50), eigenvectors of
  S_W^-1 S_B), perceptron ((4.54)-(4.55), cycling through the data).
* Section 4.2, probabilistic generative models:
  gaussian_generative_fit (maximum likelihood (4.75)-(4.80), shared covariance) and
  generative_posterior (the logistic sigmoid (4.57)-(4.67) for two classes, the softmax
  (4.62)-(4.65) for K classes).
* Section 4.3, probabilistic discriminative models:
  logistic_loglik, logistic_grad_hess ((4.90)-(4.97)), irls_logistic (Newton-Raphson, the weighted
  least-squares form (4.99)-(4.101), optional L2 term), softmax and multiclass_irls
  ((4.104)-(4.110)), probit_irls (Fisher scoring for the inverse probit link (4.114)-(4.116)).
* Section 4.4, Laplace approximation: laplace_1d ((4.125)-(4.127)), laplace_logistic (the
  Gaussian posterior (4.142)-(4.143), the evidence (4.137) and the BIC (4.139)).
* Section 4.5.2, the predictive distribution: bayes_predictive ((4.150)-(4.155), the probit
  approximation sigma(a) ~ Phi(lambda a) with lambda^2 = pi/8 and kappa(sigma^2)).
* JSRandom (imported from ch11_sampling.py): the seeded generator DK.rng of the browser demos.
  The data sets below (blobs, make_two_class, make_three_class, make_fisher_data) draw from
  JSRandom in a fixed order, so linear-classifiers-demo.html and logistic-regression-demo.html
  reproduce the same points number for number.

Figures: 4.1, 4.2, 4.3, 4.4, 4.5, 4.6, 4.7, 4.9, 4.10, 4.11, 4.12, 4.13, 4.14. Run
`python3 ch04_linear_classification.py --list`, `python3 ch04_linear_classification.py` (all
figures, under a minute) or `python3 ch04_linear_classification.py --fig 4.4 4.6`. Files go to
../img/ch04/. Not reproduced: Figure 4.8 (a photograph of the Mark 1 perceptron).

The book states the settings of its synthetic data only in the figures. What we chose (and what
we did not copy) is listed here:
* Figure 4.4: two Gaussian classes in the plane, 40 points each. Red (x): mean (-1.3, 1.3),
  standard deviations (1.2, 1.0), correlation 0.3; blue (o): mean (1.0, -1.6), standard
  deviations (0.9, 1.2), correlation 0.7. The right panel adds 10 blue points around (7.7, -6)
  (standard deviations 0.5, correlation 0). The least-squares boundary is the line where the two
  fitted outputs y_1(x) = y_2(x) agree; logistic regression is fitted by IRLS (4.99).
* Figure 4.5: three Gaussian classes, 40 points each: red (-2.5, 2.3), sd (1.3, 1.2), rho 0.3;
  green (-0.4, 0), sd (1.0, 1.0), rho 0.8 (a strip along the diagonal between the other two);
  blue (2.0, -2.2), sd (1.1, 1.0), rho 0.3. The multiclass logistic regression of the right
  panel uses IRLS (4.109)-(4.110) with a tiny L2 term (alpha = 1e-4).
* Figure 4.6: two classes with the same covariance (sd (1.3, 0.45), correlation 0.3), blue mean
  (1.0, 3.0), red mean (3.8, 1.5), 100 points each. The left panel projects onto the line joining
  the class means, the right panel onto the Fisher direction (4.30).
* Figure 4.7: the eight-point-like data set is read off the book: five red points (+1) at (0.23, 0.72),
  (0.45, 0.30), (0.97, 0.90), (0.37, -0.49), (0.66, -0.90) and five blue points (-1) at (-0.75, 0.95),
  (-0.77, 0.61), (-0.55, 0.64), (-0.40, 0.73), (-0.10, -0.20); features phi = x without a bias,
  initial w = (-0.42, -0.90), points visited in this order, continuing cyclically after the
  last update. The sequence of updates then is the book's two (the figure's panels (a)-(d)).
* Figure 4.9: sigma(a) and Phi(lambda a) with lambda^2 = pi/8 (an exact reproduction).
* Figure 4.10: two Gaussian classes with means (-0.55, 0) and (0.55, 0), the shared covariance
  diag(0.4^2, 0.5^2) and equal priors (FIG10); surfaces are drawn with matplotlib's 3-D axes. The
  left surface is p(x|C_k) p(C_k) coloured by the posterior.
* Figure 4.11: three Gaussian classes: red (-1, -1) and green (-1, 1) share the covariance
  diag(0.3^2, 0.2^2); blue (1.5, 0) has diag(0.12^2, 0.32^2); equal priors; contours at 1, 2, 3
  standard deviations (in the right panel the RGB colour is the posterior vector).
* Figure 4.12: red cluster N((0, 0), 0.3^2 I) and two blue clusters around (-1.1, -0.95) and
  (0.95, 1.05); two Gaussian basis functions with centres (-1, -1) and (0, 0), width 0.78;
  logistic regression in the (phi_1, phi_2) plane (alpha = 0.01).
* Figure 4.13 is a schematic in the book (a mixture of two Gaussians and its cumulative
  distribution); we compute it (weights 0.7, 0.3; means 1.5, 3.0; sd 0.4, 0.3).
* Figure 4.14: exactly the book's p(z) proportional to exp(-z^2/2) sigma(20 z + 4). The mode is
  z_0 = 0.0775 and the exact curvature -d^2 ln f/dz^2 = 1 + 400 sigma(1 - sigma) = A = 2.54 (checked by
  finite differences in checks_ch04.py), so the Gaussian has peak sqrt(A / 2 pi) = 0.64. The red
  curve printed in the book is narrower (peak about 0.76, which corresponds to A of about 3.6);
  we plot the exact Laplace approximation, and the yellow density agrees with the book's.
* Figures 4.1, 4.2, 4.3 are geometric drawings in the book; they are drawn here from computed
  discriminants (planes, one-versus-the-rest, one-versus-one, the convex decision regions).
"""
import math

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from scipy import integrate
from scipy.special import expit, log_ndtr, ndtr
from scipy.stats import norm

import common
from common import RED, BLUE, GREEN, MAGENTA, GREY
from ch11_sampling import JSRandom  # noqa: F401  (re-exported: the demos' generator)

PINK, LIGHTGREEN, LIGHTBLUE = "#fde4e4", "#e4fbe4", "#e4e4fb"
LAMBDA2 = math.pi / 8.0   # lambda^2 in the probit approximation (4.153)


# =============================================================================================
# data sets (shared with the browser demos through JSRandom)
# =============================================================================================
def gauss2(rng, n, mean, sd, rho):
    """n draws from a bivariate normal with standard deviations sd, correlation rho.

    The Cholesky factor is L = [[s1, 0], [rho s2, s2 sqrt(1 - rho^2)]] and each point is
    mean + L (z1, z2) with z1 and z2 consecutive rng.norm() values: the demos do the same.
    """
    s1, s2 = sd
    out = np.empty((n, 2))
    for i in range(n):
        z1, z2 = rng.norm(), rng.norm()
        out[i] = (mean[0] + s1 * z1, mean[1] + rho * s2 * z1 + s2 * math.sqrt(1 - rho * rho) * z2)
    return out


#: settings of the synthetic sets: (mean, sd, rho) for each class
TWO_CLASS = {"red": ((-1.3, 1.3), (1.2, 1.0), 0.3), "blue": ((1.0, -1.6), (0.9, 1.2), 0.7)}
OUTLIERS = ((7.7, -6.0), (0.5, 0.5), 0.0)
THREE_CLASS = [((-2.5, 2.3), (1.3, 1.2), 0.3), ((-0.4, 0.0), (1.0, 1.0), 0.8), ((2.0, -2.2), (1.1, 1.0), 0.3)]
FISHER_SET = [((1.0, 3.0), (1.3, 0.45), 0.3), ((3.8, 1.5), (1.3, 0.45), 0.3)]   # blue, red
SEED_44, SEED_45, SEED_46 = 441, 451, 461
ALPHA_44 = 1e-3   # tiny L2 term for the logistic fit of Figure 4.4 (the classes are linearly separable)


def make_two_class(seed=SEED_44, n=40, outliers=0):
    """Figure 4.4 data: X (N x 2), class labels y in {0 (red), 1 (blue)}. outliers extra blue points."""
    rng = JSRandom(seed)
    Xr = gauss2(rng, n, *TWO_CLASS["red"])
    Xb = gauss2(rng, n, *TWO_CLASS["blue"])
    X, y = np.vstack([Xr, Xb]), np.r_[np.zeros(n, int), np.ones(n, int)]
    if outliers:
        Xo = gauss2(rng, outliers, *OUTLIERS)
        X, y = np.vstack([X, Xo]), np.r_[y, np.ones(outliers, int)]
    return X, y


def make_three_class(seed=SEED_45, n=40):
    """Figure 4.5 data: three classes 0 (red), 1 (green), 2 (blue)."""
    rng = JSRandom(seed)
    X = np.vstack([gauss2(rng, n, *c) for c in THREE_CLASS])
    return X, np.repeat(np.arange(3), n)


def make_fisher_data(seed=SEED_46, n=100):
    """Figure 4.6 data: class 0 (blue) and class 1 (red) with the same covariance."""
    rng = JSRandom(seed)
    X = np.vstack([gauss2(rng, n, *c) for c in FISHER_SET])
    return X, np.r_[np.zeros(n, int), np.ones(n, int)]


def blobs(rng, n, sep, sd=1.0):
    """Two overlapping classes, as in logistic-regression-demo.html: class 0 around -sep (1, -0.5),
    class 1 around +sep (1, -0.5), isotropic standard deviation sd; n points each."""
    X = np.empty((2 * n, 2))
    for k in (0, 1):
        m = (2 * k - 1) * sep
        for i in range(n):
            X[k * n + i] = (m * 1.0 + sd * rng.norm(), -0.5 * m + sd * rng.norm())
    return X, np.r_[np.zeros(n, int), np.ones(n, int)]


PERCEPTRON_RED = [(0.23, 0.72), (0.45, 0.30), (0.97, 0.90), (0.37, -0.49), (0.66, -0.90)]
PERCEPTRON_BLUE = [(-0.75, 0.95), (-0.77, 0.61), (-0.55, 0.64), (-0.40, 0.73), (-0.10, -0.20)]
PERCEPTRON_W0 = (-0.42, -0.90)


def design(X, bias=True):
    """Feature matrix with a leading column of ones (the bias w_0)."""
    X = np.atleast_2d(X)
    return np.hstack([np.ones((len(X), 1)), X]) if bias else X


def one_hot(y, K=None):
    K = K or int(y.max()) + 1
    T = np.zeros((len(y), K))
    T[np.arange(len(y)), y] = 1.0
    return T


# =============================================================================================
# 4.1 discriminant functions
# =============================================================================================
def least_squares_classifier(Phi, T):
    """Section 4.1.3: W = (Phi^T Phi)^-1 Phi^T T = Phi^+ T (4.16); y(x) = W^T phi(x) (4.14)-(4.15)."""
    return np.linalg.pinv(Phi) @ T


def class_means_scatter(X, y):
    """Class means m_k, within-class scatter S_W (4.28)/(4.40), between-class scatter S_B (4.46)."""
    K, D = int(y.max()) + 1, X.shape[1]
    m = X.mean(0)
    mk = np.array([X[y == k].mean(0) for k in range(K)])
    SW, SB = np.zeros((D, D)), np.zeros((D, D))
    for k in range(K):
        d = X[y == k] - mk[k]
        SW += d.T @ d
        SB += (y == k).sum() * np.outer(mk[k] - m, mk[k] - m)
    return mk, SW, SB


def fisher_two_class(X, y):
    """Section 4.1.4: w proportional to S_W^-1 (m_2 - m_1) (4.30); returns the unit vector, m_1, m_2, S_W."""
    mk, SW, _ = class_means_scatter(X, y)
    w = np.linalg.solve(SW, mk[1] - mk[0])
    return w / np.linalg.norm(w), mk, SW


def fisher_criterion(w, SW, SB):
    """J(w) = (w^T S_B w) / (w^T S_W w) (4.26)."""
    return float(w @ SB @ w / (w @ SW @ w))


def fisher_multiclass(X, y, dim=None):
    """Section 4.1.6: the columns of W are the leading eigenvectors of S_W^-1 S_B ((4.50) and the text
    after (4.50)); at most K - 1 of them have non-zero eigenvalue. Returns W (D x dim), eigenvalues."""
    _, SW, SB = class_means_scatter(X, y)
    vals, vecs = np.linalg.eig(np.linalg.solve(SW, SB))
    idx = np.argsort(-vals.real)
    vals, vecs = vals.real[idx], vecs.real[:, idx]
    dim = dim or int(y.max())
    return vecs[:, :dim], vals


def perceptron(Phi, t, w0=None, max_epochs=1000, order=None):
    """Section 4.1.7: the perceptron criterion E_P(w) = -sum_{misclassified} w^T phi_n t_n (4.54), minimised by
    stochastic gradient descent with learning rate 1: w <- w + phi_n t_n (4.55) for a misclassified n.

    The points are visited in the given order (default 0, 1, ...), cyclically and continuing after the
    last update, as in the book. t in {-1, +1}. Returns w and a history of (w before, n, w after).
    """
    N = len(t)
    w = np.zeros(Phi.shape[1]) if w0 is None else np.array(w0, float)
    hist, n, clean = [], 0, 0
    for _ in range(max_epochs * N):
        if t[n] * (Phi[n] @ w) <= 0:
            hist.append((w.copy(), n, w + t[n] * Phi[n]))
            w = w + t[n] * Phi[n]
            clean = 0
        else:
            clean += 1
            if clean >= N:
                break
        n = (n + 1) % N
    return w, hist


# =============================================================================================
# 4.2 probabilistic generative models
# =============================================================================================
def gaussian_generative_fit(X, y):
    """Section 4.2.2 and 4.2.3: class-conditional Gaussians with a SHARED covariance, maximum likelihood.
    pi_k = N_k / N (4.75-4.76), mu_k = class mean (4.77-4.78), Sigma = sum_k (N_k / N) S_k (4.79-4.80)."""
    K = int(y.max()) + 1
    N = len(y)
    pi = np.array([(y == k).mean() for k in range(K)])
    mu = np.array([X[y == k].mean(0) for k in range(K)])
    Sigma = sum(((X[y == k] - mu[k]).T @ (X[y == k] - mu[k])) for k in range(K)) / N
    return pi, mu, Sigma


def generative_posterior_params(pi, mu, Sigma):
    """Weights of the posterior. Two classes: p(C_1 | x) = sigma(w^T x + w_0) with
    w = Sigma^-1 (mu_1 - mu_2) (4.66) and w_0 = -1/2 mu_1^T Sigma^-1 mu_1 + 1/2 mu_2^T Sigma^-1 mu_2
    + ln p(C_1)/p(C_2) (4.67), the two classes being indexed 0 (= C_1 here is class 0).
    General K: w_k = Sigma^-1 mu_k (4.65), w_k0 = -1/2 mu_k^T Sigma^-1 mu_k + ln pi_k (4.65); softmax (4.62)."""
    Si = np.linalg.inv(Sigma)
    W = (Si @ mu.T).T
    w0 = -0.5 * np.einsum("kd,de,ke->k", mu, Si, mu) + np.log(pi)
    return W, w0


def softmax(A, axis=-1):
    """Normalised exponential (4.104)."""
    A = A - A.max(axis, keepdims=True)
    E = np.exp(A)
    return E / E.sum(axis, keepdims=True)


# =============================================================================================
# 4.3 probabilistic discriminative models
# =============================================================================================
def ln_sigmoid(a):
    return -np.logaddexp(0.0, -a)


def logistic_loglik(w, Phi, t, alpha=0.0):
    """ln p(t | w) = sum_n {t_n ln y_n + (1 - t_n) ln(1 - y_n)} = -E(w) (4.89-4.90), minus alpha/2 |w|^2 if alpha > 0."""
    a = Phi @ w
    return float(np.sum(t * ln_sigmoid(a) + (1 - t) * ln_sigmoid(-a)) - 0.5 * alpha * w @ w)


def logistic_grad_hess(w, Phi, t, alpha=0.0):
    """Gradient (4.91)/(4.96) and Hessian (4.97) of the NEGATIVE log likelihood (plus alpha/2 |w|^2):
    grad E = Phi^T (y - t) + alpha w, H = Phi^T R Phi + alpha I with R = diag(y_n (1 - y_n))."""
    y = expit(Phi @ w)
    R = y * (1 - y)
    return Phi.T @ (y - t) + alpha * w, (Phi * R[:, None]).T @ Phi + alpha * np.eye(Phi.shape[1])


def irls_logistic(Phi, t, alpha=0.0, w0=None, max_iter=100, tol=1e-10, store=True):
    """Iterative reweighted least squares for logistic regression (Section 4.3.3).

    Newton-Raphson step w <- w - H^-1 grad E (4.92) with the logistic gradient and Hessian (4.96)-(4.97);
    equivalently the weighted least-squares problem w = (Phi^T R Phi)^-1 Phi^T R z with
    z = Phi w_old - R^-1 (y - t) (4.99)-(4.101). alpha > 0 adds the Gaussian prior N(0, alpha^-1 I)
    (Section 4.5.1), which keeps w finite when the classes are linearly separable.

    Returns w and a dict with the path of w, log likelihood and iterations used.
    """
    w = np.zeros(Phi.shape[1]) if w0 is None else np.array(w0, float)
    path, ll = [w.copy()], [logistic_loglik(w, Phi, t, alpha)]
    for it in range(max_iter):
        g, H = logistic_grad_hess(w, Phi, t, alpha)
        step = np.linalg.solve(H, g)
        w = w - step
        path.append(w.copy())
        ll.append(logistic_loglik(w, Phi, t, alpha))
        if np.abs(step).max() < tol:
            break
    return w, {"path": np.array(path), "loglik": np.array(ll), "iters": len(path) - 1}


def multiclass_grad_hess(W, Phi, T, alpha=0.0):
    """Multiclass logistic regression, W is M x K. Gradient (4.109): grad_{w_j} E = sum_n (y_nj - t_nj) phi_n;
    Hessian blocks (4.110): H_jk = sum_n y_nk (I_kj - y_nj) phi_n phi_n^T (the leading minus sign in the
    book's (4.110) is an erratum). Returns the gradient flattened column by column (class by class) and the (MK x MK) Hessian."""
    M, K = W.shape
    Y = softmax(Phi @ W)
    G = Phi.T @ (Y - T) + alpha * W
    H = np.zeros((K * M, K * M))
    for j in range(K):
        for k in range(K):
            r = Y[:, k] * ((k == j) - Y[:, j])
            H[j * M:(j + 1) * M, k * M:(k + 1) * M] = (Phi * r[:, None]).T @ Phi
    H += alpha * np.eye(K * M)
    return G.T.ravel(), H


def multiclass_loglik(W, Phi, T, alpha=0.0):
    A = Phi @ W
    lse = np.log(np.exp(A - A.max(1, keepdims=True)).sum(1)) + A.max(1)
    return float(np.sum(T * (A - lse[:, None])) - 0.5 * alpha * np.sum(W * W))


def multiclass_irls(Phi, T, alpha=1e-4, max_iter=100, tol=1e-9):
    """Newton-Raphson for multiclass logistic regression (Section 4.3.4, (4.109)-(4.110)); returns W (M x K)."""
    M, K = Phi.shape[1], T.shape[1]
    W = np.zeros((M, K))
    for _ in range(max_iter):
        g, H = multiclass_grad_hess(W, Phi, T, alpha)
        step = np.linalg.solve(H, g)
        W = W - step.reshape(K, M).T
        if np.abs(step).max() < tol:
            break
    return W


def probit_loglik(w, Phi, t, alpha=0.0):
    a = Phi @ w
    return float(np.sum(t * log_ndtr(a) + (1 - t) * log_ndtr(-a)) - 0.5 * alpha * w @ w)


def probit_grad(w, Phi, t, alpha=0.0):
    """Gradient of the negative log likelihood for y_n = Phi(a_n) (4.114): sum_n (y_n - t_n) phi(a_n) / (y_n (1 - y_n)) phi_n,
    using dy/da = N(a | 0, 1) (Section 4.3.5)."""
    a = Phi @ w
    y = np.clip(ndtr(a), 1e-300, 1 - 1e-16)
    d = norm.pdf(a)
    return Phi.T @ ((y - t) * d / (y * (1 - y))) + alpha * w


def probit_irls(Phi, t, alpha=0.0, w0=None, max_iter=100, tol=1e-10):
    """IRLS for probit regression: Fisher scoring, i.e. Newton's method with the expected Hessian
    Phi^T diag(N(a_n)^2 / (y_n (1 - y_n))) Phi (Section 4.3.5: 'the same IRLS applies with a different R')."""
    w = np.zeros(Phi.shape[1]) if w0 is None else np.array(w0, float)
    ll = [probit_loglik(w, Phi, t, alpha)]
    for it in range(max_iter):
        a = Phi @ w
        y = np.clip(ndtr(a), 1e-300, 1 - 1e-16)
        W = norm.pdf(a) ** 2 / (y * (1 - y))
        H = (Phi * W[:, None]).T @ Phi + alpha * np.eye(Phi.shape[1])
        step = np.linalg.solve(H, probit_grad(w, Phi, t, alpha))
        w = w - step
        ll.append(probit_loglik(w, Phi, t, alpha))
        if np.abs(step).max() < tol:
            break
    return w, {"loglik": np.array(ll), "iters": len(ll) - 1}


# =============================================================================================
# 4.4 Laplace approximation, 4.5 Bayesian logistic regression
# =============================================================================================
def laplace_1d(ln_f, dln_f, d2ln_f, z_init=0.0, max_iter=100):
    """Section 4.4: for p(z) = f(z)/Z find a mode z_0 (dln f = 0 by Newton) and A = -d^2 ln f / dz^2 at z_0 (4.127);
    then q(z) = N(z | z_0, A^-1) (4.128) and Z ~ f(z_0) sqrt(2 pi / A) (4.126). Returns z_0, A, ln Z."""
    z = z_init
    for _ in range(max_iter):
        step = dln_f(z) / d2ln_f(z)
        z -= step
        if abs(step) < 1e-13:
            break
    A = -d2ln_f(z)
    return z, A, ln_f(z) + 0.5 * math.log(2 * math.pi / A)


def fig414_functions(a=20.0, b=4.0):
    """ln f, its derivatives for f(z) = exp(-z^2/2) sigma(a z + b) (Figure 4.14 has a = 20, b = 4)."""
    ln_f = lambda z: -0.5 * z * z + ln_sigmoid(a * z + b)
    dln_f = lambda z: -z + a * expit(-(a * z + b))
    d2ln_f = lambda z: -1.0 - a * a * expit(a * z + b) * expit(-(a * z + b))
    return ln_f, dln_f, d2ln_f


def laplace_logistic(Phi, t, alpha, w_init=None):
    """Section 4.5.1: Gaussian prior N(w | 0, alpha^-1 I); the posterior mode w_MAP maximises (4.142) (by IRLS)
    and q(w) = N(w | w_MAP, S_N) with S_N^-1 = alpha I + Phi^T R Phi (4.143, erratum: S_N should read S_N^-1).
    Also returns the Laplace evidence (4.137) ln p(D) ~ ln p(D | w_MAP) + ln p(w_MAP) + M/2 ln 2 pi - 1/2 ln |A|."""
    M = Phi.shape[1]
    w, info = irls_logistic(Phi, t, alpha=alpha, w0=w_init)
    _, H = logistic_grad_hess(w, Phi, t, alpha)
    S = np.linalg.inv(H)
    ln_prior = 0.5 * M * math.log(alpha / (2 * math.pi)) - 0.5 * alpha * w @ w
    logev = logistic_loglik(w, Phi, t) + ln_prior + 0.5 * M * math.log(2 * math.pi) - 0.5 * np.linalg.slogdet(H)[1]
    return w, S, logev


def bic_logistic(Phi, t):
    """BIC (4.139): ln p(D) ~ ln p(D | w_ML) - M/2 ln N (here ln p(D | w_ML) is the maximum log likelihood)."""
    w, _ = irls_logistic(Phi, t, alpha=1e-10)
    return logistic_loglik(w, Phi, t) - 0.5 * Phi.shape[1] * math.log(len(t))


def kappa(s2):
    """kappa(sigma^2) = (1 + pi sigma^2 / 8)^(-1/2) (4.154)."""
    return 1.0 / np.sqrt(1.0 + math.pi * s2 / 8.0)


def bayes_predictive(Phi_new, w_map, S):
    """Section 4.5.2: p(C_1 | phi, t) ~ sigma(kappa(sigma_a^2) mu_a) (4.155) with mu_a = w_MAP^T phi (4.149) and
    sigma_a^2 = phi^T S_N phi (4.150), from the convolution of a sigmoid with a Gaussian approximated by the
    probit approximation (4.151)-(4.153): sigma(a) ~ Phi(lambda a), lambda^2 = pi/8."""
    mu = Phi_new @ w_map
    s2 = np.einsum("nd,de,ne->n", Phi_new, S, Phi_new)
    return expit(kappa(s2) * mu)


# =============================================================================================
# plotting helpers
# =============================================================================================
def scatter_classes(ax, X, y, ms=22, lw=1.3):
    """Red crosses, green plus signs, blue open circles (as in the book), by class index 0, 1, 2."""
    style = [dict(marker="x", color=RED, s=ms, linewidths=lw), dict(marker="o", facecolors="none", edgecolors=BLUE, s=ms, linewidths=lw),
             dict(marker="+", color=GREEN, s=ms * 1.4, linewidths=lw * 1.2)]
    for k in sorted(set(y)):
        st = style[k] if len(set(y)) == 2 else [style[0], style[2], style[1]][k]
        ax.scatter(X[y == k, 0], X[y == k, 1], zorder=3, **st)


def grid2d(xlim, ylim, n=300):
    xs, ys = np.linspace(*xlim, n), np.linspace(*ylim, n)
    XX, YY = np.meshgrid(xs, ys)
    return XX, YY, np.c_[XX.ravel(), YY.ravel()]


def regions(ax, XX, YY, cls, K):
    """Pale background colours per class with black boundaries."""
    cols = [PINK, LIGHTGREEN, LIGHTBLUE]
    cmap = ListedColormap(cols[:K])
    ax.contourf(XX, YY, cls, levels=np.arange(-0.5, K, 1.0), cmap=cmap, zorder=0, antialiased=True)
    ax.contour(XX, YY, cls, levels=np.arange(0.5, K - 1, 1.0), colors="k", linewidths=0.9, zorder=1)


def newaxes(ax, xlim, ylim, ticks=None):
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.set_aspect("equal")
    if ticks:
        ax.set_xticks(ticks[0])
        ax.set_yticks(ticks[1])


# =============================================================================================
# figures
# =============================================================================================
def fig4_1(out_dir=None):
    """Figure 4.1: geometry of a linear discriminant function in two dimensions (computed from w, w_0, x)."""
    fig, ax = plt.subplots(figsize=(4.8, 3.9))
    w, w0 = np.array([0.55, 0.85]), -0.8
    x = np.array([2.6, 1.9])
    nw = np.linalg.norm(w)
    xperp = x - (w @ x + w0) / nw ** 2 * w          # orthogonal projection onto y = 0 (4.6)
    d = np.array([-w[1], w[0]]) / nw
    p0 = -w0 / nw ** 2 * w                           # the point of the plane closest to the origin, at distance -w0/|w| (4.5)
    line = p0[:, None] + d[:, None] * np.array([-2.9, 2.2])
    ax.plot(line[0], line[1], color=RED, lw=1.8)
    arrow = dict(arrowstyle="-|>", mutation_scale=11, lw=1.5)
    ax.annotate("", xy=w / nw * 1.6, xytext=(0, 0), arrowprops=dict(arrow, color=GREEN))
    ax.annotate("", xy=x, xytext=(0, 0), arrowprops=dict(arrow, color=BLUE))
    ax.plot([x[0], xperp[0]], [x[1], xperp[1]], "--", color=BLUE, lw=1.4)
    ax.plot([0, p0[0]], [0, p0[1]], color="k", lw=1.2)
    ax.annotate("", xy=(2.8, 0), xytext=(-0.3, 0), arrowprops=dict(arrowstyle="-|>", color="k", lw=0.9))
    ax.annotate("", xy=(0, 3.2), xytext=(0, -0.3), arrowprops=dict(arrowstyle="-|>", color="k", lw=0.9))
    ax.text(2.85, -0.05, "$x_1$", fontsize=11, va="top")
    ax.text(-0.08, 3.25, "$x_2$", fontsize=11, ha="right")
    ax.text(-0.05, 1.05, r"$\mathbf{w}$", color=GREEN, fontsize=11, ha="right")
    ax.text(x[0] + 0.05, x[1] + 0.08, r"$\mathbf{x}$", color=BLUE, fontsize=11)
    ax.text(xperp[0] - 0.1, xperp[1] + 0.22, r"$\mathbf{x}_\perp$", fontsize=11, ha="right")
    ax.text(0.5 * (x[0] + xperp[0]) + 0.1, 0.5 * (x[1] + xperp[1]) + 0.05, r"$y(\mathbf{x})/\|\mathbf{w}\|$ = %.2f" % ((w @ x + w0) / nw),
            color=BLUE, fontsize=9, va="center")
    ax.annotate(r"$-w_0/\|\mathbf{w}\|$ = %.2f" % (-w0 / nw), xy=0.5 * p0, xytext=(0.7, -0.75), fontsize=9, va="center",
                arrowprops=dict(arrowstyle="-", color="k", lw=0.6))
    ax.text(-1.9, 2.75, "$y>0$", fontsize=10)
    ax.text(-1.9, 2.45, "$y=0$", fontsize=10, color=RED)
    ax.text(-1.9, 2.15, "$y<0$", fontsize=10)
    ax.text(1.0, 2.9, r"$\mathcal{R}_1$", fontsize=11)
    ax.text(-1.0, 0.6, r"$\mathcal{R}_2$", fontsize=11)
    ax.set_xlim(-2.1, 3.3)
    ax.set_ylim(-1.2, 3.5)
    ax.set_aspect("equal")
    ax.axis("off")
    ax.set_title(r"$y(\mathbf{x}) = \mathbf{w}^{\mathrm{T}}\mathbf{x} + w_0$, $\mathbf{w} = (%.2f, %.2f)$, $w_0 = %.1f$" % (w[0], w[1], w0), fontsize=9.5)
    common.savefig(fig, 4, "fig4-1", out_dir)


def _centroid_labels(ax, XX, YY, winner, K, fontsize=12):
    for k in range(K):
        m = winner == k
        if m.any():
            ax.text(XX[m].mean(), YY[m].mean(), r"$\mathcal{R}_%d$" % (k + 1), fontsize=fontsize, ha="center", va="center")


def _ambiguous_panel(ax, kind):
    """One-versus-the-rest (kind 'ovr': discriminants C1 vs not C1 and C2 vs not C2) or one-versus-one ('ovo': three
    pairwise discriminants, majority vote) and the ambiguous region (green)."""
    xlim, ylim = (-3, 3), (-3, 3)
    XX, YY, P = grid2d(xlim, ylim, 500)
    x1, x2 = P[:, 0], P[:, 1]
    if kind == "ovr":
        y1 = 0.9 * x1 + 0.5 * x2 - 0.1
        y2 = -0.6 * x1 + 0.8 * x2 + 0.1
        amb = (y1 > 0) & (y2 > 0)                            # both claim the point
        winner = np.where(amb, -1, np.where(y1 > 0, 0, np.where(y2 > 0, 1, 2)))   # R3 = neither claims it
        lines = [y1, y2]
    else:
        V = np.array([[0.0, 1.3], [-1.3, -0.9], [1.3, -0.9]])
        def inside(a, b, c):                                  # signed distance to the edge a-b, positive on the side of c
            n = np.array([-(b - a)[1], (b - a)[0]])
            n = n / np.linalg.norm(n) * np.sign(n @ (c - a))
            return (P - a) @ n
        y12 = inside(V[1], V[2], V[0])                        # C1 beats C2 inside the triangle
        y23 = inside(V[0], V[1], V[2])                        # C2 beats C3 inside
        y31 = inside(V[0], V[2], V[1])                        # C3 beats C1 inside
        votes = np.stack([(y12 > 0).astype(int) + (y31 <= 0), (y12 <= 0).astype(int) + (y23 > 0), (y23 <= 0).astype(int) + (y31 > 0)], 1)
        amb = (votes == 1).all(1)
        winner = np.where(amb, -1, votes.argmax(1))
        lines = [y12, y23, y31]
    W = winner.reshape(XX.shape)
    ax.contourf(XX, YY, (W == -1).astype(float), levels=[0.5, 1.5], colors=["#4dff4d"], zorder=0)
    for L in lines:
        ax.contour(XX, YY, L.reshape(XX.shape), levels=[0], colors=RED, linewidths=1.8)
    _centroid_labels(ax, XX, YY, W, 3)
    ax.text(XX[W == -1].mean(), YY[W == -1].mean(), "?", fontsize=13, ha="center", va="center")
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.set_aspect("equal")
    ax.axis("off")


def fig4_2(out_dir=None):
    """Figure 4.2: ambiguous regions of one-versus-the-rest (left) and one-versus-one (right) discriminants."""
    fig, axes = common.panel_grid(1, 2, size=(3.6, 3.7))
    _ambiguous_panel(axes[0, 0], "ovr")
    _ambiguous_panel(axes[0, 1], "ovo")
    axes[0, 0].set_title("(a) one versus the rest:\ngreen = claimed by both $\\mathcal{C}_1$ and $\\mathcal{C}_2$", fontsize=8.5)
    axes[0, 1].set_title("(b) one versus one:\ngreen = one vote for each class", fontsize=8.5)
    common.savefig(fig, 4, "fig4-2", out_dir)


def fig4_3(out_dir=None):
    """Figure 4.3: decision regions of y_k = w_k^T x + w_k0 are singly connected and convex (computed argmax)."""
    fig, ax = plt.subplots(figsize=(4.4, 3.6))
    W = np.array([[1.0, 0.2], [-0.6, 0.9], [-0.4, -1.1]])
    b = np.array([0.0, 0.1, -0.1])
    XX, YY, P = grid2d((-3, 3), (-2.5, 2.5), 400)
    cls = (P @ W.T + b).argmax(1).reshape(XX.shape)
    Yk = P @ W.T + b
    for kk in range(3):
        other = np.delete(Yk, kk, axis=1).max(1)
        ax.contour(XX, YY, (Yk[:, kk] - other).reshape(XX.shape), levels=[0], colors=RED, linewidths=2.2)
    for s, pos in [("i", (-2.6, 1.0)), ("j", (0.2, 1.9)), ("k", (-0.3, -1.5))]:
        ax.text(*pos, r"$\mathcal{R}_%s$" % s, fontsize=13, ha="center")
    # pick A and B in the same region k so that the segment stays inside R_k (convexity)
    k = 2
    cand = [(-2.5, -0.9), (0.5, -2.0)]
    xa, xb = np.array(cand[0]), np.array(cand[1])
    assert (np.array([xa, xb]) @ W.T + b).argmax(1).tolist() == [k, k]
    xh = 0.4 * xa + 0.6 * xb
    ax.plot([xa[0], xb[0]], [xa[1], xb[1]], color=BLUE, lw=2.2)
    for p, s in [(xa, r"$\mathbf{x}_A$"), (xb, r"$\mathbf{x}_B$"), (xh, r"$\widehat{\mathbf{x}}$")]:
        ax.plot(*p, "o", color=BLUE, ms=6)
        ax.text(p[0], p[1] - 0.45, s, fontsize=11, ha="center", va="top")
    ax.set_xlim(-3, 3)
    ax.set_ylim(-2.5, 2.5)
    ax.set_aspect("equal")
    ax.axis("off")
    ax.set_title(r"$\widehat{\mathbf{x}} = 0.4\,\mathbf{x}_A + 0.6\,\mathbf{x}_B$ stays in $\mathcal{R}_k$ because $y_k$ is linear", fontsize=8.5)
    common.savefig(fig, 4, "fig4-3", out_dir)


def fig4_4(out_dir=None):
    """Figure 4.4: least squares (magenta) and logistic regression (green) on two Gaussian classes, without and with outliers."""
    fig, axes = common.panel_grid(1, 2, size=(3.6, 3.7))
    xlim, ylim = (-4, 9), (-9, 4)
    XX, YY, P = grid2d(xlim, ylim, 400)
    for ax, nout, lab in [(axes[0, 0], 0, "a"), (axes[0, 1], 10, "b")]:
        X, y = make_two_class(outliers=nout)
        Phi = design(X)
        W = least_squares_classifier(Phi, one_hot(y))
        f = (design(P) @ W)
        ax.contour(XX, YY, (f[:, 0] - f[:, 1]).reshape(XX.shape), levels=[0], colors=MAGENTA, linewidths=1.8, zorder=2)
        w, _ = irls_logistic(Phi, y.astype(float), alpha=ALPHA_44)  # class 1 (blue) is t = 1; the classes are separable, so a tiny L2 term
        ax.contour(XX, YY, (design(P) @ w).reshape(XX.shape), levels=[0], colors=GREEN, linewidths=1.8, zorder=2)
        scatter_classes(ax, X, y)
        newaxes(ax, xlim, ylim, ([-4, -2, 0, 2, 4, 6, 8], [-8, -6, -4, -2, 0, 2, 4]))
        common.panel_label(ax, lab, "lower left")
        err = lambda f_: int(((f_[:, 0] - f_[:, 1] < 0).astype(int) != y).sum())
        print("  Fig 4.4(%s): N = %d, least-squares errors %d, logistic errors %d, |w_logistic| = %.2f" %
              (lab, len(y), err(Phi @ W), int(((Phi @ w > 0).astype(int) != y).sum()), np.linalg.norm(w)))
    common.savefig(fig, 4, "fig4-4", out_dir)


def fig4_5(out_dir=None):
    """Figure 4.5: three classes; least squares (left) leaves the green class a sliver, logistic regression (right) separates them."""
    fig, axes = common.panel_grid(1, 2, size=(3.7, 3.7))
    xlim, ylim = (-6, 6), (-6, 6)
    XX, YY, P = grid2d(xlim, ylim, 400)
    X, y = make_three_class()
    Phi, T = design(X), one_hot(y)
    W_ls = least_squares_classifier(Phi, T)
    W_lg = multiclass_irls(Phi, T, alpha=1e-4)
    for ax, W, lab in [(axes[0, 0], W_ls, "a"), (axes[0, 1], W_lg, "b")]:
        cls = (design(P) @ W).argmax(1).reshape(XX.shape)
        regions(ax, XX, YY, cls, 3)
        scatter_classes(ax, X, y)
        newaxes(ax, xlim, ylim, ([-6, -4, -2, 0, 2, 4, 6], [-6, -4, -2, 0, 2, 4, 6]))
        common.panel_label(ax, lab, "lower left")
        acc = ((Phi @ W).argmax(1) == y).mean()
        print("  Fig 4.5(%s): training accuracy %.3f, area share of green region %.3f" % (lab, acc, (cls == 1).mean()))
    common.savefig(fig, 4, "fig4-5", out_dir)


def _proj_hist(ax, X, y, u, bins=18, height=1.0, gap=1.6, colors=(BLUE, RED)):
    """Draw, on a line below/left of the data, the histograms of the projections u^T x of the two classes
    (bars stand perpendicular to the line, as in the book)."""
    u = u / np.linalg.norm(u)
    n = np.array([-u[1], u[0]])
    if n @ (X.mean(0)) < 0:
        n = -n
    c = (X @ n).min() - gap
    s = X @ u
    edges = np.linspace(s.min(), s.max(), bins + 1)
    cmax = max(np.histogram(s[y == k], edges)[0].max() for k in (0, 1))
    P = lambda sv, h: sv * u + (c + h) * n
    ax.plot(*np.c_[P(edges[0] - 0.2, 0), P(edges[-1] + 0.2, 0)], "k--", lw=0.9, zorder=1)
    for k in (0, 1):
        cnt = np.histogram(s[y == k], edges)[0]
        for i, ci in enumerate(cnt):
            if ci == 0:
                continue
            h = height * ci / cmax
            poly = np.array([P(edges[i], 0), P(edges[i + 1], 0), P(edges[i + 1], h), P(edges[i], h)])
            ax.fill(poly[:, 0], poly[:, 1], facecolor=colors[k], edgecolor="k", lw=0.5, alpha=0.9, zorder=2)
    return P, c, n


def fig4_6(out_dir=None):
    """Figure 4.6: projection onto the line joining the class means (overlap) and onto the Fisher direction (good separation)."""
    X, y = make_fisher_data()
    wF, mk, SW = fisher_two_class(X, y)
    _, _, SB = class_means_scatter(X, y)
    d = (mk[1] - mk[0]) / np.linalg.norm(mk[1] - mk[0])
    fig, axes = common.panel_grid(1, 2, size=(3.7, 3.0))
    for ax, u, lab in [(axes[0, 0], d, "a"), (axes[0, 1], wF, "b")]:
        ax.scatter(X[y == 0, 0], X[y == 0, 1], s=3, color=BLUE, zorder=3)
        ax.scatter(X[y == 1, 0], X[y == 1, 1], s=3, color=RED, zorder=3)
        P, c, n = _proj_hist(ax, X, y, u, height=0.9, gap=1.1)
        for m in mk:
            ax.plot(*m, "+", color="k", ms=7, zorder=4)
            foot = (m @ u) * u + c * n
            ax.plot([m[0], foot[0]], [m[1], foot[1]], color=GREEN, lw=0.9, zorder=2)
        ax.plot(*np.c_[mk[0], mk[1]], color=GREEN, lw=0.9, zorder=2)
        newaxes(ax, (-3.5, 7.5), (-3.5, 4.5), ([-2, 2, 6], [-2, 0, 2, 4]))
        common.panel_label(ax, lab, "lower right")
        print("  Fig 4.6(%s): J(w) = %.3f for u = (%.3f, %.3f)" % (lab, fisher_criterion(u, SW, SB), u[0], u[1]))
    common.savefig(fig, 4, "fig4-6", out_dir)


def fig4_7(out_dir=None):
    """Figure 4.7: the perceptron learning algorithm on ten points (two updates, panels a to d)."""
    Xr, Xb = np.array(PERCEPTRON_RED), np.array(PERCEPTRON_BLUE)
    X = np.vstack([Xr, Xb])
    t = np.r_[np.ones(len(Xr)), -np.ones(len(Xb))]
    w, hist = perceptron(X, t, w0=PERCEPTRON_W0)
    print("  Fig 4.7: %d updates, final w = (%.2f, %.2f)" % (len(hist), w[0], w[1]))
    fig, axes = common.panel_grid(2, 2, size=(3.3, 3.3))
    seq = []   # (w to draw, index to circle, update vector or None)
    for (w_before, n, w_after) in hist:
        seq.append((w_before, n))
    seq.append((w, None))
    panels = [(0, True), (1, False), (1, True), (2, False)]   # (index in seq, show the update arrow)
    letters = "abcd"
    for ax, (k, show), lab in zip(axes.ravel(), panels, letters):
        wk, n = seq[min(k, len(seq) - 1)]
        ax.scatter(Xr[:, 0], Xr[:, 1], s=22, color=RED, zorder=3)
        ax.scatter(Xb[:, 0], Xb[:, 1], s=22, color=BLUE, zorder=3)
        d = np.array([-wk[1], wk[0]]) / np.linalg.norm(wk)
        ax.plot([-3 * d[0], 3 * d[0]], [-3 * d[1], 3 * d[1]], "k", lw=1.5)
        sc = 1.0
        ax.annotate("", xy=wk * sc, xytext=(0, 0), arrowprops=dict(arrowstyle="-|>", color="k", lw=1.6, mutation_scale=12), zorder=4)
        if show and n is not None:
            ax.add_patch(plt.Circle(X[n], 0.06, fill=False, ec="#00ee00", lw=1.8, zorder=5))
            ax.annotate("", xy=X[n] * t[n], xytext=(0, 0), arrowprops=dict(arrowstyle="-|>", color=RED, lw=1.3, mutation_scale=10), zorder=4)
            ax.annotate("", xy=wk * sc + X[n] * t[n], xytext=wk * sc, arrowprops=dict(arrowstyle="-|>", color=RED, lw=1.3, mutation_scale=10), zorder=4)
        newaxes(ax, (-1, 1), (-1, 1), ([-1, -0.5, 0, 0.5, 1], [-1, -0.5, 0, 0.5, 1]))
        common.panel_label(ax, lab, "lower left")
    common.savefig(fig, 4, "fig4-7", out_dir)


def fig4_9(out_dir=None):
    """Figure 4.9: the logistic sigmoid (red) and the scaled inverse probit function Phi(lambda a), lambda^2 = pi/8 (dashed blue)."""
    a = np.linspace(-8, 8, 800)
    fig, ax = plt.subplots(figsize=(4.4, 3.0))
    ax.axhline(0.5, color="k", ls=":", lw=0.7)
    ax.axvline(0, color="k", ls=":", lw=0.7)
    ax.plot(a, expit(a), color=RED, lw=1.8, label=r"$\sigma(a)$")
    ax.plot(a, ndtr(math.sqrt(LAMBDA2) * a), "--", color=BLUE, lw=1.8, label=r"$\Phi(\lambda a)$, $\lambda^2=\pi/8$")
    ax.set_xlim(-8, 8)
    ax.set_ylim(0, 1)
    ax.set_xticks([-5, 0, 5])
    ax.set_yticks([0, 0.5, 1])
    ax.legend(loc="lower right", frameon=False)
    print("  Fig 4.9: max |sigma(a) - Phi(lambda a)| = %.4f; slopes at 0: %.5f and %.5f" %
          (np.abs(expit(a) - ndtr(math.sqrt(LAMBDA2) * a)).max(), 0.25, math.sqrt(LAMBDA2) / math.sqrt(2 * math.pi)))
    common.savefig(fig, 4, "fig4-9", out_dir)


FIG10 = dict(mu=((-0.55, 0.0), (0.55, 0.0)), sd=(0.4, 0.5), prior=0.5)


def fig4_10(out_dir=None):
    """Figure 4.10: class-conditional densities of two classes with a shared covariance (left) and the posterior p(C_1|x), a sigmoid of a linear function (right)."""
    mu = np.array(FIG10["mu"])
    Sigma = np.diag(np.array(FIG10["sd"]) ** 2)
    pi = np.array([FIG10["prior"], 1 - FIG10["prior"]])
    xs = np.linspace(-1, 1, 70)
    XX, YY = np.meshgrid(xs, xs)
    P = np.c_[XX.ravel(), YY.ravel()]
    dens = [pi[k] * np.exp(common.gauss_logpdf(P, mu[k], Sigma)).reshape(XX.shape) for k in range(2)]
    post = dens[0] / (dens[0] + dens[1])
    W, w0 = generative_posterior_params(pi, mu, Sigma)
    w, b = W[0] - W[1], w0[0] - w0[1]
    assert np.allclose(post.ravel(), expit(P @ w + b))
    fig = plt.figure(figsize=(7.4, 3.3))
    ax1, ax2 = fig.add_subplot(1, 2, 1, projection="3d"), fig.add_subplot(1, 2, 2, projection="3d")
    tot = dens[0] + dens[1]
    col = lambda r: np.stack([r, 0 * r + 0.15 + 0.0, 1 - r, 0 * r + 1.0], -1) * np.array([1, 1, 1, 1])
    rgba = np.stack([post, 0.05 + 0 * post, 1 - post, 0 * post + 1.0], -1)
    fade = np.clip(tot / tot.max(), 0, 1)[..., None]
    rgba1 = rgba.copy()
    rgba1[..., :3] = 1 - fade * (1 - rgba1[..., :3])
    ax1.plot_surface(XX, YY, tot, facecolors=rgba1, rstride=1, cstride=1, linewidth=0, antialiased=False, shade=False, rasterized=True)
    ax2.plot_surface(XX, YY, post, facecolors=rgba, rstride=1, cstride=1, linewidth=0, antialiased=False, shade=False, rasterized=True)
    for ax in (ax1, ax2):
        ax.set_xlim(-1, 1)
        ax.set_ylim(-1, 1)
        ax.view_init(24, -60)
        ax.set_xlabel("$x_1$")
        ax.set_ylabel("$x_2$")
    ax1.set_zlim(0, 0.45)
    ax2.set_zlim(0, 1)
    ax1.set_title(r"$p(\mathbf{x}\mid\mathcal{C}_k)\,p(\mathcal{C}_k)$, red $k=1$, blue $k=2$", fontsize=9)
    ax2.set_title(r"$p(\mathcal{C}_1\mid\mathbf{x}) = \sigma(\mathbf{w}^{\mathrm{T}}\mathbf{x}+w_0)$", fontsize=9)
    print("  Fig 4.10: w = (%.2f, %.2f), w_0 = %.2f (4.66)-(4.67); posterior equals sigma(w^T x + w_0) on the grid" % (w[0], w[1], b))
    common.savefig(fig, 4, "fig4-10", out_dir)


FIG11 = dict(mu=((-1.0, -1.0), (-1.0, 1.0), (1.5, 0.0)), cov=(np.diag([0.3 ** 2, 0.2 ** 2]), np.diag([0.3 ** 2, 0.2 ** 2]), np.diag([0.12 ** 2, 0.32 ** 2])))


def fig4_11(out_dir=None):
    """Figure 4.11: three Gaussian classes (red and green share a covariance) and the posterior probabilities as an RGB image; boundaries between red and green are linear, the others quadratic."""
    mu, cov = np.array(FIG11["mu"]), FIG11["cov"]
    fig, axes = common.panel_grid(1, 2, size=(3.7, 3.7))
    ax = axes[0, 0]
    for k, c in enumerate([RED, GREEN, BLUE]):
        for r in (1, 2, 3):
            ax.plot(*common.ellipse_xy(mu[k], cov[k], nsd=r), color=c, lw=1.0)
    newaxes(ax, (-2.5, 2.5), (-2.5, 2.5), ([-2, -1, 0, 1, 2], [-2.5, -2, -1.5, -1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5]))
    common.panel_label(ax, "a", "lower right")
    XX, YY, P = grid2d((-2.5, 2.5), (-2.5, 2.5), 400)
    L = np.stack([common.gauss_logpdf(P, mu[k], cov[k]) for k in range(3)], 1)
    post = softmax(L).reshape(XX.shape + (3,))
    ax = axes[0, 1]
    ax.imshow(post[::-1], extent=(-2.5, 2.5, -2.5, 2.5), aspect="equal", interpolation="bilinear")
    ax.contour(XX, YY, post.argmax(2), levels=[0.5, 1.5], colors="w", linewidths=1.4)
    newaxes(ax, (-2.5, 2.5), (-2.5, 2.5), ([-2, -1, 0, 1, 2], [-2.5, -2, -1.5, -1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5]))
    common.panel_label(ax, "b", "lower right")
    common.savefig(fig, 4, "fig4-11", out_dir)


FIG12 = dict(centres=((-1.0, -1.0), (0.0, 0.0)), width=0.78)


def fig4_12(out_dir=None):
    """Figure 4.12: two Gaussian basis functions turn a non-linearly-separable problem into a linearly separable one."""
    rng = JSRandom(412)
    red = gauss2(rng, 60, (0.0, 0.0), (0.3, 0.3), 0.0)
    b1 = gauss2(rng, 30, (-1.1, -0.95), (0.2, 0.2), 0.0)
    b2 = gauss2(rng, 30, (0.95, 1.05), (0.2, 0.2), 0.0)
    X = np.vstack([red, b1, b2])
    y = np.r_[np.zeros(60, int), np.ones(60, int)]       # red = 0, blue = 1
    c, s = np.array(FIG12["centres"]), FIG12["width"]
    feat = lambda Z: np.exp(-((Z[:, None, :] - c[None]) ** 2).sum(-1) / (2 * s * s))
    F = feat(X)
    w, _ = irls_logistic(design(F), y.astype(float), alpha=0.01)
    print("  Fig 4.12: training errors in feature space %d" % int(((design(F) @ w > 0).astype(int) != y).sum()))
    fig, axes = common.panel_grid(1, 2, size=(3.7, 3.7))
    ax = axes[0, 0]
    XX, YY, P = grid2d((-1.5, 1.5), (-1.5, 1.5), 300)
    ax.contour(XX, YY, (design(feat(P)) @ w).reshape(XX.shape), levels=[0], colors="k", linewidths=1.3)
    for ci in c:
        ax.plot(*ci, "+", color="#00cc00", ms=10, mew=2)
        ax.add_patch(plt.Circle(ci, 1.5 * s, fill=False, ec="#00cc00", lw=1.4))
    ax.scatter(*X[y == 0].T, s=6, color=RED, zorder=3)
    ax.scatter(*X[y == 1].T, s=6, color=BLUE, zorder=3)
    newaxes(ax, (-1.5, 1.5), (-1.5, 1.5), ([-1, 0, 1], [-1, 0, 1]))
    ax.set_xlabel("$x_1$")
    ax.set_ylabel("$x_2$", rotation=0, labelpad=8)
    common.panel_label(ax, "a", "lower right")
    ax = axes[0, 1]
    ax.scatter(*F[y == 0].T, s=6, color=RED, zorder=3)
    ax.scatter(*F[y == 1].T, s=6, color=BLUE, zorder=3)
    p1 = np.linspace(-0.1, 1.1, 2)
    ax.plot(p1, -(w[0] + w[1] * p1) / w[2], "k", lw=1.5)
    ax.set_xlim(-0.1, 1.1)
    ax.set_ylim(-0.1, 1.1)
    ax.set_xlabel(r"$\phi_1$")
    ax.set_ylabel(r"$\phi_2$", rotation=0, labelpad=8)
    ax.set_xticks([0, 0.5, 1])
    ax.set_yticks([0, 0.5, 1])
    ax.set_aspect("equal")
    common.panel_label(ax, "b", "lower right")
    common.savefig(fig, 4, "fig4-12", out_dir)


def fig4_13(out_dir=None):
    """Figure 4.13: a mixture of two Gaussians p(theta) (blue) and its cumulative distribution f(a) (red); f(a) is the area under p up to a."""
    th = np.linspace(0, 4, 800)
    comp = [(0.7, 1.5, 0.4), (0.3, 3.0, 0.3)]
    p = sum(w * norm.pdf(th, m, s) for w, m, s in comp)
    F = sum(w * norm.cdf(th, m, s) for w, m, s in comp)
    a = 2.8
    fig, ax = plt.subplots(figsize=(4.4, 3.2))
    ax.fill_between(th, p, where=th <= a, color="#c8f7c8", zorder=0)
    ax.plot(th, p, color=BLUE, lw=1.8)
    ax.plot(th, F, color=RED, lw=1.8)
    ax.axvline(a, color="#00dd00", lw=1.8)
    ax.set_xlim(0, 4)
    ax.set_ylim(0, 1)
    ax.set_xticks([0, 1, 2, 3, 4])
    ax.set_xlabel(r"$\theta$")
    print("  Fig 4.13: f(%.1f) = %.3f = area of p up to a (numerical %.3f)" % (a, F[np.searchsorted(th, a)], np.trapz(p[th <= a], th[th <= a])))
    common.savefig(fig, 4, "fig4-13", out_dir)


def fig4_14(out_dir=None):
    """Figure 4.14: Laplace approximation to p(z) proportional to exp(-z^2/2) sigma(20 z + 4): p(z) (yellow), Gaussian q(z) at the mode (red), and their negative logarithms."""
    ln_f, dln_f, d2ln_f = fig414_functions(20.0, 4.0)
    z0, A, lnZ_lap = laplace_1d(ln_f, dln_f, d2ln_f, 0.0)
    zs = np.linspace(-2, 4, 1500)
    f = np.exp(ln_f(zs))
    Z, _ = integrate.quad(lambda z: math.exp(ln_f(z)), -10, 10, limit=200, points=[z0])
    p = f / Z
    q = norm.pdf(zs, z0, 1 / math.sqrt(A))
    print("  Fig 4.14: mode z0 = %.4f, A = %.3f, ln Z exact %.4f, Laplace %.4f" % (z0, A, math.log(Z), lnZ_lap))
    fig, axes = common.panel_grid(1, 2, size=(3.8, 3.0))
    ax = axes[0, 0]
    ax.fill_between(zs, p, color="#e5cc55", zorder=1)
    ax.plot(zs, q, color=RED, lw=1.6, zorder=2)
    ax.set_xlim(-2, 4)
    ax.set_ylim(0, 0.8)
    ax.set_xticks([-2, -1, 0, 1, 2, 3, 4])
    common.panel_label(ax, "a", "upper right")
    ax = axes[0, 1]
    with np.errstate(divide="ignore"):
        ax.plot(zs, -np.log(p), color="#c9a800", lw=1.6)
        ax.plot(zs, -np.log(q), color=RED, lw=1.6)
    ax.set_xlim(-2, 4)
    ax.set_ylim(0, 40)
    ax.set_xticks([-2, -1, 0, 1, 2, 3, 4])
    common.panel_label(ax, "b", "upper right")
    common.savefig(fig, 4, "fig4-14", out_dir)


FIGURES = {"4.1": fig4_1, "4.2": fig4_2, "4.3": fig4_3, "4.4": fig4_4, "4.5": fig4_5, "4.6": fig4_6, "4.7": fig4_7,
           "4.9": fig4_9, "4.10": fig4_10, "4.11": fig4_11, "4.12": fig4_12, "4.13": fig4_13, "4.14": fig4_14}

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