"""Free-knot regression splines via reversible-jump MCMC -- from scratch.

DiMatteo, I., Genovese, C. R. & Kass, R. E. (2001), "Bayesian curve-fitting with free-knot
splines", Biometrika 88, 1055-1071.  Generalises Congdon (2005) Example 3.9 -- a cubic
spline whose knot *locations* are unknown but whose *number* is fixed at five -- to the
unknown-dimension case where the number of knots is itself inferred.

Model.  y_i = s(x_i) + eps_i, eps ~ N(0, sigma^2), where s is a cubic spline in the
truncated-power basis with k knots xi_1 < ... < xi_k:

    s(x) = b0 + b1 x + b2 x^2 + b3 x^3 + sum_j b_{3+j} (x - xi_j)_+^3.

Both k and the knot positions are unknown.  As in the change-point sampler, the regression
coefficients and error variance are given a conjugate (Zellner g-prior) treatment and
integrated out, leaving a closed-form marginal likelihood for a knot configuration -- so the
reversible-jump BIRTH (add a knot), DEATH (remove a knot) and MOVE (relocate a knot) moves
change dimension with NO Jacobian.  Knots live on a fixed candidate grid; a Poisson(lambda)
prior on their number supplies the parsimony penalty.  x is rescaled to [0,1] for numerical
stability.  g = n (unit-information).
"""
import numpy as np


def _design(xs, knots):
    """Cubic truncated-power spline design at rescaled locations xs with given knots."""
    cols = [np.ones_like(xs), xs, xs ** 2, xs ** 3]
    for k in knots:
        cols.append(np.where(xs > k, (xs - k) ** 3, 0.0))
    return np.column_stack(cols)


def _logml(xs, y, knots, g):
    """g-prior marginal log-likelihood of a knot configuration (coefficients + sigma^2 integrated out)."""
    n = len(y)
    X = _design(xs, knots)
    Xc = X[:, 1:] - X[:, 1:].mean(0)                    # non-intercept columns, centered
    yc = y - y.mean()
    beta, *_ = np.linalg.lstsq(Xc, yc, rcond=None)
    rss = yc - Xc @ beta
    R2 = 1.0 - (rss @ rss) / (yc @ yc)
    p = Xc.shape[1]                                     # 3 (cubic) + k (knots)
    return 0.5 * (n - 1 - p) * np.log1p(g) - 0.5 * (n - 1) * np.log1p(g * (1.0 - R2))


def _fit(xs, y, knots, g):
    """Posterior-mean fitted values under the g-prior (shrinkage g/(1+g) on the OLS coefficients)."""
    X = _design(xs, knots)
    m = X.mean(0)
    Xc = X[:, 1:] - m[1:]; yc = y - y.mean()
    beta, *_ = np.linalg.lstsq(Xc, yc, rcond=None)
    return y.mean() + Xc @ (g / (1 + g) * beta)


def rj_spline(x, y, ncand=80, lam=4.0, kmax=25, g=None, niter=80000, burn=25000, seed=0):
    """Reversible-jump sampler over cubic-spline knot configurations.

    Returns the posterior over the number of knots, the knot-location probabilities (on the
    candidate grid) and the posterior-mean fitted curve.
    """
    rng = np.random.default_rng(seed)
    x = np.asarray(x, float); y = np.asarray(y, float)
    n = len(y)
    if g is None: g = float(n)
    xlo, xhi = x.min(), x.max()
    xs = (x - xlo) / (xhi - xlo)                         # rescale to [0,1]
    cand = np.linspace(0.02, 0.98, ncand)               # candidate knot positions (interior grid)

    knots = []                                           # indices into cand
    inset = np.zeros(ncand, bool)
    cur = _logml(xs, y, cand[knots], g)
    Ktr = np.zeros(niter - burn, int); loc = np.zeros(ncand); fit = np.zeros(n); nkeep = 0

    for it in range(niter):
        k = len(knots); u = rng.random()
        if u < 1/3 and k < kmax:                         # BIRTH
            empt = np.where(~inset)[0]
            c = int(empt[rng.integers(len(empt))])
            nw = sorted(knots + [c]); nl = _logml(xs, y, cand[nw], g)
            if np.log(rng.random()) < (nl - cur) + np.log(lam / (k + 1)):
                knots = nw; inset[c] = True; cur = nl
        elif u < 2/3 and k > 0:                          # DEATH
            c = knots[rng.integers(k)]
            nw = [z for z in knots if z != c]; nl = _logml(xs, y, cand[nw], g)
            if np.log(rng.random()) < (nl - cur) + np.log(k / lam):
                knots = nw; inset[c] = False; cur = nl
        elif k > 0:                                       # MOVE
            c = knots[rng.integers(k)]; empt = np.where(~inset)[0]
            cn = int(empt[rng.integers(len(empt))])
            nw = sorted([z for z in knots if z != c] + [cn]); nl = _logml(xs, y, cand[nw], g)
            if np.log(rng.random()) < (nl - cur):
                inset[c] = False; inset[cn] = True; knots = nw; cur = nl

        if it >= burn:
            j = it - burn
            Ktr[j] = len(knots)
            for c in knots: loc[c] += 1
            fit += _fit(xs, y, cand[knots], g); nkeep += 1

    return dict(K=Ktr, kpost=np.bincount(Ktr, minlength=kmax + 1) / nkeep,
                loc=loc / nkeep, cand=cand * (xhi - xlo) + xlo, fit=fit / nkeep, x=x, g=g, lam=lam)