"""
gp.py -- GAUSSIAN-PROCESS REGRESSION (from scratch).

Backs the notebooks in  "Gaussian-Process Regression".

A Gaussian process is a prior over FUNCTIONS: any finite set of function values is jointly
Gaussian, with covariance given by a kernel k(x,x'). Regression with a GP prior and Gaussian
noise has everything in closed form. With y = f(X) + eps, eps ~ N(0, sigma_n^2):

    f ~ GP(0, k),   [f(X), f(X*)] jointly Gaussian
    posterior mean at X*:  m* = K(X*,X) [K(X,X)+sigma_n^2 I]^{-1} y
    posterior cov  at X*:  C* = K(X*,X*) - K(X*,X)[K(X,X)+sigma_n^2 I]^{-1} K(X,X*)

and the kernel HYPERPARAMETERS are learned by maximising the log marginal likelihood

    log p(y) = -1/2 y^T [K+sigma_n^2 I]^{-1} y - 1/2 log|K+sigma_n^2 I| - n/2 log 2pi,

which trades data fit against model complexity automatically -- no cross-validation needed.
Everything is computed stably through the Cholesky factor of K + sigma_n^2 I.

The kernel is where prior structure lives: an RBF/Matern kernel encodes smoothness and a
length-scale; a periodic kernel encodes a cycle; SUMS and PRODUCTS of kernels build structured
priors (the Mauna Loa CO2 model = long trend + periodic seasonal + short-term wiggle + noise).
This module is 1-D (a single input), which covers both notebook datasets.
"""

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


# --------------------------------------------------------------------------- #
#  kernel primitives (1-D)                                                      #
# --------------------------------------------------------------------------- #

def _D(xa, xb):
    xa = np.asarray(xa, float).ravel(); xb = np.asarray(xb, float).ravel()
    return xa[:, None] - xb[None, :]

def k_rbf(xa, xb, ell, eta):
    """squared-exponential (infinitely smooth), length-scale ell, variance eta."""
    return eta * np.exp(-0.5 * (_D(xa, xb) / ell) ** 2)

def k_matern52(xa, xb, ell, eta):
    """Matern-5/2 (twice differentiable -- rougher than RBF)."""
    r = np.abs(_D(xa, xb)) / ell
    return eta * (1 + np.sqrt(5) * r + 5.0 / 3.0 * r ** 2) * np.exp(-np.sqrt(5) * r)

def k_periodic(xa, xb, ell, period, eta):
    """exactly periodic kernel (Mackay), period `period`, smoothness ell."""
    return eta * np.exp(-2.0 * np.sin(np.pi * np.abs(_D(xa, xb)) / period) ** 2 / ell ** 2)

def k_linear(xa, xb, c, eta):
    xa = np.asarray(xa, float).ravel(); xb = np.asarray(xb, float).ravel()
    return eta * np.outer(xa - c, xb - c)


# --------------------------------------------------------------------------- #
#  fit (maximise log marginal likelihood) and predict                           #
# --------------------------------------------------------------------------- #

def _neg_lml(theta, X, y, build):
    kfun, nv = build(theta)
    n = len(X); K = kfun(X, X) + (nv + 1e-8) * np.eye(n)
    try:
        L = cholesky(K, lower=True)
    except np.linalg.LinAlgError:
        return 1e25
    alpha = cho_solve((L, True), y)
    return 0.5 * y @ alpha + np.log(np.diag(L)).sum() + 0.5 * n * np.log(2 * np.pi)

def fit_gp(X, y, build, theta0, bounds=None, restarts=1, rng=None):
    """maximise the log marginal likelihood over hyperparameters `theta` (usually log-scale).
    `build(theta)` returns (kfun, noise_var), where kfun(Xa,Xb) is the signal covariance.
    Returns (theta_hat, log_marginal_likelihood)."""
    best = None
    for r in range(restarts):
        t0 = np.asarray(theta0, float) if r == 0 else np.asarray(theta0, float) + 0.4 * rng.standard_normal(len(theta0))
        res = minimize(_neg_lml, t0, args=(X, y, build), method="L-BFGS-B", bounds=bounds)
        if best is None or res.fun < best.fun:
            best = res
    return best.x, -best.fun

def predict_gp(theta, X, y, Xs, build):
    """posterior predictive mean and covariance of the latent function at Xs."""
    kfun, nv = build(theta); n = len(X)
    K = kfun(X, X) + (nv + 1e-8) * np.eye(n); L = cholesky(K, lower=True)
    alpha = cho_solve((L, True), y)
    Ks = kfun(Xs, X); mean = Ks @ alpha
    v = cho_solve((L, True), Ks.T); cov = kfun(Xs, Xs) - Ks @ v
    return mean, cov

def sample_gp(mean, cov, n, rng):
    """draw n function samples from a Gaussian with the given mean and covariance."""
    m = np.asarray(mean, float); C = np.asarray(cov, float)
    L = cholesky(C + 1e-9 * np.eye(len(m)), lower=True)
    return m[:, None] + L @ rng.standard_normal((len(m), n))


# --------------------------------------------------------------------------- #
#  simulation for validation                                                    #
# --------------------------------------------------------------------------- #

def simulate_regression(f, x, noise_sd, rng):
    """y = f(x) + Normal(0, noise_sd) -- a known function to recover."""
    x = np.asarray(x, float)
    return f(x) + noise_sd * rng.standard_normal(len(x))