"""Bishop (2006), Chapter 13: Sequential Data.

Reproduces the numerical figures of Chapter 13 of C. M. Bishop, Pattern Recognition and
Machine Learning (Springer, 2006), pp. 605-652, written from scratch with numpy, scipy and
matplotlib. Every algorithm is a small function meant to be read next to the lecture notes
(ch13-sequential-data.html):

* Section 13.2, hidden Markov models:
  hmm_sample (ancestral sampling, Section 13.2 and Figure 13.8),
  emission_logprob (Gaussian emissions (13.20)-(13.21) or a discrete table (13.22)),
  forward_backward_unscaled (the alpha-beta recursions (13.36)-(13.43)),
  forward_backward (the scaled recursions (13.59)-(13.65), with ln p(X) from (13.63)),
  viterbi (max-sum on the lattice, (13.68)-(13.71), with the errata to (13.68), (13.70), (13.71)),
  baum_welch (EM, E step (13.33)-(13.43), M step (13.18)-(13.23); several sequences;
  Gaussian or discrete emissions; a left-to-right model keeps its zero transitions).
* Section 13.3, linear dynamical systems:
  lds_sample ((13.75)-(13.77)), kalman_filter ((13.88)-(13.97), with V0 renamed P0 as in the
  errata), rts_smoother ((13.100)-(13.104)), lds_em (E step (13.105)-(13.107), M step
  (13.110)-(13.116) with the errata to (13.106), (13.114) and (13.116)).
* Section 13.3.4, particle_filter ((13.117)-(13.119): weights, mixture representation of the
  predictive distribution, resampling from that mixture).
* JSRandom (imported from ch11_sampling.py): the seeded generator DK.rng of the browser demos.
  Passing JSRandom(seed) to hmm_sample, lds_sample, pf_simulate and particle_filter reproduces
  the demos' sequences and particles number for number.

Figures: 13.1, 13.8, 13.11, 13.12, 13.13, 13.16, 13.22, 13.23. Run
`python3 ch13_sequential_data.py --list`, `python3 ch13_sequential_data.py` (all figures,
under a minute) or `python3 ch13_sequential_data.py --fig 13.8 13.22`. Files go to ../img/ch13/.

What differs from the book, and the settings chosen where the book gives none:
* Figure 13.1: the book's recording of the spoken words "Bayes' theorem" is not available.
  We synthesise a speech-like signal with a source-filter model (a glottal pulse train with a
  falling pitch, three formant resonators whose frequencies move through the vowels of
  "Bayes' theorem", and filtered noise for the fricatives z and th) and show its spectrogram.
  It is a SUBSTITUTE, labelled as such; it makes the book's point (successive spectra are
  strongly correlated), which a third panel quantifies.
* Figure 13.8: the emission means and covariances are read off the book's contours:
  means (0.31, 0.72), (0.30, 0.20), (0.70, 0.47); standard deviations along the principal
  axes (0.10, 0.035) at +45 degrees, (0.085, 0.035) horizontal, (0.10, 0.04) at -45 degrees;
  contours at 2/3, 4/3 and 2 standard deviations. Transition matrix: 0.9 on the diagonal and
  0.05 elsewhere, pi uniform, as in the caption. The 50 points use JSRandom(SEED_FIG8), so
  preset 1 of hmm-demo.html draws the same sequence.
* Figure 13.11: the book's 45 on-line handwritten digits are not available. We SYNTHESISE 45
  pen trajectories of a digit '2' (an arc, a cusp or small loop, and a base stroke, with
  random proportions, slant, size and jitter), turn each into line segments of fixed length
  whose angles are quantised to 16 values (as in the book), and fit a 16-state left-to-right
  HMM with a 16 x 16 emission table by 25 iterations of Baum-Welch. A left-to-right chain
  never leaves its last state, so to draw a sample we run the chain until it enters state 16
  and then stay there for the average time the training sequences spend in that state
  (from the posterior gamma); every other step is sampled from the fitted model.
* Figures 13.12, 13.13 and 13.16 are lattice diagrams in the book. We draw the same lattices
  for a small 3-state HMM with discrete emissions (LATTICE_* below) and write the actual
  values of alpha, beta and omega on the nodes. The example is chosen so that the Viterbi path
  differs from the sequence of individually most probable states.
* Figure 13.22: the book's object and noise levels are not given. We use a constant-velocity
  model in the plane (state: position and velocity; white-noise acceleration with spectral
  density q = 0.05; position measurements with standard deviation 0.8), 12 time steps, and
  add a second panel with the Rauch-Tung-Striebel smoother.
* Figure 13.23 is a sketch in the book. We compute it from a run of the particle filter with
  L = 12 particles on the univariate nonlinear growth model of Kitagawa (1996) and Gordon,
  Salmond and Smith (1993), whose observation x = z^2/20 + noise makes p(x|z) bimodal in z,
  as in the book's sketch; a second panel shows the whole run with L = 500 particles.
"""
import math

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from scipy import signal

import common
from ch11_sampling import JSRandom  # noqa: F401  (re-exported: the demos' generator)

LOG2PI = math.log(2 * math.pi)


def _rot(deg):
    t = math.radians(deg)
    return np.array([[math.cos(t), -math.sin(t)], [math.sin(t), math.cos(t)]])


def _cov(sd_major, sd_minor, deg):
    R = _rot(deg)
    return R @ np.diag([sd_major ** 2, sd_minor ** 2]) @ R.T


# The three-state HMM of Figure 13.8 (and of hmm-demo.html).
FIG8_PI = np.ones(3) / 3
FIG8_A = np.full((3, 3), 0.05) + 0.85 * np.eye(3)
FIG8_MU = np.array([[0.31, 0.72], [0.30, 0.20], [0.70, 0.47]])
FIG8_SIGMA = np.array([_cov(0.10, 0.035, 45), _cov(0.085, 0.035, 0), _cov(0.10, 0.04, -45)])
SEED_FIG8 = 8

# The small discrete HMM drawn on the lattices of Figures 13.12, 13.13 and 13.16.
LATTICE_PI = np.array([0.5, 0.3, 0.2])
LATTICE_A = np.array([[0.6, 0.3, 0.1], [0.2, 0.6, 0.2], [0.1, 0.3, 0.6]])
LATTICE_B = np.array([[0.7, 0.2, 0.1], [0.1, 0.6, 0.3], [0.2, 0.2, 0.6]])   # p(x = symbol | state)
LATTICE_X = np.array([0, 1, 0, 2, 2])   # Viterbi path 1 1 1 3 3, most probable states 1 2 1 3 3

# The constant-velocity tracking model of Figure 13.22 (and kalman-demo.html).
def tracking_model(q=0.05, r=0.8, dt=1.0):
    """Constant-velocity model in the plane: z = (position, velocity), x = position + noise.

    A moves the position by dt times the velocity (13.75); Gamma is the covariance of
    white-noise acceleration with spectral density q; C picks the position (13.76) and
    Sigma = r^2 I. Returns a dict with A, Gamma, C, Sigma, mu0, P0 (prior (13.77)).
    """
    I2 = np.eye(2)
    A = np.block([[I2, dt * I2], [0 * I2, I2]])
    Gamma = q * np.block([[dt ** 3 / 3 * I2, dt ** 2 / 2 * I2], [dt ** 2 / 2 * I2, dt * I2]])
    C = np.hstack([I2, 0 * I2])
    return dict(A=A, Gamma=Gamma, C=C, Sigma=r ** 2 * I2,
                mu0=np.array([0.0, 0.0, 1.0, 0.6]), P0=np.diag([0.5, 0.5, 0.1, 0.1]))


# =========================================================================== random numbers
def _u(rng):
    return rng.u() if isinstance(rng, JSRandom) else float(rng.random())


def _n(rng):
    return rng.norm() if isinstance(rng, JSRandom) else float(rng.standard_normal())


def _categorical(rng, p):
    """Index k with probability p[k]: the first k whose cumulative sum exceeds a uniform u
    (the same rule as the demos, so JSRandom reproduces them)."""
    u = _u(rng)
    c = np.cumsum(p)
    k = int(np.searchsorted(c, u, side="right"))
    return min(k, len(p) - 1)


def _mvn(rng, mu, Sigma):
    """mu + L e with L the lower Cholesky factor and e standard normal (as ML.mvn)."""
    L = np.linalg.cholesky(Sigma)
    e = np.array([_n(rng) for _ in range(len(mu))])
    return mu + L @ e


