"""Bishop (2006), Chapter 14: Combining models.

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

Algorithms (small readable functions, each citing the book's equations):
    committee_errors                       Section 14.2, E_AV and E_COM, (14.7)-(14.11)
    fit_stump, adaboost, adaboost_f        Section 14.3, AdaBoost with decision stumps, (14.15)-(14.21)
    exponential_error, ... (four losses)   Section 14.3.2, the error functions of Figures 14.3, 14.4
    grow_tree, prune_tree, pruning_path    Section 14.4, CART trees and the pruning criterion (14.29)-(14.33)
    bagging, fit_polynomial                Section 14.2, bootstrap committees (14.7)
    mixture_linear_regression_em           Section 14.5.1, EM for a mixture of linear regressions, (14.34)-(14.44)
    irls_logistic, mixture_logistic_em     Section 14.5.2, EM with IRLS in the M step, (14.45)-(14.52)
    mixture_of_experts_em                  Section 14.5.3, linear experts with a softmax gate, (14.53)

What could not be reproduced exactly, and the choices made:
* Figure 14.2: the book uses 30 points of its two-class synthetic data (Appendix A, Figure A.7) but
  does not publish them. BOOST_X, BOOST_T below are those 30 points read off the book's figure
  (circle centres located in the image, accurate to about 0.005). The base learners are decision
  stumps with thresholds half-way between neighbouring data values; equal weighted errors are
  broken by taking the last best stump in the order (x1 before x2, increasing threshold, sign +1
  before -1). With these points and this rule our stumps for m = 1, 2, 3, 6, 10 and 150 are exactly
  the book's dashed lines (x1 = 0.59, x1 = -0.007, x2 = -1.72, x2 = 0.087, x2 = -0.867, x1 = 0.417);
  the first-best rule swaps the first two stumps and then drifts away from the book. Circle
  radius is proportional to the weight w_n^(m) used to train the latest stump, scaled so that the
  largest circle has a fixed size (at m = 1 all weights are equal and the circles are drawn at half
  that size, as in the book). The ensemble boundary is drawn exactly, as the edges between the
  axis-aligned cells on which sign(f_m) is constant.
* Figures 14.3 and 14.4 are exact (the cross-entropy is divided by ln 2 as in the text, p. 662).
* Figures 14.5 and 14.6 are schematic in the book. Here the partition and the tree are drawn from
  one tree object (the Node class used by the CART code), with theta_1 = 0.4, theta_2 = 0.3,
  theta_3 = 0.7, theta_4 = 0.8 on the unit square, the proportions of the book's drawing.
* Figures 14.8 and 14.9: the book's 20 points have x = linspace(-1, 1, 20) exactly; their targets t
  (TWO_LINES_T) and the two initial lines (MIX_W0) are read off the book's figure. The text says
  beta was initialised to the reciprocal of the variance of the targets; we use the variance with
  the N - 1 divisor, with which our state after 30 EM iterations matches the book's centre panel
  (the change near iteration 30 is so fast that the divisor N gives the same picture one iteration
  earlier). After 50 iterations the log likelihood is -3.08 (book: -3.0) and a single linear
  regression has -27.66 (book: -27.6).
* Figure 14.10: the book's data are not available. We draw 100 points from each of four Gaussian
  clusters with standard deviation 0.5: class "red" around (-1.5, 0) and (1.5, 0), class "blue"
  around (0, 1.5) and (0, -1.5), which is what the book's figure shows. Panel (a) is shaded with the
  exact class probability of this generating model, which is sharper than the book's shading (the
  book's colours fit unit-variance clusters, p(red|x) = cosh(1.5 x1)/(cosh(1.5 x1) + cosh(1.5 x2)),
  although its points have a standard deviation close to 0.5). A single logistic regression gives
  probabilities close to 1/2 everywhere, so the direction of its p = 1/2 line is decided by sampling
  noise; with seed 25 it is nearly horizontal, as in the book. The mixture of two logistic
  regressions is started from two vertical boundaries (x1 = -1 and x1 = 1); a start with horizontal
  boundaries finds the mirror-image solution (a horizontal red band) of similar likelihood.
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

import common

# The book's colours: light green for the boosting boundary, pure green rings for Fig 14.8.
BOOST_GREEN = "#7be05a"
RING_GREEN = "#21c421"
WHITE_MAGENTA = LinearSegmentedColormap.from_list("white_magenta", ["#ffffff", "#ff00ff"])

# =========================================================================== data
# Figure 14.2: 30 points read off the book's figure (13 of class t = -1, blue; 17 of class t = +1, red).
BOOST_X = np.array([
    [-1.351, 1.769], [-1.245, 1.009], [-0.078, 1.002], [-0.052, 0.991], [-0.213, 0.000],
    [-0.222, -0.441], [-0.366, -0.666], [-0.369, -1.257], [0.504, -0.288], [0.581, -0.785],
    [0.563, -1.640], [1.094, -1.451], [1.602, -0.030],
    [0.038, -1.436], [0.256, 0.949], [0.330, -1.874], [0.599, -1.481], [0.628, 1.058],
    [0.705, -2.624], [0.835, -1.933], [0.867, -1.063], [0.893, 1.903], [0.973, -0.949],
    [1.055, 0.601], [1.179, 1.526], [1.183, -1.800], [1.209, 0.423], [1.286, -1.084],
    [1.702, -0.646], [2.424, 0.174]])
BOOST_T = np.r_[-np.ones(13), np.ones(17)].astype(int)
BOOST_LIMITS = (-1.53, 2.60, -2.82, 2.11)          # the book's axis box

# Figures 14.8, 14.9: x = linspace(-1, 1, 20); targets read off the book's figure.
TWO_LINES_X = np.linspace(-1.0, 1.0, 20)
TWO_LINES_T = np.array([-1.133, -0.856, -1.071, -0.870, -0.958, 0.863, 0.637, 1.024, 0.715, 1.138,
                        0.815, 0.775, 0.976, 0.801, 1.098, -0.946, -1.249, -1.233, -0.991, -0.914])
# initial weight vectors (bias, slope) of the blue and the red line in Figure 14.8 (left)
MIX_W0 = np.array([[0.0434, 0.0136], [-0.1427, 0.1023]])


def two_lines_data():
    """The 20 points of Figures 14.8 and 14.9: design matrix Phi = [1, x], targets t, inputs x."""
    x = TWO_LINES_X
    return np.c_[np.ones_like(x), x], TWO_LINES_T.copy(), x


def four_clusters(n_per=100, sd=0.5, seed=25):
    """Two-class data like Figure 14.10: class 1 ("blue") around (0, +-1.5), class 0 ("red")
    around (+-1.5, 0); n_per points per cluster with standard deviation sd. Returns X, t."""
    rng = np.random.default_rng(seed)
    centres = [(-1.5, 0.0), (1.5, 0.0), (0.0, 1.5), (0.0, -1.5)]
    X = np.vstack([np.array(c) + sd * rng.standard_normal((n_per, 2)) for c in centres])
    t = np.r_[np.zeros(2 * n_per), np.ones(2 * n_per)].astype(int)
    return X, t


def four_clusters_prob_blue(X, sd=0.5):
    """Exact p(t = blue | x) for the generating model of four_clusters (equal priors and weights)."""
    a = 1.5 / sd ** 2
    # each class density is proportional to exp(-|x|^2/(2 sd^2)) cosh(a x_i); the common factor cancels
    lr = np.logaddexp(a * X[:, 0], -a * X[:, 0])
    lb = np.logaddexp(a * X[:, 1], -a * X[:, 1])
    return common.sigmoid(lb - lr)


def sinusoidal_data(n=30, noise=0.3, seed=0):
    """Appendix A regression data: x uniform on (0, 1), t = sin(2 pi x) + Gaussian noise (sd 0.3)."""
    rng = np.random.default_rng(seed)
    x = rng.uniform(0, 1, n)
    return x, np.sin(2 * np.pi * x) + noise * rng.standard_normal(n)


# =========================================================================== Section 14.2: committees
def committee_errors(Y, h):
    """E_AV (14.10) and E_COM (14.11) for M model outputs Y (M x G) on a grid of G inputs.

    The expectation over x is the average over the grid (uniform input density), and the error of
    model m is eps_m(x) = y_m(x) - h(x) (14.8). E_COM = E_AV - (1/M) sum_m E_x[(y_m - y_COM)^2]
    (the spread of the members about their average), so E_COM <= E_AV always (Exercise 14.3).
    """
    eps = Y - h[None, :]
    e_av = np.mean(np.mean(eps ** 2, axis=1))
    e_com = np.mean(np.mean(eps, axis=0) ** 2)
    return e_av, e_com


def legendre_design(x, degree):
    """Legendre polynomials P_0..P_degree of z = 2x - 1 (well conditioned on (0, 1))."""
    return np.polynomial.legendre.legvander(2.0 * np.asarray(x) - 1.0, degree)


def fit_polynomial(x, t, degree, ridge=1e-8):
    """Least-squares polynomial of the given degree; returns a function x -> y(x).

    The tiny ridge only matters when a bootstrap sample has fewer distinct x than coefficients.
    """
    A = legendre_design(x, degree)
    c = np.linalg.solve(A.T @ A + ridge * len(x) * np.eye(degree + 1), A.T @ t)
    return lambda xn: legendre_design(xn, degree) @ c


def bagging(x, t, fit, M, rng, n_draw=None):
    """Bootstrap aggregation (Breiman 1996), Section 14.2: fit M models y_m to bootstrap samples.

    fit(x_b, t_b) must return a prediction function. Returns the list of M functions and the
    bootstrap index sets. The committee prediction is their average, y_COM = (1/M) sum_m y_m (14.7).
    """
    n = len(x)
    models, idx = [], []
    for _ in range(M):
        b = rng.integers(0, n, n_draw or n)
        models.append(fit(x[b], t[b]))
        idx.append(b)
    return models, idx


# =========================================================================== Section 14.3: boosting
def stump_predict(stump, X):
    """Decision stump y(x) = s if x_d > theta else -s, with stump = (d, theta, s)."""
    d, theta, s = stump
    return np.where(X[:, d] > theta, s, -s)


def fit_stump(X, t, w):
    """The decision stump minimising the weighted error J_m = sum_n w_n I(y(x_n) != t_n) (14.15).

    Thresholds are half-way between neighbouring distinct values of each input; for each threshold
    both signs s = +1, -1 are tried. Equal errors are common with few points; ties go to the last
    best stump in the order (input x1 then x2, increasing threshold, s = +1 then -1), the rule that
    reproduces the book's sequence of stumps in Figure 14.2. Returns ((d, theta, s), J).
    """
    best, best_err = None, np.inf
    for d in range(X.shape[1]):
        order = np.argsort(X[:, d], kind="stable")
        xs, ts, ws = X[order, d], t[order], w[order]
        # with s = +1, a point left of the threshold is predicted -1: it is an error iff t = +1
        left_err = np.cumsum(ws * (ts == 1))                 # errors among points 0..i (left)
        right_err = np.sum(ws * (ts == -1)) - np.cumsum(ws * (ts == -1))   # errors among i+1..N-1
        err_plus = left_err + right_err
        err_minus = ws.sum() - err_plus
        for i in range(len(xs) - 1):
            if xs[i + 1] <= xs[i]:
                continue
            theta = 0.5 * (xs[i] + xs[i + 1])
            for s, e in ((1, err_plus[i]), (-1, err_minus[i])):
                if e <= best_err + 1e-12:
                    best, best_err = (d, theta, s), e
    return best, best_err


def adaboost(X, t, M):
    """AdaBoost (Freund and Schapire 1996) with decision stumps, as on pp. 658-659.

    t in {-1, +1}. For m = 1..M: fit the stump minimising (14.15) with weights w^(m);
    eps_m = weighted error rate (14.16); alpha_m = ln((1 - eps_m)/eps_m) (14.17);
    w^(m+1)_n = w^(m)_n exp(alpha_m I(y_m(x_n) != t_n)) (14.18), rescaled to sum to 1 (only ratios
    of weights matter in (14.16)). Returns a dict with the stumps, eps, alpha and the weights
    w^(1), ..., w^(M) used to train each stump (each normalised to sum to 1).
    """
    n = len(t)
    w = np.full(n, 1.0 / n)
    out = {"stumps": [], "eps": [], "alpha": [], "weights": []}
    for _ in range(M):
        stump, _ = fit_stump(X, t, w)
        miss = stump_predict(stump, X) != t
        eps = w[miss].sum() / w.sum()
        eps = min(max(eps, 1e-12), 1 - 1e-12)
        alpha = np.log((1 - eps) / eps)
        out["stumps"].append(stump)
        out["eps"].append(eps)
        out["alpha"].append(alpha)
        out["weights"].append(w.copy())
        w = w * np.exp(alpha * miss)
        w = w / w.sum()
    out["eps"], out["alpha"] = np.array(out["eps"]), np.array(out["alpha"])
    out["weights"] = np.array(out["weights"])
    return out


def adaboost_f(model, X, m):
    """f_m(x) = (1/2) sum_{l <= m} alpha_l y_l(x) (14.21); the classifier is sign f_m(x) (14.19)."""
    f = np.zeros(len(X))
    for stump, a in zip(model["stumps"][:m], model["alpha"][:m]):
        f += 0.5 * a * stump_predict(stump, X)
    return f


def adaboost_curves(model, X, t):
    """Training error of sign(f_m), the exponential error (1/N) sum_n exp(-t_n f_m(x_n)) (14.20) and
    the products prod_l 2 sqrt(eps_l (1 - eps_l)) and exp(-2 sum_l (1/2 - eps_l)^2), for m = 1..M."""
    M = len(model["alpha"])
    f = np.zeros(len(t))
    train, expo = np.zeros(M), np.zeros(M)
    for m in range(M):
        f += 0.5 * model["alpha"][m] * stump_predict(model["stumps"][m], X)
        train[m] = np.mean(np.sign(f) != t)
        expo[m] = np.mean(np.exp(-t * f))
    eps = model["eps"]
    prod = np.cumprod(2 * np.sqrt(eps * (1 - eps)))
    loose = np.exp(-2 * np.cumsum((0.5 - eps) ** 2))
    return train, expo, prod, loose


def exponential_error(z):
    """E(z) = exp(-z), z = t y(x) (14.20)."""
    return np.exp(-z)


def cross_entropy_error(z):
    """E(z) = ln(1 + exp(-z)) / ln 2, the cross-entropy for t in {-1, 1} rescaled to pass through (0, 1)."""
    return np.logaddexp(0, -z) / np.log(2)


def hinge_error(z):
    """E(z) = max(0, 1 - z), the hinge error of support vector machines (Section 7.1.2)."""
    return np.maximum(0, 1 - z)


def misclassification_error(z):
    """E(z) = I(z < 0)."""
    return (np.asarray(z) < 0).astype(float)


# =========================================================================== Section 14.4: trees
class Node:
    """A node of a binary decision tree.

    An internal node sends x to `left` if x[feature] <= threshold and to `right` otherwise.
    Every node keeps what it would predict as a leaf: `value` (the mean target (14.29) for
    regression, the class proportions p_tau,k for classification), the number of training points
    `n` and `Q`, the node's contribution to the residual error if it were a leaf ((14.30), or N_tau
    times an impurity (14.32), (14.33) for classification). `name` labels leaves in drawings.
    """

    def __init__(self, value=None, n=0, Q=0.0, feature=None, threshold=None, left=None, right=None,
                 name=None, depth=0):
        self.value, self.n, self.Q = value, n, Q
        self.feature, self.threshold, self.left, self.right = feature, threshold, left, right
        self.name, self.depth = name, depth

    @property
    def is_leaf(self):
        return self.left is None

    def leaves(self):
        return [self] if self.is_leaf else self.left.leaves() + self.right.leaves()

    def internal(self):
        return [] if self.is_leaf else [self] + self.left.internal() + self.right.internal()

    def copy(self):
        c = Node(self.value, self.n, self.Q, self.feature, self.threshold, None, None, self.name, self.depth)
        if not self.is_leaf:
            c.left, c.right = self.left.copy(), self.right.copy()
        return c


def gini(p):
    """Gini index (14.33): sum_k p_k (1 - p_k)."""
    p = np.asarray(p, float)
    return float(np.sum(p * (1 - p)))


def cross_entropy_impurity(p):
    """Cross-entropy (14.32) with the minus sign of the errata: -sum_k p_k ln p_k (0 ln 0 = 0)."""
    p = np.asarray(p, float)
    q = p[p > 0]
    return float(-np.sum(q * np.log(q)))


def misclassification_impurity(p):
    """Misclassification rate of a region that predicts its majority class: 1 - max_k p_k."""
    return float(1 - np.max(p))


IMPURITY = {"gini": gini, "entropy": cross_entropy_impurity, "misclass": misclassification_impurity}


def _node_stats(t, kind, criterion, n_classes):
    """value, Q of a node holding targets t."""
    if kind == "regression":
        value = float(np.mean(t))
        return value, float(np.sum((t - value) ** 2))           # (14.29), (14.30)
    p = np.bincount(t, minlength=n_classes) / len(t)
    return p, len(t) * IMPURITY[criterion](p)


def _best_split(X, t, kind, criterion, n_classes, min_leaf):
    """Exhaustive search over inputs and thresholds (half-way between neighbouring distinct values)
    for the split with the smallest Q_left + Q_right. Returns (Q, feature, threshold) or None."""
    best = None
    n = len(t)
    for d in range(X.shape[1]):
        order = np.argsort(X[:, d], kind="stable")
        xs, ts = X[order, d], t[order]
        if kind == "regression":
            c1, c2 = np.cumsum(ts), np.cumsum(ts ** 2)
        for i in range(min_leaf - 1, n - min_leaf):
            if xs[i + 1] <= xs[i]:
                continue
            nl, nr = i + 1, n - i - 1
            if kind == "regression":
                ql = c2[i] - c1[i] ** 2 / nl
                qr = (c2[-1] - c2[i]) - (c1[-1] - c1[i]) ** 2 / nr
            else:
                ql = _node_stats(ts[:nl], kind, criterion, n_classes)[1]
                qr = _node_stats(ts[nl:], kind, criterion, n_classes)[1]
            q = ql + qr
            if best is None or q < best[0] - 1e-12:
                best = (q, d, 0.5 * (xs[i] + xs[i + 1]))
    return best


def grow_tree(X, t, kind="regression", criterion="gini", max_depth=8, min_leaf=1, n_classes=None):
    """Grow a CART tree greedily (Section 14.4, pp. 664-665).

    Starting from the root (the whole input space), each region is split on the input and threshold
    that give the smallest residual error of the two new leaves: the sum of squares (14.30) for
    regression, or N_tau times the Gini index (14.33) or cross-entropy (14.32) for classification
    (weighting each region by its number of points, as in CART). Growth stops at max_depth, when a
    leaf would get fewer than min_leaf points, or when no split lowers the error.
    """
    X = np.asarray(X, float).reshape(len(t), -1)
    if kind != "regression":
        t = np.asarray(t, int)
        n_classes = n_classes or int(t.max()) + 1

    def build(idx, depth):
        value, Q = _node_stats(t[idx], kind, criterion, n_classes)
        node = Node(value, len(idx), Q, depth=depth)
        if depth >= max_depth or len(idx) < 2 * min_leaf or Q <= 1e-12:
            return node
        split = _best_split(X[idx], t[idx], kind, criterion, n_classes, min_leaf)
        if split is None or split[0] >= Q - 1e-12:
            return node
        _, d, theta = split
        go_left = X[idx, d] <= theta
        node.feature, node.threshold = d, theta
        node.left = build(idx[go_left], depth + 1)
        node.right = build(idx[~go_left], depth + 1)
        return node

    return build(np.arange(len(t)), 0)


def tree_predict(tree, X):
    """Prediction of a tree: the leaf value (regression) or the leaf's majority class."""
    X = np.asarray(X, float)
    X = X.reshape(len(X), -1)
    out = []
    for x in X:
        node = tree
        while not node.is_leaf:
            node = node.left if x[node.feature] <= node.threshold else node.right
        out.append(node.value if np.isscalar(node.value) else int(np.argmax(node.value)))
    return np.array(out)


