"""Numerical checks of the derivations in Bishop (2006), Chapter 13 (Sequential Data).

Each check prints PASS or FAIL with the numbers behind it; the last line counts them.
Run from this folder: python3 checks_ch13.py (about half a minute). The algorithms checked are
the ones in ch13_sequential_data.py, which also produce the figures: the HMM recursions are
compared with brute-force enumeration of every state path, the Kalman filter and smoother with
the conditionals of the joint Gaussian of all latent and observed variables.
"""
import itertools
import math

import numpy as np

import common
import ch13_sequential_data as m

RESULTS = []


def check(name, ok, detail=""):
    RESULTS.append(bool(ok))
    print("%s  %s%s" % ("PASS" if ok else "FAIL", name, ("  [" + detail + "]") if detail else ""))


def random_hmm(rng, K, D=2, M=None):
    pi = rng.dirichlet(np.ones(K))
    A = rng.dirichlet(np.ones(K), size=K)
    if M is not None:
        return dict(pi=pi, A=A, B=rng.dirichlet(np.ones(M), size=K))
    mus = rng.normal(0, 1.2, size=(K, D))
    S = []
    for _ in range(K):
        Q = rng.normal(size=(D, D))
        S.append(Q @ Q.T * 0.3 + 0.3 * np.eye(D))
    return dict(pi=pi, A=A, mus=mus, Sigmas=np.array(S))


def logE_of(p, X):
    return m.emission_logprob(X, B=p["B"]) if "B" in p else m.emission_logprob(X, p["mus"], p["Sigmas"])


def brute_force(pi, A, logE):
    """p(X), gamma and xi by summing p(X, Z) over all K^N paths (13.6)."""
    N, K = logE.shape
    pX, gam, xi = 0.0, np.zeros((N, K)), np.zeros((N - 1, K, K))
    best, bestz = -np.inf, None
    for z in itertools.product(range(K), repeat=N):
        lj = m.log_joint(pi, A, logE, z)
        pj = math.exp(lj)
        pX += pj
        for n in range(N):
            gam[n, z[n]] += pj
        for n in range(1, N):
            xi[n - 1, z[n - 1], z[n]] += pj
        if lj > best:
            best, bestz = lj, np.array(z)
    return pX, gam / pX, xi / pX, best, bestz


# ----------------------------------------------------------------------------- 13.2.2 and 13.2.4
print("== Forward-backward against brute-force enumeration (Sections 13.2.2, 13.2.4)")
rng = np.random.default_rng(1)
cases = [("discrete lattice example of Figs 13.12, 13.13", dict(pi=m.LATTICE_PI, A=m.LATTICE_A, B=m.LATTICE_B), m.LATTICE_X)]
pg = random_hmm(rng, 3)
_, Xg = m.hmm_sample(rng, pg["pi"], pg["A"], 7, pg["mus"], pg["Sigmas"])
cases.append(("Gaussian emissions, K = 3, N = 7", pg, Xg))
pd = random_hmm(rng, 4, M=5)
_, Xd = m.hmm_sample(rng, pd["pi"], pd["A"], 6, B=pd["B"])
cases.append(("discrete emissions, K = 4, N = 6", pd, Xd))
for name, p, X in cases:
    lE = logE_of(p, X)
    pX, gam, xi, best, bestz = brute_force(p["pi"], p["A"], lE)
    u = m.forward_backward_unscaled(p["pi"], p["A"], lE)
    s = m.forward_backward(p["pi"], p["A"], lE)
    check("unscaled alpha-beta (13.36)-(13.43): p(X), gamma, xi equal enumeration (%s)" % name,
          abs(u["pX"] / pX - 1) < 1e-10 and np.allclose(u["gamma"], gam, atol=1e-12) and np.allclose(u["xi"], xi, atol=1e-12),
          "p(X) = %.6g, max |gamma error| %.1e, max |xi error| %.1e" % (pX, np.abs(u["gamma"] - gam).max(), np.abs(u["xi"] - xi).max()))
    ok_alpha = np.allclose(s["alpha"], u["alpha"] / u["alpha"].sum(1, keepdims=True), atol=1e-12)
    ok_c = np.allclose(np.exp(np.cumsum(s["log_c"])), u["alpha"].sum(1), rtol=1e-10)
    check("scaled recursions (13.59)-(13.65) agree with the unscaled ones: alpha_hat = p(z_n|x_1..x_n), prod_{m<=n} c_m = p(x_1..x_n), gamma, xi, ln p(X) (%s)" % name,
          ok_alpha and ok_c and np.allclose(s["gamma"], u["gamma"], atol=1e-12) and np.allclose(s["xi"], u["xi"], atol=1e-12)
          and abs(s["loglik"] - math.log(pX)) < 1e-10,
          "ln p(X): scaled %.10f, enumeration %.10f" % (s["loglik"], math.log(pX)))
    path, vbest, _, _ = m.viterbi(p["pi"], p["A"], lE)
    check("Viterbi (13.68)-(13.71) returns the brute-force argmax path and max ln p(X, Z) (%s)" % name,
          np.array_equal(path, bestz) and abs(vbest - best) < 1e-10, "path %s, ln p(X, Z) = %.6f" % ("".join(str(k + 1) for k in path), vbest))
    check("beta_hat(z_n) = p(x_{n+1..N}|z_n) / p(x_{n+1..N}|x_{1..n}) (13.60), and sum_k gamma = 1, sum xi = gamma",
          np.allclose(s["beta"], u["beta"] / np.exp(s["loglik"] - np.cumsum(s["log_c"]))[:, None], rtol=1e-9)
          and np.allclose(s["gamma"].sum(1), 1) and np.allclose(s["xi"].sum(2), s["gamma"][:-1]) and np.allclose(s["xi"].sum(1), s["gamma"][1:]))

