"""Numerical checks of the derivations in Bishop (2006), Chapter 8 (Graphical Models).

Each check prints PASS or FAIL with the numbers behind it; the last line counts them.
Run from this folder: python3 checks_ch08.py (about half a minute). The algorithms checked
are the ones in ch08_graphical_models.py, which also produce the figures. Every claim is
compared with brute-force enumeration of the joint distribution, or with Monte Carlo using
fixed seeds and tolerances of about four standard errors.
"""
import itertools
import math

import numpy as np

import ch08_graphical_models as m

RESULTS = []
TOL_INDEP = 1e-12     # conditional mutual information below this counts as independence


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_dag(rng, n, p=0.4, max_par=3):
    names = ["v%d" % i for i in range(n)]
    par = {}
    for i, v in enumerate(names):
        cand = [names[j] for j in range(i) if rng.random() < p]
        par[v] = tuple(cand[:max_par])
    return par


FIG2 = {"x1": (), "x2": (), "x3": (), "x4": ("x1", "x2", "x3"), "x5": ("x1", "x3"), "x6": ("x4",), "x7": ("x4", "x5")}

# ----------------------------------------------------------------------------- 8.1
print("== Bayesian networks: factorisation and ancestral sampling (Section 8.1)")
rng = np.random.default_rng(0)
card = {"x1": 2, "x2": 3, "x3": 2, "x4": 3, "x5": 2, "x6": 2, "x7": 4}
cp = m.random_cpts(FIG2, card, rng)
order = m.topological_order(FIG2)
P = m.joint_table(order, FIG2, cp)
check("(8.4)/(8.5): the product of the conditionals of Figure 8.2 sums to 1 (mixed numbers of states)",
      abs(P.sum() - 1) < 1e-12 and P.min() >= 0, "sum = %.15f, %d states" % (P.sum(), P.size))
# Exercise 8.1: summing out a leaf leaves the factorisation of the remaining graph
sub = {k: v for k, v in FIG2.items() if k not in ("x6", "x7")}
P_sub = m.joint_table(m.topological_order(sub), sub, {k: cp[k] for k in sub})
check("Exercise 8.1: summing out the leaves x6, x7 leaves the factorisation of the subgraph",
      np.allclose(P.sum(axis=(5, 6)), P_sub, atol=1e-15), "max diff %.1e" % np.abs(P.sum(axis=(5, 6)) - P_sub).max())
marg_ok = all(np.allclose(P.sum(axis=tuple(j for j in range(7) if j != order.index(v))), cp[v]) for v in ("x1", "x2", "x3"))
check("parentless nodes: the marginal p(x_k) of the joint is the table p(x_k)", marg_ok)
S = m.ancestral_sample(FIG2, cp, np.random.default_rng(1), 400000)
counts = np.zeros(P.shape)
np.add.at(counts, tuple(S[v] for v in order), 1)
freq = counts / counts.sum()
z = np.abs(freq - P) / np.sqrt(P * (1 - P) / 400000 + 1e-300)
chi2 = (((counts - 400000 * P) ** 2) / (400000 * P)).sum()
dof = P.size - 1
check("ancestral sampling (Section 8.1.2): frequencies of all %d configurations match p(x) (chi-squared)" % P.size,
      abs(chi2 - dof) < 5 * math.sqrt(2 * dof), "chi2 = %.0f on %d d.o.f., max |z| = %.1f" % (chi2, dof, z.max()))

print("== Linear-Gaussian models (Section 8.1.4)")
rng = np.random.default_rng(2)
D = 6
W = np.tril(rng.normal(0, 0.8, (D, D)), -1) * (rng.random((D, D)) < 0.6)
b = rng.normal(0, 1, D)
v = rng.uniform(0.3, 1.5, D)
mu, Sg = m.linear_gaussian_moments(W, b, v)
A = np.linalg.inv(np.eye(D) - W)
check("recursions (8.15), (8.16) equal the closed form mu = (I - W)^-1 b, Sigma = (I - W)^-1 diag(v) (I - W)^-T",
      np.allclose(mu, A @ b) and np.allclose(Sg, A @ np.diag(v) @ A.T), "max diff %.1e" % np.abs(Sg - A @ np.diag(v) @ A.T).max())
