"""
pspline.py -- BAYESIAN PENALISED SPLINES (P-splines) and ADDITIVE MODELS (from scratch).

Backs the notebooks in  "Bayesian Penalised Splines & Additive Models".

Two routes to a flexible curve. The variable-selection arc's free-knot notebook chooses the
NUMBER and LOCATION of a few knots by trans-dimensional MCMC. The PENALISED-spline route
(Eilers & Marx 1996; Bayesian version: Lang & Brezger 2004) does the opposite: lay down MANY
equally spaced B-spline basis functions -- far more than needed -- and control smoothness not by
the knots but by a PENALTY on the coefficients. As a Bayesian prior this penalty is a random walk:

    y_i = sum_j B_j(x_i) beta_j + eps_i,   eps ~ N(0, sigma^2)
    beta_j - 2 beta_{j-1} + beta_{j-2} ~ N(0, tau^2)          (2nd-order random-walk prior)

i.e. beta ~ N(0, tau^2 (D'D)^{-1}) with D the 2nd-difference matrix. The smoothing parameter is
the ratio sigma^2/tau^2, and -- crucially -- it is INFERRED from the data through tau^2's own
prior, so there is nothing to cross-validate. The full conditionals are all conjugate (a Gaussian
Markov random field):

    beta   ~ N( (B'B/sigma^2 + D'D/tau^2)^{-1} B'y/sigma^2,  (B'B/sigma^2 + D'D/tau^2)^{-1} )
    sigma^2 ~ InvGamma(a + n/2,      b + 1/2 ||y - B beta||^2)
    tau^2   ~ InvGamma(a + rank(D)/2, b + 1/2 beta' D'D beta)

ADDITIVE MODELS (GAMs) stack one penalised spline per covariate, y = beta0 + sum_k f_k(x_k),
each with its own smoothing variance, and are sampled by the same Gibbs sweep applied to partial
residuals (backfitting inside Gibbs). A P-spline is also exactly a Gaussian process with a
particular kernel, so this connects straight back to the GP-regression notebook.
"""

import numpy as np
from scipy.interpolate import BSpline
from scipy.linalg import cholesky, cho_solve, solve_triangular


# --------------------------------------------------------------------------- #
#  B-spline design matrix (equally spaced knots) + difference penalty           #
# --------------------------------------------------------------------------- #

def make_knots(xl, xr, ndx, deg):
    dx = (xr - xl) / ndx
    return xl + dx * np.arange(-deg, ndx + deg + 1)

def bspline_design(x, knots, deg):
    nb = len(knots) - deg - 1; B = np.empty((len(x), nb))
    for j in range(nb):
        c = np.zeros(nb); c[j] = 1.0
        B[:, j] = BSpline(knots, c, deg, extrapolate=True)(x)
    return B

def diff_penalty(nb, order=2):
    return np.diff(np.eye(nb), order, axis=0)          # (nb-order, nb); K = D'D


def _draw_beta(BtB, Bty, K, sig2, tau2, rng):
    nb = BtB.shape[0]
    Prec = BtB / sig2 + K / tau2 + 1e-8 * np.eye(nb)
    L = cholesky(Prec, lower=True)
    mean = cho_solve((L, True), Bty / sig2)
    return mean + solve_triangular(L.T, rng.standard_normal(nb))


# --------------------------------------------------------------------------- #
#  single penalised spline                                                      #
# --------------------------------------------------------------------------- #