def pruning_criterion(tree, lam):
    """C(T) = sum_tau Q_tau(T) + lambda |T| (14.31), summing over the leaves of T."""
    leaves = tree.leaves()
    return sum(l.Q for l in leaves) + lam * len(leaves)


def prune_tree(tree, lam):
    """The subtree T of `tree` that minimises C(T) = sum_tau Q_tau(T) + lambda |T| (14.31).

    Bottom-up dynamic programming: a node is collapsed into a leaf when Q(node) + lambda is no
    larger than the best cost of its two pruned subtrees (ties collapse, giving the smallest
    optimal subtree). This is exact because C(T) adds up over the two subtrees.
    """
    def best(node):
        if node.is_leaf:
            return node.copy(), node.Q + lam
        left, cl = best(node.left)
        right, cr = best(node.right)
        if node.Q + lam <= cl + cr + 1e-12:
            return Node(node.value, node.n, node.Q, name=node.name, depth=node.depth), node.Q + lam
        keep = Node(node.value, node.n, node.Q, node.feature, node.threshold, left, right, node.name, node.depth)
        return keep, cl + cr

    return best(tree)[0]


def pruning_path(tree):
    """Weakest-link pruning (Breiman et al. 1984): the nested subtrees T_0 = tree > T_1 > ... > root
    and the values lambda_0 = 0 < lambda_1 < ... at which each becomes optimal for (14.31).

    At each stage the internal node t with the smallest g(t) = (Q(t) - Q(T_t)) / (|T_t| - 1) is
    collapsed, where Q(T_t) sums Q over the leaves below t. Returns a list of (lambda_k, T_k).
    """
    T = tree.copy()
    path = [(0.0, T.copy())]
    while not T.is_leaf:
        g = []
        for node in T.internal():
            leaves = node.leaves()
            g.append(((node.Q - sum(l.Q for l in leaves)) / (len(leaves) - 1), node))
        gmin = min(v for v, _ in g)
        for v, node in g:
            if v <= gmin + 1e-12 and not node.is_leaf:
                node.left = node.right = None
                node.feature = node.threshold = None
        path.append((max(gmin, 0.0), T.copy()))
    return path


