"""Numerical checks of the derivations in Bishop (2006), Chapter 1 (Introduction).

Each check prints PASS or FAIL with the numbers behind it; the last line counts them. Run from this
folder: python3 checks_ch01.py (a few seconds). The functions checked are those of
ch01_introduction.py, which also produce the figures.
"""
import numpy as np
from scipy import integrate, stats
from scipy.optimize import minimize

import ch01_introduction 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(0)

# ----------------------------------------------------------------------------- Section 1.1
print("== Polynomial curve fitting (Section 1.1)")
x, t = m.js_noisy_sine(10, m.DATA_SEED)
for M in (1, 3, 5):
    Phi = m.design_matrix(x, M)
    w_ne = np.linalg.solve(Phi.T @ Phi, Phi.T @ t)          # the normal equations
    w_ls = m.fit_poly(x, t, M)
    check("normal equations (Phi^T Phi) w = Phi^T t agree with least squares (M = %d)" % M,
          np.allclose(w_ne, w_ls, rtol=1e-6, atol=1e-8), "max |diff| %.1e" % np.abs(w_ne - w_ls).max())
w9 = m.fit_poly(x, t, 9)
check("M = 9, N = 10: the polynomial interpolates every training point (E_RMS = 0)", m.rms_error(x, t, w9) < 1e-8,
      "E_RMS = %.1e" % m.rms_error(x, t, w9))
Phi = m.design_matrix(x, 9)
check("the normal equations are ill-conditioned for M = 9 (cond(Phi^T Phi) > 1e10), which is why fit_poly uses an SVD",
      np.linalg.cond(Phi.T @ Phi) > 1e10, "cond = %.1e" % np.linalg.cond(Phi.T @ Phi))
# the sum-of-squares error gradient vanishes at the least-squares solution
M = 3
w = m.fit_poly(x, t, M)
grad = m.design_matrix(x, M).T @ (m.design_matrix(x, M) @ w - t)
check("gradient of E(w) (1.2) is zero at the least-squares solution", np.abs(grad).max() < 1e-9, "max |grad| %.1e" % np.abs(grad).max())
# ridge
lam = np.exp(-4.0)
M = 9
Phi = m.design_matrix(x, M)
w_r = m.fit_poly(x, t, M, lam)
w_closed = np.linalg.solve(lam * np.eye(M + 1) + Phi.T @ Phi, Phi.T @ t)
check("ridge solution w = (lam I + Phi^T Phi)^{-1} Phi^T t equals the stacked least squares of fit_poly (ln lambda = -4)",
      np.allclose(w_r, w_closed, rtol=1e-8, atol=1e-10), "max |diff| %.1e" % np.abs(w_r - w_closed).max())
E_reg = lambda w_: m.sse(x, t, w_) + 0.5 * lam * w_ @ w_
opt = minimize(E_reg, np.zeros(M + 1), jac=lambda w_: Phi.T @ (Phi @ w_ - t) + lam * w_, method="BFGS", options=dict(gtol=1e-11, maxiter=2000))
check("regularised error (1.4) at the ridge solution is no larger than at a numerical minimiser",
      E_reg(w_r) <= opt.fun + 1e-9, "E(w*) = %.10f, BFGS %.10f" % (E_reg(w_r), opt.fun))
check("ridge shrinks the weights: ||w|| decreases as lambda grows", all(np.linalg.norm(m.fit_poly(x, t, 9, np.exp(a))) > np.linalg.norm(m.fit_poly(x, t, 9, np.exp(b)))
      for a, b in [(-18, -10), (-10, -2), (-2, 2)]))

# ----------------------------------------------------------------------------- Sections 1.2.5, 1.2.6
print("== Probabilistic and Bayesian curve fitting (Sections 1.2.5, 1.2.6)")
alpha, beta, M = m.ALPHA, m.BETA, 9
w_map = m.map_poly(x, t, M, alpha, beta)
neg_log_post = lambda w_: 0.5 * beta * np.sum((m.design_matrix(x, M) @ w_ - t) ** 2) + 0.5 * alpha * w_ @ w_
opt = minimize(neg_log_post, np.zeros(M + 1), method="BFGS", options=dict(gtol=1e-10, maxiter=5000))
check("MAP = regularised least squares with lambda = alpha/beta (1.67): beta E + alpha/2 w^T w is minimised by map_poly",
      neg_log_post(w_map) <= opt.fun + 1e-8 and np.allclose(w_map, m.fit_poly(x, t, M, alpha / beta)),
      "objective at map_poly %.8f, BFGS %.8f" % (neg_log_post(w_map), opt.fun))
