"""Numerical checks of the derivations of Bishop (2006), Chapter 14 (Combining models).

Run `python3 checks_ch14.py`; every line reports PASS or FAIL, and the exit status is 1 if any
check fails. The checks use the algorithms of ch14_combining_models.py.
"""
import itertools
import sys

import numpy as np
from scipy.optimize import minimize_scalar

import ch14_combining_models 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 ""))


# --------------------------------------------------------------------------- Section 14.2
def committees():
    G, M = 2000, 8
    x = (np.arange(G) + 0.5) / G
    h = np.sin(2 * np.pi * x)
    # errors with zero mean and zero cross-correlation over x: sin(2 pi m x), m = 1..M
    eps = np.array([0.3 * np.sin(2 * np.pi * (m + 1) * x) for m in range(M)])
    e_av, e_com = c.committee_errors(h + eps, h)
    check("(14.14) E_COM = E_AV / M for uncorrelated zero-mean errors", abs(e_com - e_av / M) < 1e-12,
          "E_AV = %.5f, E_COM = %.6f, E_AV/M = %.6f" % (e_av, e_com, e_av / M))
    rng = np.random.default_rng(0)
    worst, ident = -np.inf, 0.0
    for _ in range(500):
        A = rng.standard_normal((M, M))
        eps = A @ rng.standard_normal((M, G)) * rng.uniform(0.1, 2)       # correlated, nonzero mean
        eps += rng.standard_normal((M, 1))
        e_av, e_com = c.committee_errors(h + eps, h)
        worst = max(worst, e_com - e_av)
        Y = h + eps
        spread = np.mean(np.mean((Y - Y.mean(0)) ** 2, axis=1))
        ident = max(ident, abs(e_av - (e_com + spread)))
    check("Exercise 14.3: E_COM <= E_AV for correlated errors (500 random committees)", worst <= 1e-12,
          "max E_COM - E_AV = %.2e" % worst)
    check("E_AV = E_COM + (1/M) sum_m E[(y_m - y_COM)^2] (why E_COM <= E_AV)", ident < 1e-9, "max error %.1e" % ident)
    # Exercise 14.5: with alpha_m >= 0 summing to 1 the committee stays between the members
    Y = rng.standard_normal((5, 200))
    a = rng.dirichlet(np.ones(5))
    yc = a @ Y
    inside = np.all((yc >= Y.min(0) - 1e-12) & (yc <= Y.max(0) + 1e-12))
    a_bad = np.r_[1.3, -0.3, 0, 0, 0]
    outside = np.any((a_bad @ Y < Y.min(0)) | (a_bad @ Y > Y.max(0)))
    check("Exercise 14.5: convex weights keep y_COM in [y_min, y_max]; a negative weight can leave it",
          inside and outside)