def fit_tree_regression(max_depth=8, min_leaf=1, lam=0.0):
    """A fit(x, t) function for bagging: a regression tree, pruned with lambda if lam > 0."""
    def fit(x, t):
        tree = grow_tree(x, t, "regression", max_depth=max_depth, min_leaf=min_leaf)
        if lam > 0:
            tree = prune_tree(tree, lam)
        return lambda xn: tree_predict(tree, xn)
    return fit


# =========================================================================== Section 14.5.1
def gaussian_mixture_loglik(t, mu, pi, beta):
    """log p(t_n | theta) for each n, where mu is N x K, pi is N x K or K, (14.34)."""
    lp = np.log(np.clip(pi, 1e-300, None)) + 0.5 * np.log(beta / (2 * np.pi)) - 0.5 * beta * (t[:, None] - mu) ** 2
    return common.logsumexp(lp, axis=1), lp


def mixture_linear_regression_em(Phi, t, W, pi, beta, n_iter):
    """EM for a mixture of K linear regressions with a common noise precision beta (Section 14.5.1).

    Model p(t | phi) = sum_k pi_k N(t | w_k^T phi, 1/beta) (14.34); W is K x M (row k = w_k).
    One iteration: E step, responsibilities gamma_nk (14.37); M step, pi_k = (1/N) sum_n gamma_nk
    (14.38), w_k = (Phi^T R_k Phi)^{-1} Phi^T R_k t with R_k = diag(gamma_nk) (14.42), and
    1/beta = (1/N) sum_n sum_k gamma_nk (t_n - w_k^T phi_n)^2 (14.44), which uses the new w_k
    (together the exact maximiser of Q). Returns a list of n_iter + 1 states theta_0, ..., theta_n,
    each a dict with W, pi, beta, the log likelihood (14.35) and the responsibilities gamma(theta).
    """
    W, pi = np.array(W, float), np.array(pi, float)
    states = []
    for it in range(n_iter + 1):
        ll, lp = gaussian_mixture_loglik(t, Phi @ W.T, pi, beta)
        gamma = np.exp(lp - ll[:, None])                                # E step (14.37)
        states.append({"W": W.copy(), "pi": pi.copy(), "beta": beta, "loglik": float(ll.sum()), "gamma": gamma})
        if it == n_iter:
            break
        pi = gamma.mean(0)                                              # (14.38)
        for k in range(len(pi)):
            R = gamma[:, k]
            W[k] = np.linalg.solve(Phi.T @ (R[:, None] * Phi), Phi.T @ (R * t))   # (14.42)
        beta = len(t) / np.sum(gamma * (t[:, None] - Phi @ W.T) ** 2)  # (14.44)
    return states