xs = np.linspace(0, 1, 11)
mean, var = m.bayes_predictive(x, t, xs, M, alpha, beta)
# direct joint-Gaussian computation: (t_1..t_N, t*) is jointly Gaussian with the prior w ~ N(0, alpha^-1 I)
Phi = m.design_matrix(x, M)
Phis = m.design_matrix(xs, M)
C = Phi @ Phi.T / alpha + np.eye(10) / beta           # cov of the training targets
Ks = Phis @ Phi.T / alpha                             # cov of t*, training targets
Kss = Phis @ Phis.T / alpha + np.eye(len(xs)) / beta
mean_d = Ks @ np.linalg.solve(C, t)
var_d = np.diag(Kss - Ks @ np.linalg.solve(C, Ks.T))
check("Bayesian predictive mean (1.70) equals the conditional mean of the joint Gaussian of (t, t*)", np.allclose(mean, mean_d, atol=1e-9),
      "max |diff| %.1e" % np.abs(mean - mean_d).max())
check("Bayesian predictive variance (1.71) equals the conditional variance of the joint Gaussian", np.allclose(var, var_d, atol=1e-9),
      "max |diff| %.1e" % np.abs(var - var_d).max())
check("predictive variance (1.71) is at least the noise variance 1/beta, with equality only without parameter uncertainty",
      np.all(var >= 1 / beta - 1e-12) and np.all(var > 1 / beta), "min s^2 = %.4f, 1/beta = %.4f" % (var.min(), 1 / beta))
mean_big, var_big = m.bayes_predictive(*m.js_noisy_sine(10000, 3), xs, 3, 1e-6, beta)
check("as N grows the predictive variance tends to 1/beta (here M = 3, N = 10000)", np.abs(var_big - 1 / beta).max() < 1e-3,
      "max |s^2 - 1/beta| = %.1e" % np.abs(var_big - 1 / beta).max())

# ----------------------------------------------------------------------------- Section 1.2.4
print("== The Gaussian (Section 1.2.4)")
mu_true, sig2 = 0.7, 2.0
N = 5
samples = rng.normal(mu_true, np.sqrt(sig2), size=(400000, N))
mu_ml = samples.mean(1)
s2_ml = ((samples - mu_ml[:, None]) ** 2).mean(1)
se = s2_ml.std() / np.sqrt(len(s2_ml))
check("E[mu_ML] = mu (1.57)", abs(mu_ml.mean() - mu_true) < 4 * mu_ml.std() / np.sqrt(len(mu_ml)), "mean %.4f vs %.4f" % (mu_ml.mean(), mu_true))
check("E[sigma^2_ML] = (N-1)/N sigma^2 (1.58), by simulation, N = 5", abs(s2_ml.mean() - m.expected_ml_variance(sig2, N)) < 4 * se,
      "simulated %.4f, theory %.4f (se %.4f)" % (s2_ml.mean(), m.expected_ml_variance(sig2, N), se))
check("the unbiased estimator N/(N-1) sigma^2_ML (1.59) has mean sigma^2", abs((s2_ml * N / (N - 1)).mean() - sig2) < 4 * se * N / (N - 1))
data = rng.normal(1.3, 0.8, 50)
nll = lambda p: -np.sum(stats.norm.logpdf(data, p[0], np.exp(p[1])))
o = minimize(nll, [0.0, 0.0], method="Nelder-Mead", options=dict(xatol=1e-10, fatol=1e-12, maxiter=5000))
mu_h, v_h = m.gauss_ml(data)
check("closed forms (1.55), (1.56) maximise the Gaussian log likelihood (1.54) (Nelder-Mead)", abs(o.x[0] - mu_h) < 1e-5 and abs(np.exp(2 * o.x[1]) - v_h) < 1e-5,
      "mu %.6f vs %.6f, sigma^2 %.6f vs %.6f" % (o.x[0], mu_h, np.exp(2 * o.x[1]), v_h))
total = integrate.quad(lambda z: m.gauss_pdf(z, 0.7, 2.0), -np.inf, np.inf)[0]
check("the Gaussian density (1.46) integrates to 1", abs(total - 1) < 1e-10)