# random Viterbi checks
ok, n_cases = True, 30
for t in range(n_cases):
    p = random_hmm(rng, 3, M=4)
    _, X = m.hmm_sample(rng, p["pi"], p["A"], 6, B=p["B"])
    lE = logE_of(p, X)
    _, _, _, best, bestz = brute_force(p["pi"], p["A"], lE)
    path, vb, _, _ = m.viterbi(p["pi"], p["A"], lE)
    ok &= np.array_equal(path, bestz) and abs(vb - best) < 1e-10
check("Viterbi equals the brute-force argmax on 30 random discrete HMMs (K = 3, N = 6)", ok)

lE = m.emission_logprob(m.LATTICE_X, B=m.LATTICE_B)
path, vb, _, _ = m.viterbi(m.LATTICE_PI, m.LATTICE_A, lE)
gp = m.forward_backward(m.LATTICE_PI, m.LATTICE_A, lE)["gamma"].argmax(1)
lj = m.log_joint(m.LATTICE_PI, m.LATTICE_A, lE, gp)
check("most probable path differs from the most probable states (lattice example, Fig 13.16): the sequence of argmax gamma has a lower joint probability",
      not np.array_equal(path, gp) and lj < vb,
      "Viterbi %s, ln p(X, Z) = %.3f; argmax gamma %s, ln p(X, Z) = %.3f" % ("".join(str(k + 1) for k in path), vb, "".join(str(k + 1) for k in gp), lj))
# a left-to-right example where the individually most probable states form an impossible path
piL = np.array([1.0, 0, 0])
AL = np.array([[0.6, 0.4, 0], [0, 0.2, 0.8], [0, 0, 1.0]])     # state 2 is short-lived
BL = np.array([[0.8, 0.2], [0.5, 0.5], [0.2, 0.8]])
found = None
for x in itertools.product(range(2), repeat=3):
    lE2 = m.emission_logprob(np.array(x), B=BL)
    fb = m.forward_backward(piL, AL, lE2)
    g2 = fb["gamma"].argmax(1)
    if not np.isfinite(m.log_joint(piL, AL, lE2, g2)):
        found = (x, g2, m.viterbi(piL, AL, lE2)[0])
        break
check("in a left-to-right HMM the per-step argmax of gamma can even be an impossible path (probability 0); Viterbi never is",
      found is not None, "x = %s: argmax gamma %s, Viterbi %s" % (found[0], "".join(str(k + 1) for k in found[1]), "".join(str(k + 1) for k in found[2])) if found else "none found")