# --------------------------------------------------------------------------- Section 14.3
def boosting():
    X, t = c.BOOST_X, c.BOOST_T
    model = c.adaboost(X, t, 150)
    book = {1: (0, 0.590), 2: (0, -0.007), 3: (1, -1.720), 6: (1, 0.087), 10: (1, -0.867), 150: (0, 0.417)}
    ok = all(model["stumps"][m - 1][0] == d and abs(model["stumps"][m - 1][1] - th) < 2e-3 for m, (d, th) in book.items())
    check("Figure 14.2: stumps at m = 1, 2, 3, 6, 10, 150 are the book's dashed lines", ok,
          ", ".join("m=%d: x%d=%.3f" % (m, model["stumps"][m - 1][0] + 1, model["stumps"][m - 1][1]) for m in book))
    # alpha_m minimises the exponential error (14.22) along the new direction (Exercise 14.6)
    f = np.zeros(len(t))
    worst_a, worst_w, worst_23 = 0.0, 0.0, 0.0
    for m in range(40):
        y = c.stump_predict(model["stumps"][m], X)
        w = np.exp(-t * f)                                  # w_n^(m) = exp(-t_n f_{m-1}(x_n)) (14.22)
        E = lambda a: np.sum(w * np.exp(-0.5 * t * a * y))
        a_ls = minimize_scalar(E, bounds=(-5, 20), method="bounded", options={"xatol": 1e-10}).x
        worst_a = max(worst_a, abs(a_ls - model["alpha"][m]))
        # (14.23): E = (e^{a/2} - e^{-a/2}) sum_n w_n I(miss) + e^{-a/2} sum_n w_n
        a = model["alpha"][m]
        miss = y != t
        worst_23 = max(worst_23, abs(E(a) - ((np.exp(a / 2) - np.exp(-a / 2)) * w[miss].sum() + np.exp(-a / 2) * w.sum())) / E(a))
        # weights used by AdaBoost are proportional to exp(-t_n f_{m-1}(x_n))
        worst_w = max(worst_w, np.max(np.abs(model["weights"][m] - w / w.sum())))
        f = f + 0.5 * a * y
    check("Exercise 14.6: alpha_m = ln((1 - eps_m)/eps_m) minimises (14.22) by line search (40 rounds)", worst_a < 1e-6,
          "max |alpha - argmin| = %.1e" % worst_a)
    check("(14.23): the split of E into correctly and wrongly classified points", worst_23 < 1e-12, "rel. error %.1e" % worst_23)
    check("(14.18) reweighting = exp(-t_n f_{m-1}(x_n)) up to normalisation (14.22), (14.26)", worst_w < 1e-12,
          "max difference %.1e" % worst_w)
    # Freund-Schapire bound, on the book's data and on a noisy random data set
    rng = np.random.default_rng(3)
    Xr = rng.standard_normal((120, 2))
    tr_ = np.where(Xr[:, 0] + 0.5 * Xr[:, 1] ** 2 + 0.6 * rng.standard_normal(120) > 0.3, 1, -1)
    for name, XX, tt in (("Figure 14.2 data", X, t), ("120 noisy random points", Xr, tr_)):
        mod = c.adaboost(XX, tt, 200)
        train, expo, prod, loose = c.adaboost_curves(mod, XX, tt)
        ok = np.all(train <= expo + 1e-12) and np.allclose(expo, prod, rtol=1e-10) and np.all(prod <= loose + 1e-12)
        check("Freund-Schapire: training error <= (1/N) sum exp(-t f_m) = prod 2 sqrt(eps(1-eps)) <= exp(-2 sum gamma^2) (%s)" % name,
              ok, "m = 200: %.3f <= %.2e = %.2e <= %.2e" % (train[-1], expo[-1], prod[-1], loose[-1]))
    tr150 = c.adaboost_curves(model, X, t)[0]
    check("Figure 14.2: training error reaches 0 (first at m = %d) and is 0 at m = 150" % (np.argmax(tr150 == 0) + 1),
          tr150[-1] == 0)
    # (14.28): the minimiser of the expected exponential error is half the log odds (Exercise 14.7)
    worst = 0.0
    for p in (0.05, 0.2, 0.5, 0.7, 0.93):
        y = minimize_scalar(lambda y: p * np.exp(-y) + (1 - p) * np.exp(y), tol=1e-12).x
        worst = max(worst, abs(y - 0.5 * np.log(p / (1 - p))))
    check("(14.28): argmin_y E[exp(-t y)] = (1/2) ln(p(t=1|x)/p(t=-1|x))", worst < 1e-6, "max error %.1e" % worst)
    # Exercise 14.8: exp(-t y) cannot be a normalised likelihood: exp(-y) + exp(y) = 2 cosh y > 1 unless y = 0
    yy = np.linspace(-3, 3, 13)
    check("Exercise 14.8: sum_t exp(-t y) = 2 cosh(y) >= 2, never 1 (not a likelihood)", np.all(np.exp(-yy) + np.exp(yy) >= 2 - 1e-12))
    # error functions of Figure 14.3 pass through (0, 1) and bound the misclassification error
    z = np.linspace(-3, 3, 601)
    ok = all(abs(fn(np.array([0.0]))[0] - 1) < 1e-12 for fn in (c.exponential_error, c.cross_entropy_error, c.hinge_error))
    ok = ok and np.all(c.exponential_error(z) >= c.misclassification_error(z)) and np.all(c.cross_entropy_error(z) >= c.misclassification_error(z))
    check("Figure 14.3: E(0) = 1 for the three smooth errors, each >= misclassification", ok)
    # Exercise 14.9: boosting with squared error fits the residuals
    rng = np.random.default_rng(1)
    tt, fprev, y = rng.standard_normal(50), rng.standard_normal(50), rng.standard_normal(50)
    a_opt = minimize_scalar(lambda a: np.sum((tt - fprev - 0.5 * a * y) ** 2)).x
    a_res = 2 * np.dot(y, tt - fprev) / np.dot(y, y)     # least-squares fit of 0.5 a y to the residuals
    check("Exercise 14.9: sequential squared-error fit = least squares on residuals t - f_(m-1)", abs(a_opt - a_res) < 1e-6)