X = m.linear_gaussian_sample(W, b, v, np.random.default_rng(3), 400000)
se_mu = np.sqrt(np.diag(Sg) / len(X))
se_S = np.sqrt((Sg ** 2 + np.outer(np.diag(Sg), np.diag(Sg))) / len(X))
check("(8.15), (8.16) against Monte Carlo (400000 ancestral samples of (8.14))",
      np.all(np.abs(X.mean(0) - mu) < 4.5 * se_mu) and np.all(np.abs(np.cov(X.T) - Sg) < 4.5 * se_S),
      "max |z| mean %.1f, cov %.1f" % ((np.abs(X.mean(0) - mu) / se_mu).max(), (np.abs(np.cov(X.T) - Sg) / se_S).max()))
w21, w32, b3 = 0.7, -1.3, np.array([0.5, -1.0, 2.0])
v3 = np.array([0.4, 0.9, 1.6])
W3 = np.array([[0, 0, 0], [w21, 0, 0], [0, w32, 0]])
mu3, S3 = m.linear_gaussian_moments(W3, b3, v3)
mu_book = [b3[0], b3[1] + w21 * b3[0], b3[2] + w32 * b3[1] + w32 * w21 * b3[0]]
S_book = np.array([[v3[0], w21 * v3[0], w32 * w21 * v3[0]],
                   [w21 * v3[0], v3[1] + w21 ** 2 * v3[0], w32 * (v3[1] + w21 ** 2 * v3[0])],
                   [w32 * w21 * v3[0], w32 * (v3[1] + w21 ** 2 * v3[0]), v3[2] + w32 ** 2 * (v3[1] + w21 ** 2 * v3[0])]])
check("Exercise 8.7: the chain x1 -> x2 -> x3 (Figure 8.14) has mean (8.17) and covariance (8.18)",
      np.allclose(mu3, mu_book) and np.allclose(S3, S_book))

# ----------------------------------------------------------------------------- 8.2
print("== The three 3-node graphs (Section 8.2.1, Figures 8.15 to 8.20), 20 random tables each")
for kind, marg_dep, cond_dep in (("tail", True, False), ("chain", True, False), ("head", False, True)):
    pos, par = m.THREE[kind]
    order = m.topological_order(par)
    mins = [np.inf, np.inf]
    maxs = [0.0, 0.0]
    for seed in range(20):
        Pk = m.joint_table(order, par, m.random_cpts(par, {x: 2 for x in par}, np.random.default_rng(100 + seed)))
        for j, C in enumerate(((), ("c",))):
            val = m.cond_mutual_info(Pk, order, ["a"], ["b"], C)
            mins[j] = min(mins[j], val)
            maxs[j] = max(maxs[j], abs(val))
    ok0 = (mins[0] > 1e-8) if marg_dep else (maxs[0] < TOL_INDEP)
    ok1 = (mins[1] > 1e-8) if cond_dep else (maxs[1] < TOL_INDEP)
    name = {"tail": "tail-to-tail (8.23)-(8.24)", "chain": "head-to-tail (8.25)-(8.26)", "head": "head-to-head (8.27)-(8.28)"}[kind]
    check("%s: a, b %s marginally" % (name, "dependent" if marg_dep else "independent"), ok0,
          "I(a;b) in [%.1e, %.1e]" % (mins[0], maxs[0]))
    check("%s: a, b %s given c" % (name, "dependent" if cond_dep else "independent"), ok1,
          "I(a;b|c) in [%.1e, %.1e]" % (mins[1], maxs[1]))
    check("%s: d_separated gives the same two answers" % name,
          m.d_separated(par, ["a"], ["b"]) != marg_dep and m.d_separated(par, ["a"], ["b"], ["c"]) != cond_dep)

