"""
gplink.py -- GP CLASSIFICATION and LOG-GAUSSIAN COX PROCESSES (from scratch).

Backs the notebooks in  "GP Classification & Log-Gaussian Cox Processes".

A Gaussian-process regression puts a GP prior on a real-valued function. Push that latent
function through a LINK and it models other things:

  * GP CLASSIFICATION -- binary/count outcomes via a logit link:
        f ~ GP(0,k),   y_i | f ~ Bernoulli(sigmoid(f(x_i)))   (or Binomial)
    the GP is a flexible, nonlinear log-odds surface with calibrated uncertainty.

  * LOG-GAUSSIAN COX PROCESS -- a point pattern whose intensity is a GP:
        log lambda(t) = f(t),   f ~ GP(0,k),   counts in a bin ~ Poisson(lambda * width)
    the smooth, positive intensity of events over time/space, with uncertainty.

Now the likelihood is non-Gaussian, so the posterior over f is not closed form. The classic
from-scratch solution is the LAPLACE APPROXIMATION (Rasmussen & Williams, Algorithm 3.1): Newton-
iterate to the posterior mode of f, then approximate the posterior by the Gaussian with that mode
and the curvature there. The mode step and the predictive variance are computed stably through the
Cholesky factor of B = I + W^{1/2} K W^{1/2}, where W is the (diagonal) negative Hessian of the
log-likelihood. The Laplace log marginal likelihood then selects the kernel hyperparameters, and
elliptical slice sampling gives an exact-posterior cross-check.

The kernels are D-dimensional (isotropic RBF over standardised inputs), covering the 2-D Pima
classification surface and the 1-D coal-mining intensity alike.
"""

import numpy as np
from scipy.linalg import cholesky, cho_solve, solve_triangular
from scipy.optimize import minimize
from scipy.special import expit


# --------------------------------------------------------------------------- #
#  kernels (D-dimensional)                                                      #
# --------------------------------------------------------------------------- #

def sqdist(Xa, Xb):
    Xa = np.atleast_2d(Xa); Xb = np.atleast_2d(Xb)
    return np.sum(Xa ** 2, 1)[:, None] + np.sum(Xb ** 2, 1)[None, :] - 2 * Xa @ Xb.T

def k_rbf(Xa, Xb, ell, eta):
    return eta * np.exp(-0.5 * sqdist(Xa, Xb) / ell ** 2)


# --------------------------------------------------------------------------- #
#  likelihoods: return (loglik, gradient d/df, W = -d2/df2 diagonal)            #
# --------------------------------------------------------------------------- #

def lik_bernoulli(f, y, extra=None):
    p = expit(f); ll = np.sum(y * np.log(p + 1e-12) + (1 - y) * np.log(1 - p + 1e-12))
    return ll, y - p, p * (1 - p)

def lik_binomial(f, y, n):
    p = expit(f); ll = np.sum(y * f - n * np.log1p(np.exp(f)))
    return ll, y - n * p, n * p * (1 - p)

def lik_poisson(f, y, e):
    lam = e * np.exp(f); ll = np.sum(y * (np.log(e + 1e-12) + f) - lam)
    return ll, y - lam, lam


# --------------------------------------------------------------------------- #
#  Laplace approximation (Rasmussen-Williams Algorithm 3.1)                      #
# --------------------------------------------------------------------------- #

def laplace_mode(K, y, likfun, extra, maxit=100, tol=1e-8):
    """find the posterior mode of the latent f and the Laplace log marginal likelihood."""
    n = len(y); f = np.zeros(n); obj_old = -np.inf; a = np.zeros(n); L = np.eye(n)
    for _ in range(maxit):
        _, g, W = likfun(f, y, extra); sW = np.sqrt(np.maximum(W, 1e-12))
        B = np.eye(n) + sW[:, None] * K * sW[None, :]
        L = cholesky(B, lower=True)
        b = W * f + g
        a = b - sW * cho_solve((L, True), sW * (K @ b))
        f = K @ a
        ll, _, _ = likfun(f, y, extra)
        obj = -0.5 * a @ f + ll - np.log(np.diag(L)).sum()
        if abs(obj - obj_old) < tol:
            break
        obj_old = obj
    return f, a, L, sW, obj

def fit_predict(X, y, Xs, likfun, extra, kernel, theta0, bounds=None):
    """fit kernel hyperparameters by maximising the Laplace log marginal likelihood, then return
    the predictive latent mean and variance at Xs. kernel(theta, Xa, Xb) -> covariance."""
    X = np.asarray(X, float); Xs = np.asarray(Xs, float)
    if X.ndim == 1: X = X[:, None]
    if Xs.ndim == 1: Xs = Xs[:, None]
    def neg_obj(theta):
        K = kernel(theta, X, X) + 1e-6 * np.eye(len(y))
        return -laplace_mode(K, y, likfun, extra)[4]
    res = minimize(neg_obj, theta0, method="L-BFGS-B", bounds=bounds)
    th = res.x; K = kernel(th, X, X) + 1e-6 * np.eye(len(y))
    f, a, L, sW, obj = laplace_mode(K, y, likfun, extra)
    Ks = kernel(th, Xs, X); fs_mean = Ks @ a
    v = solve_triangular(L, sW[:, None] * Ks.T, lower=True)
    fs_var = np.maximum(np.diag(kernel(th, Xs, Xs)) - np.sum(v ** 2, 0), 0.0)
    return dict(theta=th, f=f, mean=fs_mean, var=fs_var, logml=obj)


# --------------------------------------------------------------------------- #
#  elliptical slice sampling -- exact-posterior cross-check                      #
# --------------------------------------------------------------------------- #

def ess_sample(K, y, likfun, extra, draws=400, burn=200, rng=None):
    """draw latent-function samples f | y by elliptical slice sampling (Murray et al. 2010)."""
    n = len(y); Lk = cholesky(K + 1e-6 * np.eye(n), lower=True)
    f = Lk @ rng.standard_normal(n); out = []
    for it in range(draws + burn):
        nu = Lk @ rng.standard_normal(n)
        ll0 = likfun(f, y, extra)[0] + np.log(rng.random())
        th = rng.uniform(0, 2 * np.pi); lo, hi = th - 2 * np.pi, th
        while True:
            fp = f * np.cos(th) + nu * np.sin(th)
            if likfun(fp, y, extra)[0] > ll0:
                f = fp; break
            if th < 0: lo = th
            else: hi = th
            th = rng.uniform(lo, hi)
        if it >= burn: out.append(f.copy())
    return np.array(out)


# --------------------------------------------------------------------------- #
#  helpers: link transforms + LGCP binning                                      #
# --------------------------------------------------------------------------- #

def predictive_prob(mean, var):
    """MacKay probit approximation to E[sigmoid(f)] under f ~ N(mean, var)."""
    return expit(mean / np.sqrt(1 + np.pi * var / 8))

def bin_counts(events, lo, hi, nbins):
    """bin event times into counts; returns bin centres, counts, and bin width (exposure)."""
    edges = np.linspace(lo, hi, nbins + 1); width = edges[1] - edges[0]
    counts, _ = np.histogram(events, bins=edges)
    centres = 0.5 * (edges[:-1] + edges[1:])
    return centres, counts.astype(float), width