# --------------------------------------------------------------------------- Section 14.4
def trees():
    rng = np.random.default_rng(2)
    tt = rng.standard_normal(25)
    q = lambda v: np.sum((tt - v) ** 2)
    best = minimize_scalar(q).x
    check("Exercise 14.10 / (14.29): the region mean minimises the sum of squares", abs(best - tt.mean()) < 1e-6
          and q(tt.mean()) <= min(q(tt.mean() + d) for d in (-1e-3, 1e-3)))
    # Exercise 14.11 (as corrected in the errata)
    A = [(300, 100), (100, 300)]
    B = [(200, 400), (200, 0)]
    def crit(tree, f):
        return sum(f(np.array(n) / sum(n)) for n in tree)
    mis = lambda tree: sum(min(n) for n in tree) / 800
    gA, gB = crit(A, c.gini), crit(B, c.gini)
    eA, eB = crit(A, c.cross_entropy_impurity), crit(B, c.cross_entropy_impurity)
    check("Exercise 14.11: equal misclassification rates (%.2f), lower Gini (%.3f < %.3f) and cross-entropy (%.3f < %.3f) for B"
          % (mis(A), gB, gA, eB, eA), mis(A) == mis(B) and gB < gA and eB < eA)
    # impurities vanish at a pure node and peak at p_k = 1/K (errata to p. 666)
    K = 3
    pure = np.eye(K)[0]
    peak = np.full(K, 1 / K)
    ps = rng.dirichlet(np.ones(K), 2000)
    ok = True
    for f in (c.gini, c.cross_entropy_impurity):
        ok = ok and abs(f(pure)) < 1e-15 and all(f(p) <= f(peak) + 1e-12 for p in ps)
    check("(14.32), (14.33): impurities are 0 for a pure region and largest at p_k = 1/K", ok)
    # growing a regression tree: every split lowers the residual sum of squares
    x, t = c.sinusoidal_data(40, seed=4)
    tree = c.grow_tree(x, t, "regression", max_depth=10, min_leaf=1)
    ok = all(n.Q >= n.left.Q + n.right.Q - 1e-12 for n in tree.internal())
    check("CART growth: each split reduces sum of Q_tau (14.30); full tree has %d leaves" % len(tree.leaves()), ok)
    pruned = c.prune_tree(tree, 0.05)

    def leaf_of(node, xv):
        while not node.is_leaf:
            node = node.left if xv <= node.threshold else node.right
        return node

    members = {}
    for xv, tv in zip(x, t):
        members.setdefault(id(leaf_of(pruned, xv)), []).append(tv)
    ok = all(abs(l.value - np.mean(members[id(l)])) < 1e-12 and l.n == len(members[id(l)]) for l in pruned.leaves())
    check("(14.29): each leaf of a pruned tree predicts the mean of its %d training targets" % len(t), ok,
          "%d leaves" % len(pruned.leaves()))
    # pruning: dynamic programming = weakest-link sequence = brute force on a small tree
    path = c.pruning_path(tree)
    lams = [p[0] for p in path]
    bad = 0
    for lam in rng.uniform(0, 1.2 * lams[-1], 300):
        Tdp = c.prune_tree(tree, lam)
        k = max(i for i, l in enumerate(lams) if l <= lam + 1e-12)
        if abs(c.pruning_criterion(Tdp, lam) - c.pruning_criterion(path[k][1], lam)) > 1e-9:
            bad += 1
    check("(14.31): exact pruning by dynamic programming agrees with weakest-link pruning (300 values of lambda)", bad == 0,
          "%d mismatches, %d subtrees on the path" % (bad, len(path)))
    small = c.grow_tree(x[:20], t[:20], "regression", max_depth=4, min_leaf=1)

    def subtrees(node):
        yield c.Node(node.value, node.n, node.Q)
        if not node.is_leaf:
            for l, r in itertools.product(list(subtrees(node.left)), list(subtrees(node.right))):
                yield c.Node(node.value, node.n, node.Q, node.feature, node.threshold, l, r)

    allsub = list(subtrees(small))
    worst = 0.0
    for lam in np.linspace(0, 3, 31):
        brute = min(c.pruning_criterion(T, lam) for T in allsub)
        worst = max(worst, c.pruning_criterion(c.prune_tree(small, lam), lam) - brute)
    check("(14.31): pruned tree minimises C(T) over all %d subtrees of a small tree" % len(allsub), worst < 1e-9,
          "max excess %.1e" % worst)
    # a classification tree separates two Gaussian classes
    X = np.r_[rng.standard_normal((60, 2)) + [1.5, 0], rng.standard_normal((60, 2)) - [1.5, 0]]
    y = np.r_[np.zeros(60, int), np.ones(60, int)]
    for crit_name in ("gini", "entropy"):
        ct = c.grow_tree(X, y, "classification", crit_name, max_depth=12)
        acc = np.mean(c.tree_predict(ct, X) == y)
        check("classification tree (%s) grown to purity fits the training data" % crit_name, acc == 1.0, "accuracy %.3f" % acc)