# =========================================================================== hidden Markov models
def hmm_sample(rng, pi, A, N, mus=None, Sigmas=None, B=None):
    """Ancestral sampling from an HMM (Bishop 2006, Section 13.2, Figure 13.8).

    z_1 ~ p(z_1|pi) (13.8), z_n ~ p(z_n|z_{n-1}, A) (13.7), x_n ~ p(x_n|z_n, phi): a Gaussian
    N(mus[k], Sigmas[k]) or, if B is given, a discrete table B[k, symbol] (13.22).
    For every n the uniform for z_n is drawn first, then the emission. Returns (z, X).
    """
    z = np.zeros(N, dtype=int)
    X = []
    for n in range(N):
        z[n] = _categorical(rng, pi if n == 0 else A[z[n - 1]])
        if B is not None:
            X.append(_categorical(rng, B[z[n]]))
        else:
            X.append(_mvn(rng, mus[z[n]], Sigmas[z[n]]))
    return z, np.array(X)


def emission_logprob(X, mus=None, Sigmas=None, B=None):
    """N x K matrix of ln p(x_n|z_n = k): Gaussian emissions, or a discrete table B (13.22)."""
    if B is not None:
        with np.errstate(divide="ignore"):
            return np.log(B[:, np.asarray(X, dtype=int)].T)
    X = np.atleast_2d(X)
    return np.column_stack([common.gauss_logpdf(X, mus[k], Sigmas[k]) for k in range(len(mus))])


def forward_backward_unscaled(pi, A, logE):
    """The alpha-beta recursions without scaling (Section 13.2.2).

    alpha(z_1) = p(z_1) p(x_1|z_1) (13.37); alpha(z_n) = p(x_n|z_n) sum alpha(z_{n-1}) p(z_n|z_{n-1})
    (13.36); beta(z_N) = 1 (13.39-13.40); beta(z_n) = sum beta(z_{n+1}) p(x_{n+1}|z_{n+1})
    p(z_{n+1}|z_n) (13.38); p(X) = sum alpha(z_N) (13.42); gamma = alpha beta / p(X) (13.33);
    xi(z_{n-1}, z_n) = alpha(z_{n-1}) p(x_n|z_n) p(z_n|z_{n-1}) beta(z_n) / p(X) (13.43).
    Underflows for long sequences: use forward_backward. Returns a dict.
    """
    E = np.exp(logE)
    N, K = E.shape
    alpha = np.zeros((N, K))
    beta = np.ones((N, K))
    alpha[0] = pi * E[0]
    for n in range(1, N):
        alpha[n] = E[n] * (alpha[n - 1] @ A)
    for n in range(N - 2, -1, -1):
        beta[n] = A @ (E[n + 1] * beta[n + 1])
    pX = alpha[-1].sum()
    gamma = alpha * beta / pX
    xi = np.array([alpha[n - 1][:, None] * A * (E[n] * beta[n])[None, :] / pX for n in range(1, N)])
    return dict(alpha=alpha, beta=beta, pX=pX, gamma=gamma, xi=xi)


def forward_backward(pi, A, logE):
    """Scaled forward-backward algorithm (Bishop 2006, Section 13.2.4, (13.59)-(13.65)).

    alpha_hat(z_n) = p(z_n|x_1..x_n) and c_n = p(x_n|x_1..x_{n-1}) (13.55)-(13.56):
      c_n alpha_hat(z_n) = p(x_n|z_n) sum_{z_{n-1}} alpha_hat(z_{n-1}) p(z_n|z_{n-1})   (13.59)
      c_{n+1} beta_hat(z_n) = sum_{z_{n+1}} beta_hat(z_{n+1}) p(x_{n+1}|z_{n+1}) p(z_{n+1}|z_n) (13.62)
      p(X) = prod_n c_n                                                               (13.63)
      gamma(z_n) = alpha_hat(z_n) beta_hat(z_n)                                         (13.64)
      xi(z_{n-1}, z_n) = c_n^{-1} alpha_hat(z_{n-1}) p(x_n|z_n) p(z_n|z_{n-1}) beta_hat(z_n) (13.65,
      with the erratum c_n^{-1} and z_{n-1}).
    logE is the N x K matrix ln p(x_n|z_n = k). Each row is shifted by its maximum m_n before
    exponentiating, so even Gaussian densities far in the tails do not underflow; the shift is
    added back to ln c_n. Returns a dict with alpha, beta, log_c, gamma, xi (N-1 x K x K), loglik.
    """
    logE = np.asarray(logE, dtype=float)
    N, K = logE.shape
    m = logE.max(1)
    E = np.exp(logE - m[:, None])
    alpha = np.zeros((N, K))
    beta = np.ones((N, K))
    c = np.zeros(N)
    a = pi * E[0]
    c[0] = a.sum()
    alpha[0] = a / c[0]
    for n in range(1, N):
        a = E[n] * (alpha[n - 1] @ A)
        c[n] = a.sum()
        alpha[n] = a / c[n]
    for n in range(N - 2, -1, -1):
        beta[n] = (A @ (E[n + 1] * beta[n + 1])) / c[n + 1]
    gamma = alpha * beta
    xi = np.array([alpha[n - 1][:, None] * A * (E[n] * beta[n])[None, :] / c[n] for n in range(1, N)]).reshape(N - 1, K, K)
    log_c = np.log(c) + m
    return dict(alpha=alpha, beta=beta, log_c=log_c, gamma=gamma, xi=xi, loglik=float(log_c.sum()))


def viterbi(pi, A, logE):
    """Most probable state sequence by max-sum (Bishop 2006, Section 13.2.5, (13.68)-(13.71)).

    omega(z_1) = ln p(z_1) + ln p(x_1|z_1)                                       (13.69)
    omega(z_{n+1}) = ln p(x_{n+1}|z_{n+1}) + max_{z_n} {ln p(z_{n+1}|z_n) + omega(z_n)} (13.68)
    so omega(z_n) = max over z_1..z_{n-1} of ln p(x_1..x_n, z_1..z_n) (13.70, with the errata:
    ln p(z_{n+1}|z_n) in (13.68), the missing ln in (13.70)). psi(k_n) records the maximising
    z_{n-1}; back-tracking k_{n-1} = psi(k_n) (13.71) recovers the path. Ties go to the lowest
    state index (as in the demo). Returns (path, max ln p(X, Z), omega, psi).
    """
    logE = np.asarray(logE, dtype=float)
    N, K = logE.shape
    with np.errstate(divide="ignore"):
        logA, logpi = np.log(A), np.log(pi)
    omega = np.zeros((N, K))
    psi = np.zeros((N, K), dtype=int)
    omega[0] = logpi + logE[0]
    for n in range(1, N):
        s = omega[n - 1][:, None] + logA          # s[j, k] = omega(z_{n-1} = j) + ln A_jk
        psi[n] = s.argmax(0)
        omega[n] = logE[n] + s.max(0)
    path = np.zeros(N, dtype=int)
    path[-1] = int(omega[-1].argmax())
    for n in range(N - 1, 0, -1):
        path[n - 1] = psi[n][path[n]]
    return path, float(omega[-1].max()), omega, psi


def log_joint(pi, A, logE, z):
    """ln p(X, Z) for a given state path z: ln p(z_1) + sum ln A + sum ln p(x_n|z_n) (13.6)."""
    with np.errstate(divide="ignore"):
        s = math.log(pi[z[0]]) if pi[z[0]] > 0 else -np.inf
        s += sum(np.log(A[z[n - 1], z[n]]) for n in range(1, len(z)))
    return float(s + logE[np.arange(len(z)), z].sum())


