"""Numerical checks of the derivations of Chapter 6 (Kernel Methods): prints PASS/FAIL lines.

Run: python3 checks_ch06.py
"""
import sys
import numpy as np
import ch06_kernel_methods as m
from common import sigmoid

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 fd_grad(f, x, h=1e-6):
    g = np.zeros_like(x)
    for i in range(len(x)):
        e = np.zeros_like(x)
        e[i] = h
        g[i] = (f(x + e) - f(x - e)) / (2 * h)
    return g


rng = np.random.default_rng(0)

# ---------------------------------------------------------------- Section 6.1: dual representation
N, M = 15, 6
Phi = rng.standard_normal((N, M))
t = rng.standard_normal(N)
lam = 0.7
w = m.ridge_primal(Phi, t, lam)
K = m.basis_kernel(Phi, Phi)
a = m.ridge_dual(K, t, lam)
check("(6.3) w = Phi^T a with a from the dual (6.8) equals the primal ridge solution",
      np.allclose(Phi.T @ a, w, atol=1e-10), "max |Phi^T a - w| = %.1e" % np.abs(Phi.T @ a - w).max())
phis = rng.standard_normal((5, M))
check("(6.9) dual prediction k(x)^T (K + lam I)^(-1) t equals w^T phi(x)",
      np.allclose(m.basis_kernel(phis, Phi) @ a, phis @ w, atol=1e-10))
check("(6.4) a_n = -(1/lam) (w^T phi(x_n) - t_n)", np.allclose(a, -(Phi @ w - t) / lam, atol=1e-10))

# ---------------------------------------------------------------- Section 6.2: valid kernels
X = rng.standard_normal((40, 3))
Z = rng.standard_normal((40, 3))


def psd(Kmat, tol=1e-9):
    return np.linalg.eigvalsh((Kmat + Kmat.T) / 2).min() > -tol


k1 = lambda A, B: m.k_gauss(A, B, 1.3)
k2 = lambda A, B: m.k_poly(A, B, 2, 1.0)
check("Gram matrices of the Gaussian (6.23), polynomial (6.12) and exponential (6.56) kernels are PSD",
      psd(k1(X, X)) and psd(k2(X, X)) and psd(m.k_exp(X, X, 2.0)),
      "min eigenvalues %.1e, %.1e, %.1e" % tuple(np.linalg.eigvalsh(f(X, X)).min() for f in (k1, k2, lambda A, B: m.k_exp(A, B, 2.0))))
f = lambda A: np.sin(A[:, 0]) + 2.0
rules = {
    "(6.13) c k1": 3.0 * k1(X, X),
    "(6.14) f(x) k1 f(x')": f(X)[:, None] * k1(X, X) * f(X)[None, :],
    "(6.15) polynomial q with non-negative coefficients": 1.0 + 2.0 * k1(X, X) + 0.5 * k1(X, X) ** 3,
    "(6.16) exp(k1)": np.exp(k1(X, X)),
    "(6.17) k1 + k2": k1(X, X) + k2(X, X),
    "(6.18) k1 k2": k1(X, X) * k2(X, X),
    "(6.19) k3(phi(x), phi(x'))": m.k_gauss(np.c_[np.sin(X[:, 0]), X[:, 1] ** 2], np.c_[np.sin(X[:, 0]), X[:, 1] ** 2], 0.8),
    "(6.20) x^T A x' with A PSD": X @ (lambda B: B @ B.T)(rng.standard_normal((3, 3))) @ X.T,
    "(6.21) ka(xa, xa') + kb(xb, xb')": m.k_gauss(X[:, :2], X[:, :2], 1.0) + m.k_poly(X[:, 2:], X[:, 2:], 3, 1.0),
    "(6.22) ka kb": m.k_gauss(X[:, :2], X[:, :2], 1.0) * m.k_poly(X[:, 2:], X[:, 2:], 3, 1.0),
}
for name, G in rules.items():
    check("closure rule %s gives a PSD Gram matrix" % name, psd(G, 1e-8 * max(1, np.abs(G).max())),
          "min eigenvalue %.1e" % np.linalg.eigvalsh(G).min())
