"""Bishop (2006), Chapter 9: Mixture Models and EM.

Reproduces the numerical figures of Chapter 9 of C. M. Bishop, Pattern Recognition
and Machine Learning (Springer, 2006), pp. 423-459, written from scratch with numpy
and matplotlib. The algorithms are small functions meant to be read next to the
lecture notes (ch09-mixture-models-em.html):

* K-means (Section 9.1): kmeans_e_step (9.2), kmeans_m_step (9.4), distortion (9.1),
  kmeans (the batch algorithm).
* EM for Gaussian mixtures (Section 9.2.2): gmm_e_step (9.23), gmm_m_step (9.24)-(9.27),
  gmm_loglik (9.28), em_gmm.
* EM for mixtures of Bernoulli distributions (Section 9.3.3): bernoulli_e_step (9.56),
  bernoulli_m_step (9.57)-(9.60), bernoulli_mixture_em.
* The lower bound L(q, theta) (9.71) and KL(q||p) (9.72) of Section 9.4 for a mixture of
  two 1-D Gaussians with one free parameter (the mix1d_* functions), used for Figure 9.14
  and by the demo em-bound-demo.html.

Figures: 9.1, 9.2, 9.3, 9.5, 9.7, 9.8, 9.10, 9.14. Run `python3 ch09_mixture_em.py --list`,
`python3 ch09_mixture_em.py` (all figures, under a minute) or
`python3 ch09_mixture_em.py --fig 9.1 9.8`. The PNGs are written to ../img/ch09/.

What differs from the book, and the choices made:
* Figures 9.1, 9.2 and 9.8: the book does not print its starting values; we read them off
  Figure 9.1(a) and Figure 9.8(a): centres (1.4, -1.0) (red) and (-1.4, 1.0) (blue), and,
  for EM, covariances 0.6 I (the circles in Figure 9.8(a) have radius 0.77) and pi_k = 1/2.
  With this start every panel matches the book's: J after the first E step is 1048.4
  (Figure 9.2 shows about 1050), and the fourth E step still moves one point from blue to
  red, as in the book's panel (h); the fifth E step changes nothing. EM reaches
  ln p(X) = -385.46 (standardised data) and is within 1e-8 of its limit by L = 20.
* Figure 9.3: the book's photographs are not available; we use two public-domain
  photographs (Grace Hopper, U.S. Navy; astronaut Eileen Collins, NASA), averaged over
  3 x 3 pixel blocks (200 x 170 and 170 x 170 pixels; the book's image has 240 x 180).
  K-means starts from K randomly chosen pixels (seeded; 3 starts, the lowest J kept) and
  runs until no assignment changes.
* Figure 9.5: the mixture of Figure 2.23 is not tabulated in the book. Fitting ellipses
  to the contours of Figure 2.23(a) gives pi = (0.5, 0.3, 0.2), means (0.2, 0.4),
  (0.5, 0.5), (0.8, 0.6) and covariances 0.016 [[1, r], [r, 1]] with r = 0.75 for the red
  and blue components and r = -0.75 for the green one; we sample 500 points from it.
* Figure 9.7: the book's picture is a sketch. We compute it for seven points: the broad
  component is the single-Gaussian fit to all seven points, the narrow one sits on the
  right-most point with sd 0.15, and the weights are 6/7 and 1/7. Panel (b), which is not
  in the book, shows ln p(X) as the narrow sd shrinks (it grows without bound), against a
  single Gaussian collapsing onto the same point (it falls without bound; compare Figure 1.14).
* Figure 9.10: MNIST training images (the first 200 of each of the digits 2, 3 and 4),
  binarised at 0.5. EM is initialised exactly as the book describes (pi_k = 1/3; mu_ki
  uniform on (0.25, 0.75), then normalised so that sum_i mu_ki = 1) and run for 10
  iterations. Logs of mu are taken after clipping to [1e-10, 1 - 1e-10], because some
  border pixels are never on. With this initialisation 10 iterations do not always
  separate the three digits (local maxima); with our seed 526 of the 600 images land
  in the component of their own digit.
* Figure 9.14: the book's curves are schematic. We compute them for a real model:
  40 points (the "two clusters" data of em-bound-demo.html) and the mixture
  0.5 N(x|0, 2^2) + 0.5 N(x|theta, 0.6^2) whose only free parameter theta is the second
  mean; theta_old = 0. ln p(X|theta) is then not concave (a second, lower maximum near
  theta = -2.1), the bound at theta_old touches it with the same slope, and one M step
  moves theta from 0 to 0.712.
"""
import math

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import patheffects
from matplotlib.colors import LinearSegmentedColormap, to_rgb

import common

# --------------------------------------------------------------------------- shared settings
BOOK_CENTRES = np.array([[1.4, -1.0], [-1.4, 1.0]])  # read off Figure 9.1(a): red, then blue
BOOK_COV = 0.6  # Figure 9.8(a): initial covariances 0.6 I (circles of radius 0.77)
LIM = (-2.5, 2.5)  # axis range of Figures 9.1 and 9.8
HALO = [patheffects.Stroke(linewidth=5.5, foreground="white"), patheffects.Normal()]
EPS = 1e-10  # clip for Bernoulli logs