# long sequence: the unscaled recursion underflows, the scaled one does not
p8 = dict(pi=m.FIG8_PI, A=m.FIG8_A, mus=m.FIG8_MU, Sigmas=m.FIG8_SIGMA)
_, X8 = m.hmm_sample(np.random.default_rng(3), p8["pi"], p8["A"], 3000, p8["mus"], p8["Sigmas"])
lE8 = logE_of(p8, X8)
with np.errstate(all="ignore"):
    u8 = m.forward_backward_unscaled(p8["pi"], p8["A"], lE8 - 40)  # shift so the unscaled values underflow
s8 = m.forward_backward(p8["pi"], p8["A"], lE8)
lalpha = np.log(p8["pi"]) + lE8[0]
for n in range(1, len(X8)):
    lalpha = lE8[n] + common.logsumexp(lalpha[:, None] + np.log(p8["A"]), axis=0)
check("N = 3000: unscaled alpha underflows to 0, the scaled recursion gives ln p(X) = sum ln c_n (13.63), equal to a log-domain forward pass",
      u8["pX"] == 0 and abs(s8["loglik"] - common.logsumexp(lalpha)) < 1e-6 * abs(s8["loglik"]),
      "ln p(X) = %.4f versus %.4f" % (s8["loglik"], common.logsumexp(lalpha)))

# ----------------------------------------------------------------------------- 13.2.1 Baum-Welch
print("== Baum-Welch (Section 13.2.1)")


def Q(p, X, gam, xi):
    """Q(theta, theta_old) of (13.17) for one sequence, with gamma and xi from theta_old."""
    with np.errstate(divide="ignore"):
        lp, lA = np.log(p["pi"]), np.log(p["A"])
    q = np.sum(gam[0] * np.where(gam[0] > 0, lp, 0)) + np.sum(xi.sum(0) * np.where(xi.sum(0) > 0, lA, 0))
    return q + np.sum(gam * logE_of(p, X))


_, X8 = m.hmm_sample(m.JSRandom(5), p8["pi"], p8["A"], 300, p8["mus"], p8["Sigmas"])
r0 = np.random.default_rng(7)
start = dict(pi=np.ones(3) / 3, A=np.full((3, 3), 0.1) + 0.7 * np.eye(3), mus=X8[r0.choice(300, 3, replace=False)],
             Sigmas=np.array([np.cov(X8.T, bias=True)] * 3))
pBW, llg = m.baum_welch([X8], start, n_iter=60)
d = np.diff(llg)
check("Gaussian emissions (Fig 13.8 model, N = 300): ln p(X) never decreases over 60 EM iterations",
      d.min() > -1e-8, "ln p(X): %.3f -> %.3f, smallest change %.2e" % (llg[0], llg[-1], d.min()))
pd = random_hmm(np.random.default_rng(8), 3, M=4)
_, Xd1 = m.hmm_sample(np.random.default_rng(9), pd["pi"], pd["A"], 200, B=pd["B"])
_, Xd2 = m.hmm_sample(np.random.default_rng(10), pd["pi"], pd["A"], 150, B=pd["B"])
pstart = random_hmm(np.random.default_rng(11), 3, M=4)
pBWd, lld = m.baum_welch([Xd1, Xd2], pstart, n_iter=60)
check("discrete emissions, two sequences: ln p(X) never decreases", np.diff(lld).min() > -1e-8,
      "%.3f -> %.3f" % (lld[0], lld[-1]))
D = m.fit_digits()
check("left-to-right HMM of Fig 13.11 (16 states, 45 synthetic digits, 25 iterations): monotone, and every forbidden transition stays exactly 0",
      np.diff(D["ll"]).min() > -1e-8 and np.all(D["p"]["A"][~(np.eye(16, dtype=bool) | np.eye(16, k=1, dtype=bool))] == 0),
      "ln p(X): %.2f -> %.2f; mean self-transition %.3f" % (D["ll"][0], D["ll"][-1], np.mean(np.diag(D["p"]["A"])[:-1])))

