"""Numerical checks for Bishop (2006), Sections 10.5 to 10.7 (ch10_local_variational_ep.py).

Run `python3 checks_ch10b.py`; every line starts with PASS or FAIL (INFO lines report
numbers without a test). The checks cover the local bounds of Section 10.5, variational
logistic regression (Section 10.6, including the corrected bound (10.164)), and
expectation propagation for the clutter problem (Section 10.7.1, including the evidence).
"""
import math
import sys

import numpy as np
from scipy import integrate, optimize

import ch10_local_variational_ep as C

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 info(text):
    print("INFO  " + text)


sig = C.common.sigmoid


# --------------------------------------------------------------------------- generator shared with the demos
def check_rng():
    r = C.DemoRNG(13)
    got = [r.u(), r.u(), r.norm(), r.norm(), r.norm(), r.u()]
    browser = [0.56632264936342835, 0.36011716164648533, 2.1978995479370087,
               0.68625593086156556, 2.5679056971854517, 0.091207450954243541]   # DK.rng(13) in Chromium
    err = max(abs(a - b) for a, b in zip(got, browser))
    check("DemoRNG reproduces DK.rng of demo-kit.js (seed 13, six draws)", err < 1e-15, "max diff %.1e" % err)


# --------------------------------------------------------------------------- Section 10.5
def check_bounds():
    x = np.linspace(-30, 30, 20001)
    worst, touch = 0.0, 0.0
    for xi in (0.05, 0.5, 1.0, 2.5, 5.0, 12.0):
        worst = max(worst, np.max(C.sigmoid_lower_bound(x, xi) - sig(x)))
        touch = max(touch, abs(C.sigmoid_lower_bound(xi, xi) - sig(xi)), abs(C.sigmoid_lower_bound(-xi, xi) - sig(-xi)))
    check("Jaakkola-Jordan bound (10.144) <= sigma(x) on [-30, 30] for six xi", worst <= 1e-15, "max excess %.1e" % worst)
    check("Jaakkola-Jordan bound is exact at x = +xi and x = -xi", touch < 1e-14, "max gap %.1e" % touch)

    worst, touch = np.inf, 0.0
    for eta in (0.05, 0.2, 0.5, 0.7, 0.95):
        worst = min(worst, np.min(C.sigmoid_upper_bound(x[np.abs(x) < 10], eta) - sig(x[np.abs(x) < 10])))
        x0 = math.log((1 - eta) / eta)
        touch = max(touch, abs(C.sigmoid_upper_bound(x0, eta) - sig(x0)))
    check("exponential bound (10.137) >= sigma(x) for five eta", worst >= -1e-15, "min excess %.1e" % worst)
    check("exponential bound touches sigma at x = ln((1 - eta)/eta)", touch < 1e-14, "max gap %.1e" % touch)

    xi = np.array([1e-6, 1e-3, 0.3, 1.0, 2.5, 7.0])
    lam_a = C.jj_lambda(xi)
    lam_b = (sig(xi) - 0.5) / (2 * xi)
    dl = np.diff(C.jj_lambda(np.linspace(0, 10, 1001)))
    check("lambda(xi) = tanh(xi/2)/(4 xi) = (sigma(xi) - 1/2)/(2 xi), lambda(0) = 1/8 (10.141), (10.150)",
          np.allclose(lam_a, lam_b, rtol=1e-9, atol=1e-12) and abs(C.jj_lambda(0.0) - 0.125) < 1e-15)
    check("lambda(xi) decreases on xi > 0, so lambda'(xi) != 0 in (10.162)", np.all(dl < 0))