# =========================================================================== K-means (Section 9.1)
def distortion(X, r, mu):
    """The K-means objective J of (9.1): sum of squared distances of each point to its centre.

    r[n] is the index k of the cluster of x_n, i.e. r_nk = 1 for k = r[n] in the book's 1-of-K notation.
    """
    return float(((X - mu[r]) ** 2).sum())


def kmeans_e_step(X, mu):
    """E step of K-means (9.2): assign every point to its nearest centre (squared Euclidean distance)."""
    d2 = ((X[:, None, :] - mu[None, :, :]) ** 2).sum(-1)
    return d2.argmin(1)


def kmeans_m_step(X, r, mu_old):
    """M step of K-means (9.4): move each centre to the mean of the points assigned to it.

    A centre that has no points keeps its old position (the book does not discuss empty clusters).
    """
    mu = mu_old.copy()
    for k in range(len(mu)):
        if np.any(r == k):
            mu[k] = X[r == k].mean(0)
    return mu


def kmeans(X, mu0, max_iter=500, record=False):
    """Batch K-means (Section 9.1) started from the centres mu0.

    Alternates the E step (9.2) and the M step (9.4) until an E step changes no assignment
    (Exercise 9.1: this happens after finitely many steps, because J never increases and
    there are finitely many assignments). Returns a dict with the final centres 'mu' and
    assignments 'r', 'n_m_steps' (number of M steps taken), 'J' (J after every half step, in
    order E, M, E, M, ..., ending with the E step that changed nothing) and, with
    record=True, 'steps': a list of ('E' or 'M', r, mu) after every half step.
    """
    mu = np.array(mu0, dtype=float)
    r, Js, steps, n_m = None, [], [], 0
    for _ in range(max_iter):
        r_new = kmeans_e_step(X, mu)
        Js.append(distortion(X, r_new, mu))
        if record:
            steps.append(("E", r_new, mu.copy()))
        if r is not None and np.array_equal(r_new, r):
            break
        r = r_new
        mu = kmeans_m_step(X, r, mu)
        n_m += 1
        Js.append(distortion(X, r, mu))
        if record:
            steps.append(("M", r, mu.copy()))
    return {"mu": mu, "r": r, "J": Js, "n_m_steps": n_m, "steps": steps}


# =========================================================================== EM for Gaussian mixtures (Section 9.2)
def gmm_log_terms(X, pi, mu, Sigma):
    """N x K array of ln pi_k + ln N(x_n | mu_k, Sigma_k): the terms inside the sum of (9.14)."""
    return np.stack([np.log(pi[k]) + common.gauss_logpdf(X, mu[k], Sigma[k]) for k in range(len(pi))], axis=1)


def gmm_loglik(X, pi, mu, Sigma):
    """Log likelihood (9.14) = (9.28): sum_n ln sum_k pi_k N(x_n | mu_k, Sigma_k)."""
    return float(common.logsumexp(gmm_log_terms(X, pi, mu, Sigma), axis=1).sum())


def gmm_e_step(X, pi, mu, Sigma):
    """E step (9.23): responsibilities gamma(z_nk) = pi_k N(x_n|mu_k, Sigma_k) / sum_j pi_j N(x_n|mu_j, Sigma_j).

    Computed on the log scale (log-sum-exp) for stability. Also returns ln p(X | pi, mu, Sigma),
    (9.28), at the parameters used.
    """
    T = gmm_log_terms(X, pi, mu, Sigma)
    lse = common.logsumexp(T, axis=1)
    return np.exp(T - lse[:, None]), float(lse.sum())


def floor_eigenvalues(S, floor):
    """Raise every eigenvalue of the symmetric matrix S to at least `floor`.

    This is the maximiser of the Gaussian M-step objective under the constraint that all
    eigenvalues of Sigma_k are at least `floor`, so EM with it still never decreases the
    log likelihood. It stops a component collapsing onto a data point (Section 9.2.1).
    """
    w, V = np.linalg.eigh(S)
    return (V * np.maximum(w, floor)) @ V.T


def gmm_m_step(X, gamma, var_floor=0.0, shared_var=None):
    """M step (9.24)-(9.27): re-estimate pi, mu and Sigma from the responsibilities gamma (N x K).

    N_k = sum_n gamma_nk (9.27); mu_k = sum_n gamma_nk x_n / N_k (9.24); Sigma_k = sum_n gamma_nk
    (x_n - mu_k)(x_n - mu_k)^T / N_k with the NEW mean (9.25); pi_k = N_k / N (9.26).
    var_floor > 0 applies floor_eigenvalues to every Sigma_k. shared_var = eps keeps every
    Sigma_k fixed at eps I (the model of Section 9.3.2 whose eps -> 0 limit is K-means).
    """
    N, K = gamma.shape
    D = X.shape[1]
    Nk = gamma.sum(0)
    mu = (gamma.T @ X) / Nk[:, None]
    Sigma = np.empty((K, D, D))
    for k in range(K):
        if shared_var is not None:
            Sigma[k] = shared_var * np.eye(D)
            continue
        d = X - mu[k]
        Sigma[k] = (gamma[:, k, None] * d).T @ d / Nk[k]
        if var_floor > 0:
            Sigma[k] = floor_eigenvalues(Sigma[k], var_floor)
    return Nk / N, mu, Sigma