# --------------------------------------------------------------------------- Section 14.5.1
def mixture_linear():
    states, Phi, t, x = c.run_two_lines_em(50)
    ll = np.array([s["loglik"] for s in states])
    check("Figure 14.8: EM never decreases the log likelihood (14.35) from the book's start", np.all(np.diff(ll) >= -1e-10),
          "ln p = %.3f -> %.3f -> %.3f at 0, 30, 50 iterations" % (ll[0], ll[30], ll[50]))
    check("Figure 14.9: converged log likelihood -3.08 (book -3.0)", abs(ll[50] + 3.0) < 0.15, "%.3f" % ll[50])
    w1, b1, ll1 = c.linear_regression_ml(Phi, t)
    check("Figure 14.9: single linear regression log likelihood -27.66 (book -27.6)", abs(ll1 + 27.6) < 0.15, "%.3f" % ll1)
    rng = np.random.default_rng(5)
    worst = 0.0
    for trial in range(20):
        K = rng.integers(2, 5)
        W0 = rng.standard_normal((K, 2))
        st = c.mixture_linear_regression_em(Phi, t, W0, np.full(K, 1 / K), rng.uniform(0.2, 5), 60)
        d = np.diff([s["loglik"] for s in st])
        worst = min(worst, d.min())
    check("EM for mixtures of K = 2..4 linear regressions never decreases ln p (20 random starts)", worst >= -1e-9,
          "smallest change %.1e" % worst)
    # the M step solves the weighted normal equations (14.40)-(14.42) and the beta equation (14.44)
    s0, s1 = states[10], states[11]
    g = s0["gamma"]
    res = [np.abs(Phi.T @ (g[:, k] * (t - Phi @ s1["W"][k]))).max() for k in range(2)]
    beta = s1["beta"]
    dQ = 0.5 * len(t) / beta - 0.5 * np.sum(g * (t[:, None] - Phi @ s1["W"].T) ** 2)   # dQ/dbeta from (14.43)
    check("(14.40)-(14.42): new w_k solve Phi^T R_k (t - Phi w_k) = 0", max(res) < 1e-10, "max residual %.1e" % max(res))
    check("(14.44): dQ/dbeta = 0 at the new beta", abs(dQ) < 1e-9, "dQ/dbeta = %.1e" % dQ)
    check("(14.38): pi_k = (1/N) sum_n gamma_nk (Lagrange multiplier, Exercise 14.14)",
          np.allclose(s1["pi"], g.mean(0)) and abs(s1["pi"].sum() - 1) < 1e-12)
    # Q increases under the M step and the new pi maximises sum gamma ln pi over the simplex
    Qf = lambda pi: np.sum(g * np.log(pi))
    worst = min(Qf(s1["pi"]) - Qf(p) for p in rng.dirichlet(np.ones(2), 500))
    check("(14.38) maximises sum_n sum_k gamma_nk ln pi_k over the simplex (500 random pi)", worst >= -1e-12)
    # Exercise 14.15: the conditional mean is sum_k pi_k w_k^T phi
    s = states[50]
    tg = np.linspace(-4, 4, 40001)
    for xv in (-0.7, 0.2):
        dens = c.predictive_density_grid(s["W"], s["pi"], s["beta"], np.array([xv]), tg)[:, 0]
        mean_num = np.sum(tg * dens) / np.sum(dens)
        mean_th = s["pi"] @ (s["W"][:, 0] + s["W"][:, 1] * xv)
        if abs(mean_num - mean_th) > 1e-6:
            break
    check("Exercise 14.15: E[t|x] = sum_k pi_k w_k^T phi(x) (numerical integration)", abs(mean_num - mean_th) < 1e-6,
          "%.6f vs %.6f" % (mean_num, mean_th))
    # Exercise 14.17, first part: a mixture of mixtures with constant coefficients is a flat mixture
    a, b = np.array([0.3, 0.7]), np.array([[0.2, 0.8], [0.5, 0.5]])
    mus = np.array([[-1.0, 0.0], [1.0, 2.0]])
    tg = np.linspace(-3, 4, 50)
    nrm = lambda tv, m: np.exp(-0.5 * (tv - m) ** 2) / np.sqrt(2 * np.pi)
    two = sum(a[i] * sum(b[i, j] * nrm(tg, mus[i, j]) for j in range(2)) for i in range(2))
    flat = sum((a[:, None] * b).ravel()[k] * nrm(tg, mus.ravel()[k]) for k in range(4))
    check("Exercise 14.17: a two-level mixture with constant weights equals a flat mixture (weights a_i b_ij)",
          np.allclose(two, flat, atol=1e-15))


