"""Numerical checks of the derivations in Bishop (2006), Chapter 2 (Probability Distributions).

Each check prints PASS or FAIL with the numbers behind it; the last line counts them.
Run from this folder: python3 checks_ch02.py (a few seconds). The functions checked are the ones in
ch02_probability_distributions.py, which also produce the figures and are mirrored in the demos' JavaScript.
"""
import math

import numpy as np
from scipy import integrate, special

import common
import ch02_probability_distributions 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 ""))


rng = np.random.default_rng(2)

# ---- normalisations by quadrature
for a, b in [(0.1, 0.1), (1, 1), (2, 3), (8, 4)]:
    v = integrate.quad(lambda u: m.beta_pdf(u, a, b), 0, 1, limit=200)[0]
    check("beta(%g, %g) integrates to 1 (2.13)" % (a, b), abs(v - 1) < 1e-6, "integral %.8f" % v)
al = np.array([2.0, 3.0, 4.5])
v = integrate.dblquad(lambda y, x: math.exp(m.dirichlet_logpdf(np.array([x, y, 1 - x - y]), al)), 0, 1, 0, lambda x: 1 - x)[0]
check("Dirichlet(2, 3, 4.5) integrates to 1 over the simplex (2.38)", abs(v - 1) < 1e-6, "integral %.8f" % v)
v = integrate.quad(lambda u: m.gauss_pdf(u, 1.3, 0.7), -20, 20)[0]
check("Gaussian integrates to 1 (1.48)", abs(v - 1) < 1e-9, "integral %.10f" % v)
S = np.array([[1.0, 0.6], [0.6, 2.0]])
v = integrate.dblquad(lambda y, x: math.exp(common.gauss_logpdf(np.array([[x, y]]), np.zeros(2), S)[0]), -12, 12, -12, 12)[0]
check("bivariate Gaussian integrates to 1 (2.43)", abs(v - 1) < 1e-6, "integral %.8f" % v)
for mm in (0.5, 5.0):
    v = integrate.quad(lambda t: m.von_mises_pdf(t, 1.0, mm), 0, 2 * np.pi)[0]
    check("von Mises m = %g integrates to 1 on [0, 2 pi) (2.179)" % mm, abs(v - 1) < 1e-9, "integral %.10f" % v)
v = integrate.quad(lambda x: m.student_pdf(x, 0.5, 2.0, 3.5), -np.inf, np.inf)[0]
check("Student t integrates to 1 (2.159)", abs(v - 1) < 1e-8, "integral %.9f" % v)
v = integrate.quad(lambda x: m.gamma_pdf(x, 4, 6), 0, np.inf)[0]
check("gamma(4, 6) integrates to 1 (2.146)", abs(v - 1) < 1e-8, "integral %.9f" % v)

# ---- moments by Monte Carlo
a, b = 2.0, 3.0
x = rng.beta(a, b, 400000)
mean, var, mode = m.beta_moments(a, b)
check("beta mean (2.15) and variance (2.16) against Monte Carlo", abs(x.mean() - mean) < 3e-3 and abs(x.var() - var) < 2e-3, "%.4f vs %.4f, %.5f vs %.5f" % (x.mean(), mean, x.var(), var))
D = rng.dirichlet(al, 400000)
mu_, C = m.dirichlet_moments(al)
check("Dirichlet mean and covariance (2.273)-(2.276) against Monte Carlo", np.abs(D.mean(0) - mu_).max() < 3e-3 and np.abs(np.cov(D.T) - C).max() < 1e-3,
      "max cov error %.2e" % np.abs(np.cov(D.T) - C).max())
vm = rng.vonmises(1.0, 2.0, 200000)
check("von Mises: E[cos(theta - theta0)] = A(m) (2.186) by Monte Carlo", abs(np.cos(vm - 1.0).mean() - m.bessel_ratio_A(2.0)) < 3e-3, "%.4f vs %.4f" % (np.cos(vm - 1.0).mean(), m.bessel_ratio_A(2.0)))

# ---- conjugate posteriors against grid posteriors
xs = rng.random(12) < 0.7
(aN, bN) = m.beta_update(2, 2, xs.astype(int))[-1]
g = np.linspace(1e-6, 1 - 1e-6, 20001)
post = g ** (2 - 1 + xs.sum()) * (1 - g) ** (2 - 1 + (~xs).sum())
post /= np.trapz(post, g)
check("beta-binomial: sequential update (2.18) equals the grid posterior", np.abs(post - m.beta_pdf(g, aN, bN)).max() < 1e-3 and (aN, bN) == (2 + xs.sum(), 2 + (~xs).sum()),
      "a_N = %d, b_N = %d, max density error %.1e" % (aN, bN, np.abs(post - m.beta_pdf(g, aN, bN)).max()))