print("== Explaining away: the fuel gauge (Section 8.2.1, Figure 8.21)")
fg = m.fuel_gauge()
check("(8.30) p(G = 0) = 0.315", abs(fg["p(G=0)"] - 0.315) < 1e-12, "%.6f" % fg["p(G=0)"])
check("(8.31) p(G = 0 | F = 0) = 0.81", abs(fg["p(G=0|F=0)"] - 0.81) < 1e-12, "%.6f" % fg["p(G=0|F=0)"])
check("(8.32) p(F = 0 | G = 0) = 0.257", round(fg["p(F=0|G=0)"], 3) == 0.257, "%.6f" % fg["p(F=0|G=0)"])
check("(8.33) p(F = 0 | G = 0, B = 0) = 0.111", round(fg["p(F=0|G=0,B=0)"], 3) == 0.111, "%.6f" % fg["p(F=0|G=0,B=0)"])
Pf = fg["joint"]
indep = np.allclose(Pf.sum(2), np.outer(Pf.sum((1, 2)), Pf.sum((0, 2))))
pBF_g0 = Pf[:, :, 0] / Pf[:, :, 0].sum()
dep = not np.allclose(pBF_g0, np.outer(pBF_g0.sum(1), pBF_g0.sum(0)))
check("B and F are independent a priori and become dependent once G = 0 is observed", indep and dep)
check("p(F = 0) = 0.1 < p(F = 0 | G = 0, B = 0) < p(F = 0 | G = 0): the battery explains the gauge away only partly",
      0.1 < fg["p(F=0|G=0,B=0)"] < fg["p(F=0|G=0)"])

print("== D-separation against brute-force independence (Section 8.2.2)")
P22 = m.joint_table(m.topological_order(m.FIG22_PAR), m.FIG22_PAR, m.fig22_cpts())
o22 = m.topological_order(m.FIG22_PAR)
i_c = m.cond_mutual_info(P22, o22, ["a"], ["b"], ["c"])
i_f = m.cond_mutual_info(P22, o22, ["a"], ["b"], ["f"])
check("Figure 8.22(a): a, b not d-separated by c and I(a;b|c) > 0; (b): d-separated by f and I(a;b|f) = 0",
      not m.d_separated(m.FIG22_PAR, "a", "b", "c") and i_c > 1e-3 and m.d_separated(m.FIG22_PAR, "a", "b", "f") and abs(i_f) < TOL_INDEP,
      "I(a;b|c) = %.4f, I(a;b|f) = %.1e" % (i_c, i_f))
rng = np.random.default_rng(5)
n_q = n_agree = n_trail = n_moral = n_sep = 0
weakest = np.inf
for g in range(40):
    n = int(rng.integers(5, 8))
    par = random_dag(rng, n)
    order = m.topological_order(par)
    Pg = m.joint_table(order, par, m.random_cpts(par, {x: 2 for x in par}, rng, alpha=0.7))
    nodes = list(par)
    for a, b_ in itertools.combinations(nodes, 2):
        rest = [x for x in nodes if x not in (a, b_)]
        for _ in range(3):
            C = [x for x in rest if rng.random() < 0.35]
            sep = m.d_separated(par, [a], [b_], C)
            cmi = m.cond_mutual_info(Pg, order, [a], [b_], C)
            n_q += 1
            n_sep += sep
            n_agree += sep == (abs(cmi) < TOL_INDEP)
            if not sep:
                weakest = min(weakest, cmi)
            n_trail += sep == all(blk for _, blk in m.trails(par, a, b_, C))
            # Lauritzen's criterion: separation in the moral graph of the ancestral set of {a, b} and C
            An = m.ancestors_of(par, [a, b_] + C)
            E = m.moralize({x: tuple(p for p in par[x]) for x in An})
            nb = {x: set() for x in An}
            for e in E:
                u, w = tuple(e)
                nb[u].add(w)
                nb[w].add(u)
            seen, stack = {a}, [a]
            while stack:
                for w in nb[stack.pop()]:
                    if w not in seen and w not in C:
                        seen.add(w)
                        stack.append(w)
            n_moral += sep == (b_ not in seen)