def check_duality():
    errs = []
    for eta in (-2.0, -1.0, -1 / math.e, -0.3, -0.05):
        res = optimize.minimize_scalar(lambda x: -(eta * x - math.exp(-x)), bounds=(-20, 40), method="bounded",
                                       options={"xatol": 1e-12})
        errs.append(abs(-res.fun - C.conjugate_exp_neg(eta)))
    check("g(eta) = max_x {eta x - exp(-x)} equals eta - eta ln(-eta) (10.129), (10.131)", max(errs) < 1e-9,
          "max error %.1e" % max(errs))
    errs = []
    for x in (-1.0, 0.0, 1.0, 2.5):
        res = optimize.minimize_scalar(lambda e: -(e * x - C.conjugate_exp_neg(e)), bounds=(-50, -1e-9), method="bounded",
                                       options={"xatol": 1e-12})
        errs.append(abs(-res.fun - math.exp(-x)))
    check("exp(-x) = max_eta {eta x - g(eta)} recovers f (10.130)", max(errs) < 1e-8, "max error %.1e" % max(errs))
    errs = []
    for eta in (0.1, 0.3, 0.5, 0.8):
        res = optimize.minimize_scalar(lambda x: eta * x - C.log_sigmoid(x), bounds=(-40, 40), method="bounded",
                                       options={"xatol": 1e-12})
        errs.append(abs(res.fun - C.binary_entropy(eta)))
    check("min_x {eta x - ln sigma(x)} is the binary entropy (10.135)", max(errs) < 1e-9, "max error %.1e" % max(errs))
    # Exercises 10.30 and 10.31: ln sigma is concave; -ln(e^{x/2} + e^{-x/2}) is convex as a function of u = x^2
    x = np.linspace(-15, 15, 3001)
    h = 1e-4
    d2 = (C.log_sigmoid(x + h) - 2 * C.log_sigmoid(x) + C.log_sigmoid(x - h)) / h ** 2
    u = np.linspace(0.01, 100, 3001)
    f = lambda uu: -np.logaddexp(np.sqrt(uu) / 2, -np.sqrt(uu) / 2)
    hu = 1e-3
    d2u = (f(u + hu) - 2 * f(u) + f(u - hu)) / hu ** 2
    check("ln sigma(x) is concave (Exercise 10.30); -ln(e^(x/2)+e^(-x/2)) is convex in x^2 (Exercise 10.31)",
          np.all(d2 < 1e-6) and np.all(d2u > -1e-7), "max d2 ln sigma %.1e, min d2 in x^2 %.1e" % (d2.max(), d2u.min()))


# --------------------------------------------------------------------------- Section 10.6
def _small_logit_problem():
    """A 2-parameter model (no bias) on 12 points, small enough for exact 2-D quadrature."""
    r = C.DemoRNG(21)
    X = np.array([[r.norm() + (1.0 if k % 2 else -1.0), r.norm()] for k in range(12)])
    t = np.array([k % 2 for k in range(12)], dtype=float)
    return X, t


def _log_integral_2d(logf, lim=12.0, n=801):
    g = np.linspace(-lim, lim, n)
    W1, W2 = np.meshgrid(g, g)
    L = logf(np.c_[W1.ravel(), W2.ravel()])
    h = g[1] - g[0]
    return C.common.logsumexp(L) + 2 * math.log(h)


def check_vlr_bound_closed_form():
    X, t = _small_logit_problem()
    alpha = 0.5
    m0, S0inv = np.zeros(2), alpha * np.eye(2)
    xi = np.linspace(0.4, 2.6, len(t))
    mN, SN = C.vlr_e_step(X, t, xi, m0, S0inv)
    closed = C.vlr_bound(xi, mN, SN, m0, S0inv)

    def log_h_prior(W):
        a = W @ X.T                                                   # (n_grid, N)
        lh = np.sum(C.log_sigmoid(xi) + a * t - (a + xi) / 2 - C.jj_lambda(xi) * (a ** 2 - xi ** 2), axis=1)
        return lh + np.log(alpha / (2 * np.pi)) - 0.5 * alpha * (W ** 2).sum(1)

    numeric = _log_integral_2d(log_h_prior)
    check("corrected bound (10.164) equals ln of the integral of h(w, xi) p(w) (10.159)", abs(closed - numeric) < 1e-6,
          "closed form %.8f, quadrature %.8f" % (closed, numeric))
    exact = _log_integral_2d(lambda W: np.sum(C.log_sigmoid((2 * t - 1) * (W @ X.T)), axis=1)
                             + np.log(alpha / (2 * np.pi)) - 0.5 * alpha * (W ** 2).sum(1))
    res = C.variational_logistic_regression(X, t, alpha=alpha, tol=1e-12)
    check("L(xi) <= ln p(t) for the initial and the optimised xi", closed <= exact and res["bound"][-1] <= exact,
          "L(xi0) = %.4f, L(xi*) = %.4f, ln p(t) = %.4f" % (closed, res["bound"][-1], exact))