def em_gmm(X, pi, mu, Sigma, n_iter=100, tol=None, var_floor=0.0, shared_var=None):
    """EM for a Gaussian mixture (the box on pp. 438-439), started from (pi, mu, Sigma).

    Returns a list `hist`: hist[0] = {'pi', 'mu', 'Sigma', 'loglik'} is the start and hist[L]
    is the state after L complete cycles (E step, then M step), with 'gamma' the
    responsibilities of the L-th E step and 'loglik' = ln p(X) at the new parameters.
    Stops after n_iter cycles, or earlier when the log likelihood rises by less than tol.
    """
    pi, mu, Sigma = np.array(pi, float), np.array(mu, float), np.array(Sigma, float)
    hist = [{"pi": pi, "mu": mu, "Sigma": Sigma, "loglik": gmm_loglik(X, pi, mu, Sigma)}]
    for _ in range(n_iter):
        gamma, _ = gmm_e_step(X, pi, mu, Sigma)
        pi, mu, Sigma = gmm_m_step(X, gamma, var_floor, shared_var)
        hist.append({"pi": pi, "mu": mu, "Sigma": Sigma, "gamma": gamma, "loglik": gmm_loglik(X, pi, mu, Sigma)})
        if tol is not None and hist[-1]["loglik"] - hist[-2]["loglik"] < tol:
            break
    return hist


def sample_gmm(n, pi, mu, Sigma, rng):
    """Ancestral sampling (Section 8.1.2): z_n ~ p(z) = pi, then x_n ~ N(mu_{z_n}, Sigma_{z_n}).

    Returns X (n x D) and the component labels z (n,).
    """
    z = rng.choice(len(pi), size=n, p=pi)
    X = np.empty((n, mu.shape[1]))
    for k in range(len(pi)):
        idx = np.nonzero(z == k)[0]
        X[idx] = rng.multivariate_normal(mu[k], Sigma[k], size=len(idx))
    return X, z


# =========================================================================== EM for Bernoulli mixtures (Section 9.3.3)
def bernoulli_log_terms(X, pi, mu):
    """N x K array of ln pi_k + ln p(x_n | mu_k), with ln p(x|mu_k) = sum_i x_i ln mu_ki + (1 - x_i) ln(1 - mu_ki) (9.48)."""
    m = np.clip(mu, EPS, 1 - EPS)
    return X @ np.log(m).T + (1 - X) @ np.log(1 - m).T + np.log(pi)


def bernoulli_e_step(X, pi, mu):
    """E step (9.56): gamma(z_nk) proportional to pi_k p(x_n | mu_k); also returns ln p(X | mu, pi) (9.51)."""
    T = bernoulli_log_terms(X, pi, mu)
    lse = common.logsumexp(T, axis=1)
    return np.exp(T - lse[:, None]), float(lse.sum())


def bernoulli_m_step(X, gamma):
    """M step: N_k = sum_n gamma_nk (9.57), mu_k = xbar_k = sum_n gamma_nk x_n / N_k (9.58), (9.59), pi_k = N_k / N (9.60)."""
    Nk = gamma.sum(0)
    return Nk / len(X), (gamma.T @ X) / Nk[:, None]


def bernoulli_mixture_em(X, pi, mu, n_iter=10):
    """n_iter cycles of EM for a mixture of Bernoulli distributions (9.47).

    Returns the final (pi, mu, gamma) and the list of log likelihoods (9.51): entry t is
    ln p(X) at the parameters after t M steps (entry 0: the start).
    """
    pi, mu = np.array(pi, float), np.array(mu, float)
    lls = []
    for _ in range(n_iter):
        gamma, ll = bernoulli_e_step(X, pi, mu)
        lls.append(ll)
        pi, mu = bernoulli_m_step(X, gamma)
    gamma, ll = bernoulli_e_step(X, pi, mu)
    lls.append(ll)
    return pi, mu, gamma, lls


# =========================================================================== the bound of Section 9.4, one free parameter
# Model: p(x | theta) = (1 - pi) N(x | m1, s1^2) + pi N(x | m2, s2^2); theta replaces either
# m2 ('mu2') or pi ('pi'), the other parameters are fixed. q_n is the probability that point n
# belongs to component 2; the posterior is q_n = gamma(z_n2), the E step of (9.13).
EM_BOUND_SCENARIOS = {
    # name: data generator (seed, N, weight of cluster 2, (mean, sd) of cluster 1, (mean, sd) of cluster 2),
    #       model (fixed values), default free parameter and start. Mirrors em-bound-demo.html.
    "two": {"seed": 11, "n": 40, "w2": 0.65, "c1": (-2.0, 0.6), "c2": (1.5, 0.6),
            "model": {"m1": 0.0, "s1": 2.0, "m2": 1.5, "s2": 0.6, "pi": 0.5}, "free": "mu2", "theta0": 0.0},
    "overlap": {"seed": 3, "n": 100, "w2": 0.6, "c1": (-0.5, 1.0), "c2": (0.5, 1.0),
                "model": {"m1": -0.5, "s1": 1.0, "m2": 0.5, "s2": 1.0, "pi": 0.6}, "free": "pi", "theta0": 0.1},
    "separated": {"seed": 11, "n": 40, "w2": 0.65, "c1": (-2.0, 0.6), "c2": (1.5, 0.6),
                  "model": {"m1": -2.0, "s1": 0.6, "m2": 1.5, "s2": 0.6, "pi": 0.65}, "free": "pi", "theta0": 0.1},
}