check("d_separated (reachability) agrees with I(a;b|C) = 0 on %d queries in 40 random DAGs with random tables" % n_q,
      n_agree == n_q, "%d agree, %d d-separated; smallest I for a non-separated pair %.1e" % (n_agree, n_sep, weakest))
check("d_separated agrees with 'every path blocked' (trails, the definition of Section 8.2.2)", n_trail == n_q, "%d / %d" % (n_trail, n_q))
check("d_separated agrees with separation in the moralised ancestral graph (Lauritzen et al. 1990)", n_moral == n_q,
      "%d / %d" % (n_moral, n_q))

print("== Markov blankets (Figures 8.26, 8.28)")
mb = m.markov_blanket(m.MB_PAR, "xi")
check("Figure 8.26: the blanket of x_i is its parents, children and co-parents (not g, not d)",
      mb == {"p1", "p2", "c1", "c2", "q1", "q2"}, ", ".join(sorted(mb)))
worst = worst38 = 0.0
for seed in range(10):
    rng = np.random.default_rng(200 + seed)
    par = m.MB_PAR if seed < 3 else random_dag(rng, 7, p=0.45)
    order = m.topological_order(par)
    cp = m.random_cpts(par, {x: 2 for x in par}, rng)
    Pg = m.joint_table(order, par, cp)
    for v in par:
        i = order.index(v)
        cond_all = Pg / Pg.sum(axis=i, keepdims=True)
        blanket = m.markov_blanket(par, v)
        keep = [order.index(u) for u in blanket] + [i]
        drop = tuple(j for j in range(len(order)) if j not in keep)
        Pm = Pg.sum(axis=drop, keepdims=True) if drop else Pg
        cond_mb = np.broadcast_to(Pm / Pm.sum(axis=i, keepdims=True), Pg.shape)
        worst = max(worst, np.abs(cond_all - cond_mb).max())
        # (8.38): p(x_i | rest) is the normalised product of the factors that contain x_i
        F = np.ones(Pg.shape)
        for u in par:
            if u == v or v in par[u]:      # the factors p(x_u | pa_u) that contain x_i
                axes = [order.index(p) for p in par[u]] + [order.index(u)]
                T = np.moveaxis(cp[u].reshape(cp[u].shape + (1,) * (len(order) - len(axes))),
                                list(range(len(axes))), axes)
                F = F * T
        worst38 = max(worst38, np.abs(F / F.sum(axis=i, keepdims=True) - cond_all).max())
check("p(x_i | all other nodes) = p(x_i | Markov blanket) for every node of 10 networks", worst < 1e-12, "max diff %.1e" % worst)
check("(8.38): p(x_i | all others) is the normalised product of the factors that contain x_i", worst38 < 1e-12, "max diff %.1e" % worst38)

# ----------------------------------------------------------------------------- 8.3
print("== Markov random fields: moralisation and the Ising model (Section 8.3)")
E = m.moralize({"x1": (), "x2": (), "x3": (), "x4": ("x1", "x2", "x3")})
check("Figure 8.33: moralising x1, x2, x3 -> x4 gives the complete graph on four nodes", len(E) == 6)
y = np.array([[1, -1, 1], [1, 1, -1], [-1, 1, 1]], float)
x = np.array([[1, 1, 1], [-1, 1, -1], [1, 1, -1]], float)
e_loop = 0.3 * x.sum() - 0.7 * sum(x[i, j] * x[k, l] for i in range(3) for j in range(3) for k in range(3) for l in range(3)
                                   if abs(i - k) + abs(j - l) == 1 and (i, j) < (k, l)) - 1.9 * (x * y).sum()