# ----------------------------------------------------------------------------- Section 1.4
print("== Curse of dimensionality (Section 1.4)")
for D in (2, 5, 20):
    eps = 0.1
    frac_numeric = 1 - (integrate.quad(lambda r: r ** (D - 1), 0, 1 - eps)[0] / integrate.quad(lambda r: r ** (D - 1), 0, 1)[0])
    check("shell volume fraction 1 - (1 - eps)^D (1.76) equals the ratio of radial integrals (D = %d, eps = 0.1)" % D,
          abs(m.shell_fraction(eps, D) - frac_numeric) < 1e-10, "%.6f" % frac_numeric)
check("shell with eps = 0.1 holds 87.8% of the volume of a 20-dimensional sphere", abs(m.shell_fraction(0.1, 20) - 0.8784) < 1e-3, "%.4f" % m.shell_fraction(0.1, 20))
for D in (1, 2, 20):
    r = np.linspace(0, 12, 120001)
    p = m.radial_density(r, D, 0.5)
    mass = integrate.trapezoid(p, r)
    check("radial Gaussian density integrates to 1 (D = %d, sigma = 1/2)" % D, abs(mass - 1) < 1e-6, "%.8f" % mass)
    Xs = rng.normal(0, 0.5, size=(200000, D))
    rs = np.linalg.norm(Xs, axis=1)
    h, edges = np.histogram(rs, bins=40, range=(0, 4), density=True)
    cen = 0.5 * (edges[1:] + edges[:-1])
    check("... and agrees with the histogram of ||x|| for x ~ N(0, (1/2)^2 I_D) (D = %d)" % D, np.abs(h - m.radial_density(cen, D, 0.5)).max() < 0.03,
          "max |diff| %.3f" % np.abs(h - m.radial_density(cen, D, 0.5)).max())
mode20 = np.linspace(0, 5, 50001)[np.argmax(m.radial_density(np.linspace(0, 5, 50001), 20, 1.0))]
check("the radial density of N(0, I_D) peaks at r = sqrt(D - 1) (D = 20: %.3f)" % np.sqrt(19), abs(mode20 - np.sqrt(19)) < 1e-3, "mode %.4f" % mode20)

# ----------------------------------------------------------------------------- Section 1.5
print("== Decision theory (Section 1.5)")
xg = np.linspace(-6, 8, 4001)
pr = np.array([0.4, 0.6])
p_x_c = np.stack([stats.norm.pdf(xg, 0.0, 1.0), stats.norm.pdf(xg, 2.0, 1.0)], 1)
joint = p_x_c * pr
post = joint / joint.sum(1, keepdims=True)
dx = xg[1] - xg[0]
dec = m.bayes_decision(post, np.array([[0, 1], [1, 0]]))
check("0-1 loss: the minimum expected loss rule picks the most probable class", np.array_equal(dec, post.argmax(1)))
# misclassification probability of a threshold rule R1 = (-inf, xb): grid search over xb
def p_err(xb):
    return joint[xg >= xb, 0].sum() * dx + joint[xg < xb, 1].sum() * dx
cand = xg[(xg > -3) & (xg < 5)]
errs = np.array([p_err(c) for c in cand])
region = np.flatnonzero(np.diff(dec) != 0)
check("the Bayes boundary (posteriors equal) minimises p(mistake) (1.78) against a grid search over thresholds",
      abs(cand[errs.argmin()] - xg[region[0]]) < 0.05 and abs(errs.min() - p_err(xg[region[0] + 1])) < 1e-3,
      "grid argmin %.3f, Bayes boundary %.3f, p(mistake) %.4f" % (cand[errs.argmin()], xg[region[0]], errs.min()))
# general loss matrix: the rule beats every other deterministic assignment of classes to grid cells, and a grid search over thresholds
L = np.array([[0.0, 8.0], [1.0, 0.0]])    # cancer-style loss: L[truth, decision]
dec_L = m.bayes_decision(post, L)
R_bayes = (joint.sum(1) * m.expected_loss(post, L, dec_L)).sum() * dx
bestR = np.inf
for a_ in cand:
    for rule in (np.where(xg < a_, 0, 1), np.where(xg < a_, 1, 0)):
        R = (joint.sum(1) * m.expected_loss(post, L, rule)).sum() * dx
        bestR = min(bestR, R)
check("loss matrix L = [[0, 8], [1, 0]]: the rule argmin_j sum_k L_kj p(C_k|x) (1.81) has expected loss no larger than any threshold rule",
      R_bayes <= bestR + 1e-9, "Bayes %.5f, best threshold rule %.5f" % (R_bayes, bestR))
check("a loss matrix that punishes deciding C2 when the truth is C1 eight times more shrinks R2 (fewer grid points go to C2)", (dec_L == 1).sum() < (dec == 1).sum(),
      "%d of %d grid points in R2, against %d with 0-1 loss" % ((dec_L == 1).sum(), len(xg), (dec == 1).sum()))