# M step = stationary point of Q (finite differences, with the simplex constraints)
fbg = m.forward_backward(start["pi"], start["A"], logE_of(start, X8))
new, _ = m.baum_welch([X8], start, n_iter=1)
g, xi = fbg["gamma"], fbg["xi"]
h = 1e-6
worst, scale = 0.0, 0.0
for key in ("pi", "A"):
    for r in range(1 if key == "pi" else 3):
        for i, j in [(0, 1), (1, 2), (0, 2)]:
            for base, acc in ((new, "w"), (start, "s")):
                pp, pm = {k: v.copy() for k, v in base.items()}, {k: v.copy() for k, v in base.items()}
                for q_, sgn in ((pp, 1), (pm, -1)):
                    if key == "pi":
                        q_["pi"][i] += sgn * h
                        q_["pi"][j] -= sgn * h
                    else:
                        q_["A"][r, i] += sgn * h
                        q_["A"][r, j] -= sgn * h
                dq = (Q(pp, X8, g, xi) - Q(pm, X8, g, xi)) / (2 * h)
                if acc == "w":
                    worst = max(worst, abs(dq))
                else:
                    scale = max(scale, abs(dq))
for k in range(3):
    for i in range(2):
        for base, acc in ((new, "w"), (start, "s")):
            pp, pm = {kk: v.copy() for kk, v in base.items()}, {kk: v.copy() for kk, v in base.items()}
            pp["mus"][k, i] += h
            pm["mus"][k, i] -= h
            dq = (Q(pp, X8, g, xi) - Q(pm, X8, g, xi)) / (2 * h)
            worst, scale = (max(worst, abs(dq)), scale) if acc == "w" else (worst, max(scale, abs(dq)))
    for (i, j) in [(0, 0), (1, 1), (0, 1)]:
        for base, acc in ((new, "w"), (start, "s")):
            pp, pm = {kk: v.copy() for kk, v in base.items()}, {kk: v.copy() for kk, v in base.items()}
            E = np.zeros((2, 2))
            E[i, j] = E[j, i] = 1e-7
            pp["Sigmas"][k] += E
            pm["Sigmas"][k] -= E
            dq = (Q(pp, X8, g, xi) - Q(pm, X8, g, xi)) / 2e-7
            worst, scale = (max(worst, abs(dq)), scale) if acc == "w" else (worst, max(scale, abs(dq)))
check("M step (13.18)-(13.21) is a stationary point of Q(theta, theta_old): all directional derivatives (pi and A along the simplex, mu, symmetric Sigma) vanish",
      worst < 1e-4 * max(scale, 1), "max |dQ| at the M-step solution %.1e, at theta_old %.1e" % (worst, scale))
qn = Q(new, X8, g, xi)
rr = np.random.default_rng(12)
better = 0
for _ in range(200):
    pp = {k: v.copy() for k, v in new.items()}
    pp["pi"] = np.abs(pp["pi"] + 0.02 * rr.normal(size=3))
    pp["pi"] /= pp["pi"].sum()
    pp["A"] = np.abs(pp["A"] + 0.02 * rr.normal(size=(3, 3)))
    pp["A"] /= pp["A"].sum(1, keepdims=True)
    pp["mus"] = pp["mus"] + 0.01 * rr.normal(size=(3, 2))
    pp["Sigmas"] = np.array([S + 0.001 * np.eye(2) * rr.normal() for S in pp["Sigmas"]])
    better += Q(pp, X8, g, xi) > qn
check("... and a maximum: 200 random feasible perturbations all lower Q", better == 0, "%d of 200 higher" % better)
fbd = m.forward_backward(pstart["pi"], pstart["A"], logE_of(pstart, Xd1))
newd, _ = m.baum_welch([Xd1], pstart, n_iter=1)
worst = 0.0
for k in range(3):
    for i, j in [(0, 1), (1, 2), (2, 3), (0, 3)]:
        pp, pm = {kk: v.copy() for kk, v in newd.items()}, {kk: v.copy() for kk, v in newd.items()}
        pp["B"][k, i] += h
        pp["B"][k, j] -= h
        pm["B"][k, i] -= h
        pm["B"][k, j] += h
        worst = max(worst, abs(Q(pp, Xd1, fbd["gamma"], fbd["xi"]) - Q(pm, Xd1, fbd["gamma"], fbd["xi"])) / (2 * h))
check("discrete M step (13.23): Q is stationary along every simplex direction of each row of the emission table", worst < 1e-4,
      "max |dQ| %.1e" % worst)