check("ising_energy is (8.42) with each neighbouring pair counted once", abs(m.ising_energy(x, y, 0.3, 0.7, 1.9) - e_loop) < 1e-12)
# Gibbs conditional and the ICM rule from the energy difference (Exercise 8.13)
okc = True
for (i, j) in [(0, 0), (1, 1), (2, 1)]:
    xp, xm = x.copy(), x.copy()
    xp[i, j], xm[i, j] = 1, -1
    dE = m.ising_energy(xm, y, 0.3, 0.7, 1.9) - m.ising_energy(xp, y, 0.3, 0.7, 1.9)
    a = 0.7 * m._nbr_sum(x, i, j) + 1.9 * y[i, j] - 0.3
    okc &= abs(dE - 2 * a) < 1e-12 and abs(1 / (1 + math.exp(-dE)) - float(m.common.sigmoid(2 * a))) < 1e-12
check("Exercise 8.13: E(x_j = -1) - E(x_j = +1) = 2 (beta sum_nbrs x_i + eta y_j - h), hence the ICM rule and the Gibbs conditional", okc)

x0, yn = m.fig30_setup()
xi, tr = m.icm(yn)
check("ICM on Figure 8.30: the energy never increases from sweep to sweep (Exercise 8.13)",
      np.all(np.diff(tr) <= 1e-9), "energies " + ", ".join("%.1f" % e for e in tr))
loc = True
for i in range(yn.shape[0]):
    for j in range(yn.shape[1]):
        a = m._nbr_sum(xi, i, j) + 2.1 * yn[i, j]
        loc &= a * xi[i, j] >= 0
check("ICM stops at a local minimum: no single-pixel change lowers the energy", loc)
xg = m.graph_cut_map(yn)
eg = m.ising_energy(xg, yn)
check("graph cut: E(graph cut) <= E(ICM) <= E(noisy image) on Figure 8.30",
      eg <= tr[-1] + 1e-9 and tr[-1] <= tr[0], "%.1f <= %.1f <= %.1f" % (eg, tr[-1], tr[0]))
ex = lambda z: 100 * np.mean(z != x0)
check("Figure 8.30 error rates: noisy about 10%, ICM < noisy, graph cut < ICM",
      9 < ex(yn) < 11 and ex(xg) < ex(xi) < ex(yn), "noisy %.2f%%, ICM %.2f%%, graph cut %.2f%%" % (ex(yn), ex(xi), ex(xg)))
xi2, _ = m.icm(yn, x0=xg)
check("the graph-cut image is a fixed point of ICM (a global minimum is also a local one)", np.array_equal(xi2, xg))
rng = np.random.default_rng(9)
nbf, okbf = 0, True
worst_gap = 0.0
for t in range(16):
    H, W_ = (3, 4) if t % 2 else (4, 4)
    yy = np.where(rng.random((H, W_)) < 0.5, 1.0, -1.0)
    h, beta, eta = [(0.0, 1.0, 2.1), (0.4, 0.8, 1.0), (-0.3, 1.5, 0.7), (0.25, 0.5, 2.0)][t % 4]
    xb, eb = m.brute_force_map(yy, h, beta, eta)
    xc = m.graph_cut_map(yy, h, beta, eta)
    ec = m.ising_energy(xc, yy, h, beta, eta)
    worst_gap = max(worst_gap, abs(ec - eb))
    okbf &= abs(ec - eb) < 1e-9
    nbf += 1
check("graph cut equals the brute-force minimum of (8.42) over all 2^16 images on %d tiny noisy images" % nbf, okbf,
      "max energy gap %.1e" % worst_gap)
check("Exercise 8.14: with beta = 0 (and h = 0) the most probable image is x = y",
      np.array_equal(m.graph_cut_map(yn, beta=0.0), yn) and np.array_equal(m.icm(yn, beta=0.0)[0], yn))
check("image_to_hex packs the test image losslessly (the demo stores it this way)",
      np.array_equal(np.where(np.unpackbits(np.frombuffer(bytes.fromhex(m.image_to_hex(x0)), np.uint8))[:x0.size] > 0, 1.0, -1.0).reshape(x0.shape), x0))