Gs = m.k_sigmoid(X, X, 1.0, -1.0)
check("the sigmoidal kernel (6.37) is NOT PSD in general (the book notes this)", not psd(Gs, 1e-9),
      "min eigenvalue %.2f" % np.linalg.eigvalsh(Gs).min())
G = k1(X, Z)
check("(6.96) Cauchy-Schwarz: k(x1, x2)^2 <= k(x1, x1) k(x2, x2) for the Gaussian kernel",
      np.all(G ** 2 <= np.outer(np.diag(k1(X, X)), np.diag(k1(Z, Z))) + 1e-12))
# (6.25): the Gaussian kernel as a normalised exponential of the linear kernel
sg = 1.3
alt = np.exp(-(X ** 2).sum(1)[:, None] / (2 * sg ** 2)) * np.exp(X @ Z.T / sg ** 2) * np.exp(-(Z ** 2).sum(1)[None, :] / (2 * sg ** 2))
check("(6.25) exp(-||x-x'||^2/2s^2) = exp(-x^Tx/2s^2) exp(x^Tx'/s^2) exp(-x'^Tx'/2s^2)", np.allclose(alt, k1(X, Z)))
# (6.12): (x^T z)^2 = phi(x)^T phi(z), phi = (x1^2, sqrt2 x1 x2, x2^2)
x2d, z2d = rng.standard_normal((6, 2)), rng.standard_normal((6, 2))
ph = lambda U: np.c_[U[:, 0] ** 2, np.sqrt(2) * U[:, 0] * U[:, 1], U[:, 1] ** 2]
check("(6.12) the homogeneous quadratic kernel (x^T z)^2 equals phi(x)^T phi(z) with phi = (x1^2, sqrt2 x1 x2, x2^2)",
      np.allclose(m.k_poly(x2d, z2d, 2, 0.0), ph(x2d) @ ph(z2d).T))

# ---------------------------------------------------------------- Section 6.3: Nadaraya-Watson
xs = np.linspace(-0.2, 1.2, 30)
Xn = np.linspace(0, 1, 10)
tn = np.sin(2 * np.pi * Xn) + 0.3 * rng.standard_normal(10)
W, mean, var = m.nadaraya_watson(xs, Xn, tn, 0.08)
check("(6.46) Nadaraya-Watson weights k(x, x_n) sum to 1 at every x (even far from the data)",
      np.allclose(W.sum(1), 1.0, atol=1e-12) and np.all(W >= 0), "max |sum - 1| = %.1e" % np.abs(W.sum(1) - 1).max())
# conditional mean and variance of the Gaussian mixture (6.48) by numerical integration over t
tg = np.linspace(-6, 6, 20001)
j = 12
dens = (W[j][:, None] * np.exp(-(tg[None, :] - tn[:, None]) ** 2 / (2 * 0.08 ** 2)) / np.sqrt(2 * np.pi * 0.08 ** 2)).sum(0)
mt = np.trapz(dens * tg, tg)
vt = np.trapz(dens * (tg - mt) ** 2, tg)
check("(6.45), (6.48) mean and variance of the Gaussian-mixture conditional agree with integration over t",
      abs(mt - mean[j]) < 1e-6 and abs(vt - var[j]) < 1e-6, "mean %.6f vs %.6f, var %.6f vs %.6f" % (mt, mean[j], vt, var[j]))

# ---------------------------------------------------------------- Section 6.4: GP regression
Xg = rng.uniform(-1, 1, (12, 1))
tg_ = np.sin(3 * Xg[:, 0]) + 0.2 * rng.standard_normal(12)
Xs = np.linspace(-1, 1, 7)[:, None]
beta = 1 / 0.2 ** 2
kern = m.GaussKernel(1.3, 6.0, D=1, tie=True)
mu, v, cov = m.gp_predict(kern, Xg, tg_, Xs, beta)
# direct conditioning of the joint Gaussian over (t, f*)
Xa = np.vstack([Xg, Xs])
Ja = kern.K(Xa)
Ctt = Ja[:12, :12] + np.eye(12) / beta
Cst = Ja[12:, :12]
mu_d = Cst @ np.linalg.solve(Ctt, tg_)
cov_d = Ja[12:, 12:] - Cst @ np.linalg.solve(Ctt, Cst.T)
check("(6.66)-(6.67) GP predictive mean and covariance equal direct conditioning of the joint Gaussian (2.81)-(2.82)",
      np.allclose(mu, mu_d, atol=1e-10) and np.allclose(cov, cov_d, atol=1e-10), "max diff %.1e" % max(np.abs(mu - mu_d).max(), np.abs(cov - cov_d).max()))