def linear_regression_ml(Phi, t):
    """Maximum likelihood linear regression: w (least squares), beta = 1/mean squared residual, log lik."""
    w = np.linalg.lstsq(Phi, t, rcond=None)[0]
    s2 = np.mean((t - Phi @ w) ** 2)
    return w, 1 / s2, -0.5 * len(t) * (np.log(2 * np.pi * s2) + 1)


# =========================================================================== Section 14.5.2
def logistic_q(Phi, t, r, w):
    """Q_k(w) = sum_n r_n [t_n ln y_n + (1 - t_n) ln(1 - y_n)], y_n = sigma(w^T phi_n): the part of (14.49)
    that depends on w_k, with responsibilities r_n = gamma_nk."""
    a = Phi @ w
    return float(np.sum(r * (t * -np.logaddexp(0, -a) + (1 - t) * -np.logaddexp(0, a))))


def irls_logistic(Phi, t, r, w, n_steps=20, ridge=1e-6, tol=1e-10):
    """Weighted logistic regression by iterative reweighted least squares (Newton-Raphson).

    Gradient (14.51): sum_n r_n (t_n - y_n) phi_n; Hessian of -Q (14.52): sum_n r_n y_n (1 - y_n) phi_n phi_n^T.
    Each Newton step w <- w + H^{-1} grad is halved until Q does not decrease (a tiny ridge keeps H
    invertible when the weighted data are separable). With r_n = 1 this is ordinary logistic
    regression (Section 4.3.3).
    """
    w = np.array(w, float)
    q = logistic_q(Phi, t, r, w)
    for _ in range(n_steps):
        y = common.sigmoid(Phi @ w)
        grad = Phi.T @ (r * (t - y))
        H = (Phi * (r * y * (1 - y))[:, None]).T @ Phi + ridge * np.eye(len(w))
        step = np.linalg.solve(H, grad)
        for _ in range(40):
            q_new = logistic_q(Phi, t, r, w + step)
            if q_new >= q - 1e-12:
                break
            step = step / 2
        w = w + step
        if q_new - q < tol * (1 + abs(q)):
            q = q_new
            break
        q = q_new
    return w