# ----------------------------------------------------------------------------- 8.4
print("== Inference on a chain (Section 8.4.1)")
rng = np.random.default_rng(11)
okz = okm = okp = True
for t in range(6):
    Ks = rng.integers(2, 5, size=6)
    psis = [rng.uniform(0.1, 2.0, (Ks[n], Ks[n + 1])) for n in range(5)]
    r = m.chain_sum_product(psis)
    fgc = m.chain_graph(psis)
    bm, Zb = fgc.brute_marginals()
    Pj = fgc.joint()
    okz &= abs(r["Z"] - Zb) < 1e-10 * Zb
    okm &= all(np.allclose(a, b_) for a, b_ in zip(r["marginals"], bm))
    for n in range(5):
        pair = Pj.sum(axis=tuple(j for j in range(6) if j not in (n, n + 1))) / Zb
        okp &= np.allclose(pair, r["pairwise"][n])
check("(8.54)-(8.57): forward/backward messages give every marginal of 6 random chains (2 to 4 states)", okm)
check("Z = sum over x_n of mu_alpha(x_n) mu_beta(x_n) for every n equals the brute-force normaliser", okz)
check("(8.58): p(x_{n-1}, x_n) = mu_alpha psi mu_beta / Z equals the brute-force pairwise marginals", okp)
r = m.chain_sum_product(m.CHAIN_PSI)
check("the same marginals from the general sum_product on the chain of the demo (CHAIN_PSI)",
      all(np.allclose(a, b_) for a, b_ in zip(r["marginals"], m.sum_product(m.chain_graph(), root=4)["marginals"])),
      "Z = %.4f" % r["Z"])

print("== The sum-product algorithm (Section 8.4.4)")
fg51 = m.fig51_graph()
res = m.sum_product(fg51, root=2)
bm, Zb = fg51.brute_marginals()
check("Figure 8.51: marginals of x1..x4 from the messages (8.74)-(8.85) equal brute force (8.86)",
      all(np.allclose(a, b_) for a, b_ in zip(res["marginals"], bm)) and abs(res["Z"] - Zb) < 1e-12,
      "Z = %g, p(x2) = (%.4f, %.4f)" % (res["Z"], *res["marginals"][1]))
msgs = {(m._key(s), m._key(d)): v for s, d, v in res["messages"]}
check("Figure 8.52: the messages of the example are (8.75) mu_fa->x2 = (3, 4), (8.79) mu_fb->x3 = (33, 26), (8.83) mu_fa->x1 = (26, 33)",
      np.allclose(msgs[("fa", "x2")], [3, 4]) and np.allclose(msgs[("fb", "x3")], [33, 26]) and np.allclose(msgs[("fa", "x1")], [26, 33]))
Pj = fg51.joint()
okf = True
for s, (vs, _) in enumerate(fg51.factors):
    okf &= np.allclose(Pj.sum(axis=tuple(j for j in range(4) if j not in vs)) / Zb, res["factor_marginals"][s])
check("(8.72): factor marginals p(x_s) = f_s(x_s) prod of incoming messages / Z", okf)
okr = all(all(np.allclose(a, b_) for a, b_ in zip(m.sum_product(fg51, root=rt)["marginals"], bm)) for rt in range(4))
check("the marginals do not depend on the choice of root", okr)
rng = np.random.default_rng(12)
okt, okev, n_tree = True, True, 0
for t in range(12):
    nv = int(rng.integers(4, 8))
    card = list(rng.integers(2, 4, size=nv))
    fgt = m.FactorGraph(card)
    # a random tree: each new variable is linked by a pairwise factor to one earlier variable
    for v_ in range(1, nv):
        a_ = int(rng.integers(0, v_))
        fgt.add((a_, v_), rng.uniform(0.1, 2, (card[a_], card[v_])))
    for v_ in range(nv):
        if rng.random() < 0.5:
            fgt.add((v_,), rng.uniform(0.1, 2, card[v_]))
    assert fgt.is_tree()
    n_tree += 1
    nv2 = nv
    sp = m.sum_product(fgt, root=int(rng.integers(0, nv2)))
    bmt, Zt = fgt.brute_marginals()
    okt &= all(np.allclose(a, b_) for a, b_ in zip(sp["marginals"], bmt)) and abs(sp["Z"] - Zt) < 1e-9 * Zt
    ev = {int(rng.integers(0, nv2)): 0, int(rng.integers(0, nv2)): 1}
    ev = {k: min(v2, fgt.card[k] - 1) for k, v2 in ev.items()}
    spe = m.sum_product(fgt, root=0, evidence=ev)
    bme, Ze = fgt.brute_marginals(ev)
    okev &= all(np.allclose(a, b_) for a, b_ in zip(spe["marginals"], bme)) and abs(spe["Z"] - Ze) < 1e-9 * Ze