amean = np.linalg.solve(Ctt, tg_)
check("(6.68) the predictive mean is sum_n a_n k(x_n, x*) with a = C^(-1) t", np.allclose(mu, kern.K(Xg, Xs).T @ amean))
# linear-basis model: GP view equals the weight-space result of Section 3.3.2
Phi = np.c_[np.ones(12), Xg[:, 0], Xg[:, 0] ** 2]
alpha_w = 2.0
SN = np.linalg.inv(alpha_w * np.eye(3) + beta * Phi.T @ Phi)
mN = beta * SN @ Phi.T @ tg_
Phis = np.c_[np.ones(7), Xs[:, 0], Xs[:, 0] ** 2]
Kb = Phi @ Phi.T / alpha_w
Cb = Kb + np.eye(12) / beta
kb = Phis @ Phi.T / alpha_w
mu_gp = kb @ np.linalg.solve(Cb, tg_)
var_gp = np.diag(Phis @ Phis.T / alpha_w - kb @ np.linalg.solve(Cb, kb.T))
check("a GP with the kernel phi(x)^T phi(x')/alpha reproduces the Bayesian linear regression predictive (3.58)-(3.59)",
      np.allclose(mu_gp, Phis @ mN, atol=1e-9) and np.allclose(var_gp, np.einsum("ij,jk,ik->i", Phis, SN, Phis), atol=1e-9))

# log marginal likelihood (6.69) and its gradient (6.70)
Xr = rng.standard_normal((25, 3))
tr_ = np.sin(Xr[:, 0]) + 0.1 * rng.standard_normal(25)
kard = m.GaussKernel(1.2, [2.0, 0.5, 0.1], D=3)
p0 = np.concatenate([kard.lp, [np.log(30.0)]])
val, g = m.gp_log_marginal(kard, p0[:4], p0[4], Xr, tr_)
gfd = fd_grad(lambda p: m.gp_log_marginal(kard, p[:4], p[4], Xr, tr_, grad=False), p0)
from scipy.stats import multivariate_normal
C0 = kard.K(Xr) + np.eye(25) / 30.0
check("(6.69) ln p(t|theta) equals the multivariate normal log density N(t | 0, C_N)",
      abs(val - multivariate_normal(np.zeros(25), C0).logpdf(tr_)) < 1e-9)
check("(6.70) gradient of ln p(t|theta) with respect to (ln theta0, ln eta_1..3, ln beta) matches finite differences",
      np.allclose(g, gfd, rtol=1e-5, atol=1e-6), "max |analytic - FD| = %.1e" % np.abs(g - gfd).max())
ke = m.ExpKernel(1.1, 2.5)
p1 = np.concatenate([ke.lp, [np.log(20.0)]])
_, g1 = m.gp_log_marginal(ke, p1[:2], p1[2], Xr[:, :1], tr_)
g1fd = fd_grad(lambda p: m.gp_log_marginal(ke, p[:2], p[2], Xr[:, :1], tr_, grad=False), p1)
check("(6.70) the same for the exponential kernel (6.56)", np.allclose(g1, g1fd, rtol=1e-5, atol=1e-6), "max diff %.1e" % np.abs(g1 - g1fd).max())
kt = m.GaussKernel(1.0, 3.0, D=3, tie=True, th2=0.5, th3=0.2)
pt = np.concatenate([kt.lp, [np.log(10.0)]])
_, gt = m.gp_log_marginal(kt, pt[:2], pt[2], Xr, tr_)
gtfd = fd_grad(lambda p: m.gp_log_marginal(kt, p[:2], p[2], Xr, tr_, grad=False), pt)
check("(6.70) the same for (6.63)-(6.72) with the constant and linear terms theta2, theta3", np.allclose(gt, gtfd, rtol=1e-5, atol=1e-6))
# optimisation: marginal likelihood increases and the gradient at the optimum is ~ 0
X1, t1 = m.ard_data(np.random.default_rng(10))
kf, bf, tr = m.gp_fit(m.GaussKernel(1.0, 1.0, D=3), X1, t1, maxiter=200)
lls = [m.gp_log_marginal(kf, p[:4], p[4], X1, t1, grad=False) for p in tr]
_, gopt = m.gp_log_marginal(kf, tr[-1][:4], tr[-1][4], X1, t1)
check("maximum marginal likelihood by L-BFGS: ln p(t|theta) never decreases along the iterates and ends at a stationary point",
      all(b >= a_ - 1e-9 for a_, b in zip(lls, lls[1:])) and np.abs(gopt[:2]).max() < 1e-2 and np.abs(gopt[4]) < 1e-2,
      "ln p: %.2f -> %.2f in %d iterations, |grad| = %.1e" % (lls[0], lls[-1], len(tr) - 1, np.abs(gopt).max()))