def mixture_logistic_loglik(Phi, t, W, pi):
    """ln p(t | theta) for the mixture of logistic regressions (14.46), and log of each term."""
    A = Phi @ W.T                                                        # N x K
    lb = np.where(t[:, None] == 1, -np.logaddexp(0, -A), -np.logaddexp(0, A))
    lp = np.log(pi) + lb
    ll = common.logsumexp(lp, axis=1)
    return float(ll.sum()), lp, ll


def mixture_logistic_em(Phi, t, W, pi, n_iter=100, irls_steps=20, tol=1e-9):
    """EM for a mixture of K logistic regressions (Section 14.5.2).

    E step: gamma_nk proportional to pi_k y_nk^t_n (1 - y_nk)^(1 - t_n) (14.48).
    M step: pi_k = (1/N) sum_n gamma_nk (14.50); each w_k by IRLS on the data weighted by gamma_nk,
    using the gradient (14.51) and Hessian (14.52). Returns W, pi and the log likelihoods (14.46).
    """
    W, pi = np.array(W, float), np.array(pi, float)
    lls = []
    for _ in range(n_iter):
        ll, lp, lln = mixture_logistic_loglik(Phi, t, W, pi)
        lls.append(ll)
        gamma = np.exp(lp - lln[:, None])                                # (14.48)
        pi = gamma.mean(0)                                               # (14.50)
        for k in range(len(pi)):
            W[k] = irls_logistic(Phi, t, gamma[:, k], W[k], irls_steps)  # (14.51), (14.52)
        if len(lls) > 1 and abs(lls[-1] - lls[-2]) < tol * (1 + abs(lls[-1])):
            break
    lls.append(mixture_logistic_loglik(Phi, t, W, pi)[0])
    return W, pi, np.array(lls)


# =========================================================================== Section 14.5.3
def softmax_rows(A):
    A = A - A.max(1, keepdims=True)
    E = np.exp(A)
    return E / E.sum(1, keepdims=True)


def gate_newton(G, gamma, V, n_steps=20, ridge=1e-6):
    """M step for a softmax gate pi_k(x) = exp(v_k^T g(x)) / sum_j exp(v_j^T g(x)).

    Maximises sum_n sum_k gamma_nk ln pi_k(x_n), a multinomial logistic regression with soft targets
    gamma (concave), by Newton steps with step halving; v_K = 0 is fixed for identifiability.
    Gradient for v_j: sum_n (gamma_nj - pi_nj) g_n; Hessian blocks -sum_n pi_nj (delta_jl - pi_nl) g_n g_n^T.
    """
    V = np.array(V, float)
    K, D = V.shape

    def obj(V):
        return float(np.sum(gamma * np.log(np.clip(softmax_rows(G @ V.T), 1e-300, None))))

    f = obj(V)
    for _ in range(n_steps):
        P = softmax_rows(G @ V.T)
        grad = ((gamma - P)[:, :K - 1].T @ G).ravel()
        H = np.zeros(((K - 1) * D, (K - 1) * D))
        for j in range(K - 1):
            for l in range(K - 1):
                c = P[:, j] * ((j == l) - P[:, l])
                H[j * D:(j + 1) * D, l * D:(l + 1) * D] = (G * c[:, None]).T @ G
        step = np.linalg.solve(H + ridge * np.eye(len(grad)), grad).reshape(K - 1, D)
        for _ in range(40):
            Vn = V.copy()
            Vn[:K - 1] += step
            fn = obj(Vn)
            if fn >= f - 1e-12:
                break
            step = step / 2
        V, gain, f = Vn, fn - f, fn
        if gain < 1e-10:
            break
    return V


def mixture_of_experts_em(Phi, G, t, W, V, beta, n_iter=200, tol=1e-10):
    """EM for a mixture of linear-regression experts with a softmax gate (Jacobs et al. 1991;
    Jordan and Jacobs 1994), the model (14.53) with p_k(t|x) = N(t | w_k^T phi(x), 1/beta) and
    pi_k(x) = softmax_k(V g(x)).

    E step: gamma_nk proportional to pi_k(x_n) N(t_n | w_k^T phi_n, 1/beta) (as (14.37) with
    input-dependent pi). M step: w_k and beta exactly as (14.42), (14.44); the gate by gate_newton.
    Returns W, V, beta and the log likelihoods, which never decrease.
    """
    W, V = np.array(W, float), np.array(V, float)
    lls = []
    for _ in range(n_iter):
        P = softmax_rows(G @ V.T)
        ll, lp = gaussian_mixture_loglik(t, Phi @ W.T, P, beta)
        lls.append(float(ll.sum()))
        gamma = np.exp(lp - ll[:, None])
        for k in range(W.shape[0]):
            R = gamma[:, k]
            W[k] = np.linalg.solve(Phi.T @ (R[:, None] * Phi) + 1e-10 * np.eye(Phi.shape[1]), Phi.T @ (R * t))
        beta = len(t) / np.sum(gamma * (t[:, None] - Phi @ W.T) ** 2)
        V = gate_newton(G, gamma, V)
        if len(lls) > 1 and abs(lls[-1] - lls[-2]) < tol * (1 + abs(lls[-1])):
            break
    P = softmax_rows(G @ V.T)
    lls.append(float(gaussian_mixture_loglik(t, Phi @ W.T, P, beta)[0].sum()))
    return W, V, beta, np.array(lls)


# =========================================================================== drawing helpers
def ensemble_boundary_segments(model, m, limits):
    """Exact decision boundary of sign f_m(x) for stumps: f_m is constant on the axis-aligned cells cut
    out by the thresholds, so the boundary is made of cell edges with different signs on each side.
    Returns a list of segments [(x0, y0), (x1, y1)]."""
    x0, x1, y0, y1 = limits
    box = [(x0, x1), (y0, y1)]
    cuts = [[x0, x1], [y0, y1]]
    for d, theta, _ in model["stumps"][:m]:
        if box[d][0] < theta < box[d][1]:
            cuts[d].append(theta)
    xs, ys = np.unique(cuts[0]), np.unique(cuts[1])
    cx, cy = 0.5 * (xs[1:] + xs[:-1]), 0.5 * (ys[1:] + ys[:-1])
    CX, CY = np.meshgrid(cx, cy, indexing="ij")
    pos = adaboost_f(model, np.c_[CX.ravel(), CY.ravel()], m).reshape(CX.shape) > 0
    segs = []
    for i in range(len(cx) - 1):                     # vertical edges x = xs[i + 1]
        for j in np.where(pos[i] != pos[i + 1])[0]:
            segs.append([(xs[i + 1], ys[j]), (xs[i + 1], ys[j + 1])])
    for j in range(len(cy) - 1):                     # horizontal edges y = ys[j + 1]
        for i in np.where(pos[:, j] != pos[:, j + 1])[0]:
            segs.append([(xs[i], ys[j + 1]), (xs[i + 1], ys[j + 1])])
    return segs