def baum_welch(seqs, params, n_iter=50, tol=None, reg=0.0, update_pi=True):
    """EM for an HMM, the Baum-Welch algorithm (Bishop 2006, Section 13.2.1).

    seqs: a list of sequences (arrays N_r x D for Gaussian emissions, integer arrays for a
    discrete table). params: dict with 'pi', 'A' and either 'mus', 'Sigmas' or 'B'.
    E step: gamma and xi of every sequence from forward_backward (13.33)-(13.43, scaled).
    M step, summing over sequences r:
      pi_k = sum_r gamma(z_{r1k}) / sum_r sum_j gamma(z_{r1j})                 (13.18)
      A_jk = sum_r sum_n xi(z_{r,n-1,j}, z_{rnk}) / sum_l of the same            (13.19)
      mu_k = sum gamma(z_nk) x_n / sum gamma(z_nk)                               (13.20)
      Sigma_k = sum gamma(z_nk) (x_n - mu_k)(x_n - mu_k)^T / sum gamma(z_nk)     (13.21)
      B_ki = sum gamma(z_nk) [x_n = i] / sum gamma(z_nk)                          (13.23)
    reg adds reg * I to every Sigma_k (0 for the exact M step). Zeros in pi and A stay zero,
    because gamma and xi vanish there, so a left-to-right model stays left-to-right.
    Returns (params, loglik) with loglik[t] = ln p(X|theta after t M steps), t = 0..n_iter.
    """
    p = {k: np.array(v, dtype=float) for k, v in params.items()}
    disc = "B" in p
    ll = []
    for it in range(n_iter + 1):
        g1 = np.zeros(len(p["pi"]))
        xs = np.zeros_like(p["A"])
        stats, total = [], 0.0
        for X in seqs:
            logE = emission_logprob(X, B=p["B"]) if disc else emission_logprob(X, p["mus"], p["Sigmas"])
            fb = forward_backward(p["pi"], p["A"], logE)
            total += fb["loglik"]
            g1 += fb["gamma"][0]
            xs += fb["xi"].sum(0)
            stats.append((X, fb["gamma"]))
        ll.append(total)
        if it == n_iter or (tol is not None and it > 0 and ll[-1] - ll[-2] < tol):
            break
        if update_pi:
            p["pi"] = g1 / g1.sum()
        with np.errstate(invalid="ignore"):
            p["A"] = np.nan_to_num(xs / xs.sum(1, keepdims=True))
        G = np.concatenate([g for _, g in stats])            # all gamma, stacked over sequences
        Nk = G.sum(0)
        if disc:
            Xall = np.concatenate([np.asarray(X, dtype=int) for X, _ in stats])
            M = p["B"].shape[1]
            p["B"] = np.array([[G[Xall == i, k].sum() for i in range(M)] for k in range(len(Nk))]) / Nk[:, None]
        else:
            Xall = np.concatenate([np.atleast_2d(X) for X, _ in stats])
            p["mus"] = (G.T @ Xall) / Nk[:, None]
            D = Xall.shape[1]
            p["Sigmas"] = np.array([((G[:, k, None] * (Xall - p["mus"][k])).T @ (Xall - p["mus"][k])) / Nk[k] + reg * np.eye(D)
                                    for k in range(len(Nk))])
    return p, np.array(ll)