# ----------------------------------------------------------------------------- 13.3 linear dynamical systems
print("== Kalman filter and RTS smoother against the joint Gaussian (Section 13.3.1)")


def joint_gaussian(md, N):
    """Mean and covariance of (z_1..z_N, x_1..x_N) under (13.75)-(13.77)."""
    A, G, C, S = md["A"], md["Gamma"], md["C"], md["Sigma"]
    Dz, Dx = A.shape[0], C.shape[0]
    mz = [md["mu0"]]
    for n in range(1, N):
        mz.append(A @ mz[-1])
    Szz = np.zeros((N * Dz, N * Dz))
    blocks = {}
    for n in range(N):
        blocks[n, n] = md["P0"] if n == 0 else A @ blocks[n - 1, n - 1] @ A.T + G
        for k in range(n - 1, -1, -1):
            blocks[n, k] = A @ blocks[n - 1, k]
    for (n, k), B in blocks.items():
        Szz[n * Dz:(n + 1) * Dz, k * Dz:(k + 1) * Dz] = B
        Szz[k * Dz:(k + 1) * Dz, n * Dz:(n + 1) * Dz] = B.T
    Cb = np.kron(np.eye(N), C)
    mu = np.concatenate([np.concatenate(mz), Cb @ np.concatenate(mz)])
    Sxx = Cb @ Szz @ Cb.T + np.kron(np.eye(N), S)
    Szx = Szz @ Cb.T
    Sig = np.block([[Szz, Szx], [Szx.T, Sxx]])
    return mu, Sig, N * Dz


def condition(mu, Sig, nz, obs_idx, xval):
    """Mean and covariance of all z given the x components obs_idx (indices into x)."""
    ix = nz + np.asarray(obs_idx, dtype=int)
    Szx = Sig[:nz][:, ix]
    Sxx = Sig[np.ix_(ix, ix)]
    Kg = Szx @ np.linalg.inv(Sxx)
    return mu[:nz] + Kg @ (xval - mu[ix]), Sig[:nz, :nz] - Kg @ Szx.T


rl = np.random.default_rng(20)
md = dict(A=np.array([[0.9, 0.3, 0.0], [-0.2, 0.8, 0.1], [0.0, 0.2, 0.7]]), Gamma=np.diag([0.3, 0.2, 0.1]) + 0.05,
          C=rl.normal(size=(2, 3)), Sigma=np.array([[0.5, 0.1], [0.1, 0.4]]), mu0=np.array([1.0, -1.0, 0.5]),
          P0=np.diag([1.0, 0.5, 0.8]))
Nl = 8
Z, X = m.lds_sample(rl, md, Nl)
mu_j, Sig_j, nz = joint_gaussian(md, Nl)
Dz, Dx = 3, 2
for label, Xu in (("all observed", X.copy()), ("x_4 and x_5 missing", X.copy())):
    if "missing" in label:
        Xu[3:5] = np.nan
    kf = m.kalman_filter(Xu, md)
    sm = m.rts_smoother(kf, md)
    errf = errs = errc = 0.0
    for n in range(Nl):
        idx = [i for k in range(n + 1) if not np.isnan(Xu[k, 0]) for i in range(k * Dx, (k + 1) * Dx)]
        vals = np.array([Xu[k, i] for k in range(n + 1) if not np.isnan(Xu[k, 0]) for i in range(Dx)])
        cm, cS = condition(mu_j, Sig_j, nz, idx, vals)
        sl = slice(n * Dz, (n + 1) * Dz)
        errf = max(errf, np.abs(kf["mu"][n] - cm[sl]).max(), np.abs(kf["V"][n] - cS[sl, sl]).max())
    idx = [i for k in range(Nl) if not np.isnan(Xu[k, 0]) for i in range(k * Dx, (k + 1) * Dx)]
    vals = np.array([Xu[k, i] for k in range(Nl) if not np.isnan(Xu[k, 0]) for i in range(Dx)])
    cm, cS = condition(mu_j, Sig_j, nz, idx, vals)
    for n in range(Nl):
        sl = slice(n * Dz, (n + 1) * Dz)
        errs = max(errs, np.abs(sm["mu"][n] - cm[sl]).max(), np.abs(sm["V"][n] - cS[sl, sl]).max())
        if n >= 1:
            slp = slice((n - 1) * Dz, n * Dz)
            errc = max(errc, np.abs(sm["J"][n - 1] @ sm["V"][n] - cS[slp, sl]).max())
    ix = nz + np.array(idx)
    lj = common.gauss_logpdf(vals[None, :], mu_j[ix], Sig_j[np.ix_(ix, ix)])[0]
    check("Kalman filter (13.89)-(13.97): mu_n, V_n equal the mean and covariance of p(z_n|x_1..x_n) from the joint Gaussian (%s)" % label,
          errf < 1e-9, "max error %.1e" % errf)
    check("RTS smoother (13.100)-(13.102): mu_hat_n, V_hat_n equal p(z_n|x_1..x_N) from the joint Gaussian (%s)" % label,
          errs < 1e-9, "max error %.1e" % errs)
    check("cov[z_{n-1}, z_n] = J_{n-1} V_hat_n (13.104), from the joint Gaussian (%s)" % label, errc < 1e-9, "max error %.1e" % errc)
    check("ln p(X) = sum ln c_n (13.63) with c_n from (13.91) equals the joint Gaussian marginal density of the observed x (%s)" % label,
          abs(kf["loglik"] - lj) < 1e-9, "%.8f versus %.8f" % (kf["loglik"], lj))