def arrow_axes(ax, xlim, ylim, xlabel, ylabel, ylabel_side="right"):
    """Book-style axes for Figures 14.3 to 14.5: two arrows through the origin, no frame."""
    for s in ("top", "right", "left", "bottom"):
        ax.spines[s].set_visible(False)
    kw = dict(arrowstyle="-|>", lw=1.6, color="k", mutation_scale=14, shrinkA=0, shrinkB=0)
    ax.annotate("", xy=(xlim[1], 0), xytext=(xlim[0], 0), arrowprops=kw, zorder=1)
    ax.annotate("", xy=(0, ylim[1]), xytext=(0, ylim[0]), arrowprops=kw, zorder=1)
    dx = 0.02 * (xlim[1] - xlim[0])
    ax.text(xlim[1] + dx, 0, xlabel, ha="left", va="center", fontsize=13)
    ax.text(dx if ylabel_side == "right" else -dx, ylim[1], ylabel, ha="left" if ylabel_side == "right" else "right",
            va="top", fontsize=13)
    ax.set_xlim(xlim[0], xlim[1] + 0.09 * (xlim[1] - xlim[0]))
    ax.set_ylim(ylim[0], ylim[1] + 0.02 * (ylim[1] - ylim[0]))
    ax.set_yticks([])


def book_tree():
    """The tree of Figures 14.5 and 14.6 on the unit square: theta_1 = 0.4, theta_2 = 0.3, theta_3 = 0.7,
    theta_4 = 0.8. Each internal node sends x_d <= theta to the left; `label` and `side` say which
    branch the book annotates."""
    A, B, C, D, E = (Node(name=s) for s in "ABCDE")
    n4 = Node(feature=0, threshold=0.8, left=C, right=D)
    n3 = Node(feature=1, threshold=0.7, left=n4, right=E)
    n2 = Node(feature=1, threshold=0.3, left=A, right=B)
    root = Node(feature=0, threshold=0.4, left=n2, right=n3)
    for node, k, side in ((root, 1, "right"), (n2, 2, "left"), (n3, 3, "right"), (n4, 4, "left")):
        node.theta_index, node.side = k, side
    return root


def draw_partition(ax, node, box, color=common.GREEN, lw=2.2):
    """Draw the axis-aligned regions of a tree inside box = (x0, x1, y0, y1); leaves are named at
    the centre of their region (Figure 14.5)."""
    x0, x1, y0, y1 = box
    if node.is_leaf:
        if node.name:
            ax.text((x0 + x1) / 2, (y0 + y1) / 2, node.name, ha="center", va="center", fontsize=17,
                    family="DejaVu Sans")
        return
    th = node.threshold
    if node.feature == 0:
        ax.plot([th, th], [y0, y1], color=color, lw=lw, solid_capstyle="butt")
        draw_partition(ax, node.left, (x0, th, y0, y1), color, lw)
        draw_partition(ax, node.right, (th, x1, y0, y1), color, lw)
    else:
        ax.plot([x0, x1], [th, th], color=color, lw=lw, solid_capstyle="butt")
        draw_partition(ax, node.left, (x0, x1, y0, th), color, lw)
        draw_partition(ax, node.right, (x0, x1, th, y1), color, lw)


def tree_layout(tree):
    """Positions for drawing a tree: leaves equally spaced on the bottom row, each internal node at the
    mean position of the leaves below it and at the height of its depth."""
    leaves = tree.leaves()
    depth = max(l.depth for l in leaves) if leaves else 0

    def set_depth(node, d):
        node.depth = d
        if not node.is_leaf:
            set_depth(node.left, d + 1)
            set_depth(node.right, d + 1)

    set_depth(tree, 0)
    leaves = tree.leaves()
    bottom = max(l.depth for l in leaves)
    pos = {}
    for i, l in enumerate(leaves):
        pos[id(l)] = (float(i), -float(bottom))
    for node in tree.internal():
        xs = [pos[id(l)][0] for l in node.leaves()]
        pos[id(node)] = (float(np.mean(xs)), -float(node.depth))
    return pos


# =========================================================================== figures
def fig14_2(out_dir=None):
    """Figure 14.2: AdaBoost with decision stumps on 30 two-class points, m = 1, 2, 3, 6, 10, 150.

    The points are BOOST_X, BOOST_T: the book's 30 points read off its figure (13 blue, t = -1;
    17 red, t = +1), see the module docstring. Circle radius = weight w_n^(m) of the latest stump.
    """
    from matplotlib.collections import LineCollection
    X, t = BOOST_X, BOOST_T
    model = adaboost(X, t, 150)
    lim = BOOST_LIMITS
    fig, axes = common.panel_grid(2, 3, size=(3.0, 2.3))
    R = 9.0                                            # largest circle radius in points
    for ax, m in zip(axes.ravel(), (1, 2, 3, 6, 10, 150)):
        w = model["weights"][m - 1]
        r = np.maximum(R * w / max(w.max(), 2.0 / len(w)), 1.3)   # tiny weights stay visible as dots
        for cls, col in ((1, common.RED), (-1, common.BLUE)):
            k = t == cls
            ax.scatter(X[k, 0], X[k, 1], s=(2 * r[k]) ** 2, facecolors="none", edgecolors=col, linewidths=1.5, zorder=3)
        segs = ensemble_boundary_segments(model, m, lim)
        ax.add_collection(LineCollection(segs, colors=BOOST_GREEN, linewidths=2.2, capstyle="projecting", zorder=4))
        d, theta, _ = model["stumps"][m - 1]
        if d == 0:
            ax.plot([theta, theta], lim[2:], "k--", lw=1.2, dashes=(6, 4), zorder=5)
        else:
            ax.plot(lim[:2], [theta, theta], "k--", lw=1.2, dashes=(6, 4), zorder=5)
        ax.text(0.96, 0.95, r"$m = %d$" % m, transform=ax.transAxes, ha="right", va="top", fontsize=12)
        ax.set_xlim(lim[:2])
        ax.set_ylim(lim[2:])
        ax.set_xticks([-1, 0, 1, 2])
        ax.set_yticks([-2, 0, 2])
    fig.tight_layout()
    tr, expo, prod, _ = adaboost_curves(model, X, t)
    print("  AdaBoost: eps_1..3 = %s, alpha_1..3 = %s; training error at m = 150: %.3f (first zero at m = %d);"
          " (1/N) sum exp(-t f) at m = 150: %.5f" % (np.round(model["eps"][:3], 4), np.round(model["alpha"][:3], 4),
                                                    tr[149], np.argmax(tr == 0) + 1, expo[149]))
    common.savefig(fig, 14, "fig14-2", out_dir)