eta = np.exp(tr[-1][1:4])
check("ARD (Section 6.4.4): eta_1 for the informative input is far larger than eta_2, eta_3", eta[0] > 50 * eta[1] and eta[0] > 50 * eta[2],
      "eta = %s" % np.array2string(eta, precision=4))

# ---------------------------------------------------------------- Section 6.4.6: Laplace approximation
Xc, tc = m.two_class_data(np.random.default_rng(3), 30)
kc = m.GaussKernel(4.0, 1.0, D=2, tie=True)
nu = 1e-4
Cc = kc.K(Xc) + nu * np.eye(len(tc))
a_star, hist, gn = m.laplace_mode(Cc, tc)
check("(6.81)/(6.84) at the Newton limit the gradient t - sigma - C^(-1) a vanishes and a* = C (t - sigma(a*))",
      np.linalg.norm(m.grad_psi(a_star, Cc, tc)) < 1e-8 and np.allclose(a_star, Cc @ (tc - sigmoid(a_star)), atol=1e-7),
      "|grad Psi| = %.1e after %d iterations" % (np.linalg.norm(m.grad_psi(a_star, Cc, tc)), len(hist) - 1))
gfd = fd_grad(lambda a_: m.psi(a_, Cc, tc), a_star + 0.1 * rng.standard_normal(len(tc)), 1e-6)
aa = a_star + 0.1 * rng.standard_normal(len(tc))
check("(6.81) the analytic gradient of Psi equals finite differences",
      np.allclose(m.grad_psi(aa, Cc, tc), fd_grad(lambda a_: m.psi(a_, Cc, tc), aa), atol=1e-5))
check("(6.83) the Newton-Raphson iteration increases Psi at every step (log-concave posterior)",
      np.all(np.diff(hist) >= -1e-10), "Psi: %.3f -> %.3f in %d steps" % (hist[0], hist[-1], len(hist) - 1))
H = np.diag((lambda s: s * (1 - s))(sigmoid(aa))) + np.linalg.inv(Cc)
check("(6.82) the Hessian -grad grad Psi = W + C^(-1) is positive definite, and the diagonal of W lies in (0, 1/4)",
      np.linalg.eigvalsh((H + H.T) / 2).min() > 0 and np.all((sigmoid(aa) * (1 - sigmoid(aa))) <= 0.25))
# the global maximum: no random point does better
best = m.psi(a_star, Cc, tc)
check("(6.80) a* is the global maximum of the log-concave Psi (random perturbations never beat it)",
      all(m.psi(a_star + 0.3 * rng.standard_normal(len(tc)), Cc, tc) < best for _ in range(200)))
# predictive (6.87)-(6.88): compare with Gaussian marginalisation by the linear-Gaussian result (2.115)
Xq = np.array([[0.2, 0.3], [1.0, -1.0]])
mean, varq, pq = m.laplace_predict(kc, Xc, tc, Xq, nu)
s = sigmoid(a_star)
Wd = np.diag(s * (1 - s))
Hm = Wd + np.linalg.inv(Cc)
ksq = kc.K(Xc, Xq)
mean2 = ksq.T @ np.linalg.solve(Cc, a_star)                    # k^T C^(-1) a*  (6.78) with a = a*
var2 = np.diag(kc.K(Xq)) + nu - np.einsum("ij,ij->j", ksq, np.linalg.solve(Cc, ksq)) + np.einsum("ij,ij->j", np.linalg.solve(Cc, ksq), np.linalg.solve(Hm, np.linalg.solve(Cc, ksq)))
check("(6.87)-(6.88) the predictive mean k^T(t - sigma) and variance c - k^T(W^(-1)+C)^(-1)k equal the convolution (2.115) of (6.78) with q(a) (6.86)",
      np.allclose(mean, mean2, atol=1e-8) and np.allclose(varq, var2, atol=1e-8), "max diff %.1e" % max(np.abs(mean - mean2).max(), np.abs(varq - var2).max()))