def check_vlr_em():
    X, t = C.logistic_data()
    Phi = C.design_matrix(X)
    res = C.variational_logistic_regression(Phi, t, alpha=C.LOGIT_ALPHA, tol=1e-12, max_iter=20000)
    b = res["bound"]
    check("EM for xi: the bound (10.164) never decreases (%d iterations, Figure 10.13 data)" % len(b),
          np.all(np.diff(b) >= -1e-10), "L from %.4f to %.4f" % (b[0], b[-1]))
    # M step (10.163) is the maximiser of Q(xi, xi_old) (10.161) for fixed q(w)
    xi_old = np.full(len(t), 1.0)
    mN, SN = C.vlr_e_step(Phi, t, xi_old, np.zeros(3), C.LOGIT_ALPHA * np.eye(3))
    xi_new = C.vlr_m_step(Phi, mN, SN)
    EwwT = SN + np.outer(mN, mN)
    errs = []
    for n in (0, 7, 30, 49):
        A = Phi[n] @ EwwT @ Phi[n]
        q = lambda s: -(C.log_sigmoid(s) - s / 2 - C.jj_lambda(s) * (A - s ** 2))
        opt = optimize.minimize_scalar(q, bounds=(1e-6, 100), method="bounded", options={"xatol": 1e-12})
        errs.append(abs(opt.x - xi_new[n]))
    check("M step (10.163) equals the numerical maximiser of Q(xi, xi_old) (10.161)", max(errs) < 1e-6,
          "max |difference| %.1e" % max(errs))
    # Exercise 10.34: at the EM fixed point the gradient of L(xi) itself vanishes
    xi = res["xi"]
    S0inv, m0 = C.LOGIT_ALPHA * np.eye(3), np.zeros(3)

    def L_of(xx):
        a, b_ = C.vlr_e_step(Phi, t, xx, m0, S0inv)
        return C.vlr_bound(xx, a, b_, m0, S0inv)

    grads = []
    for n in (0, 12, 33, 49):
        e = np.zeros(len(xi))
        e[n] = 1e-5
        grads.append(abs(L_of(xi + e) - L_of(xi - e)) / 2e-5)
    check("at the EM fixed point dL/dxi_n = 0: direct maximisation gives the same xi (Exercise 10.34)",
          max(grads) < 1e-6, "max |dL/dxi_n| %.1e" % max(grads))


def check_vlr_alpha():
    X, t = C.logistic_data()
    Phi = C.design_matrix(X)
    res = C.variational_logistic_regression_alpha(Phi, t, tol=1e-12, max_iter=20000)
    b = res["bound"]
    check("Section 10.6.3: the bound L~(q, xi) never decreases over the cycles", np.all(np.diff(b) >= -1e-10),
          "E[alpha] = %.4f after %d cycles" % (res["Ealpha"][-1], len(b)))
    # with a very concentrated Gam(alpha | a0, b0), E[alpha] = a0 / b0 and the fixed-alpha answer is recovered
    a_fix = 0.2
    big = C.variational_logistic_regression_alpha(Phi, t, a0=1e8 * a_fix, b0=1e8, tol=1e-13, max_iter=20000)
    fix = C.variational_logistic_regression(Phi, t, alpha=a_fix, tol=1e-13, max_iter=20000)
    d = max(np.max(np.abs(big["mN"] - fix["mN"])), np.max(np.abs(big["SN"] - fix["SN"])))
    check("(10.174) to (10.176) reduce to (10.156) to (10.158) when q(alpha) is a point mass", d < 1e-5,
          "max difference in m_N, S_N %.1e" % d)


# --------------------------------------------------------------------------- Section 10.7: KL(p||q) and EP
def check_moment_matching():
    z, p, dz = C._toy_grid()
    mean = np.sum(p * z) * dz
    var = np.sum(p * (z - mean) ** 2) * dz

    def kl_pq(par):
        mu, ls = par
        lq = -0.5 * (z - mu) ** 2 * np.exp(-2 * ls) - ls - 0.5 * C.LN2PI
        return -np.sum(p * lq) * dz                      # KL(p||q) up to the entropy of p

    res = optimize.minimize(kl_pq, x0=[0.0, 0.0], method="BFGS", options={"gtol": 1e-12})
    check("minimising KL(p||q) over Gaussians matches the mean and variance of p (10.187)",
          abs(res.x[0] - mean) < 1e-6 and abs(math.exp(2 * res.x[1]) - var) < 1e-6,
          "q: %.6f, %.6f; p: %.6f, %.6f" % (res.x[0], math.exp(2 * res.x[1]), mean, var))
    m_ep, v_ep = C.ep_toy()
    check("EP on exp(-z^2/2) sigma(20z + 4) returns the moments of p (Figure 10.14)",
          abs(m_ep - mean) < 1e-10 and abs(v_ep - var) < 1e-10)
    z0, v_lap = C.laplace_toy()
    h = 1e-4
    hess = (C.toy_log_density(z0 + h) - 2 * C.toy_log_density(z0) + C.toy_log_density(z0 - h)) / h ** 2
    grad = (C.toy_log_density(z0 + h) - C.toy_log_density(z0 - h)) / (2 * h)
    check("Laplace for the toy density: zero gradient at z0 and variance = -1 / (d2 ln p) by finite differences",
          abs(grad) < 1e-6 and abs(v_lap + 1 / hess) < 1e-6, "z0 = %.6f, var = %.6f, finite-difference var %.6f" % (z0, v_lap, -1 / hess))