chunk = m.beta_update(2, 2, xs[:6].astype(int))[-1]
seq = m.beta_update(*chunk, xs[6:].astype(int))[-1]
check("... and one batch of 12 equals two batches of 6 (sequential updating, Section 2.1.1)", seq == (aN, bN))
counts = np.array([3, 0, 5])
dN = m.dirichlet_update(al, counts)
g1, g2 = np.meshgrid(np.linspace(0.002, 0.998, 500), np.linspace(0.002, 0.998, 500))
inside = g1 + g2 < 0.998
pts = np.stack([g1, g2, 1 - g1 - g2], -1)
lp = (np.log(pts[inside]) * (al - 1 + counts)).sum(-1)
lq = m.dirichlet_logpdf(pts[inside], dN)
check("Dirichlet-multinomial: posterior (2.41) proportional to prior x likelihood on a grid", np.ptp(lp - lq) < 1e-9, "spread of ln ratio %.1e" % np.ptp(lp - lq))
xg = rng.normal(0.8, 0.3, 7)
muN, vN = m.gauss_mean_posterior(xg, 0.09, 0.0, 0.2)
grid = np.linspace(-3, 4, 40001)
lpst = -0.5 * (grid - 0.0) ** 2 / 0.2 - 0.5 * ((xg[:, None] - grid[None, :]) ** 2).sum(0) / 0.09
pp = np.exp(lpst - lpst.max()); pp /= np.trapz(pp, grid)
gm = np.trapz(grid * pp, grid); gv = np.trapz((grid - gm) ** 2 * pp, grid)
check("Gaussian mean, known variance: posterior mean (2.141) and variance (2.142) equal the grid values", abs(gm - muN) < 1e-6 and abs(gv - vN) < 1e-6, "%.6f vs %.6f, %.6f vs %.6f" % (gm, muN, gv, vN))
check("... the posterior mean is a convex combination of prior mean and ML estimate (2.141)", abs(muN - (0.09 / (7 * 0.2 + 0.09) * 0.0 + 7 * 0.2 / (7 * 0.2 + 0.09) * xg.mean())) < 1e-12)
aN, bN = m.gauss_precision_posterior(xg, 0.8, 2.0, 1.5)
lam = np.linspace(1e-4, 60, 200001)
lpl = (2 - 1) * np.log(lam) - 1.5 * lam + 0.5 * len(xg) * np.log(lam) - 0.5 * lam * ((xg - 0.8) ** 2).sum()
pl = np.exp(lpl - lpl.max()); pl /= np.trapz(pl, lam)
check("Gamma prior on the precision: posterior (2.150), (2.151) equals the grid posterior", np.abs(pl - m.gamma_pdf(lam, aN, bN)).max() < 1e-4, "a_N = %.1f, b_N = %.4f" % (aN, bN))
muN2, betaN, aN2, bN2 = m.normal_gamma_posterior(xg, 0.3, 2.0, 5.0, 6.0)
M, L = np.meshgrid(np.linspace(-1, 2.5, 700), np.linspace(0.02, 6, 700))
lpost = (np.log(m.normal_gamma_pdf(M, L, 0.3, 2.0, 5.0, 6.0)) + 0.5 * len(xg) * np.log(L)
         - 0.5 * L * ((xg[:, None, None] - M[None]) ** 2).sum(0))
lclosed = np.log(m.normal_gamma_pdf(M, L, muN2, betaN, aN2, bN2))
check("normal-gamma: posterior parameters (Exercise 2.44) match prior x likelihood up to a constant", np.ptp(lpost - lclosed) < 1e-8, "spread of ln ratio %.1e" % np.ptp(lpost - lclosed))

# ---- Gaussian conditional / marginal against numerics
Q = rng.normal(size=(4, 4)); Sig = Q @ Q.T + 0.5 * np.eye(4); mu = rng.normal(size=4)
a_, b_ = [0, 2], [1, 3]
xb = rng.normal(size=2)
cm, cS = m.gauss_conditional(mu, Sig, a_, b_, xb)
cm2, cS2 = m.gauss_conditional_precision(mu, Sig, a_, b_, xb)
check("Gaussian conditional: covariance form (2.81), (2.82) equals precision form (2.73), (2.75)", np.allclose(cm, cm2) and np.allclose(cS, cS2), "max diff %.1e" % max(np.abs(cm - cm2).max(), np.abs(cS - cS2).max()))
# numerical conditioning by weighting: sample a very large set, keep points near xb
Z = rng.multivariate_normal(mu, Sig, 2000000)
near = np.abs(Z[:, b_] - xb).max(1) < 0.15
sel = Z[near][:, a_]
check("... and against rejection sampling of the joint (points with x_b near the value)", np.abs(sel.mean(0) - cm).max() < 0.15 and np.abs(np.cov(sel.T) - cS).max() < 0.4,
      "%d points kept; mean error %.3f, covariance error %.3f" % (near.sum(), np.abs(sel.mean(0) - cm).max(), np.abs(np.cov(sel.T) - cS).max()))