# Exercises 13.27 and 13.28: limits of the Kalman filter
m1 = dict(A=np.eye(1), Gamma=np.eye(1) * 1.0, C=np.eye(1), Sigma=np.eye(1) * 1e-8, mu0=np.zeros(1), P0=np.eye(1))
xs = np.random.default_rng(21).normal(size=(10, 1)) * 3
kf1 = m.kalman_filter(xs, m1)
check("Exercise 13.27: with measurement noise much smaller than the state noise, mu_n = x_n", np.abs(kf1["mu"] - xs).max() < 1e-6,
      "max |mu_n - x_n| = %.1e" % np.abs(kf1["mu"] - xs).max())
m2 = dict(A=np.eye(1), Gamma=np.zeros((1, 1)), C=np.eye(1), Sigma=np.eye(1), mu0=np.zeros(1), P0=np.eye(1) * 1e10)
kf2 = m.kalman_filter(xs, m2)
check("Exercise 13.28: with Gamma = 0, A = C = I and a flat prior, mu_n is the average of x_1..x_n and V_n = Sigma/n",
      np.allclose(kf2["mu"][:, 0], np.cumsum(xs[:, 0]) / np.arange(1, 11), atol=1e-6) and np.allclose(kf2["V"][:, 0, 0], 1 / np.arange(1, 11), atol=1e-6))

print("== EM for the linear dynamical system (Section 13.3.2)")
rl = np.random.default_rng(30)
mt = dict(A=np.array([[0.95, 0.2], [-0.2, 0.9]]), Gamma=np.diag([0.1, 0.1]), C=np.array([[1.0, 0.0], [0.5, 1.0], [0.0, 0.7]]),
          Sigma=np.diag([0.3, 0.2, 0.4]), mu0=np.zeros(2), P0=np.eye(2))
Zt, Xt = m.lds_sample(rl, mt, 200)
m0 = dict(A=np.eye(2) * 0.5, Gamma=np.eye(2), C=rl.normal(size=(3, 2)), Sigma=np.eye(3), mu0=np.zeros(2), P0=np.eye(2))
mfit, lll = m.lds_em(Xt, m0, n_iter=100)
check("LDS EM (13.110)-(13.116, with the errata): ln p(X) never decreases over 100 iterations", np.diff(lll).min() > -1e-8,
      "ln p(X): %.2f -> %.2f (true parameters: %.2f); smallest change %.1e" % (lll[0], lll[-1], m.kalman_filter(Xt, mt)["loglik"], np.diff(lll).min()))
ev_t = np.sort(np.abs(np.linalg.eigvals(mt["A"])))
ev_f = np.sort(np.abs(np.linalg.eigvals(mfit["A"])))
check("the fitted A has eigenvalue moduli close to the true ones (A is identified only up to a change of basis of z)",
      np.abs(ev_t - ev_f).max() < 0.08, "true %s, fitted %s" % (np.round(ev_t, 3), np.round(ev_f, 3)))