def _tilted_moments(x_n, m_cav, v_cav, w, a):
    th = np.linspace(m_cav - 40 * math.sqrt(v_cav), m_cav + 40 * math.sqrt(v_cav), 400001)
    f = (1 - w) * np.exp(-0.5 * (x_n - th) ** 2) / math.sqrt(2 * math.pi) \
        + w * math.exp(-0.5 * x_n ** 2 / a) / math.sqrt(2 * math.pi * a)
    q = np.exp(-0.5 * (th - m_cav) ** 2 / v_cav) / math.sqrt(2 * math.pi * v_cav)
    g = f * q
    Z = integrate.simpson(g, x=th)
    m = integrate.simpson(g * th, x=th) / Z
    v = integrate.simpson(g * (th - m) ** 2, x=th) / Z
    return Z, m, v


def check_ep_update():
    w, a = C.CLUTTER_W, C.CLUTTER_A
    r = C.DemoRNG(5)
    worst, worst_site, neg = 0.0, 0.0, 0
    for _ in range(12):
        x_n, m_cav, v_cav = 6 * r.u() - 1, 6 * r.u() - 1, 0.02 + 2 * r.u()
        tau_n, m_n, ln_s, m_new, v_new, ln_Z = C.site_example(x_n, m_cav, v_cav, w, a)
        Z, m, v = _tilted_moments(x_n, m_cav, v_cav, w, a)
        worst = max(worst, abs(math.exp(ln_Z) - Z) / Z, abs(m_new - m), abs(v_new - v) / v)
        th = np.linspace(-6, 8, 301)
        site = np.exp(ln_s - 0.5 * tau_n * (th - m_n) ** 2)
        ratio = math.exp(ln_Z) * np.exp(-0.5 * (th - m_new) ** 2 / v_new) / math.sqrt(v_new) \
            / (np.exp(-0.5 * (th - m_cav) ** 2 / v_cav) / math.sqrt(v_cav))
        worst_site = max(worst_site, np.max(np.abs(site - ratio) / ratio))
        neg += tau_n < 0
    check("EP update: Z_n (10.216), m^new (10.217), v^new (10.218) equal the tilted moments by quadrature",
          worst < 1e-8, "12 random cavities, max relative error %.1e" % worst)
    check("site (10.220) to (10.222) equals Z_n q^new / q^\\n for every theta (also when v_n < 0)",
          worst_site < 1e-9, "max relative error %.1e, %d of 12 sites with negative variance" % (worst_site, neg))


def check_ep_evidence():
    r = C.DemoRNG(8)
    worst, printed = 0.0, 0.0
    for D in (1, 2):
        x = np.array([[2 + r.norm() for _ in range(D)] for _ in range(7)])
        res = C.ep_clutter(x, w=1e-300, max_sweeps=3)
        N, b = len(x), C.CLUTTER_B
        Cov = b * np.ones((N, N)) + np.eye(N)
        L = np.linalg.cholesky(Cov)
        exact = 0.0
        for j in range(D):
            zz = np.linalg.solve(L, x[:, j])
            exact += -0.5 * N * C.LN2PI - np.log(np.diag(L)).sum() - 0.5 * zz @ zz
        worst = max(worst, abs(res["ln_evidence"] - exact))
        # the book's printed (10.222) with (10.223) as printed (no prior factor)
        v_n = 1.0 / res["tau"]
        m = np.atleast_1d(res["m"])
        B = m @ m / res["v"] - np.sum(res["tau"] * (res["msite"] ** 2).sum(1))
        book = 0.5 * D * math.log(2 * math.pi * res["v"]) + 0.5 * B + np.sum(res["ln_s"] - 0.5 * D * np.log(2 * math.pi * v_n))
        printed = max(printed, abs(book - exact))
    check("EP evidence with the prior factor (module notes) is exact when w = 0 (conjugate model), D = 1 and 2",
          worst < 1e-9, "max error %.1e" % worst)
    info("the book's (10.222) and (10.223) taken literally miss the exact evidence by up to %.2f nats" % printed)