def fig14_3(out_dir=None):
    """Figure 14.3: exponential, rescaled cross-entropy, hinge and misclassification errors against z = t y."""
    fig, ax = plt.subplots(figsize=(4.4, 3.4))
    z = np.linspace(-2, 2, 801)
    top = 3.1
    ze = z[z >= -np.log(top)]
    ax.plot(ze, exponential_error(ze), color=common.GREEN, lw=2.2, label="exponential")
    ax.plot(z, cross_entropy_error(z), color=common.RED, lw=2.2, label=r"cross-entropy $/\ln 2$")
    ax.plot(z, hinge_error(z), color=common.BLUE, lw=2.2, label="hinge", zorder=5)
    ax.plot([-2, 0], [1, 1], color="k", lw=2.2, label="misclassification")
    arrow_axes(ax, (-2.2, 2.2), (0, top + 0.15), r"$z$", r"$E(z)$")
    ax.set_xticks([-2, -1, 0, 1, 2])
    ax.tick_params(axis="x", length=0, pad=6, labelsize=11)
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=2, frameon=False, fontsize=9)
    common.savefig(fig, 14, "fig14-3", out_dir)


def fig14_4(out_dir=None):
    """Figure 14.4: squared error (green) and absolute error (red) against z = y - t."""
    fig, ax = plt.subplots(figsize=(4.4, 2.6))
    zs = np.linspace(-np.sqrt(2), np.sqrt(2), 401)
    z = np.linspace(-2, 2, 401)
    ax.plot(z, np.abs(z), color=common.RED, lw=2.2, label="absolute error $|z|$")
    ax.plot(zs, zs ** 2, color=common.GREEN, lw=2.2, label="squared error $z^2$")
    arrow_axes(ax, (-2.2, 2.2), (0, 2.15), r"$z$", r"$E(z)$", ylabel_side="left")
    ax.set_xticks([-1, 0, 1])
    ax.tick_params(axis="x", length=0, pad=6, labelsize=11)
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.14), ncol=2, frameon=False, fontsize=9)
    common.savefig(fig, 14, "fig14-4", out_dir)


def fig14_5(out_dir=None):
    """Figure 14.5: a two-dimensional input space split into five regions A to E by axis-aligned thresholds."""
    tree = book_tree()
    fig, ax = plt.subplots(figsize=(4.2, 3.6))
    draw_partition(ax, tree, (0, 1, 0, 1))
    arrow_axes(ax, (0, 1.0), (0, 1.0), r"$x_1$", "")
    ax.text(-0.02, 1.02, r"$x_2$", ha="right", va="top", fontsize=13)
    ax.set_xticks([0.4, 0.8])
    ax.set_xticklabels([r"$\theta_1$", r"$\theta_4$"], fontsize=13)
    ax.set_yticks([0.3, 0.7])
    ax.set_yticklabels([r"$\theta_2$", r"$\theta_3$"], fontsize=13)
    ax.tick_params(length=4, width=1.2)
    ax.set_xlim(-0.01, 1.06)
    ax.set_ylim(-0.01, 1.06)
    common.savefig(fig, 14, "fig14-5", out_dir)


def fig14_6(out_dir=None):
    """Figure 14.6: the binary tree of the partition in Figure 14.5."""
    tree = book_tree()
    pos = tree_layout(tree)
    fig, ax = plt.subplots(figsize=(4.6, 3.3))
    for node in tree.internal():
        for child, side in ((node.left, "left"), (node.right, "right")):
            (xa, ya), (xb, yb) = pos[id(node)], pos[id(child)]
            ax.plot([xa, xb], [ya, yb], color=common.GREEN, lw=2.2, zorder=1)
            if getattr(node, "side", None) == side:
                var = r"x_%d" % (node.feature + 1)
                rel = r"\leq" if side == "left" else ">"
                lab = r"$%s %s \theta_%d$" % (var, rel, node.theta_index)
                xm, ym = (xa + xb) / 2, (ya + yb) / 2
                ax.text(xm + (0.12 if side == "right" else -0.12), ym + 0.08, lab, fontsize=13,
                        ha="left" if side == "right" else "right", va="center")
    for node in [tree] + tree.internal()[1:] + tree.leaves():
        x, y = pos[id(node)]
        ax.plot(x, y, "o", ms=17, color=common.GREEN, zorder=2)
    for leaf in tree.leaves():
        x, y = pos[id(leaf)]
        ax.text(x, y - 0.42, leaf.name, ha="center", va="top", fontsize=16, family="DejaVu Sans")
    ax.set_xlim(-0.4, len(tree.leaves()) - 0.6)
    ax.set_ylim(-3.85, 0.3)
    ax.axis("off")
    common.savefig(fig, 14, "fig14-6", out_dir)


def run_two_lines_em(n_iter=50):
    """EM of Figure 14.8 from the book's start: lines MIX_W0, pi = (1/2, 1/2), beta = 1/var(t)."""
    Phi, t, x = two_lines_data()
    beta0 = 1.0 / np.var(t, ddof=1)
    return mixture_linear_regression_em(Phi, t, MIX_W0, [0.5, 0.5], beta0, n_iter), Phi, t, x


def fig14_8(out_dir=None):
    """Figure 14.8: EM for a mixture of two linear regressions: start, 30 and 50 iterations, responsibilities."""
    states, Phi, t, x = run_two_lines_em(50)
    fig, axes = common.panel_grid(2, 3, size=(3.0, 2.25))
    xx = np.array([-1.0, 1.0])
    for j, it in enumerate((0, 30, 50)):
        s = states[it]
        ax = axes[0, j]
        ax.scatter(x, t, s=46, facecolors="none", edgecolors=RING_GREEN, linewidths=1.6, zorder=3)
        ax.plot(xx, s["W"][0, 0] + s["W"][0, 1] * xx, color=common.BLUE, lw=2.0)
        ax.plot(xx, s["W"][1, 0] + s["W"][1, 1] * xx, color=common.RED, lw=2.0)
        ax.set_xlim(-1.1, 1.1)
        ax.set_ylim(-1.6, 1.6)
        ax.set_xticks([-1, -0.5, 0, 0.5, 1])
        ax.set_yticks([-1.5, -1, -0.5, 0, 0.5, 1, 1.5])
        ax.set_title(("initial configuration", "after 30 EM iterations", "after 50 EM iterations")[j], fontsize=10)
        ax = axes[1, j]
        g = s["gamma"][:, 1]                       # responsibility of the red line
        ax.vlines(x, 0, g, color=common.RED, lw=2.2)
        ax.vlines(x, g, 1, color=common.BLUE, lw=2.2)
        ax.hlines(g, x - 0.035, x + 0.035, color="k", lw=1.4, zorder=3)
        ax.set_xlim(-1.1, 1.1)
        ax.set_ylim(0, 1)
        ax.set_xticks([-1, -0.5, 0, 0.5, 1])
        ax.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1])
        print("  iteration %2d: blue line %.3f %+.3f x, red line %.3f %+.3f x, pi = (%.3f, %.3f), beta = %.2f,"
              " log likelihood %.3f" % (it, s["W"][0, 0], s["W"][0, 1], s["W"][1, 0], s["W"][1, 1], s["pi"][0],
                                        s["pi"][1], s["beta"], s["loglik"]))
    fig.tight_layout()
    common.savefig(fig, 14, "fig14-8", out_dir)


