"""
kriging.py -- GEOSTATISTICS: variograms and kriging for point-referenced data (from scratch).

Backs the notebooks in  "Geostatistics -- Variograms and Kriging".

Areal models (CAR/SAR) live on regions with a neighbour graph. GEOSTATISTICS instead handles
POINT-REFERENCED data -- measurements at continuous coordinates (soil samples, weather stations,
house locations) -- and asks two things: how does similarity decay with distance, and what is the
value at an UNMEASURED location?

The VARIOGRAM answers the first. The empirical semivariogram

    gamma(h) = (1 / 2 N(h)) * sum over pairs at distance ~ h of (z_i - z_j)^2

rises from a NUGGET (micro-scale noise) toward a SILL (the overall variance) over a RANGE (the
distance beyond which points are effectively uncorrelated). Fitting a model (exponential, Matern,
spherical) to it gives the covariance C(h) = sill - gamma(h) used for prediction.

KRIGING answers the second: it is the best linear unbiased predictor (BLUP) at a new location, a
weighted average of the data with weights from the covariance. ORDINARY kriging assumes an unknown
constant mean (a Lagrange constraint forces the weights to sum to one); UNIVERSAL kriging lets the
mean be a regression on covariates (a spatial trend). Each prediction comes with a KRIGING VARIANCE
-- the map of uncertainty, largest far from data.

The key identity for this arc: KRIGING IS GAUSSIAN-PROCESS REGRESSION. C(h) is the GP covariance
kernel, the kriging predictor is the GP posterior mean, and the kriging variance is the GP posterior
variance. The 1-D GP notebooks in "Gaussian-Process Regression" are the same machinery; here it
is 2-D and spatial, and the variogram is just how geostatisticians estimate the kernel.
"""

import numpy as np
from scipy.spatial.distance import cdist, pdist, squareform
from scipy.optimize import least_squares


# --------------------------------------------------------------------------- #
#  variogram                                                                    #
# --------------------------------------------------------------------------- #

def empirical_variogram(coords, z, nbins=15, maxdist=None):
    """binned empirical semivariogram. Returns (bin centres, semivariance, pair counts)."""
    z = np.asarray(z, float); D = squareform(pdist(coords)); iu = np.triu_indices(len(z), 1)
    h = D[iu]; g = 0.5 * (z[iu[0]] - z[iu[1]]) ** 2
    if maxdist is None:
        maxdist = h.max() / 2
    edges = np.linspace(0, maxdist, nbins + 1); cen = 0.5 * (edges[:-1] + edges[1:])
    gam = np.full(nbins, np.nan); cnt = np.zeros(nbins)
    for b in range(nbins):
        m = (h > edges[b]) & (h <= edges[b + 1])
        if m.sum() > 0:
            gam[b] = g[m].mean(); cnt[b] = m.sum()
    ok = ~np.isnan(gam)
    return cen[ok], gam[ok], cnt[ok]


def vgm_exponential(h, nugget, psill, rng):
    """exponential semivariogram: nugget at h>0, rising to nugget+psill over ~3*rng."""
    return nugget + psill * (1 - np.exp(-h / rng))


def fit_variogram(cen, gam, cnt):
    """weighted least-squares fit of an exponential variogram (nugget, partial sill, range)."""
    w = np.sqrt(cnt)
    def resid(p):
        nug, ps, rg = np.abs(p)
        return w * (vgm_exponential(cen, nug, ps, rg) - gam)
    p0 = [gam.min(), gam.max() - gam.min(), cen.max() / 3]
    r = least_squares(resid, p0, method="lm")
    nug, ps, rg = np.abs(r.x)
    return dict(nugget=nug, psill=ps, rng=rg, sill=nug + ps)


def _cov(D, v):
    """covariance matrix from an exponential variogram: C(h)=psill*exp(-h/rng); C(0)=psill+nugget."""
    C = v["psill"] * np.exp(-D / v["rng"])
    C[D == 0] = v["psill"] + v["nugget"]
    return C


# --------------------------------------------------------------------------- #
#  kriging                                                                       #
# --------------------------------------------------------------------------- #

def ordinary_kriging(coords, z, grid, v):
    """ordinary kriging (unknown constant mean). Returns prediction and kriging-variance arrays."""
    z = np.asarray(z, float); n = len(z)
    K = _cov(squareform(pdist(coords)), v)
    A = np.zeros((n + 1, n + 1)); A[:n, :n] = K; A[:n, n] = 1; A[n, :n] = 1
    Ai = np.linalg.inv(A)
    k0 = _cov(cdist(coords, grid), v)                              # (n, m)
    b = np.vstack([k0, np.ones((1, grid.shape[0]))])               # (n+1, m)
    W = Ai @ b                                                     # weights + Lagrange mult
    pred = W[:n].T @ z
    var = (v["psill"] + v["nugget"]) - np.sum(b * W, axis=0)
    return pred, np.maximum(var, 0)


def universal_kriging(coords, z, Xd, grid, Xg, v):
    """universal kriging with a linear trend. Xd:(n,p) trend basis at data (incl. intercept column),
    Xg:(m,p) at the grid. Returns prediction and kriging variance."""
    z = np.asarray(z, float); n, p = Xd.shape
    K = _cov(squareform(pdist(coords)), v)
    A = np.zeros((n + p, n + p)); A[:n, :n] = K; A[:n, n:] = Xd; A[n:, :n] = Xd.T
    Ai = np.linalg.inv(A)
    k0 = _cov(cdist(coords, grid), v)
    b = np.vstack([k0, Xg.T])                                      # (n+p, m)
    W = Ai @ b
    pred = W[:n].T @ z
    var = (v["psill"] + v["nugget"]) - np.sum(b * W, axis=0)
    return pred, np.maximum(var, 0)