def check_ep_clutter():
    x = C.clutter_data(1, 3.0)
    ex, ep = C.exact_clutter(x), C.ep_clutter(x)
    check("EP with a single data point is exact (mean, variance, evidence)",
          abs(ep["m"] - ex["mean"]) < 1e-9 and abs(ep["v"] - ex["var"]) < 1e-7 and abs(ep["ln_evidence"] - ex["ln_evidence"]) < 1e-9)
    rows, errs = [], []
    for N in (20, 100, 500):
        e_m, e_z = [], []
        for seed in (2, 3, 11, 13):
            x = C.clutter_data(N, 3.0, seed=seed)
            ex, ep = C.exact_clutter(x), C.ep_clutter(x)
            e_m.append(abs(ep["m"] - ex["mean"]))
            e_z.append(abs(ep["ln_evidence"] - ex["ln_evidence"]))
        errs.append((np.mean(e_m), np.mean(e_z)))
        rows.append("N=%d: %.1e, %.1e" % (N, errs[-1][0], errs[-1][1]))
    ok = all(errs[k + 1][0] < errs[k][0] and errs[k + 1][1] < errs[k][1] for k in range(2))
    check("EP posterior mean and ln p(D) approach the exact values as N grows (mean errors over 4 data sets)", ok,
          "; ".join(rows))


def check_vb_laplace_exact():
    x = C.clutter_data(200, 3.0)
    ex = C.exact_clutter(x)
    ex2 = C.exact_clutter(x, n=40001)
    quad, _ = integrate.quad(lambda th: math.exp(C.clutter_log_joint(th, x) - ex["ln_evidence"]),
                             ex["mean"] - 3, ex["mean"] + 3, points=[ex["mean"]], limit=200, epsabs=1e-13)
    check("quadrature for the exact posterior: stable under grid refinement and equal to adaptive quad",
          abs(ex["ln_evidence"] - ex2["ln_evidence"]) < 1e-10 and abs(ex["mean"] - ex2["mean"]) < 1e-10 and abs(quad - 1) < 1e-9,
          "ln p(D) = %.8f" % ex["ln_evidence"])
    vb = C.vb_clutter(x, tol=1e-13)
    L = np.array([h[3] for h in vb["history"]])
    check("VB for the clutter problem: the bound never decreases and stays below ln p(D)",
          np.all(np.diff(L) >= -1e-10) and L[-1] <= ex["ln_evidence"],
          "gap ln p(D) - L = %.4f" % (ex["ln_evidence"] - L[-1]))
    lap = C.laplace_clutter(x)
    h = 1e-3          # ln p(theta, D) is about -465: smaller steps lose digits to rounding
    f = lambda th: C.clutter_log_joint(np.array([th]), x)[0]
    g = (f(lap["m"] + h) - f(lap["m"] - h)) / (2 * h)
    H = (f(lap["m"] + h) - 2 * f(lap["m"]) + f(lap["m"] - h)) / h ** 2
    check("Laplace: zero gradient at the mode, variance = -1/H by finite differences",
          abs(g) < 1e-5 and abs(lap["v"] + 1 / H) < 1e-5 * lap["v"],
          "mode %.6f, relative variance difference %.1e" % (lap["m"], abs(lap["v"] + 1 / H) / lap["v"]))
    ep = C.ep_clutter(x)
    info("N = 200 (Figure 10.17 data): exact mean %.6f, var %.6f, ln p(D) %.4f" % (ex["mean"], ex["var"], ex["ln_evidence"]))
    for name, m, v, lz in (("EP", ep["m"], ep["v"], ep["ln_evidence"]), ("VB", vb["m"], vb["v"], vb["bound"]),
                           ("Laplace", lap["m"], lap["v"], lap["ln_evidence"])):
        info("  %-7s mean error %+.2e, var %.6f, ln p(D) error %+.2e" % (name, m - ex["mean"], v, lz - ex["ln_evidence"]))
    check("on Figure 10.17's data EP beats Laplace and VB in both the mean and ln p(D)",
          abs(ep["m"] - ex["mean"]) < min(abs(lap["m"] - ex["mean"]), abs(vb["m"] - ex["mean"]))
          and abs(ep["ln_evidence"] - ex["ln_evidence"]) < min(abs(lap["ln_evidence"] - ex["ln_evidence"]),
                                                                 abs(vb["bound"] - ex["ln_evidence"])))


if __name__ == "__main__":
    for f in (check_rng, check_bounds, check_duality, check_vlr_bound_closed_form, check_vlr_em, check_vlr_alpha,
              check_moment_matching, check_ep_update, check_ep_evidence, check_ep_clutter, check_vb_laplace_exact):
        f()
    n_fail = RESULTS.count(False)
    print("\n%d checks: %d PASS, %d FAIL" % (len(RESULTS), RESULTS.count(True), n_fail))
    sys.exit(1 if n_fail else 0)