check("sum-product equals brute force on %d random tree factor graphs (2 or 3 states, unary and pairwise factors)" % n_tree, okt)
check("Section 8.4.4, observed variables: clamping by indicator functions gives p(x_i | v) and Z = p~(v)", okev)
# a factor over three variables (Figure 8.46 style): x1, x2 -> f(x1, x2, x3) -> x3, with x4 hanging off x3
fg3 = m.FactorGraph([2, 3, 2, 3])
rng = np.random.default_rng(13)
fg3.add((0, 1, 2), rng.uniform(0.1, 2, (2, 3, 2)))
fg3.add((2, 3), rng.uniform(0.1, 2, (2, 3)))
fg3.add((1,), rng.uniform(0.1, 2, 3))
sp3, (bm3, Z3) = m.sum_product(fg3, root=3), fg3.brute_marginals()
check("sum-product with a factor of three variables (8.66) equals brute force",
      all(np.allclose(a, b_) for a, b_ in zip(sp3["marginals"], bm3)) and abs(sp3["Z"] - Z3) < 1e-10)
# normalised factors: Z is the probability of the evidence
fgB = m.FactorGraph([2, 2, 2])
fgB.add((0,), [0.1, 0.9])
fgB.add((1,), [0.1, 0.9])
Jf = m.fuel_gauge()["joint"]
fgB.add((0, 1, 2), Jf / Jf.sum(axis=2, keepdims=True))          # p(G | B, F)
spB = m.sum_product(fgB, root=2, evidence={2: 0})
check("fuel gauge as a factor graph: Z = p(G = 0) = 0.315 and p(F = 0 | G = 0) = 0.257 by sum-product",
      abs(spB["Z"] - 0.315) < 1e-12 and round(spB["marginals"][1][0], 3) == 0.257,
      "Z = %.4f, p(F=0|G=0) = %.4f" % (spB["Z"], spB["marginals"][1][0]))

print("== The max-sum algorithm (Section 8.4.5)")
rng = np.random.default_rng(14)
okc, okt2 = True, True
for t in range(10):
    Ks = rng.integers(2, 5, size=6)
    psis = [rng.uniform(0.1, 3.0, (Ks[n], Ks[n + 1])) for n in range(5)]
    rc = m.chain_max_sum(psis)
    Pj = m.chain_graph(psis).joint()
    best = np.unravel_index(Pj.argmax(), Pj.shape)
    okc &= list(best) == rc["path"] and abs(math.log(Pj.max()) - rc["logmax"]) < 1e-10
    xm, lm = m.max_sum(m.chain_graph(psis), root=int(rng.integers(0, 6)))
    okt2 &= list(best) == xm and abs(math.log(Pj.max()) - lm) < 1e-10
check("chain_max_sum (8.99)-(8.102): the back-tracked path is the brute-force arg max on 10 random chains", okc)
check("max_sum on the chain factor graph, any root, gives the same configuration and ln p_max (8.97)", okt2)
rc = m.chain_max_sum(m.CHAIN_PSI)
Pj = m.chain_graph().joint()
check("Figure 8.53 chain (CHAIN_PSI): path %s, ln max %.4f, equal to brute force" % ([p + 1 for p in rc["path"]], rc["logmax"]),
      list(np.unravel_index(Pj.argmax(), Pj.shape)) == rc["path"] and abs(math.log(Pj.max()) - rc["logmax"]) < 1e-12)