def _lognorm(x, m, s):
    return -0.5 * np.log(2 * np.pi * s * s) - (x - m) ** 2 / (2 * s * s)


def mix1d_terms(x, theta, model, free):
    """Per point: a_n = ln[(1 - pi) N(x_n|m1, s1^2)] and b_n = ln[pi N(x_n|m2, s2^2)], with theta as the free parameter."""
    pi, m2 = model["pi"], model["m2"]
    if free == "mu2":
        m2 = theta
    else:
        pi = theta
    return np.log(1 - pi) + _lognorm(x, model["m1"], model["s1"]), np.log(pi) + _lognorm(x, m2, model["s2"])


def mix1d_loglik(x, theta, model, free):
    """ln p(X | theta) = sum_n ln(exp(a_n) + exp(b_n)), (9.14) for this model."""
    a, b = mix1d_terms(x, theta, model, free)
    return float(np.logaddexp(a, b).sum())


def mix1d_posterior(x, theta, model, free):
    """E step: q_n = p(z_n = 2 | x_n, theta) = exp(b_n) / (exp(a_n) + exp(b_n)) (9.13)."""
    a, b = mix1d_terms(x, theta, model, free)
    return 1.0 / (1.0 + np.exp(a - b))


def _xlogx(q):
    q = np.asarray(q, float)
    return np.where(q > 0, q * np.log(np.where(q > 0, q, 1.0)), 0.0)


def mix1d_bound(x, theta, q, model, free):
    """Lower bound (9.71): L(q, theta) = sum_n sum_z q(z_n) ln[p(x_n, z_n | theta) / q(z_n)]."""
    a, b = mix1d_terms(x, theta, model, free)
    return float(((1 - q) * a + q * b - _xlogx(q) - _xlogx(1 - q)).sum())


def mix1d_kl(x, theta, q, model, free):
    """KL(q || p(Z|X, theta)) of (9.72), summed over the points (the posterior factorises, (9.75))."""
    g = mix1d_posterior(x, theta, model, free)
    return float((_xlogx(q) - q * np.log(g) + _xlogx(1 - q) - (1 - q) * np.log(1 - g)).sum())


def mix1d_m_step(x, q, model, free):
    """M step: the theta maximising L(q, theta) (a weighted mean for mu2, (9.17); the mean of q for pi, (9.22))."""
    return float((q * x).sum() / q.sum()) if free == "mu2" else float(q.mean())


def mix1d_em(x, theta0, model, free, tol=1e-10, max_iter=5000):
    """Iterate E and M steps from theta0 until theta moves by less than tol; returns the list of thetas."""
    th = [float(theta0)]
    for _ in range(max_iter):
        th.append(mix1d_m_step(x, mix1d_posterior(x, th[-1], model, free), model, free))
        if abs(th[-1] - th[-2]) < tol:
            break
    return th


# =========================================================================== the demos' random numbers
class DemoRNG:
    """The seeded generator DK.rng of courses/demo-kit.js (mulberry32 + Box-Muller), ported to Python.

    The demos kmeans-em-demo.html and em-bound-demo.html draw their random starts and data
    with it, so the numbers below reproduce what a student sees in the browser.
    """

    def __init__(self, seed):
        self.a = seed & 0xFFFFFFFF
        self.spare = None

    def u(self):
        self.a = (self.a + 0x6D2B79F5) & 0xFFFFFFFF
        a = self.a
        t = ((a ^ (a >> 15)) * (1 | a)) & 0xFFFFFFFF
        t = ((t + (((t ^ (t >> 7)) * (61 | t)) & 0xFFFFFFFF)) & 0xFFFFFFFF) ^ t
        return ((t ^ (t >> 14)) & 0xFFFFFFFF) / 4294967296.0

    def norm(self):
        if self.spare is not None:
            s, self.spare = self.spare, None
            return s
        x = 0.0
        while x == 0.0:
            x = self.u()
        r, th = math.sqrt(-2 * math.log(x)), 2 * math.pi * self.u()
        self.spare = r * math.sin(th)
        return r * math.cos(th)


def em_bound_data(name):
    """The 1-D data set `name` of EM_BOUND_SCENARIOS, generated exactly as em-bound-demo.html does."""
    s = EM_BOUND_SCENARIOS[name]
    rng, x = DemoRNG(s["seed"]), []
    for _ in range(s["n"]):
        second = rng.u() < s["w2"]
        m, sd = s["c2"] if second else s["c1"]
        x.append(m + sd * rng.norm())
    return np.array(x)


