"""Numerical checks of the derivations in Bishop (2006), Chapter 3 (Linear Models for Regression).

Each check prints PASS or FAIL with the numbers behind it; the last line counts them. Run from this
folder: python3 checks_ch03.py (a few seconds). The functions checked are those of
ch03_linear_regression.py, which also produce the figures.
"""
import math

import numpy as np
from scipy import integrate

import common
import ch03_linear_regression as m
from ch11_sampling import JSRandom

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)
x, t = m.sine_data(JSRandom(5), 25, 0.3)
Phi = m.gauss_design(x, 9, 0.1)
N, M = Phi.shape

# --- Section 3.1: normal equations, projection, ridge, sequential learning
w, beta_ml = m.ml_weights(Phi, t)
res = Phi.T @ (t - Phi @ w)
check("normal equations (3.14)-(3.15): Phi^T (t - Phi w_ML) = 0", np.abs(res).max() < 1e-9, "max |gradient| = %.1e" % np.abs(res).max())
P = Phi @ np.linalg.pinv(Phi)
check("projection property (Fig 3.2): y = Phi w_ML = P t with P symmetric and idempotent", np.allclose(P @ P, P, atol=1e-9) and np.allclose(P, P.T)
      and np.allclose(Phi @ w, P @ t), "|P^2 - P| = %.1e" % np.abs(P @ P - P).max())
check("(3.21): 1 / beta_ML = mean squared residual", abs(1 / beta_ml - np.mean((t - Phi @ w) ** 2)) < 1e-12)
lam = 0.7
wr = m.ridge_weights(Phi, t, lam)
grad = -Phi.T @ (t - Phi @ wr) + lam * wr
check("ridge (3.28) is the stationary point of (3.27)", np.abs(grad).max() < 1e-9, "max |gradient| = %.1e" % np.abs(grad).max())
# sequential learning with a small step converges to the least-squares solution
w_seq = np.zeros(M)
Phi4 = Phi[:, :4]
w4, _ = m.ml_weights(Phi4, t)
w_seq = np.zeros(4)
for sweep in range(30000):
    n = sweep % N
    w_seq = m.lms_step(w_seq, Phi4[n], t[n], 0.02)
check("sequential LMS (3.23) cycles toward the least-squares solution", np.abs(w_seq - w4).max() < 0.25 * max(1, np.abs(w4).max()),
      "max difference %.3f (the constant step leaves an O(eta) ripple)" % np.abs(w_seq - w4).max())

# --- Section 3.3: Bayesian posterior against joint-Gaussian conditioning
alpha, beta = 2.0, 25.0
mN, SN = m.posterior(Phi, t, alpha, beta)
# joint of (w, t): w ~ N(0, I/alpha), t = Phi w + noise
Sww = np.eye(M) / alpha
Stt = Phi @ Sww @ Phi.T + np.eye(N) / beta
Swt = Sww @ Phi.T
cond_mean = Swt @ np.linalg.solve(Stt, t)
cond_cov = Sww - Swt @ np.linalg.solve(Stt, Swt.T)
check("posterior (3.53)-(3.54) equals conditioning the joint Gaussian of (w, t)", np.allclose(mN, cond_mean, atol=1e-9) and np.allclose(SN, cond_cov, atol=1e-9),
      "max |m_N diff| = %.1e, max |S_N diff| = %.1e" % (np.abs(mN - cond_mean).max(), np.abs(SN - cond_cov).max()))
# sequential updating: the posterior after n points is the prior of the next
m1, S1 = m.posterior(Phi[:10], t[:10], alpha, beta)
S2inv = np.linalg.inv(S1) + beta * Phi[10:].T @ Phi[10:]
m2 = np.linalg.solve(S2inv, np.linalg.inv(S1) @ m1 + beta * Phi[10:].T @ t[10:])
check("sequential Bayes (Fig 3.7): posterior of the first 10 as prior for the other 15 gives the full posterior", np.allclose(m2, mN) and np.allclose(np.linalg.inv(S2inv), SN))
check("m_N tends to the ML solution as alpha -> 0 (3.53)", np.allclose(m.posterior(Phi, t, 1e-9, beta)[0], w, atol=1e-5))

# --- predictive variance (3.59) against Monte Carlo and the joint Gaussian
xn = np.array([0.13, 0.5, 0.9])
Pn = m.gauss_design(xn, 9, 0.1)
mu, var = m.predictive(Pn, mN, SN, beta)
W = np.random.default_rng(1).multivariate_normal(mN, SN, size=400000)
tt = W @ Pn.T + np.random.default_rng(2).normal(size=(400000, 3)) / math.sqrt(beta)
check("predictive (3.58)-(3.59): mean and variance 1/beta + phi^T S_N phi agree with Monte Carlo", np.allclose(tt.mean(0), mu, atol=0.01) and np.allclose(tt.var(0), var, rtol=0.02),
      "var formula %s, Monte Carlo %s" % (np.round(var, 4), np.round(tt.var(0), 4)))