# evidence (6.90) by brute force: the Laplace approximation is the Gaussian integral of exp(Psi) around the mode
ev, gev = m.laplace_evidence(kc, kc.lp, Xc, tc, nu)
sign, ld = np.linalg.slogdet(Hm)
direct = m.psi(a_star, Cc, tc) - 0.5 * np.linalg.slogdet(Cc)[1] - 0.5 * len(tc) * np.log(2 * np.pi) - 0.5 * ld + 0.5 * len(tc) * np.log(2 * np.pi)
check("(6.90) Laplace log evidence = Psi(a*) - 1/2 ln|W + C^(-1)| + N/2 ln 2 pi (evaluated directly)", abs(ev - direct) < 1e-8, "ln p(t|theta) = %.4f" % ev)
gfd = fd_grad(lambda p: m.laplace_evidence(kc, p, Xc, tc, nu, grad=False), kc.lp.copy(), 1e-5)
check("(6.91)-(6.94) gradient of the Laplace evidence (explicit and a*-dependent parts) matches finite differences",
      np.allclose(gev, gfd, rtol=1e-4, atol=1e-5), "analytic %s, FD %s" % (np.array2string(gev, precision=5), np.array2string(gfd, precision=5)))
kc2 = m.GaussKernel(2.0, [1.5, 0.4], D=2)
lp2 = kc2.lp.copy()
ev2, gev2 = m.laplace_evidence(kc2, lp2, Xc, tc, nu)
gfd2 = fd_grad(lambda p: m.laplace_evidence(kc2, p, Xc, tc, nu, grad=False), lp2, 1e-5)
check("... also for an ARD kernel with one length scale per input", np.allclose(gev2, gfd2, rtol=1e-4, atol=1e-5), "max diff %.1e" % np.abs(gev2 - gfd2).max())
# the a*-dependent term (6.92) matters: (6.91) alone is not the gradient
Ci_a = np.linalg.solve(Cc, a_star)
Rm = np.sqrt(np.diag(Wd))[:, None] * np.linalg.inv(np.eye(len(tc)) + np.sqrt(np.diag(Wd)) @ Cc @ np.sqrt(np.diag(Wd))) * np.sqrt(np.diag(Wd))[None, :]
expl = np.array([0.5 * Ci_a @ dC @ Ci_a - 0.5 * np.trace(Rm @ dC) for dC in kc.dK(Xc)])
check("the a*-dependent term (6.92) is needed: the explicit part (6.91) alone differs from the gradient", np.abs(expl - gfd).max() > 1e-3,
      "explicit %s vs full %s" % (np.array2string(expl, precision=4), np.array2string(gfd, precision=4)))
# evidence maximisation
kfit, evfit = m.gpc_fit(kc, Xc, tc, nu)
check("GP classifier: maximising the Laplace evidence improves it", evfit >= ev - 1e-9, "ln p(t|theta): %.3f -> %.3f" % (ev, evfit))
# prediction: probit approximation (4.153) close to the numerical sigmoid-Gaussian integral
u = np.random.default_rng(1).standard_normal(400000)
mn, vr = 0.8, 2.0
exact = sigmoid(mn + np.sqrt(vr) * u).mean()
approx = float(sigmoid(mn / np.sqrt(1 + np.pi * vr / 8)))
check("(4.153) the probit approximation of int sigma(a) N(a|m, s^2) da is accurate", abs(exact - approx) < 0.01, "%.4f vs %.4f" % (exact, approx))

print("\n%d checks: %d PASS, %d FAIL" % (len(RESULTS), sum(RESULTS), len(RESULTS) - sum(RESULTS)))
sys.exit(0 if all(RESULTS) else 1)