def demo_random_start(N, K, seed):
    """Indices of the K data points used as random starting centres in kmeans-em-demo.html."""
    rng, idx = DemoRNG(seed), list(range(N))
    for i in range(K):
        j = i + int(math.floor(rng.u() * (N - i)))
        idx[i], idx[j] = idx[j], idx[i]
    return idx[:K]


def demo_poor_start(K):
    """The demo's 'poor start': K centres evenly spaced from (1.4, -1) to (-1.4, 1) (the book's two for K = 2)."""
    if K == 1:
        return np.zeros((1, 2))
    t = np.linspace(0, 1, K)[:, None]
    return BOOK_CENTRES[0] + t * (BOOK_CENTRES[1] - BOOK_CENTRES[0])


# =========================================================================== plotting helpers
def _mix_colours(gamma, colours):
    """Colour each point by its responsibilities: proportions gamma_nk of the colours (Figures 9.5c, 9.8)."""
    C = np.array([to_rgb(c) for c in colours])
    return np.clip(gamma @ C, 0, 1)


def _scatter(ax, X, c, s=17):
    ax.scatter(X[:, 0], X[:, 1], s=s, c=c, edgecolors="none", zorder=2)


def _cross(ax, p, colour, size=11):
    ax.plot(p[0], p[1], marker="x", ms=size, mew=2.8, color=colour, path_effects=HALO, zorder=5, linestyle="none")


def _square(ax, lim=LIM, ticks=(-2, 0, 2)):
    ax.set_xlim(lim)
    ax.set_ylim(lim)
    ax.set_xticks(ticks, ["%g" % t for t in ticks])
    ax.set_yticks(ticks, ["%g" % t for t in ticks])
    ax.set_aspect("equal")
    ax.tick_params(direction="in", top=True, right=True)


def _bisector(ax, m1, m2, colour=common.MAGENTA):
    """The perpendicular bisector of two centres: the K = 2 decision boundary of K-means (Figure 9.1)."""
    mid, d = (m1 + m2) / 2, m1 - m2
    u = np.array([-d[1], d[0]]) / np.linalg.norm(d)
    p = mid[:, None] + u[:, None] * np.array([-10, 10])
    ax.plot(p[0], p[1], color=colour, lw=2, zorder=3)


def _ellipse(ax, mu, S, colour, lw=2.4):
    """1-standard-deviation contour of N(mu, S), drawn with a white halo as in Figure 9.8."""
    x, y = common.ellipse_xy(mu, S, 1.0)
    ax.plot(x, y, color=colour, lw=lw, zorder=4,
            path_effects=[patheffects.Stroke(linewidth=lw + 3, foreground="white"), patheffects.Normal()])


# =========================================================================== Figures 9.1 and 9.2
def fig9_1(out_dir=None):
    """Figure 9.1: K-means on the re-scaled Old Faithful data, K = 2, from the book's poor start (nine panels)."""
    X = common.old_faithful(standardize=True)
    res = kmeans(X, BOOK_CENTRES, record=True)
    cols = np.array([common.RED, common.BLUE])
    fig, axes = common.panel_grid(3, 3, size=(2.35, 2.35))
    states = [None] + res["steps"][:8]  # (a) start, (b) E1, (c) M1, ..., (i) M4
    for i, ax in enumerate(axes.flat):
        st = states[i]
        if st is None:
            _scatter(ax, X, common.GREEN)
            mu = BOOK_CENTRES
        else:
            kind, r, mu = st
            _scatter(ax, X, cols[r])
            if kind == "E":
                _bisector(ax, mu[0], mu[1])
        for k in range(2):
            _cross(ax, mu[k], cols[k])
        _square(ax)
        common.panel_label(ax, "abcdefghi"[i])
    fig.subplots_adjust(wspace=0.12, hspace=0.12)
    common.savefig(fig, 9, "fig9-1", out_dir)
    print("  K-means: %d M steps until an E step changes nothing; final centres %s"
          % (res["n_m_steps"], np.round(res["mu"], 3).tolist()))


def fig9_2(out_dir=None):
    """Figure 9.2: the distortion J of (9.1) after each E step (blue) and M step (red) of Figure 9.1."""
    X = common.old_faithful(standardize=True)
    J = kmeans(X, BOOK_CENTRES)["J"][:8]
    t = 0.5 * np.arange(1, 9)
    fig, ax = plt.subplots(figsize=(4.4, 3.0))
    ax.plot(t, J, color=common.GREEN, lw=2, zorder=1)
    for i in range(8):
        ax.plot(t[i], J[i], "o", ms=9, mew=2, mfc="white", mec=common.BLUE if i % 2 == 0 else common.RED, zorder=2)
    ax.set_xlim(0, 4.5)
    ax.set_ylim(0, 1100)
    ax.set_xticks([1, 2, 3, 4])
    ax.set_yticks([0, 500, 1000])
    ax.set_ylabel("$J$", rotation=0, labelpad=12)
    ax.tick_params(direction="in", top=True, right=True)
    common.savefig(fig, 9, "fig9-2", out_dir)
    print("  J after E1, M1, E2, ...:", np.round(J, 2).tolist())