Sxt = Pn @ Sww @ Phi.T
c_mean = Sxt @ np.linalg.solve(Stt, t)
c_var = np.diag(Pn @ Sww @ Pn.T + np.eye(3) / beta - Sxt @ np.linalg.solve(Stt, Sxt.T))
check("... and with the conditional of the joint Gaussian of (t_new, t)", np.allclose(c_mean, mu, atol=1e-9) and np.allclose(c_var, var, atol=1e-9))
Phi_far = m.gauss_design(np.array([0.5, 5.0]), 9, 0.1)
_, vfar = m.predictive(Phi_far, mN, SN, beta)
check("far from the basis centres only the bias basis function is left: variance 1/beta + (S_N)_00 (Section 3.3.2)", abs(vfar[1] - 1 / beta - SN[0, 0]) < 1e-6 and vfar[1] > vfar[0],
      "far %.5f = 1/beta + S_00 = %.5f" % (vfar[1], 1 / beta + SN[0, 0]))

# --- equivalent kernel
K = m.equivalent_kernel(Phi, Phi, SN, beta)
check("(3.60)-(3.61): the posterior mean at the training inputs is K t", np.allclose(K @ t, Phi @ mN, atol=1e-9))
xk = np.linspace(-1, 1, 200)
mu_k = np.linspace(-1, 1, 11)
Pk = m.design_matrix(xk, "gauss", mu=mu_k, s=0.2)
_, SNk = m.posterior(Pk, np.zeros(200), 1e-8, 11.1)
Kk = m.equivalent_kernel(Pk, Pk, SNk, 11.1)
inner = np.linspace(-0.5, 0.5, 5)
sums = [Kk[int(np.argmin(abs(xk - xi))), :].sum() for xi in inner]
check("(3.64): the equivalent-kernel weights sum to 1 when a constant is in the span (alpha -> 0, interior x)", np.allclose(sums, 1, atol=2e-3),
      "row sums %s" % np.round(sums, 4))
check("(3.62): k(x, x') = beta phi(x)^T S_N phi(x') is symmetric", np.allclose(Kk, Kk.T))
# S_N relation: cov(y(x), y(x')) = beta^-1 k(x, x')   (3.63)
check("(3.63): cov[y(x), y(x')] = phi^T S_N phi' = k(x, x') / beta", np.allclose(Pk @ SNk @ Pk.T, Kk / 11.1))
check("(3.62)-(3.54) relation: S_N^-1 = alpha I + beta Phi^T Phi, so beta Phi S_N Phi^T = Phi (alpha/beta I + Phi^T Phi)^-1 Phi^T",
      np.allclose(m.equivalent_kernel(Phi, Phi, SN, beta), Phi @ np.linalg.solve((alpha / beta) * np.eye(M) + Phi.T @ Phi, Phi.T)))

# --- Section 3.2: bias-variance identity
xg = np.linspace(0, 1, 200)
h = np.sin(2 * np.pi * xg)
Pg = m.gauss_design(xg, 24, 0.1)
DS = [m.sine_data(JSRandom(1000 + l), 25, 0.3) for l in range(100)]
fits = np.array([Pg @ m.ridge_weights(m.gauss_design(xd, 24, 0.1), td, math.exp(-0.31)) for xd, td in DS])
b2, v = m.bias_variance(fits, h)
loss = np.mean((fits - h) ** 2)   # E_D[(y(x; D) - h(x))^2] averaged over x, no noise
check("bias-variance (3.37)-(3.46): E_D[(y - h)^2] = bias^2 + variance exactly", abs(loss - (b2 + v)) < 1e-12, "loss %.6f = %.6f + %.6f" % (loss, b2, v))
# with noisy test targets: expected loss = bias^2 + variance + sigma^2 (Monte Carlo)
sd = 0.3
tn = h[None, :, None] + sd * np.random.default_rng(3).normal(size=(100, 200, 400))   # 400 noise draws per (data set, x)
exp_loss = np.mean((fits[:, :, None] - tn) ** 2)
check("expected loss = (bias)^2 + variance + noise (3.37), noise term sigma^2 = 0.09", abs(exp_loss - (b2 + v + sd ** 2)) < 6e-4,
      "Monte Carlo %.5f, formula %.5f" % (exp_loss, b2 + v + sd ** 2))

# --- Section 3.4: evidence (3.86) against numerical integration for a small model
xs = np.array([0.1, 0.5, 0.9, 0.3])
ts = np.array([0.4, -0.2, 0.6, 0.1])
P2 = np.column_stack([np.ones(4), xs])
a2, b2_ = 1.5, 8.0
log_ev = m.log_evidence(P2, ts, a2, b2_)