mm, mS = m.gauss_marginal(mu, Sig, a_)
check("Gaussian marginal (2.92), (2.93) against sample mean and covariance", np.abs(Z[:, a_].mean(0) - mm).max() < 5e-3 and np.abs(np.cov(Z[:, a_].T) - mS).max() < 2e-2)
# the conditional density from the definition p(x_a, x_b)/p(x_b) at a point
xa = rng.normal(size=2)
lj = common.gauss_logpdf(np.r_[xa[0], xb[0], xa[1], xb[1]][None, :], mu, Sig)[0]
lm = common.gauss_logpdf(xb[None, :], mu[b_], Sig[np.ix_(b_, b_)])[0]
lc = common.gauss_logpdf(xa[None, :], cm, cS)[0]
check("p(x_a | x_b) = p(x_a, x_b) / p(x_b) at a point", abs(lj - lm - lc) < 1e-10, "difference %.1e" % (lj - lm - lc))

# ---- sequential ML (Robbins-Monro)
xr = rng.normal(0.8, 0.3, 5000)
tr = m.robbins_monro_mean(xr)
check("sequential ML (2.126) reproduces the batch mean exactly with a_N = 1/N", abs(tr[-1] - xr.mean()) < 1e-12 and abs(tr[99] - xr[:100].mean()) < 1e-12)
tr2 = m.robbins_monro_mean(xr, a_seq=lambda n: 2.0 / (n + 5))
check("Robbins-Monro (2.129)-(2.135) with a_N = 2/(N+5) (sum a diverges, sum a^2 converges) converges to the true mean", abs(tr2[-1] - 0.8) < 0.02, "final %.4f" % tr2[-1])

# ---- Student t as a scale mixture and ML by EM
xx = np.array([-2.0, 0.0, 0.7, 3.0])
for nu in (0.5, 3.0, 20.0):
    mix = m.student_mixture_pdf(xx, 0.4, 1.7, nu, n=20000)
    check("Student t (2.159) equals the gamma scale mixture of Gaussians (2.158), nu = %g" % nu, np.allclose(mix, m.student_pdf(xx, 0.4, 1.7, nu), rtol=2e-3), "max rel diff %.1e" % np.max(np.abs(mix / m.student_pdf(xx, 0.4, 1.7, nu) - 1)))
big = m.student_pdf(np.linspace(-4, 4, 9), 0, 1, 1e6)
check("Student t tends to the Gaussian as nu -> infinity", np.allclose(big, m.gauss_pdf(np.linspace(-4, 4, 9), 0, 1), atol=1e-6))
xt = 1.0 + 0.8 * rng.standard_t(3.0, 3000)
mu_t, lam_t, nu_t, ll = m.student_t_mle(xt)
check("t maximum likelihood by EM recovers (mu, scale, nu) = (1, 0.8, 3) from 3000 points", abs(mu_t - 1) < 0.05 and abs(1 / math.sqrt(lam_t) - 0.8) < 0.05 and abs(nu_t - 3) < 0.5,
      "mu %.3f, scale %.3f, nu %.2f" % (mu_t, 1 / math.sqrt(lam_t), nu_t))
lls = []
xm = rng.normal(size=40)
for it in (1, 5, 20, 100):
    r = m.student_t_mle(xm, iters=it); lls.append(r[3])
check("EM for the t distribution never decreases the likelihood", all(b >= a - 1e-9 for a, b in zip(lls, lls[1:])), "ln L at 1, 5, 20, 100 iterations: %s" % np.round(lls, 4))

# ---- von Mises ML
th = rng.vonmises(2.0, 4.0, 20000)
t0, mhat = m.von_mises_mle(th)
check("von Mises ML (2.185), (2.186) recovers theta_0 = 2, m = 4", abs(t0 - 2) < 0.02 and abs(mhat - 4) < 0.15, "theta_0 %.3f, m %.3f" % (t0, mhat))
opt = max(np.linspace(1.8, 2.2, 401), key=lambda t: np.log(m.von_mises_pdf(th, t, mhat)).sum())
check("... theta_0 maximises the log likelihood on a grid", abs(opt - t0) < 2e-3)