# =========================================================================== Figure 9.3
def _block_mean(img, f=3):
    H, W = (img.shape[0] // f) * f, (img.shape[1] // f) * f
    return img[:H, :W].reshape(H // f, f, W // f, f, 3).mean((1, 3))


def segment_image(img, K, rng, n_starts=3):
    """K-means on the RGB vectors of an image (Section 9.1.1): every pixel is a point in [0, 1]^3.

    Each start uses K randomly chosen pixels as centres (as suggested on p. 427); the run with the
    lowest J is kept. Returns the image with each pixel replaced by its centre, and J.
    """
    X = img.reshape(-1, 3)
    best = None
    for _ in range(n_starts):
        res = kmeans(X, X[rng.choice(len(X), K, replace=False)])
        if best is None or res["J"][-1] < best["J"][-1]:
            best = res
    return best["mu"][best["r"]].reshape(img.shape), best["J"][-1]


def fig9_3(out_dir=None):
    """Figure 9.3: image segmentation and compression by K-means on RGB pixels, K = 2, 3, 10 and the original."""
    rng = np.random.default_rng(9)
    photos = [(n, _block_mean(im)) for n, im in common.photographs()]
    Ks = (2, 3, 10)
    fig, axes = plt.subplots(len(photos), 4, figsize=(8.0, 2.35 * len(photos)), squeeze=False)
    for i, (name, img) in enumerate(photos):
        N = img.shape[0] * img.shape[1]
        for j, K in enumerate(Ks):
            seg, J = segment_image(img, K, rng)
            axes[i, j].imshow(seg, interpolation="nearest")
            bits = 24 * K + N * math.ceil(math.log2(K))
            print("  %s (%d x %d = %d pixels), K = %d: J = %.1f, %d bits = %.1f%% of %d"
                  % (name, img.shape[0], img.shape[1], N, K, J, bits, 100 * bits / (24 * N), 24 * N))
        axes[i, 3].imshow(img, interpolation="nearest")
    for j, title in enumerate(["$K = 2$", "$K = 3$", "$K = 10$", "Original image"]):
        axes[0, j].set_title(title)
    for ax in axes.flat:
        ax.set_axis_off()
    fig.subplots_adjust(wspace=0.04, hspace=0.04)
    common.savefig(fig, 9, "fig9-3", out_dir)


# =========================================================================== Figure 9.5
FIG223 = {  # the mixture of Figure 2.23, read off that figure (see the module docstring)
    "pi": np.array([0.5, 0.3, 0.2]),
    "mu": np.array([[0.2, 0.4], [0.5, 0.5], [0.8, 0.6]]),
    "Sigma": 0.016 * np.array([[[1, 0.75], [0.75, 1]], [[1, -0.75], [-0.75, 1]], [[1, 0.75], [0.75, 1]]]),
}


def fig9_5(out_dir=None):
    """Figure 9.5: 500 points from the three-Gaussian mixture of Figure 2.23: (a) complete data, (b) incomplete, (c) responsibilities."""
    rng = np.random.default_rng(2)
    p = FIG223
    X, z = sample_gmm(500, p["pi"], p["mu"], p["Sigma"], rng)
    gamma, _ = gmm_e_step(X, p["pi"], p["mu"], p["Sigma"])
    cols = [common.RED, common.GREEN, common.BLUE]
    fig, axes = common.panel_grid(1, 3, size=(2.6, 2.6))
    _scatter(axes[0, 0], X, np.array(cols)[z], s=14)
    _scatter(axes[0, 1], X, common.MAGENTA, s=14)
    _scatter(axes[0, 2], X, _mix_colours(gamma, cols), s=14)
    for ax, lab in zip(axes.flat, "abc"):
        _square(ax, (-0.2, 1.2), (0, 0.5, 1))
        common.panel_label(ax, lab)
    fig.subplots_adjust(wspace=0.18)
    common.savefig(fig, 9, "fig9-5", out_dir)
    print("  component counts", np.bincount(z, minlength=3).tolist())


# =========================================================================== Figure 9.7
def fig9_7(out_dir=None):
    """Figure 9.7: the singularity: a component collapsing onto a data point makes ln p(X) unbounded."""
    x = np.array([1.3, 1.7, 2.5, 3.8, 5.1, 6.8, 8.0])
    m, s = x.mean(), x.std()  # broad component: the single-Gaussian fit to all points
    w = np.array([6 / 7, 1 / 7])
    sj = 0.15  # sd of the narrow component, centred on the right-most point x[-1]
    dens = lambda t, sd: w[0] * np.exp(_lognorm(t, m, s)) + w[1] * np.exp(_lognorm(t, x[-1], sd))
    fig, (ax, bx) = plt.subplots(1, 2, figsize=(9.2, 3.3), gridspec_kw={"width_ratios": [1.25, 1]})
    t = np.linspace(0, 9.5, 2000)
    ax.plot(t, w[0] * np.exp(_lognorm(t, m, s)), color="k", lw=1.6)
    ax.plot(t, dens(t, sj), color=common.RED, lw=2.2)
    ax.vlines(x, 0, dens(x, sj), color=common.GREEN, lw=2)
    ax.plot(x, dens(x, sj), "o", color=common.BLUE, ms=7, zorder=3)
    ax.plot(x, 0 * x, "o", color="k", ms=6, clip_on=False, zorder=3)
    ax.set_xlim(0, 9.5)
    ax.set_ylim(0, None)
    ax.set_xticks([])
    ax.set_yticks([])
    for sp in ("top", "right"):
        ax.spines[sp].set_visible(False)
    ax.set_xlabel("$x$", loc="right")
    ax.set_ylabel("$p(x)$", rotation=0, loc="top", labelpad=14)
    common.panel_label(ax, "a", "upper right")
    # (b) ln p(X) as the narrow sd shrinks, against a single Gaussian collapsing onto the same point
    sds = np.logspace(-4, 1, 300)
    mix = [np.log(dens(x, sd)).sum() for sd in sds]
    single = [_lognorm(x, x[-1], sd).sum() for sd in sds]
    bx.semilogx(sds, mix, color=common.RED, lw=2.2, label="two components: $\\mu_j = x_7$")
    bx.semilogx(sds, single, color="k", lw=1.8, ls="--", label="one Gaussian: $\\mu = x_7$")
    bx.set_ylim(-45, 0)
    bx.invert_xaxis()
    bx.set_xlabel("$\\sigma_j$ (sd of the collapsing component)")
    bx.set_ylabel("$\\ln p(\\mathbf{X})$")
    bx.legend(loc="lower right", frameon=False)
    bx.tick_params(direction="in", top=True, right=True, which="both")
    common.panel_label(bx, "b", "upper right")
    fig.subplots_adjust(wspace=0.32)
    common.savefig(fig, 9, "fig9-7", out_dir)
    print("  broad component N(%.3f, %.3f^2); ln p(X) at sigma_j = 1e-1, 1e-2, 1e-4: %s"
          % (m, s, [round(float(np.log(dens(x, v)).sum()), 3) for v in (1e-1, 1e-2, 1e-4)]))


# =========================================================================== Figure 9.8
def fig9_8(out_dir=None):
    """Figure 9.8: EM for a two-component Gaussian mixture on the re-scaled Old Faithful data (same start as Fig. 9.1)."""
    X = common.old_faithful(standardize=True)
    Sigma0 = np.stack([BOOK_COV * np.eye(2)] * 2)
    hist = em_gmm(X, [0.5, 0.5], BOOK_CENTRES, Sigma0, n_iter=20)
    cols = [common.RED, common.BLUE]
    fig, axes = common.panel_grid(2, 3, size=(2.45, 2.55))
    panels = [("a", 0, None), ("b", 0, 1), ("c", 1, 1), ("d", 2, 2), ("e", 5, 5), ("f", 20, 20)]
    for ax, (lab, L, gL) in zip(axes.flat, panels):
        st = hist[L]
        _scatter(ax, X, common.GREEN if gL is None else _mix_colours(hist[gL]["gamma"], cols))
        for k in range(2):
            _ellipse(ax, st["mu"][k], st["Sigma"][k], cols[k])
        if lab in "cdef":
            ax.text(0.07, 0.93, "$L = %d$" % L, transform=ax.transAxes, va="top", fontsize=11)
        _square(ax)
        ax.set_xlabel("(%s)" % lab, labelpad=2)
    fig.subplots_adjust(wspace=0.15, hspace=0.3)
    common.savefig(fig, 9, "fig9-8", out_dir)
    print("  ln p(X) at the start and after L = 1, 2, 5, 20:", [round(hist[L]["loglik"], 4) for L in (0, 1, 2, 5, 20)])
    print("  final pi %s, means %s" % (np.round(hist[20]["pi"], 4).tolist(), np.round(hist[20]["mu"], 4).tolist()))


# =========================================================================== Figure 9.10
def digits_234(n_each=200):
    """The first n_each MNIST training images of each of the digits 2, 3 and 4, binarised at 0.5 (Section 9.3.3)."""
    X, y = common.mnist(train=True, digits=(2, 3, 4))
    keep = np.concatenate([np.nonzero(y == d)[0][:n_each] for d in (2, 3, 4)])
    return (X[keep] > 0.5).astype(float), y[keep]


def fig9_10(out_dir=None):
    """Figure 9.10: a mixture of K = 3 Bernoulli distributions fitted by 10 EM iterations to 600 binarised digits 2, 3, 4."""
    X, y = digits_234()
    rng = np.random.default_rng(10)
    K, D = 3, X.shape[1]
    mu0 = rng.uniform(0.25, 0.75, size=(K, D))
    mu0 /= mu0.sum(1, keepdims=True)  # "normalized to satisfy the constraint that sum_j mu_kj = 1" (p. 447)
    pi, mu, gamma, lls = bernoulli_mixture_em(X, np.full(K, 1.0 / K), mu0, n_iter=10)
    lab = gamma.argmax(1)
    major = [np.bincount(y[lab == k], minlength=5)[2:].argmax() + 2 if np.any(lab == k) else 9 for k in range(K)]
    order = np.argsort(major)
    agree = sum(int(np.sum((lab == k) & (y == major[k]))) for k in range(K))
    cmap = LinearSegmentedColormap.from_list("ink", ["#ffffff", "#1a1aa8"])
    fig = plt.figure(figsize=(7.6, 3.4))
    examples = np.random.default_rng(1).choice(len(X), 5, replace=False)
    for i, n in enumerate(examples):
        ax = fig.add_axes([0.01 + 0.198 * i, 0.53, 0.18, 0.44])
        ax.imshow(X[n].reshape(28, 28), cmap=cmap, vmin=0, vmax=1, interpolation="nearest")
    imgs = [mu[k] for k in order] + [X.mean(0)]
    for i, im in enumerate(imgs):  # each mean image scaled to its own maximum, as in the book
        ax = fig.add_axes([0.01 + 0.198 * i + (0.198 if i == 3 else 0), 0.03, 0.18, 0.44])
        ax.imshow(im.reshape(28, 28), cmap=cmap, vmin=0, vmax=im.max(), interpolation="nearest")
    for ax in fig.axes:
        ax.set_xticks([])
        ax.set_yticks([])
        for sp in ax.spines.values():
            sp.set_linewidth(1.2)
    common.savefig(fig, 9, "fig9-10", out_dir)
    print("  ln p(X) after 0..10 M steps:", np.round(lls, 1).tolist())
    print("  pi (in display order) %s; components = digits %s; %d of 600 images in the cluster of their digit"
          % (np.round(pi[order], 3).tolist(), [int(major[k]) for k in order], agree))


# =========================================================================== Figure 9.14
def fig9_14(out_dir=None):
    """Figure 9.14: EM in parameter space: ln p(X|theta) (red), the bound at theta_old (blue) and at theta_new (green)."""
    name = "two"
    sc = EM_BOUND_SCENARIOS[name]
    x, model, free = em_bound_data(name), sc["model"], "mu2"
    th_old = sc["theta0"]
    q_old = mix1d_posterior(x, th_old, model, free)
    th_new = mix1d_m_step(x, q_old, model, free)
    q_new = mix1d_posterior(x, th_new, model, free)
    T = np.linspace(-4, 4, 801)
    ll = np.array([mix1d_loglik(x, t, model, free) for t in T])
    B_old = np.array([mix1d_bound(x, t, q_old, model, free) for t in T])
    B_new = np.array([mix1d_bound(x, t, q_new, model, free) for t in T])
    lo, hi = ll.min() - 8, ll.max() + 4
    fig, ax = plt.subplots(figsize=(5.6, 3.8))
    ax.plot(T, ll, color=common.RED, lw=2.2, zorder=3)
    ax.plot(T, B_old, color=common.BLUE, lw=2.2, zorder=2)
    ax.plot(T, B_new, color=common.GREEN, lw=2.2, zorder=2)
    for t, txt in ((th_old, "$\\theta^{\\rm old}$"), (th_new, "$\\theta^{\\rm new}$")):
        ax.plot([t, t], [lo, mix1d_loglik(x, t, model, free)], ls=(0, (1.5, 3)), color="k", lw=1.6)
        ax.text(t, lo - 0.035 * (hi - lo), txt, ha="center", va="top", fontsize=12)
    ax.plot(x, np.full_like(x, lo + 0.012 * (hi - lo)), "|", color=common.GREY, ms=7)
    ax.text(2.25, mix1d_loglik(x, 2.25, model, free) + 2.0, "$\\ln p(\\mathbf{X}|\\theta)$", color=common.RED, fontsize=12)
    ax.text(1.25, mix1d_bound(x, 1.25, q_old, model, free) - 1.5, "$\\mathcal{L}(q, \\theta)$", color=common.BLUE, fontsize=12, ha="left", va="top")
    ax.set_xlim(-4, 4)
    ax.set_ylim(lo, hi)
    ax.set_xticks([-4, -2, 2, 4])
    ax.set_xlabel("$\\theta$ = mean of the second component", labelpad=18)
    ax.set_ylabel("log likelihood")
    ax.tick_params(direction="in", top=True, right=True)
    common.savefig(fig, 9, "fig9-14", out_dir)
    th = mix1d_em(x, th_old, model, free)
    print("  theta_old = %.4f, theta_new = %.4f; ln p: %.4f -> %.4f; bound at theta_new %.4f"
          % (th_old, th_new, mix1d_loglik(x, th_old, model, free), mix1d_loglik(x, th_new, model, free),
             mix1d_bound(x, th_new, q_old, model, free)))
    print("  EM from theta_old converges to theta = %.5f (ln p = %.4f) in %d iterations"
          % (th[-1], mix1d_loglik(x, th[-1], model, free), len(th) - 1))


FIGURES = {
    "9.1": fig9_1,
    "9.2": fig9_2,
    "9.3": fig9_3,
    "9.5": fig9_5,
    "9.7": fig9_7,
    "9.8": fig9_8,
    "9.10": fig9_10,
    "9.14": fig9_14,
}

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