# the loss integral equals the sum over classes: E[L] = sum_{k,j} int_{R_j} L_kj p(x, C_k) dx (1.80)
R_direct = sum(L[k, j] * joint[dec_L == j, k].sum() * dx for k in range(2) for j in range(2))
check("expected loss E[L] = sum_k sum_j int_{R_j} L_kj p(x, C_k) dx (1.80) equals the integral of the conditional expected loss", abs(R_direct - R_bayes) < 1e-9,
      "%.6f vs %.6f" % (R_direct, R_bayes))
th = 0.8
rej = m.reject_decision(post, th)
check("reject option: -1 exactly where max_k p(C_k|x) <= theta, otherwise the most probable class", np.array_equal(rej == -1, post.max(1) <= th) and np.array_equal(rej[rej >= 0], post.argmax(1)[rej >= 0]))
dw = m.decide_with_reject(post, np.array([[0, 1], [1, 0]]), 1 - th)
check("decide_with_reject (smallest expected loss above lambda = 1 - theta) reduces to the posterior threshold rule for 0-1 loss",
      np.array_equal(dw == -1, post.max(1) < th - 1e-12) and np.array_equal(dw[dw >= 0], post.argmax(1)[dw >= 0]))
# reject with loss lambda: theta = 1 - lambda; cost minimised
lamb = 0.2
cost = lambda tt: ((joint.sum(1) * (1 - post.max(1)))[post.max(1) > tt]).sum() * dx + lamb * (joint.sum(1)[post.max(1) <= tt]).sum() * dx
thetas = np.linspace(0.5, 1.0, 501)
cs = [cost(tt) for tt in thetas]
check("rejecting where max_k p(C_k|x) <= 1 - lambda minimises error + lambda * P(reject) (Exercise 1.24 style, lambda = 0.2)",
      abs(thetas[int(np.argmin(cs))] - (1 - lamb)) < 0.02, "best theta %.3f" % thetas[int(np.argmin(cs))])
# conditional mean minimises squared loss
Xn = np.linspace(0, 1, 5)
tt = np.linspace(-3, 3, 6001)
ok, detail = True, ""
for x0 in Xn:
    pt = stats.norm.pdf(tt, np.sin(2 * np.pi * x0), 0.3)
    pt /= pt.sum()
    ys = np.linspace(-1.5, 1.5, 3001)
    EL = np.array([(pt * (y_ - tt) ** 2).sum() for y_ in ys])
    ok &= abs(ys[EL.argmin()] - np.sin(2 * np.pi * x0)) < 2e-3
check("the conditional mean E[t|x] minimises the expected squared loss (1.89), grid search over y", ok)
ys = np.linspace(-1.5, 1.5, 3001)
pt = 0.7 * stats.norm.pdf(tt, -1, 0.4) + 0.3 * stats.norm.pdf(tt, 1.5, 0.4)
pt /= pt.sum()
EL1 = np.array([(pt * np.abs(y_ - tt)).sum() for y_ in ys])
cdf = np.cumsum(pt)
check("the conditional median minimises the expected absolute loss (q = 1), not the mean, for a bimodal p(t|x)", abs(ys[EL1.argmin()] - tt[np.searchsorted(cdf, 0.5)]) < 0.05
      and abs(ys[EL1.argmin()] - (pt * tt).sum()) > 0.1, "argmin %.3f, median %.3f, mean %.3f" % (ys[EL1.argmin()], tt[np.searchsorted(cdf, 0.5)], (pt * tt).sum()))

# ----------------------------------------------------------------------------- Section 1.6
print("== Information theory (Section 1.6)")
M_ = 12
check("the uniform distribution maximises the discrete entropy: H = ln M (1.94)", abs(m.entropy(np.ones(M_) / M_) - np.log(M_)) < 1e-12)
ok = True
for _ in range(2000):
    p = rng.dirichlet(rng.uniform(0.2, 3) * np.ones(M_))
    ok &= m.entropy(p) <= np.log(M_) + 1e-12