# ---- exponential family identities
mu0, s2 = 0.7, 1.3
eta = m.exp_family_gaussian_natural(mu0, s2)
eps = 1e-6
grad = np.array([(m.exp_family_log_g(eta + eps * e) - m.exp_family_log_g(eta - eps * e)) / (2 * eps) for e in np.eye(2)])
check("exponential family: -grad ln g(eta) = E[u(x)] = (mu, mu^2 + sigma^2) (2.226)", np.allclose(-grad, [mu0, mu0 ** 2 + s2], atol=1e-6), "%s" % np.round(-grad, 6))
check("exponential family: g(eta) int h(x) exp(eta^T u(x)) dx = 1 (2.195)", abs(math.exp(m.exp_family_log_g(eta)) * integrate.quad(lambda x: (2 * np.pi) ** -0.5 * math.exp(eta[0] * x + eta[1] * x * x), -30, 30)[0] - 1) < 1e-9)

# ---- nonparametric estimators
xk = m.fig224_sample(50, np.random.default_rng(24))
gg = np.linspace(-1.5, 2.5, 80001)
for h in (0.02, 0.07, 0.2):
    v = np.trapz(m.kde_gaussian(xk, gg, h), gg)
    check("kernel density estimate (2.250) integrates to 1, h = %g" % h, abs(v - 1) < 1e-6, "integral %.8f" % v)
for d in (0.04, 0.25):
    gh = np.linspace(0, 1, 100001)[:-1]
    v = np.trapz(m.histogram_density(xk, gh, d), gh)
    check("histogram estimate (2.241) integrates to 1, Delta = %g" % d, abs(v - 1) < 2e-3, "integral %.5f" % v)
x1 = rng.normal(size=20000)
gk = np.linspace(-1.0, 1.0, 5)
check("K-NN density (2.246) with K = 200 approximates the N(0, 1) density", np.abs(m.knn_density(x1, gk, 200) - m.gauss_pdf(gk, 0, 1)).max() < 0.03, "max error %.3f" % np.abs(m.knn_density(x1, gk, 200) - m.gauss_pdf(gk, 0, 1)).max())
x_small = np.array([0.1, 0.4, 0.5, 0.9])
check("K-NN density K = 2 at x = 0.45 equals K / (N V) with V = 2 r_K", abs(m.knn_density(x_small, np.array([0.45]), 2)[0] - 2 / (4 * 2 * 0.05)) < 1e-12, "value %.3f" % m.knn_density(x_small, np.array([0.45]), 2)[0])
v = np.trapz(m.knn_density(x_small, np.linspace(-200, 200, 400001), 2), np.linspace(-200, 200, 400001))
check("... but it is NOT normalised (tails decay like 1/|x|; integral diverges, Section 2.5.2)", v > 1.0, "integral over [-200, 200] = %.2f" % v)
# K-NN classifier posterior p(C_k | x) = K_k / K, Bayes rule on the counts
Xc = np.vstack([rng.normal(0, 1, (300, 2)), rng.normal(1.5, 1, (300, 2))]); yc = np.r_[np.zeros(300, int), np.ones(300, int)]
q = np.array([[0.7, 0.7]])
_, votes = m.knn_classify(Xc, yc, q, 15)
dd = np.argsort(((Xc - q) ** 2).sum(1))[:15]
check("K-NN classifier: the vote counts are K_k / K of the nearest neighbours (2.256)", votes.sum() == 15 and votes[0, 1] == yc[dd].sum(), "votes %s" % votes.round(0).tolist())

# Cover-Hart: asymptotic 1-NN error is between the Bayes error R* and 2 R* (1 - R*)
n_tr, n_te = 6000, 6000
sep = 1.0
def draw(n):
    yy = rng.integers(0, 2, n)
    return rng.normal(0, 1, n) + sep * yy, yy
xtr, ytr = draw(n_tr); xte, yte = draw(n_te)
order = np.argsort(xtr); xs_ = xtr[order]; ys_ = ytr[order]
pos = np.clip(np.searchsorted(xs_, xte), 1, len(xs_) - 1)
left, right = pos - 1, pos
nn = np.where(np.abs(xte - xs_[left]) < np.abs(xte - xs_[right]), left, right)
err = np.mean(ys_[nn] != yte)
Rstar = 0.5 * math.erfc((sep / 2) / math.sqrt(2))      # Phi(-sep/2) for equal priors
check("Cover and Hart: Bayes error R* <= 1-NN error <= 2 R*(1 - R*) (Section 2.5.2)", Rstar - 0.01 < err < 2 * Rstar * (1 - Rstar) + 0.01,
      "R* = %.4f, 1-NN error = %.4f, upper bound %.4f" % (Rstar, err, 2 * Rstar * (1 - Rstar)))

# JS generator is reproducible
a1 = m.JSRandom(5); a2 = m.JSRandom(5)
check("JSRandom reproduces itself (the demos' generator)", [a1.norm() for _ in range(5)] == [a2.norm() for _ in range(5)])

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