# --------------------------------------------------------------------------- Sections 14.5.2, 14.5.3
def fd_grad(f, w, h=1e-6):
    g = np.zeros_like(w)
    for i in range(len(w)):
        e = np.zeros_like(w)
        e[i] = h
        g[i] = (f(w + e) - f(w - e)) / (2 * h)
    return g


def mixture_logistic():
    X, t, Phi, w1, W, pi, lls = c.fit_fig14_10()
    check("Figure 14.10: EM for the mixture of logistic regressions never decreases ln p (14.46)",
          np.all(np.diff(lls) >= -1e-9), "ln p from %.2f to %.2f; single logistic regression %.2f" % (
              lls[0], lls[-1], c.mixture_logistic_loglik(Phi, t, w1[None, :], np.ones(1))[0]))
    rng = np.random.default_rng(6)
    r = rng.uniform(0, 1, len(t))
    w = 0.3 * rng.standard_normal(3)
    Q = lambda v: c.logistic_q(Phi, t, r, v)
    y = c.common.sigmoid(Phi @ w)
    grad = Phi.T @ (r * (t - y))                                       # (14.51)
    H = (Phi * (r * y * (1 - y))[:, None]).T @ Phi                     # (14.52)
    g_fd = fd_grad(Q, w)
    H_fd = -np.array([fd_grad(lambda v: (Phi.T @ (r * (t - c.common.sigmoid(Phi @ v))))[i], w) for i in range(3)])
    check("(14.51): IRLS gradient of Q_k = finite differences", np.allclose(grad, g_fd, rtol=1e-6, atol=1e-6),
          "max error %.1e" % np.abs(grad - g_fd).max())
    check("(14.52): IRLS Hessian H_k = -grad grad Q_k = finite differences", np.allclose(H, H_fd, rtol=1e-6, atol=1e-5),
          "max error %.1e" % np.abs(H - H_fd).max())
    check("(14.52): H_k is positive semidefinite, so Q_k is concave and IRLS finds its maximum",
          np.linalg.eigvalsh(H).min() > -1e-12)
    w_hat = c.irls_logistic(Phi, t, r, np.zeros(3), n_steps=50)
    check("IRLS converges to a stationary point of Q_k", np.abs(Phi.T @ (r * (t - c.common.sigmoid(Phi @ w_hat)))).max() < 1e-6)
    # mixture of experts with a softmax gate
    Phi2, tt, x = c.two_lines_data()
    G = np.c_[np.ones_like(x), x]
    W0 = 0.1 * rng.standard_normal((3, 2))
    We, V, beta, lls = c.mixture_of_experts_em(Phi2, G, tt, W0, np.zeros((3, 2)), 1 / np.var(tt, ddof=1), n_iter=400)
    check("Mixture of experts (3 linear experts, softmax gate on (1, x)): EM never decreases ln p", np.all(np.diff(lls) >= -1e-9),
          "ln p from %.2f to %.3f (mixture with constant pi: -3.08)" % (lls[0], lls[-1]))
    gamma = rng.dirichlet(np.ones(3), len(x))
    Vr = rng.standard_normal((3, 2))
    Vr[2] = 0
    obj = lambda v: np.sum(gamma * np.log(c.softmax_rows(G @ np.vstack([v.reshape(2, 2), np.zeros((1, 2))]).T)))
    P = c.softmax_rows(G @ Vr.T)
    grad = ((gamma - P)[:, :2].T @ G).ravel()
    check("Softmax gate: gradient sum_n (gamma_nj - pi_nj) g_n = finite differences", np.allclose(grad, fd_grad(obj, Vr[:2].ravel()), atol=1e-6))


def main():
    print("Checks for Chapter 14 (Combining models)")
    for f in (committees, boosting, trees, mixture_linear, mixture_logistic):
        print("--", f.__name__)
        f()
    n, k = len(RESULTS), sum(RESULTS)
    print("%d checks: %d PASS, %d FAIL" % (n, k, n - k))
    return 0 if k == n else 1


if __name__ == "__main__":
    sys.exit(main())