def integrand(w1, w0):
    w_ = np.array([w0, w1])
    ll = 0.5 * 4 * math.log(b2_ / (2 * math.pi)) - 0.5 * b2_ * np.sum((ts - P2 @ w_) ** 2)
    lp = math.log(a2 / (2 * math.pi)) - 0.5 * a2 * w_ @ w_
    return math.exp(ll + lp)


val = integrate.dblquad(integrand, -12, 12, -12, 12, epsabs=1e-12, epsrel=1e-10)[0]
check("evidence (3.86) equals the integral of likelihood x prior over w (2-parameter model)", abs(math.log(val) - log_ev) < 1e-6,
      "ln integral %.8f, (3.86) %.8f" % (math.log(val), log_ev))
Sy = np.eye(4) / b2_ + P2 @ P2.T / a2
lml = -0.5 * (4 * math.log(2 * math.pi) + np.linalg.slogdet(Sy)[1] + ts @ np.linalg.solve(Sy, ts))
check("... and equals the marginal t ~ N(0, beta^-1 I + alpha^-1 Phi Phi^T)", abs(lml - log_ev) < 1e-9, "%.8f vs %.8f" % (lml, log_ev))
x14, t14 = m.sine_data(JSRandom(m.SEED + 40), 10, 0.3, grid=True)
ev14 = [m.log_evidence(m.design_matrix(x14, "poly", M=k), t14, 5e-3, 1 / 0.09) for k in range(9)]
check("the model evidence over polynomial order peaks at M = 3 for the Figure 3.14 data", int(np.argmax(ev14)) == 3, "ln evidence %s" % np.round(ev14, 2))

# --- Section 3.5: alpha/beta fixed point against direct maximisation
from scipy import optimize
xe, te, Pe = m._evidence_data()
al, be, ga, trace = m.evidence_iterate(Pe, te, 1.0, 1.0, iters=2000)
opt = optimize.minimize(lambda p: -m.log_evidence(Pe, te, math.exp(p[0]), math.exp(p[1])), [0.0, 0.0], method="Nelder-Mead",
                        options=dict(xatol=1e-10, fatol=1e-13, maxiter=5000))
ao, bo = np.exp(opt.x)
check("evidence approximation: the (3.92), (3.95) fixed point maximises ln p(t|alpha, beta) directly", abs(al - ao) / ao < 1e-4 and abs(be - bo) / bo < 1e-4,
      "fixed point alpha %.5f beta %.4f; direct alpha %.5f beta %.4f" % (al, be, ao, bo))
evs = [r[3] for r in trace]
check("the iteration increases the evidence monotonically (after the first step)", all(b_ >= a_ - 1e-9 for a_, b_ in zip(evs[1:], evs[2:])),
      "ln evidence %.4f -> %.4f in %d iterations" % (evs[0], evs[-1], len(evs)))
# gamma = sum lambda_i / (alpha + lambda_i), and gamma = M - alpha tr(S_N)
lams = np.linalg.eigvalsh(be * Pe.T @ Pe)
SNe = np.linalg.inv(al * np.eye(Pe.shape[1]) + be * Pe.T @ Pe)
g1 = np.sum(lams / (al + lams))
check("(3.91): gamma = sum lambda_i / (alpha + lambda_i) = M - alpha tr(S_N)", abs(g1 - m.gamma_eff(Pe, al, be)) < 1e-12 and abs(g1 - (Pe.shape[1] - al * np.trace(SNe))) < 1e-9,
      "gamma = %.4f (M = %d)" % (g1, Pe.shape[1]))
# derivative of the log evidence in alpha: 0 = M/(2 alpha) - m^T m / 2 - tr(A^-1)/2  (3.89 before multiplying by 2 alpha)
mNe, _ = m.posterior(Pe, te, al, be)
dlog = lambda a_: (m.log_evidence(Pe, te, a_ * (1 + 1e-6), be) - m.log_evidence(Pe, te, a_ * (1 - 1e-6), be)) / (2e-6 * a_)
check("d ln p / d alpha = 0 at the fixed point, and gamma = alpha m_N^T m_N there (3.90)", abs(dlog(al)) < 1e-4 and abs(ga - al * mNe @ mNe) < 1e-8,
      "d ln p/d alpha = %.1e; gamma %.5f vs alpha m^T m %.5f" % (dlog(al), ga, al * mNe @ mNe))
check("the effective number of parameters runs from 0 (alpha large) to M (alpha small)", m.gamma_eff(Pe, 1e9, be) < 1e-3 and abs(m.gamma_eff(Pe, 1e-9, be) - Pe.shape[1]) < 1e-3)
check("(3.95): 1/beta = sum (t - m^T phi)^2 / (N - gamma) at the fixed point", abs(1 / be - np.sum((te - Pe @ mNe) ** 2) / (len(te) - ga)) < 1e-9)

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