def left_to_right_init(seqs, K, M=None, stay=None, smooth=0.5):
    """A starting point for a left-to-right HMM: cut every sequence into K equal parts.

    Transitions only k -> k and k -> k+1 (Figure 13.10), pi = (1, 0, ..., 0) (Section 13.2);
    the stay probability matches the average part length. Discrete emissions (M symbols) get
    smoothed counts; Gaussian emissions the part means and a shared covariance.
    """
    seg = [np.minimum((np.arange(len(X)) * K) // len(X), K - 1) for X in seqs]
    L = np.mean([len(X) for X in seqs]) / K
    s = stay if stay is not None else 1 - 1 / max(L, 1.5)
    A = np.diag(np.full(K, s)) + np.diag(np.full(K - 1, 1 - s), 1)
    A[-1, -1] = 1.0
    pi = np.zeros(K)
    pi[0] = 1.0
    if M is not None:
        B = np.full((K, M), smooth)
        for X, sg in zip(seqs, seg):
            np.add.at(B, (sg, np.asarray(X, dtype=int)), 1.0)
        return dict(pi=pi, A=A, B=B / B.sum(1, keepdims=True))
    Xall = np.concatenate([np.atleast_2d(X) for X in seqs])
    sall = np.concatenate(seg)
    mus = np.array([Xall[sall == k].mean(0) for k in range(K)])
    S = np.cov(Xall.T, bias=True) * 0.5
    return dict(pi=pi, A=A, mus=mus, Sigmas=np.array([S] * K))


# =========================================================================== linear dynamical systems
def lds_sample(rng, m, N):
    """Sample (Z, X) from a linear dynamical system (13.75)-(13.77).

    z_1 ~ N(mu0, P0), z_n = A z_{n-1} + w_n with w_n ~ N(0, Gamma), x_n = C z_n + v_n with
    v_n ~ N(0, Sigma). For each n the state noise is drawn before the measurement noise.
    """
    D = len(m["mu0"])
    Z, X = [], []
    for n in range(N):
        z = _mvn(rng, m["mu0"], m["P0"]) if n == 0 else m["A"] @ Z[-1] + _mvn(rng, np.zeros(D), m["Gamma"])
        Z.append(z)
        X.append(m["C"] @ z + _mvn(rng, np.zeros(m["C"].shape[0]), m["Sigma"]))
    return np.array(Z), np.array(X)


def kalman_filter(X, m):
    """The Kalman filter (Bishop 2006, Section 13.3.1, (13.88)-(13.97)).

    alpha_hat(z_n) = N(z_n|mu_n, V_n) (13.84). With the predicted covariance
      P_{n-1} = A V_{n-1} A^T + Gamma                                    (13.88)
    and the gain K_n = P_{n-1} C^T (C P_{n-1} C^T + Sigma)^{-1}         (13.92),
      mu_n = A mu_{n-1} + K_n (x_n - C A mu_{n-1})                       (13.89)
      V_n = (I - K_n C) P_{n-1}                                           (13.90)
      c_n = N(x_n | C A mu_{n-1}, C P_{n-1} C^T + Sigma)                  (13.91)
    started from mu0 and P0 (13.94)-(13.97; the book's V0 is P0 in the errata). A row of X
    that contains NaN is a missing observation: no update, mu_n and V_n are the prediction and
    c_n = 1. Returns a dict: mu, V (filtered), mu_pred, P_pred (the mean and covariance of
    p(z_n|x_1..x_{n-1}): P_pred[n] is the book's P_{n-1}), K (gains), log_c, loglik.
    """
    A, G, C, S = m["A"], m["Gamma"], m["C"], m["Sigma"]
    N, D = len(X), len(m["mu0"])
    mu = np.zeros((N, D))
    V = np.zeros((N, D, D))
    mp = np.zeros((N, D))
    Pp = np.zeros((N, D, D))
    Ks = np.zeros((N, D, C.shape[0]))
    log_c = np.zeros(N)
    I = np.eye(D)
    for n in range(N):
        if n == 0:
            mp[n], Pp[n] = m["mu0"], m["P0"]
        else:
            mp[n], Pp[n] = A @ mu[n - 1], A @ V[n - 1] @ A.T + G
        x = X[n]
        if np.any(np.isnan(x)):
            mu[n], V[n] = mp[n], Pp[n]
            continue
        Sx = C @ Pp[n] @ C.T + S
        K = Pp[n] @ C.T @ np.linalg.inv(Sx)
        r = x - C @ mp[n]
        mu[n] = mp[n] + K @ r
        V[n] = (I - K @ C) @ Pp[n]
        V[n] = 0.5 * (V[n] + V[n].T)
        Ks[n] = K
        log_c[n] = common.gauss_logpdf(x[None, :], C @ mp[n], Sx)[0]
    return dict(mu=mu, V=V, mu_pred=mp, P_pred=Pp, K=Ks, log_c=log_c, loglik=float(log_c.sum()))


def rts_smoother(kf, m):
    """Rauch-Tung-Striebel smoother (Bishop 2006, Section 13.3.1, (13.98)-(13.104)).

    gamma(z_n) = N(z_n|mu_hat_n, V_hat_n) (13.98), from n = N backwards with
      J_n = V_n A^T (P_n)^{-1}                                  (13.102)
      mu_hat_n = mu_n + J_n (mu_hat_{n+1} - A mu_n)             (13.100, erratum mu_n)
      V_hat_n = V_n + J_n (V_hat_{n+1} - P_n) J_n^T              (13.101)
    where P_n = A V_n A^T + Gamma is the predicted covariance of z_{n+1}. The cross-covariance
    of the pair is cov[z_{n-1}, z_n] = J_{n-1} V_hat_n (13.104), i.e.
    E[z_n z_{n-1}^T] = V_hat_n J_{n-1}^T + mu_hat_n mu_hat_{n-1}^T (13.106 with its erratum).
    kf is the output of kalman_filter. Returns a dict mu, V, J, Ezz (E[z_n z_n^T]) and
    Ezz1 (Ezz1[n] = E[z_n z_{n-1}^T] for n >= 1).
    """
    A = m["A"]
    mu, V, Pp = kf["mu"], kf["V"], kf["P_pred"]
    N, D = mu.shape
    mh = mu.copy()
    Vh = V.copy()
    J = np.zeros((N, D, D))
    for n in range(N - 2, -1, -1):
        J[n] = V[n] @ A.T @ np.linalg.inv(Pp[n + 1])
        mh[n] = mu[n] + J[n] @ (mh[n + 1] - A @ mu[n])
        Vh[n] = V[n] + J[n] @ (Vh[n + 1] - Pp[n + 1]) @ J[n].T
        Vh[n] = 0.5 * (Vh[n] + Vh[n].T)
    Ezz = Vh + np.einsum("ni,nj->nij", mh, mh)
    Ezz1 = np.zeros((N, D, D))
    for n in range(1, N):
        Ezz1[n] = Vh[n] @ J[n - 1].T + np.outer(mh[n], mh[n - 1])
    return dict(mu=mh, V=Vh, J=J, Ezz=Ezz, Ezz1=Ezz1)


def lds_em(X, m, n_iter=50, fixed=()):
    """EM for a linear dynamical system (Bishop 2006, Section 13.3.2).

    E step: Kalman filter and RTS smoother give E[z_n] = mu_hat_n (13.105),
    E[z_n z_{n-1}^T] = V_hat_n J_{n-1}^T + mu_hat_n mu_hat_{n-1}^T (13.106, erratum) and
    E[z_n z_n^T] = V_hat_n + mu_hat_n mu_hat_n^T (13.107). M step:
      mu0 = E[z_1], P0 = E[z_1 z_1^T] - E[z_1] E[z_1]^T                       (13.110)-(13.111)
      A = (sum_{n>=2} E[z_n z_{n-1}^T]) (sum_{n>=2} E[z_{n-1} z_{n-1}^T])^{-1}  (13.113)
      Gamma = 1/(N-1) sum_{n>=2} {E[z_n z_n^T] - A E[z_{n-1} z_n^T]
              - E[z_n z_{n-1}^T] A^T + A E[z_{n-1} z_{n-1}^T] A^T}             (13.114, erratum A^T)
      C = (sum x_n E[z_n]^T) (sum E[z_n z_n^T])^{-1}                           (13.115)
      Sigma = 1/N sum {x_n x_n^T - C E[z_n] x_n^T - x_n E[z_n]^T C^T + C E[z_n z_n^T] C^T}
                                                                             (13.116, erratum)
    (A and C are computed first and used in Gamma and Sigma.) Names in `fixed` are not
    updated. Returns (model, loglik) with loglik[t] = ln p(X|theta after t M steps).
    """
    m = {k: np.array(v, dtype=float) for k, v in m.items()}
    N = len(X)
    ll = []
    for it in range(n_iter + 1):
        kf = kalman_filter(X, m)
        ll.append(kf["loglik"])
        if it == n_iter:
            break
        sm = rts_smoother(kf, m)
        Ez, Ezz, Ezz1 = sm["mu"], sm["Ezz"], sm["Ezz1"]
        new = {}
        new["mu0"] = Ez[0]
        new["P0"] = Ezz[0] - np.outer(Ez[0], Ez[0])
        S10 = Ezz1[1:].sum(0)                     # sum E[z_n z_{n-1}^T]
        S00 = Ezz[:-1].sum(0)                     # sum E[z_{n-1} z_{n-1}^T]
        S11 = Ezz[1:].sum(0)                      # sum E[z_n z_n^T], n >= 2
        A = S10 @ np.linalg.inv(S00) if "A" not in fixed else m["A"]
        new["A"] = A
        G = (S11 - A @ S10.T - S10 @ A.T + A @ S00 @ A.T) / (N - 1)
        new["Gamma"] = 0.5 * (G + G.T)
        Sxz = X.T @ Ez
        C = Sxz @ np.linalg.inv(Ezz.sum(0)) if "C" not in fixed else m["C"]
        new["C"] = C
        Sg = (X.T @ X - C @ Sxz.T - Sxz @ C.T + C @ Ezz.sum(0) @ C.T) / N
        new["Sigma"] = 0.5 * (Sg + Sg.T)
        for k, v in new.items():
            if k not in fixed:
                m[k] = v
    return m, np.array(ll)


# =========================================================================== particle filter
# The univariate nonlinear growth model (Kitagawa 1996; Gordon, Salmond and Smith 1993).
def growth_f(z, n):
    """Mean of p(z_n|z_{n-1}) in the nonlinear model: z/2 + 25 z/(1 + z^2) + 8 cos(1.2 n)."""
    return 0.5 * z + 25 * z / (1 + z * z) + 8 * np.cos(1.2 * n)


PF_MODELS = {
    # z_1 ~ N(m0, s0^2); z_n = f(z_{n-1}, n) + N(0, q); x_n = h(z_n) + N(0, r)
    "linear": dict(m0=0.0, s0=1.0, q=0.5, r=1.0, a=0.95),
    "growth": dict(m0=0.0, s0=math.sqrt(5.0), q=10.0, r=1.0),
}


def pf_mean(model, z, n):
    return model["a"] * z if "a" in model else growth_f(z, n)


def pf_obs(model, z):
    return z if "a" in model else z * z / 20.0


def pf_simulate(rng, model, N):
    """Simulate (z, x) of length N; for each n the state noise is drawn before the measurement noise."""
    z = np.zeros(N)
    x = np.zeros(N)
    for n in range(N):
        z[n] = model["m0"] + model["s0"] * _n(rng) if n == 0 else pf_mean(model, z[n - 1], n + 1) + math.sqrt(model["q"]) * _n(rng)
        x[n] = pf_obs(model, z[n]) + math.sqrt(model["r"]) * _n(rng)
    return z, x


def particle_filter(rng, x, model, L, resample=True):
    """Bootstrap particle filter (Bishop 2006, Section 13.3.4, (13.117)-(13.119)).

    p(z_n|X_n) is represented by particles z_n^(l) with weights
      w_n^(l) = p(x_n|z_n^(l)) / sum_m p(x_n|z_n^(m))                       (13.118)
    so E[f(z_n)] ~ sum_l w_n^(l) f(z_n^(l))                                  (13.117), and the
    predictive distribution is the mixture p(z_{n+1}|X_n) ~ sum_l w_n^(l) p(z_{n+1}|z_n^(l)) (13.119,
    erratum: approximately equal). The next particles are L draws from that mixture: choose a
    component l with probability w_n^(l) (one uniform), then draw from p(z_{n+1}|z_n^(l)) (one
    normal), particle by particle, which is the demos' order of random numbers. With
    resample=False every particle keeps its own parent and its weight is multiplied by the
    likelihood (sequential importance sampling, which degenerates). Returns a dict: particles
    and weights (N x L), parents (N x L), mean, var, ess = 1 / sum w^2.
    """
    N = len(x)
    P = np.zeros((N, L))
    W = np.zeros((N, L))
    par = np.zeros((N, L), dtype=int)
    for n in range(N):
        if n == 0:
            for l in range(L):
                P[0, l] = model["m0"] + model["s0"] * _n(rng)
            logw = np.zeros(L)
        else:
            c = np.cumsum(W[n - 1])
            for l in range(L):
                j = min(int(np.searchsorted(c, _u(rng), side="right")), L - 1) if resample else l
                par[n, l] = j
                P[n, l] = pf_mean(model, P[n - 1, j], n + 1) + math.sqrt(model["q"]) * _n(rng)
            logw = np.zeros(L) if resample else np.log(np.maximum(W[n - 1], 1e-300))
        logw = logw - 0.5 * (x[n] - pf_obs(model, P[n])) ** 2 / model["r"]
        w = np.exp(logw - logw.max())
        W[n] = w / w.sum()
    mean = (W * P).sum(1)
    var = (W * (P - mean[:, None]) ** 2).sum(1)
    return dict(particles=P, weights=W, parents=par, mean=mean, var=var, ess=1.0 / (W ** 2).sum(1))


def pf_kalman(x, model):
    """The exact Kalman filter for the linear PF model (1-D), via kalman_filter."""
    m = dict(A=np.array([[model["a"]]]), Gamma=np.array([[model["q"]]]), C=np.array([[1.0]]),
             Sigma=np.array([[model["r"]]]), mu0=np.array([model["m0"]]), P0=np.array([[model["s0"] ** 2]]))
    kf = kalman_filter(np.asarray(x, float)[:, None], m)
    return kf["mu"][:, 0], kf["V"][:, 0, 0]


# =========================================================================== Figure 13.1
def synth_speech(fs=8000, seed=0):
    """A speech-like signal for 'Bayes' theorem' by source-filter synthesis (SUBSTITUTE data).

    Voiced segments: a pulse train at a falling pitch passed through three formant resonators
    (two-pole filters, bandwidths 80, 100, 150 Hz) whose frequencies glide through each vowel;
    fricatives: band-pass filtered noise. Returns (signal, fs, segments) with segments a list
    of (label, start s, end s).
    """
    rng = np.random.default_rng(seed)
    # label, duration (s), voicing amplitude, noise amplitude, noise band (Hz), formants start, end
    plan = [("", 0.08, 0, 0, None, (500, 1500, 2500), (500, 1500, 2500)),
            ("b", 0.03, 0.4, 0.3, (100, 1500), (300, 900, 2200), (450, 1600, 2400)),
            ("ey", 0.20, 1.0, 0, None, (550, 1800, 2500), (400, 2250, 2800)),
            ("z", 0.12, 0.25, 0.5, (3000, 3900), (300, 1800, 2600), (300, 1800, 2600)),
            ("th", 0.11, 0, 0.35, (1200, 3900), (400, 1700, 2600), (400, 1700, 2600)),
            ("ee", 0.17, 1.0, 0, None, (300, 2300, 3000), (420, 1900, 2650)),
            ("r", 0.08, 0.8, 0, None, (450, 1250, 1650), (480, 1300, 1700)),
            ("e", 0.10, 0.9, 0, None, (520, 1500, 2500), (500, 1450, 2450)),
            ("m", 0.13, 0.5, 0, None, (260, 1000, 2200), (250, 950, 2200)),
            ("", 0.08, 0, 0, None, (500, 1500, 2500), (500, 1500, 2500))]
    hop = 80
    voiced, noise, segs, t0 = [], [], [], 0.0
    phase = 0.0
    zi = [np.zeros(2) for _ in range(3)]
    filters = {}
    total = sum(p[1] for p in plan)
    for lab, dur, av, an, band, F0, F1 in plan:
        nb = max(1, int(round(dur * fs)) // hop)
        v = np.zeros(nb * hop)
        w = np.zeros(nb * hop)
        if band is not None and band not in filters:
            filters[band] = signal.butter(4, [band[0] / (fs / 2), band[1] / (fs / 2)], btype="band")
        for k in range(nb):
            frac = (k + 0.5) / nb
            t = t0 + (k + 0.5) * hop / fs
            f0 = 135 - 35 * t / total + 8 * math.sin(2 * math.pi * t / 0.6)
            src = np.zeros(hop)
            for i in range(hop):
                phase += f0 / fs
                if phase >= 1:
                    phase -= 1
                    src[i] = 1.0
            y = av * src
            for j in range(3):
                F = F0[j] + (F1[j] - F0[j]) * frac
                r = math.exp(-math.pi * (80, 100, 150)[j] / fs)
                c = math.cos(2 * math.pi * F / fs)
                gain = (1 - r) * math.sqrt(1 - 2 * r * math.cos(4 * math.pi * F / fs) + r * r)   # about unit peak gain
                y, zi[j] = signal.lfilter([gain], [1, -2 * r * c, r * r], y, zi=zi[j])
            v[k * hop:(k + 1) * hop] = y
            if an > 0:
                b, a = filters[band]
                w[k * hop:(k + 1) * hop] = an * signal.lfilter(b, a, rng.standard_normal(hop + 64))[64:]
        segs.append((lab, t0, t0 + nb * hop / fs))
        t0 += nb * hop / fs
        voiced.append(v)
        noise.append(w)
    v, w = np.concatenate(voiced), np.concatenate(noise)
    rms = lambda a: math.sqrt(np.mean(a[np.abs(a) > 0] ** 2))
    s = v / rms(v) + 0.45 * w / rms(w)
    s += 0.003 * rng.standard_normal(len(s))
    return s / np.abs(s).max(), fs, segs


def spectrogram(s, fs, win=256, hop=80):
    """Log-magnitude short-time Fourier transform: frames of `win` samples (Hann), hop `hop`."""
    w = np.hanning(win)
    frames = np.array([s[i:i + win] * w for i in range(0, len(s) - win, hop)])
    S = np.abs(np.fft.rfft(frames, axis=1))
    return 20 * np.log10(S + 1e-4), np.arange(len(frames)) * hop / fs + win / (2 * fs), np.fft.rfftfreq(win, 1 / fs)


def fig13_1(out_dir=None):
    """Figure 13.1: spectrogram of a SYNTHETIC speech-like signal (substitute for the book's recording)."""
    s, fs, segs = synth_speech()
    Sdb, tt, ff = spectrogram(s, fs)
    fig = plt.figure(figsize=(10.5, 3.9))
    gs = fig.add_gridspec(2, 2, width_ratios=[2.4, 1], height_ratios=[1, 2.4], hspace=0.08, wspace=0.28)
    a = fig.add_subplot(gs[0, 0])
    tw = np.arange(len(s)) / fs
    a.plot(tw, s, color="0.25", lw=0.4, rasterized=True)
    for lab, t0, t1 in segs:
        if lab:
            a.axvline(t0, color="0.75", lw=0.6)
            a.text((t0 + t1) / 2, 1.05, lab, ha="center", va="bottom", fontsize=8.5, style="italic")
    a.set_xlim(0, tw[-1])
    a.set_ylim(-1.1, 1.1)
    a.set_xticks([])
    a.set_yticks([])
    a.set_title("synthetic substitute for the spoken words \"Bayes' theorem\"", fontsize=9.5, pad=14)
    b = fig.add_subplot(gs[1, 0])
    vmax = Sdb.max()
    b.imshow(Sdb.T, origin="lower", aspect="auto", extent=[tt[0], tt[-1], 0, ff[-1] / 1000], cmap="magma",
             vmin=vmax - 60, vmax=vmax, rasterized=True)
    b.set_xlim(0, tw[-1])
    b.set_xlabel("time (s); one spectral vector every 10 ms")
    b.set_ylabel("frequency (kHz)")
    c = fig.add_subplot(gs[:, 1])
    energy = Sdb.max(1)
    keep = energy > vmax - 40
    Z = Sdb[:, ff < 4000]
    Z = (Z - Z.mean(0)) / Z.std(0)
    lags = np.arange(1, 31)
    cor = []
    for k in lags:
        idx = np.where(keep[:-k] & keep[k:])[0]
        cor.append(np.mean([np.corrcoef(Z[i], Z[i + k])[0, 1] for i in idx]))
    c.plot(lags * 10, cor, "o-", color=common.BLUE, ms=3, lw=1.3)
    c.axhline(0, color="0.6", lw=0.7)
    c.set_xlabel("lag (ms)")
    c.set_ylabel("correlation of spectral vectors")
    c.set_title("successive spectra are\nstrongly correlated", fontsize=9.5)
    c.set_ylim(-0.1, 1.0)
    common.savefig(fig, 13, "fig13-1", out_dir)
    return dict(lag1=cor[0], lag10=cor[9])


# =========================================================================== Figure 13.8
def fig13_8(out_dir=None):
    """Figure 13.8: sampling from a 3-state HMM with 2-D Gaussian emissions (contours, 50 samples)."""
    cols = (common.RED, common.GREEN, common.BLUE)
    z, X = hmm_sample(JSRandom(SEED_FIG8), FIG8_PI, FIG8_A, 50, FIG8_MU, FIG8_SIGMA)
    fig, ax = common.panel_grid(1, 2, size=(3.6, 3.5))
    a, b = ax[0]
    for k in range(3):
        for nsd in (2 / 3, 4 / 3, 2):
            ex, ey = common.ellipse_xy(FIG8_MU[k], FIG8_SIGMA[k], nsd)
            a.plot(ex, ey, color=cols[k], lw=2.2)
    for k, (dx, dy) in enumerate([(-0.0, -0.21), (0.27, -0.0), (0.12, -0.24)]):
        a.text(FIG8_MU[k][0] + dx, FIG8_MU[k][1] + dy, "$k = %d$" % (k + 1), ha="center", va="center", fontsize=11)
    b.plot(X[:, 0], X[:, 1], "-", color="k", lw=1.1, zorder=1)
    b.scatter(X[:, 0], X[:, 1], c=[cols[k] for k in z], s=38, zorder=2, edgecolors="none")
    for p, lab in ((a, "a"), (b, "b")):
        p.set_xlim(0, 1)
        p.set_ylim(0, 1)
        p.set_aspect("equal")
        p.set_xticks([0, 0.5, 1])
        p.set_yticks([0, 0.5, 1])
        common.panel_label(p, lab)
    b.set_title("50 samples, $A_{kk} = 0.9$, $A_{jk} = 0.05$", fontsize=9.5)
    a.set_title("emission densities $p(\\mathbf{x}|\\mathbf{z})$", fontsize=9.5)
    fig.tight_layout()
    common.savefig(fig, 13, "fig13-8", out_dir)
    return dict(z=z, X=X, switches=int((np.diff(z) != 0).sum()))


# =========================================================================== Figure 13.11
def synth_two(rng):
    """One SYNTHETIC pen trajectory of a handwritten '2' in the unit square (substitute data).

    An arc from the top left over to the right and down to a cusp (or, with probability 0.3,
    a small loop) at the bottom left, then a nearly straight base stroke to the bottom right.
    Proportions, slant, size, position and a smooth pen jitter vary from example to example.
    """
    cx, cy = 0.45 + 0.05 * rng.standard_normal(), 0.66 + 0.04 * rng.standard_normal()
    R = 0.24 * math.exp(0.12 * rng.standard_normal())
    a0, a1 = math.radians(165 + 12 * rng.standard_normal()), math.radians(-25 + 10 * rng.standard_normal())
    t = np.linspace(a0, a1, 40)
    arc = np.column_stack([cx + R * np.cos(t), cy + 0.85 * R * np.sin(t)])
    cusp = np.array([0.12 + 0.06 * rng.standard_normal(), 0.12 + 0.04 * rng.standard_normal()])
    s = np.linspace(0, 1, 30)[1:]
    bend = 0.06 * rng.standard_normal()
    diag = arc[-1] + s[:, None] * (cusp - arc[-1]) + bend * np.sin(np.pi * s)[:, None] * np.array([1, -1])
    parts = [arc, diag]
    if rng.random() < 0.3:
        u = np.linspace(0, 2 * np.pi, 20)[1:]
        rl = 0.05 + 0.02 * rng.random()
        loop = cusp + rl * np.column_stack([np.sin(u) * 1.2, (np.cos(u) - 1) * -0.8])
        parts.append(loop)
    end = np.array([0.82 + 0.08 * rng.standard_normal(), cusp[1] + 0.03 * rng.standard_normal()])
    s = np.linspace(0, 1, 30)[1:]
    base = cusp + s[:, None] * (end - cusp)
    parts.append(base)
    P = np.vstack(parts)
    slant = 0.15 * rng.standard_normal()
    P[:, 0] += slant * (P[:, 1] - 0.4)
    k = np.arange(len(P))
    jit = np.column_stack([np.convolve(rng.standard_normal(len(P) + 10), np.ones(11) / 11, "valid")[:len(P)] for _ in range(2)])
    P += 0.025 * jit
    sc = math.exp(0.1 * rng.standard_normal())
    P = (P - 0.5) * sc + 0.5
    return P


def to_angles(P, step=0.035, n_angles=16):
    """Resample a pen trajectory into segments of fixed length `step` and quantise their angles
    to n_angles values (the book's representation of on-line digits). Returns integer symbols."""
    d = np.r_[0, np.cumsum(np.hypot(*np.diff(P, axis=0).T))]
    s = np.arange(0, d[-1], step)
    Q = np.column_stack([np.interp(s, d, P[:, 0]), np.interp(s, d, P[:, 1])])
    ang = np.arctan2(*np.diff(Q, axis=0).T[::-1])
    return np.mod(np.round(ang / (2 * np.pi / n_angles)), n_angles).astype(int)


def from_angles(sym, step=0.035, n_angles=16, start=(0.0, 0.0)):
    th = sym * 2 * np.pi / n_angles
    return np.vstack([start, np.array(start) + np.cumsum(step * np.column_stack([np.cos(th), np.sin(th)]), axis=0)])


def sample_left_to_right(rng, p, dwell_last):
    """Generate symbols from a left-to-right HMM: run the chain from state 1 until it reaches the
    last state, then stay there for dwell_last steps; one symbol per step from p['B']."""
    K = len(p["pi"])
    k, out, stay = 0, [], 0
    while True:
        out.append(_categorical(rng, p["B"][k]))
        if k == K - 1:
            stay += 1
            if stay >= dwell_last:
                break
        else:
            k = _categorical(rng, p["A"][k])
    return np.array(out)


def fit_digits(n_train=45, K=16, n_iter=25, seed=11):
    """Synthesise n_train digits, fit the left-to-right HMM by Baum-Welch; returns everything."""
    rng = np.random.default_rng(seed)
    pens = [synth_two(rng) for _ in range(n_train)]
    seqs = [to_angles(P) for P in pens]
    p0 = left_to_right_init(seqs, K, M=16)
    p, ll = baum_welch(seqs, p0, n_iter=n_iter, update_pi=False)
    dwell = np.mean([forward_backward(p["pi"], p["A"], emission_logprob(X, B=p["B"]))["gamma"][:, -1].sum() for X in seqs])
    return dict(pens=pens, seqs=seqs, p=p, ll=ll, dwell=dwell)


def fig13_11(out_dir=None):
    """Figure 13.11: left-to-right HMM for on-line digits, on SYNTHETIC '2's (training examples and samples)."""
    D = fit_digits()
    p = D["p"]
    rng = np.random.default_rng(5)
    fig, ax = plt.subplots(2, 4, figsize=(10.4, 5.2), gridspec_kw=dict(width_ratios=[1, 1, 1, 1.25]))
    for j, i in enumerate([0, 1, 2]):
        P = D["pens"][i]
        ax[0, j].plot(P[:, 0], P[:, 1], color=common.RED, lw=2.4, solid_capstyle="round")
    for j in range(3):
        sym = sample_left_to_right(rng, p, int(round(D["dwell"])))
        Q = from_angles(sym)
        Q = Q - (Q.min(0) + Q.max(0)) / 2 + 0.5
        ax[1, j].plot(Q[:, 0], Q[:, 1], color=common.RED, lw=2.4, solid_capstyle="round")
    for a in ax[:, :3].ravel():
        a.set_xlim(-0.05, 1.05)
        a.set_ylim(-0.05, 1.05)
        a.set_aspect("equal")
        a.set_xticks([])
        a.set_yticks([])
    ax[0, 0].set_ylabel("synthetic training digits", fontsize=9.5)
    ax[1, 0].set_ylabel("samples from the fitted HMM", fontsize=9.5)
    h = ax[0, 3]
    im = h.imshow(p["B"], aspect="auto", cmap="Greys", origin="upper", rasterized=True)
    h.set_xlabel("pen direction (16 angles)")
    h.set_ylabel("state $k$")
    h.set_xticks([0, 4, 8, 12])
    h.set_xticklabels(["0", "90", "180", "270"])
    h.set_title("emission table, 16 x 16", fontsize=9.5)
    g = ax[1, 3]
    g.plot(np.arange(len(D["ll"])), D["ll"], "o-", color=common.BLUE, ms=3, lw=1.2)
    g.set_xlabel("EM iteration")
    g.set_ylabel(r"$\ln p(\mathbf{X}|\theta)$")
    g.set_title("Baum-Welch, 45 sequences", fontsize=9.5)
    fig.suptitle("Left-to-right HMM, 16 states, fitted to 45 SYNTHETIC on-line digits '2' (the book's data are not available)",
                 fontsize=10, y=0.995)
    fig.tight_layout()
    common.savefig(fig, 13, "fig13-11", out_dir)
    return D


# =========================================================================== lattice figures
def _lattice_axes(N, K, figsize=(8.6, 3.7)):
    fig, ax = plt.subplots(figsize=figsize)
    ax.set_xlim(0.3, N + 0.75)
    ax.set_ylim(K + 0.95, -0.15)
    ax.axis("off")
    for k in range(K):
        ax.text(0.35, k + 1, "$k = %d$" % (k + 1), ha="left", va="center", fontsize=10)
    return fig, ax


def _node(ax, n, k, txt, fc="white", ec="k", lw=1.3, fs=8.2):
    ax.add_patch(plt.Circle((n, k), 0.13, fc=fc, ec=ec, lw=lw, zorder=3))
    ax.text(n + 0.17, k - 0.2, txt, fontsize=fs, ha="left", va="center", zorder=4,
            bbox=dict(boxstyle="round,pad=0.12", fc="white", ec="none", alpha=0.85))


def _edge(ax, a, b, color="0.75", lw=0.8, txt=None, arrow=True, zorder=1):
    ax.annotate("", xy=b, xytext=a, zorder=zorder,
                arrowprops=dict(arrowstyle="-|>" if arrow else "-", color=color, lw=lw, shrinkA=8, shrinkB=8, mutation_scale=9))
    if txt:
        mx, my = a[0] + 0.5 * (b[0] - a[0]), a[1] + 0.5 * (b[1] - a[1])
        ax.text(mx, my, txt, fontsize=8, color=color, ha="center", va="center",
                bbox=dict(boxstyle="round,pad=0.1", fc="white", ec="none"))


def lattice_example():
    logE = emission_logprob(LATTICE_X, B=LATTICE_B)
    fb = forward_backward_unscaled(LATTICE_PI, LATTICE_A, logE)
    return logE, fb


def fig13_12(out_dir=None):
    """Figure 13.12: the forward recursion (13.36) on the lattice, with the alpha values of a small example."""
    logE, fb = lattice_example()
    al = fb["alpha"]
    N, K = al.shape
    fig, ax = _lattice_axes(N, K)
    n0 = 3  # the highlighted step (1-based column index)
    for n in range(1, N):
        for j in range(K):
            for k in range(K):
                hl = (n + 1 == n0 and k == 0)
                _edge(ax, (n, j + 1), (n + 1, k + 1), color=common.BLUE if hl else "0.82", lw=1.6 if hl else 0.7,
                      txt=("$A_{%d1} = %.1f$" % (j + 1, LATTICE_A[j, 0])) if hl else None, zorder=2 if hl else 1)
    for n in range(N):
        ax.text(n + 1, 0.12, "$n = %d$\n$x_n = %d$" % (n + 1, LATTICE_X[n]), ha="center", va="center", fontsize=9)
        for k in range(K):
            hl = (n + 1 == n0 and k == 0)
            _node(ax, n + 1, k + 1, "%.4g" % al[n, k], fc="#cfe0ff" if hl else "white", ec=common.BLUE if hl else "k")
    E = np.exp(logE)
    terms = " + ".join("%.4g(%.1f)" % (al[n0 - 2, j], LATTICE_A[j, 0]) for j in range(K))
    ax.text((N + 1) / 2 + 0.2, K + 0.75,
            r"$\alpha(z_{%d,1}) = p(x_%d|z_{%d,1})\sum_j \alpha(z_{%d,j})A_{j1}$ = %.1f [%s] = %.4g" %
            (n0, n0, n0, n0 - 1, E[n0 - 1, 0], terms, al[n0 - 1, 0]), ha="center", va="center", fontsize=9, color=common.BLUE)
    ax.text(N + 0.7, K + 0.36, r"$p(\mathbf{X}) = \sum_k \alpha(z_{%dk})$ = %.4g" % (N, fb["pX"]), ha="right", va="center", fontsize=9)
    ax.set_title(r"Forward recursion (13.36): the value next to each node is $\alpha(z_{nk}) = p(x_1,\dots,x_n, z_{nk}=1)$", fontsize=10)
    common.savefig(fig, 13, "fig13-12", out_dir)
    return fb


def fig13_13(out_dir=None):
    """Figure 13.13: the backward recursion (13.38) on the lattice, with the beta values of the same example."""
    logE, fb = lattice_example()
    be = fb["beta"]
    E = np.exp(logE)
    N, K = be.shape
    fig, ax = _lattice_axes(N, K)
    n0 = 3
    for n in range(1, N):
        for j in range(K):
            for k in range(K):
                hl = (n == n0 and j == 0)
                _edge(ax, (n + 1, k + 1), (n, j + 1), color=common.RED if hl else "0.82", lw=1.6 if hl else 0.7,
                      txt=("$A_{1%d} = %.1f$" % (k + 1, LATTICE_A[0, k])) if hl else None, zorder=2 if hl else 1)
    for n in range(N):
        ax.text(n + 1, 0.12, "$n = %d$\n$x_n = %d$" % (n + 1, LATTICE_X[n]), ha="center", va="center", fontsize=9)
        for k in range(K):
            hl = (n + 1 == n0 and k == 0)
            _node(ax, n + 1, k + 1, "%.4g" % be[n, k], fc="#ffd9d9" if hl else "white", ec=common.RED if hl else "k")
    terms = " + ".join("%.4g(%.1f)(%.1f)" % (be[n0, k], E[n0, k], LATTICE_A[0, k]) for k in range(K))
    ax.text((N + 1) / 2 + 0.2, K + 0.75,
            r"$\beta(z_{%d,1}) = \sum_k \beta(z_{%d,k})\,p(x_%d|z_{%d,k})\,A_{1k}$ = %s = %.4g" %
            (n0, n0 + 1, n0 + 1, n0 + 1, terms, be[n0 - 1, 0]), ha="center", va="center", fontsize=8.6, color=common.RED)
    ax.set_title(r"Backward recursion (13.38): the value next to each node is $\beta(z_{nk}) = p(x_{n+1},\dots,x_N|z_{nk}=1)$", fontsize=10)
    common.savefig(fig, 13, "fig13-13", out_dir)
    return fb


def fig13_16(out_dir=None):
    """Figure 13.16: the Viterbi lattice with omega values, back-pointers and the most probable path."""
    logE, fb = lattice_example()
    path, best, omega, psi = viterbi(LATTICE_PI, LATTICE_A, logE)
    N, K = omega.shape
    fig, ax = _lattice_axes(N, K, figsize=(8.6, 3.9))
    for n in range(1, N):
        for j in range(K):
            for k in range(K):
                _edge(ax, (n, j + 1), (n + 1, k + 1), color="0.88", lw=0.6, arrow=False)
        for k in range(K):
            j = psi[n][k]
            _edge(ax, (n, j + 1), (n + 1, k + 1), color="0.45", lw=1.1)
    for n in range(1, N):
        _edge(ax, (n, path[n - 1] + 1), (n + 1, path[n] + 1), color=common.RED, lw=2.6, zorder=2)
    gpath = fb["gamma"].argmax(1)
    for n in range(N):
        ax.text(n + 1, 0.12, "$n = %d$\n$x_n = %d$" % (n + 1, LATTICE_X[n]), ha="center", va="center", fontsize=9)
        for k in range(K):
            on = path[n] == k
            _node(ax, n + 1, k + 1, "%.3f" % omega[n, k], fc="#ffd9d9" if on else "white", ec=common.RED if on else "k")
        g = gpath[n]
        ax.plot(n + 1, g + 1, marker="s", ms=17, mfc="none", mec=common.GREEN, mew=2.0, zorder=5)
    ax.text((N + 1) / 2 + 0.2, K + 0.75,
            "red: Viterbi path %s, ln p(X, Z) = %.3f;   green squares: most probable state at each n, %s, ln p(X, Z) = %.3f" %
            ("".join(str(k + 1) for k in path), best, "".join(str(k + 1) for k in gpath),
             log_joint(LATTICE_PI, LATTICE_A, logE, gpath)), ha="center", va="center", fontsize=8.6)
    ax.set_title(r"Viterbi (13.68)-(13.71): node values $\omega(z_{nk})$; dark arrows: best predecessor $\psi$", fontsize=10)
    common.savefig(fig, 13, "fig13-16", out_dir)
    return dict(path=path, gpath=gpath, best=best)


# =========================================================================== Figure 13.22
SEED_FIG22 = 22


def fig13_22(out_dir=None):
    """Figure 13.22: tracking a moving object with a linear dynamical system (Kalman filter, RTS smoother)."""
    m = tracking_model()
    Z, X = lds_sample(JSRandom(SEED_FIG22), m, 12)
    kf = kalman_filter(X, m)
    sm = rts_smoother(kf, m)
    fig, ax = common.panel_grid(1, 2, size=(4.6, 3.6))
    for a, est, lab, title in ((ax[0, 0], kf, "a", "Kalman filter: $p(\\mathbf{z}_n|\\mathbf{x}_1,\\dots,\\mathbf{x}_n)$"),
                               (ax[0, 1], sm, "b", "RTS smoother: $p(\\mathbf{z}_n|\\mathbf{x}_1,\\dots,\\mathbf{x}_N)$")):
        a.plot(Z[:, 0], Z[:, 1], "--", color=common.BLUE, lw=1.1, dashes=(5, 3))
        a.plot(Z[:, 0], Z[:, 1], "o", color=common.BLUE, ms=5, label="true position")
        a.plot(X[:, 0], X[:, 1], "--", color=common.GREEN, lw=1.1, dashes=(5, 3))
        a.plot(X[:, 0], X[:, 1], "o", color=common.GREEN, ms=5, label="measurement")
        col = common.RED if lab == "a" else common.MAGENTA
        a.plot(est["mu"][:, 0], est["mu"][:, 1], "--", color=col, lw=1.0, dashes=(5, 3))
        a.plot(est["mu"][:, 0], est["mu"][:, 1], "x", color=col, ms=8, mew=2.2, label="posterior mean")
        for n in range(len(Z)):
            ex, ey = common.ellipse_xy(est["mu"][n, :2], est["V"][n][:2, :2], 1.0)
            a.plot(ex, ey, color=col, lw=1.6)
        a.set_aspect("equal")
        a.set_title(title, fontsize=9.5)
        a.tick_params(labelsize=8)
        common.panel_label(a, lab, "lower right")
    ax[0, 0].legend(fontsize=7.5, loc="upper left", frameon=False)
    lim = (min(ax[0, 0].get_xlim()[0], ax[0, 1].get_xlim()[0]), max(ax[0, 0].get_xlim()[1], ax[0, 1].get_xlim()[1]))
    ylim = (min(ax[0, 0].get_ylim()[0], ax[0, 1].get_ylim()[0]), max(ax[0, 0].get_ylim()[1], ax[0, 1].get_ylim()[1]))
    for a in ax[0]:
        a.set_xlim(lim)
        a.set_ylim(ylim)
    rm = lambda E: float(np.sqrt(np.mean(np.sum((E - Z[:, :2]) ** 2, 1))))
    fig.suptitle("Constant-velocity model, q = 0.05, measurement sd 0.8; ellipses: one standard deviation.\n"
                 "RMS error: measurements %.2f, filter %.2f, smoother %.2f" % (rm(X), rm(kf["mu"][:, :2]), rm(sm["mu"][:, :2])),
                 fontsize=8.6, y=1.0)
    fig.tight_layout()
    common.savefig(fig, 13, "fig13-22", out_dir)
    return dict(rmse_obs=rm(X), rmse_filter=rm(kf["mu"][:, :2]), rmse_smoother=rm(sm["mu"][:, :2]))


# =========================================================================== Figure 13.23
SEED_FIG23 = 23


def fig13_23(out_dir=None):
    """Figure 13.23: one step of the particle filter for a 1-D latent variable, from a real run (L = 12)."""
    model = PF_MODELS["growth"]
    z, x = pf_simulate(JSRandom(SEED_FIG23), model, 50)
    L = 12
    pf = particle_filter(JSRandom(SEED_FIG23 + 1000), x, model, L)
    # pick the step whose new likelihood is clearly bimodal and whose particles cover both modes
    best, n0 = -1, 1
    for n in range(3, 40):
        P = pf["particles"][n + 1]
        sc = min((P > 0).sum(), (P < 0).sum()) if x[n + 1] > 3 else -1
        if sc > best:
            best, n0 = sc, n
    P0, W0 = pf["particles"][n0], pf["weights"][n0]
    P1, W1, par = pf["particles"][n0 + 1], pf["weights"][n0 + 1], pf["parents"][n0 + 1]
    lo = min(P0.min(), P1.min()) - 3
    hi = max(P0.max(), P1.max()) + 3
    fig = plt.figure(figsize=(11.0, 4.3))
    gs = fig.add_gridspec(1, 2, width_ratios=[1.5, 1], wspace=0.12)
    a = fig.add_subplot(gs[0])
    rows = [3.0, 2.0, 1.0, 0.0]
    labels = [r"$p(z_n|\mathbf{X}_n)$", r"$p(z_{n+1}|\mathbf{X}_n)$", r"$p(x_{n+1}|z_{n+1})$", r"$p(z_{n+1}|\mathbf{X}_{n+1})$"]
    for y, lab in zip(rows, labels):
        a.annotate("", xy=(hi, y), xytext=(lo, y), arrowprops=dict(arrowstyle="-|>", color="k", lw=1.3))
        a.text(lo - 0.4, y, lab, ha="right", va="center", fontsize=11)
    area = lambda w: 1400.0 * w + 12.0          # marker area in points^2, proportional to the weight
    kw = dict(fc="#b3b3ff", ec=common.BLUE, lw=1.4, zorder=3)
    a.scatter(P0, np.full(L, rows[0]), s=area(W0), **kw)
    a.scatter(P1, np.full(L, rows[1]), s=area(np.full(L, 1.0 / L)), **kw)
    a.scatter(P1, np.full(L, rows[3]), s=area(W1), **kw)
    for l in range(L):
        a.annotate("", xy=(P1[l], rows[1]), xytext=(P0[par[l]], rows[0]),
                   arrowprops=dict(arrowstyle="-|>", color=common.BLUE, lw=1.0, mutation_scale=8, shrinkA=7, shrinkB=6), zorder=2)
        a.annotate("", xy=(P1[l], rows[3]), xytext=(P1[l], rows[1]),
                   arrowprops=dict(arrowstyle="-|>", color=common.BLUE, lw=0.9, mutation_scale=8, shrinkA=6,
                                   shrinkB=math.sqrt(area(W1[l]) / math.pi) + 1), zorder=2)
    span = hi - lo
    zz = np.linspace(lo, hi, 600)
    lik = np.exp(-0.5 * (x[n0 + 1] - zz ** 2 / 20) ** 2 / model["r"])
    a.plot(zz, rows[2] + 0.75 * lik, color=common.RED, lw=2)
    a.set_xlim(lo - span * 0.42, hi + 0.5)
    a.set_ylim(-0.6, 3.6)
    a.axis("off")
    a.text(hi, -0.45, "$z$", fontsize=12, ha="right")
    a.set_title("(a) one step, n = %d, of a run with L = %d particles\n(circle area: weight; $x_{n+1}$ = %.2f)"
                % (n0 + 1, L, x[n0 + 1]), fontsize=9.5)
    b = fig.add_subplot(gs[1])
    pf2 = particle_filter(JSRandom(SEED_FIG23 + 2000), x, model, 500)
    nn = np.repeat(np.arange(1, len(x) + 1), 500)
    b.scatter(nn, pf2["particles"].ravel(), s=40 * pf2["weights"].ravel() * 500 / 10 + 0.3, color=common.BLUE, alpha=0.15,
              edgecolors="none", rasterized=True, label="particles (size: weight)")
    b.plot(np.arange(1, len(x) + 1), z, "k-", lw=1.4, label="true $z_n$")
    b.plot(np.arange(1, len(x) + 1), pf2["mean"], "--", color=common.RED, lw=1.3, label="posterior mean (13.117)")
    b.set_xlabel("$n$")
    b.set_ylabel("$z_n$")
    b.set_ylim(-28, 28)
    b.legend(fontsize=7.5, loc="lower left", frameon=True, framealpha=0.85)
    b.set_title("(b) the whole run, L = 500:\n$p(z_n|\\mathbf{X}_n)$ is often bimodal", fontsize=9.5)
    common.savefig(fig, 13, "fig13-23", out_dir)
    return dict(n=n0, ess=pf["ess"], mean_ess_500=float(pf2["ess"].mean()))


FIGURES = {
    "13.1": fig13_1, "13.8": fig13_8, "13.11": fig13_11, "13.12": fig13_12, "13.13": fig13_13,
    "13.16": fig13_16, "13.22": fig13_22, "13.23": fig13_23,
}

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