# The M step for C and Sigma maximises Q: perturbing C away from the M-step value lowers the expected complete-data log likelihood
kf = m.kalman_filter(Xt, m0)
sm = m.rts_smoother(kf, m0)
m1s, _ = m.lds_em(Xt, m0, n_iter=1)


def Qc(Cm, Sm):
    Ez, Ezz = sm["mu"], sm["Ezz"]
    Si = np.linalg.inv(Sm)
    tot = 0.0
    for n in range(len(Xt)):
        x = Xt[n]
        E = np.outer(x, x) - Cm @ np.outer(Ez[n], x) - np.outer(x, Ez[n]) @ Cm.T + Cm @ Ezz[n] @ Cm.T
        tot += -0.5 * np.trace(Si @ E)
    return tot - 0.5 * len(Xt) * np.linalg.slogdet(Sm)[1]


q0 = Qc(m1s["C"], m1s["Sigma"])
worst = 0.0
for i in range(3):
    for j in range(2):
        E = np.zeros((3, 2))
        E[i, j] = 1e-6
        worst = max(worst, abs(Qc(m1s["C"] + E, m1s["Sigma"]) - Qc(m1s["C"] - E, m1s["Sigma"])) / 2e-6)
    for j in range(3):
        E = np.zeros((3, 3))
        E[i, j] = E[j, i] = 1e-7
        worst = max(worst, abs(Qc(m1s["C"], m1s["Sigma"] + E) - Qc(m1s["C"], m1s["Sigma"] - E)) / 2e-7)
check("M step for C and Sigma (13.115), (13.116 corrected): Q is stationary in every entry (finite differences)", worst < 1e-3,
      "max |dQ| %.1e" % worst)

# ----------------------------------------------------------------------------- 13.3.4
print("== Particle filter (Section 13.3.4)")
mod = m.PF_MODELS["linear"]
zz, xx = m.pf_simulate(np.random.default_rng(40), mod, 60)
km, kv = m.pf_kalman(xx, mod)
rms = []
for L in (30, 300, 3000, 30000):
    errs = [np.sqrt(np.mean((m.particle_filter(np.random.default_rng(41 + r), xx, mod, L)["mean"] - km) ** 2)) for r in range(3)]
    rms.append(float(np.mean(errs)))
slope = np.polyfit(np.log([30, 300, 3000, 30000]), np.log(rms), 1)[0]
check("linear-Gaussian model: the particle-filter mean (13.117) approaches the Kalman mean as L grows, at about the Monte Carlo rate L^(-1/2)",
      all(a > b for a, b in zip(rms, rms[1:])) and -0.65 < slope < -0.35,
      "RMS difference from the Kalman mean for L = 30, 300, 3000, 30000: %s; slope %.2f" % (", ".join("%.4f" % r for r in rms), slope))
pf = m.particle_filter(np.random.default_rng(42), xx, mod, 30000)
check("... and its variance approaches the Kalman variance V_n", np.abs(pf["var"] - kv).max() < 0.03,
      "max |var - V_n| = %.4f (V_n about %.3f)" % (np.abs(pf["var"] - kv).max(), kv[-1]))
g = m.PF_MODELS["growth"]
zg, xg = m.pf_simulate(np.random.default_rng(43), g, 50)
p1 = m.particle_filter(np.random.default_rng(44), xg, g, 500)
p0 = m.particle_filter(np.random.default_rng(44), xg, g, 500, resample=False)
check("without resampling the weights degenerate (sequential importance sampling): the effective sample size collapses",
      p0["ess"][-10:].mean() < 2 and p1["ess"].mean() > 50, "mean ESS over the last 10 steps: %.2f without, %.1f with resampling (L = 500)"
      % (p0["ess"][-10:].mean(), p1["ess"][-10:].mean()))
js1 = m.particle_filter(m.JSRandom(7), xg, g, 50)
js2 = m.particle_filter(m.JSRandom(7), xg, g, 50)
check("JSRandom makes a particle-filter run reproducible (the demos' generator)", np.array_equal(js1["particles"], js2["particles"]))

print("\n%d checks: %d PASS, %d FAIL" % (len(RESULTS), sum(RESULTS), len(RESULTS) - sum(RESULTS)))