okm = True
for t in range(10):
    fgt = m.FactorGraph([3, 2, 3, 2, 3])
    for (a_, b_) in [(0, 1), (1, 2), (1, 3), (3, 4)]:
        fgt.add((a_, b_), rng.uniform(0.1, 3, (fgt.card[a_], fgt.card[b_])))
    fgt.add((2,), rng.uniform(0.1, 3, 3))
    ev = {4: int(rng.integers(0, 3))} if t % 2 else None
    xm, lm = m.max_sum(fgt, root=int(rng.integers(0, 5)), evidence=ev)
    Pj = fgt.joint(ev)
    okm &= list(np.unravel_index(Pj.argmax(), Pj.shape)) == xm and abs(math.log(Pj.max()) - lm) < 1e-10
check("max_sum with back-tracking on 10 random trees (half with an observed variable) equals brute force", okm)
T81 = np.array([[0.3, 0.3], [0.4, 0.0]])      # Table 8.1: rows x = 0, 1; columns y = 0, 1
fgT = m.FactorGraph([2, 2])
fgT.add((0, 1), T81)
xT, lT = m.max_sum(fgT, root=0)
mT = m.sum_product(fgT)["marginals"]
check("Table 8.1: max_sum finds (x, y) = (1, 0) with p = 0.4, while the marginal modes give (0, 0) with p = 0.3",
      xT == [1, 0] and abs(math.exp(lT) - 0.4) < 1e-12 and [int(np.argmax(p)) for p in mT] == [0, 0],
      "p(x) = (%.1f, %.1f), p(y) = (%.1f, %.1f)" % (*mT[0], *mT[1]))

print("== Loopy belief propagation (Section 8.4.7)")
hist = m.loopy_bp(fg51, iters=10)
err_tree = max(np.abs(a - b_).max() for a, b_ in zip(hist[-1], bm))
check("flooding schedule on the tree of Figure 8.51: exact marginals after a few iterations", err_tree < 1e-12, "max error %.1e" % err_tree)
rng = np.random.default_rng(15)
okl = True
for t in range(5):
    fgt = m.FactorGraph([2, 3, 2, 3, 2, 2])
    for (a_, b_) in [(0, 1), (1, 2), (1, 3), (3, 4), (3, 5)]:
        fgt.add((a_, b_), rng.uniform(0.1, 3, (fgt.card[a_], fgt.card[b_])))
    h_ = m.loopy_bp(fgt, iters=12)
    bt, _ = fgt.brute_marginals()
    okl &= max(np.abs(a - b_).max() for a, b_ in zip(h_[-1], bt)) < 1e-12
check("loopy_bp is exact on 5 random trees (diameter <= 6 iterations suffice)", okl)
fgc = m.cycle_graph()
hc = m.loopy_bp(fgc, iters=200, tol=1e-13)
bc, _ = fgc.brute_marginals()
err_c = max(np.abs(a - b_).max() for a, b_ in zip(hc[-1], bc))
check("on a cycle of five binary variables loopy BP converges to beliefs that are close to, but not, the marginals",
      1e-4 < err_c < 0.15 and len(hc) < 200, "max |belief - marginal| = %.4f after %d iterations" % (err_c, len(hc)))
errs = []
for J in (0.2, 0.5, 1.0, 1.5):
    fgj = m.cycle_graph(J)
    hj = m.loopy_bp(fgj, iters=500, tol=1e-13)
    bj, _ = fgj.brute_marginals()
    errs.append(max(np.abs(a - b_).max() for a, b_ in zip(hj[-1], bj)))
check("the loopy-BP error grows with the coupling strength J around the loop", all(np.diff(errs) > 0),
      ", ".join("J = %.1f: %.4f" % (J, e) for J, e in zip((0.2, 0.5, 1.0, 1.5), errs)))

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