def predictive_density_grid(W, pi, beta, xg, tg):
    """p(t | x) of a mixture of linear regressions (14.34) on a grid (rows t, columns x)."""
    mu = W[:, 0][None, None, :] + W[:, 1][None, None, :] * xg[None, :, None]
    dens = np.sum(pi * np.sqrt(beta / (2 * np.pi)) * np.exp(-0.5 * beta * (tg[:, None, None] - mu) ** 2), axis=2)
    return dens


def fig14_9(out_dir=None):
    """Figure 14.9: predictive density p(t|x) of the converged mixture and of a single linear regression."""
    states, Phi, t, x = run_two_lines_em(50)
    s = states[-1]
    w1, beta1, ll1 = linear_regression_ml(Phi, t)
    xg, tg = np.linspace(-1, 1, 300), np.linspace(-1.5, 1.5, 300)
    fig, axes = common.panel_grid(1, 2, size=(3.6, 2.8))
    for ax, (W, pi, beta, ll, lines) in zip(axes[0], (
            (s["W"], s["pi"], s["beta"], s["loglik"], ((s["W"][0], common.BLUE), (s["W"][1], common.RED))),
            (w1[None, :], np.ones(1), beta1, ll1, ((w1, common.RED),)))):
        dens = predictive_density_grid(W, pi, beta, xg, tg)
        ax.imshow(dens / dens.max(), extent=(-1, 1, -1.5, 1.5), origin="lower", cmap=WHITE_MAGENTA, aspect="auto",
                  vmin=0, vmax=1, interpolation="bilinear")
        for w, col in lines:
            ax.plot([-1, 1], [w[0] - w[1], w[0] + w[1]], color=col, lw=1.8)
        ax.scatter(x, t, s=46, facecolors="none", edgecolors=RING_GREEN, linewidths=1.6, zorder=3)
        ax.text(-1.05, 1.55, r"$\ln p(\mathbf{t}) = %.1f$" % ll, ha="left", va="top", fontsize=9)
        ax.set_xlim(-1.1, 1.1)
        ax.set_ylim(-1.6, 1.6)
        ax.set_xticks([-1, -0.5, 0, 0.5, 1])
        ax.set_yticks([-1.5, -1, -0.5, 0, 0.5, 1, 1.5])
    fig.tight_layout()
    print("  log likelihood: mixture %.3f (book -3.0), single linear regression %.3f (book -27.6)" % (s["loglik"], ll1))
    common.savefig(fig, 14, "fig14-9", out_dir)


def fit_fig14_10(seed=25):
    """Data, single logistic regression and the two-component mixture of Figure 14.10."""
    X, t = four_clusters(seed=seed)
    Phi = np.c_[np.ones(len(X)), X]
    w1 = irls_logistic(Phi, t, np.ones(len(t)), np.zeros(3), n_steps=50)
    W0 = np.array([[1.0, 1.0, 0.0], [1.0, -1.0, 0.0]])      # boundaries x1 = -1 and x1 = +1
    W, pi, lls = mixture_logistic_em(Phi, t, W0, [0.5, 0.5], n_iter=500)
    return X, t, Phi, w1, W, pi, lls


def fig14_10(out_dir=None):
    """Figure 14.10: true class probabilities, a single logistic regression and a mixture of two."""
    X, t, Phi, w1, W, pi, lls = fit_fig14_10()
    g = np.linspace(-3, 3, 241)
    GX, GY = np.meshgrid(g, g)
    G = np.c_[GX.ravel(), GY.ravel()]
    Pg = np.c_[np.ones(len(G)), G]
    p_true = four_clusters_prob_blue(G).reshape(GX.shape)
    p_one = common.sigmoid(Pg @ w1).reshape(GX.shape)
    p_mix = (common.sigmoid(Pg @ W.T) @ pi).reshape(GX.shape)
    fig, axes = common.panel_grid(1, 3, size=(2.75, 2.75))
    ll1 = mixture_logistic_loglik(Phi, t, w1[None, :], np.ones(1))[0]
    for ax, p, lines in zip(axes[0], (p_true, p_one, p_mix), ([], [w1], [W[0], W[1]])):
        img = np.stack([0.5 + 0.5 * (1 - p), 0.5 * np.ones_like(p), 0.5 + 0.5 * p], axis=-1)   # red ... blue
        ax.imshow(img, extent=(-3, 3, -3, 3), origin="lower", interpolation="bilinear")
        ax.scatter(X[t == 0, 0], X[t == 0, 1], s=11, facecolors="none", edgecolors="#ff0000", linewidths=0.8)
        ax.scatter(X[t == 1, 0], X[t == 1, 1], s=11, facecolors="none", edgecolors="#0000ff", linewidths=0.8)
        for w in lines:                     # the line w^T phi = 0, where that model gives probability 1/2
            if abs(w[2]) > abs(w[1]):
                xs = np.array([-3, 3.0])
                ax.plot(xs, -(w[0] + w[1] * xs) / w[2], "k", lw=1.0)
            else:
                ys = np.array([-3, 3.0])
                ax.plot(-(w[0] + w[2] * ys) / w[1], ys, "k", lw=1.0)
        ax.set_xlim(-3, 3)
        ax.set_ylim(-3, 3)
        ax.set_xticks([-2, 0, 2])
        ax.set_yticks([-3, -2, -1, 0, 1, 2, 3])
    fig.tight_layout()
    print("  log likelihood: single logistic regression %.2f, mixture of two %.2f (pi = %s, %d EM iterations)"
          % (ll1, lls[-1], np.round(pi, 3), len(lls) - 1))
    common.savefig(fig, 14, "fig14-10", out_dir)


FIGURES = {
    "14.2": fig14_2,
    "14.3": fig14_3,
    "14.4": fig14_4,
    "14.5": fig14_5,
    "14.6": fig14_6,
    "14.8": fig14_8,
    "14.9": fig14_9,
    "14.10": fig14_10,
}

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