check("H[p] <= ln M for 2000 random distributions on 12 states", ok)
# constrained optimisation: maximise H subject to normalisation only
res = minimize(lambda z: -m.entropy(np.exp(z) / np.exp(z).sum()), rng.normal(size=M_), method="BFGS")
pp = np.exp(res.x) / np.exp(res.x).sum()
check("numerical maximiser of H over the simplex is uniform", np.abs(pp - 1 / M_).max() < 1e-4, "max |p - 1/M| %.1e" % np.abs(pp - 1 / M_).max())
# Gaussian maximises differential entropy for given variance: compare with other densities of variance 1
h_gauss = 0.5 * (1 + np.log(2 * np.pi))
h_lap = 1 + np.log(2 / np.sqrt(2))                        # Laplace with variance 1: b = 1/sqrt 2, h = 1 + ln 2b
h_unif = np.log(np.sqrt(12))                              # uniform with variance 1
check("Gaussian maximises differential entropy for fixed variance (1.110): h = 1/2 ln(2 pi e sigma^2) beats Laplace and uniform with variance 1",
      h_gauss > h_lap and h_gauss > h_unif, "%.4f > %.4f, %.4f" % (h_gauss, h_lap, h_unif))
hn = integrate.quad(lambda z: -stats.norm.pdf(z, 0, 1.7) * stats.norm.logpdf(z, 0, 1.7), -20, 20)[0]
check("numerical differential entropy of N(0, 1.7^2) equals 1/2 ln(2 pi e sigma^2) (1.110)", abs(hn - 0.5 * np.log(2 * np.pi * np.e * 1.7 ** 2)) < 1e-9)
ok_kl, kmin = True, np.inf
for _ in range(2000):
    p = rng.dirichlet(np.ones(M_)); q = rng.dirichlet(np.ones(M_))
    k = m.kl_discrete(p, q)
    ok_kl &= k >= -1e-14
    kmin = min(kmin, k)
check("KL(p||q) >= 0 for 2000 random pairs of distributions (1.115)", ok_kl, "smallest %.2e" % kmin)
p = rng.dirichlet(np.ones(M_))
check("KL(p||p) = 0", abs(m.kl_discrete(p, p)) < 1e-15)
check("KL is not symmetric", abs(m.kl_discrete(p, rng.dirichlet(np.ones(M_))) - m.kl_discrete(rng.dirichlet(np.ones(M_)), p)) > 1e-3)
for (m1, v1, m2, v2) in [(0, 1, 1, 2), (0.5, 0.3, -1, 1.5), (2, 0.5, 2, 0.5)]:
    num = integrate.quad(lambda z: stats.norm.pdf(z, m1, np.sqrt(v1)) * (stats.norm.logpdf(z, m1, np.sqrt(v1)) - stats.norm.logpdf(z, m2, np.sqrt(v2))), -30, 30)[0]
    check("Gaussian KL closed form 1/2 [ln(v2/v1) + (v1 + (m1-m2)^2)/v2 - 1] equals numerical integration (%g, %g | %g, %g)" % (m1, v1, m2, v2),
          abs(m.kl_gauss(m1, v1, m2, v2) - num) < 1e-9, "%.8f vs %.8f" % (m.kl_gauss(m1, v1, m2, v2), num))
# entropy as expected information and mutual information I[x,y] = KL(p(x,y) || p(x)p(y)) = H[x] + H[y] - H[x,y]
P = rng.dirichlet(np.ones(12)).reshape(3, 4)
I_kl = m.kl_discrete(P.ravel(), np.outer(P.sum(1), P.sum(0)).ravel())
I_h = m.entropy(P.sum(1)) + m.entropy(P.sum(0)) - m.entropy(P.ravel())
check("mutual information I[x, y] = KL(p(x,y) || p(x) p(y)) = H[x] + H[y] - H[x, y] (1.120)", abs(I_kl - I_h) < 1e-12 and I_kl >= 0, "%.6f" % I_kl)
H1, H2 = m.entropy(m.discretised_gaussian_hist(1.42)), m.entropy(m.discretised_gaussian_hist(5.49))
check("Figure 1.30: the two 30-bin histograms have H = 1.77 and 3.09, below ln 30 = 3.40", abs(H1 - 1.77) < 0.005 and abs(H2 - 3.09) < 0.005 and H2 < np.log(30),
      "%.3f, %.3f, ln 30 = %.3f" % (H1, H2, np.log(30)))
f = lambda z: z * np.log(z)
a_, b_ = 0.4, 2.4
lam_ = rng.uniform(0, 1, 1000)
check("x ln x is convex: f(lam a + (1-lam) b) <= lam f(a) + (1-lam) f(b) (Fig 1.31, (1.114))", np.all(f(lam_ * a_ + (1 - lam_) * b_) <= lam_ * f(a_) + (1 - lam_) * f(b_) + 1e-14))

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