def pspline_gibbs(x, y, rng, ndx=25, deg=3, penord=2, draws=2500, burn=1000,
                  grid=None, a0=1e-3, b0=1e-3):
    x = np.asarray(x, float); y = np.asarray(y, float); n = len(y)
    xl, xr = x.min(), x.max(); knots = make_knots(xl, xr, ndx, deg)
    B = bspline_design(x, knots, deg); nb = B.shape[1]
    K = diff_penalty(nb, penord); K = K.T @ K
    if grid is None: grid = np.linspace(xl, xr, 200)
    Bg = bspline_design(grid, knots, deg)
    BtB = B.T @ B; Bty = B.T @ y; rankD = nb - penord
    tau2 = 1.0; sig2 = float(np.var(y)); FIT = np.empty((draws, len(grid)))
    TAU = np.empty(draws)
    for it in range(draws + burn):
        beta = _draw_beta(BtB, Bty, K, sig2, tau2, rng)
        resid = y - B @ beta
        sig2 = 1.0 / rng.gamma(a0 + n / 2, 1.0 / (b0 + 0.5 * resid @ resid))
        tau2 = 1.0 / rng.gamma(a0 + rankD / 2, 1.0 / (b0 + 0.5 * beta @ K @ beta))
        if it >= burn:
            FIT[it - burn] = Bg @ beta; TAU[it - burn] = tau2
    return dict(grid=grid, fit=FIT, tau2=TAU, sig2=sig2)


# --------------------------------------------------------------------------- #
#  additive model (GAM) -- one penalised spline per covariate                   #
# --------------------------------------------------------------------------- #

def gam_gibbs(Xcols, y, rng, ndx=15, deg=3, penord=2, draws=2500, burn=1000,
              grids=None, a0=1e-3, b0=1e-3):
    """Xcols: list of 1-D covariate arrays. Returns per-term partial-effect draws on `grids`
    (centred so each smooth sums to zero; the intercept carries the level)."""
    y = np.asarray(y, float); n = len(y); p = len(Xcols)
    knots = [make_knots(x.min(), x.max(), ndx, deg) for x in Xcols]
    Braw = [bspline_design(np.asarray(x, float), knots[k], deg) for k, x in enumerate(Xcols)]
    cmean = [b.mean(0) for b in Braw]
    B = [Braw[k] - cmean[k] for k in range(p)]                  # centred bases (identifiability)
    K = [diff_penalty(b.shape[1], penord) for b in B]; K = [d.T @ d for d in K]
    BtB = [B[k].T @ B[k] for k in range(p)]
    if grids is None:
        grids = [np.linspace(x.min(), x.max(), 200) for x in Xcols]
    Bg = [bspline_design(grids[k], knots[k], deg) - cmean[k] for k in range(p)]
    beta = [np.zeros(b.shape[1]) for b in B]; b0i = float(y.mean())
    tau2 = [1.0] * p; sig2 = float(np.var(y)); rankD = [b.shape[1] - penord for b in B]
    PART = [np.empty((draws, len(grids[k]))) for k in range(p)]
    B0 = np.empty(draws)
    for it in range(draws + burn):
        fitsum = sum(B[k] @ beta[k] for k in range(p))
        b0i = (y - fitsum).mean() + rng.standard_normal() * np.sqrt(sig2 / n)
        for k in range(p):
            r = y - b0i - sum(B[j] @ beta[j] for j in range(p) if j != k)
            beta[k] = _draw_beta(BtB[k], B[k].T @ r, K[k], sig2, tau2[k], rng)
            tau2[k] = 1.0 / rng.gamma(a0 + rankD[k] / 2, 1.0 / (b0 + 0.5 * beta[k] @ K[k] @ beta[k]))
        resid = y - b0i - sum(B[k] @ beta[k] for k in range(p))
        sig2 = 1.0 / rng.gamma(a0 + n / 2, 1.0 / (b0 + 0.5 * resid @ resid))
        if it >= burn:
            B0[it - burn] = b0i
            for k in range(p):
                PART[k][it - burn] = Bg[k] @ beta[k]
    return dict(partial=PART, grids=grids, intercept=B0, sig2=sig2)


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

def simulate_additive(n, funcs, ranges, noise_sd, rng, intercept=0.0):
    """y = intercept + sum_k f_k(x_k) + noise; funcs = list of callables, ranges = list of (lo,hi)."""
    Xcols = [rng.uniform(lo, hi, n) for (lo, hi) in ranges]
    mu = intercept + sum(f(x) for f, x in zip(funcs, Xcols))
    y = mu + noise_sd * rng.standard_normal(n)
    return Xcols, y