"""
pointproc.py -- SPATIAL POINT PROCESSES: CSR testing, inhomogeneous Poisson, LGCP (from scratch).

Backs the notebooks in  "Spatial Point Processes -- CSR, Intensity, and Cox Processes".

Areal and geostatistical models take the locations as given and model a value there. A POINT PROCESS
models the LOCATIONS THEMSELVES -- where events occur (trees, crimes, disease cases, retail sites).
The questions are different: is the pattern random, clustered, or regular? what drives the local
INTENSITY (expected points per unit area)? is there clustering left after covariates?

The benchmark is COMPLETE SPATIAL RANDOMNESS (CSR) -- a homogeneous Poisson process, constant
intensity, points independent. RIPLEY'S K counts, for each distance r, the average number of other
points within r of a typical point, scaled by intensity; under CSR K(r) = pi r^2, so the centred
L-function L(r) - r is zero. Above the CSR envelope means CLUSTERING, below means INHIBITION.

For intensity we use the standard computational device: overlay a fine GRID, count points per cell,
and treat the counts as Poisson with mean (cell area) x intensity. Then:

  INHOMOGENEOUS POISSON:  n_c ~ Poisson( area_c * exp(x_c' beta) )
      the intensity is a log-linear function of covariates -- an ordinary Poisson regression on the
      grid (the Berman-Turner / Baddeley quadrature approximation to the point-process likelihood).

  LOG-GAUSSIAN COX PROCESS (LGCP):  n_c ~ Poisson( area_c * exp(x_c' beta + phi_c) ),  phi ~ CAR
      adds a latent Gaussian spatial field phi to the log-intensity, capturing clustering BEYOND the
      covariates (a doubly-stochastic Poisson). On a grid this is exactly the areal Poisson-CAR
      (BYM) model of the areal notebooks -- so the LGCP is disease-mapping applied to a point pattern,
      and the same latent-field machinery as the coal-mining LGCP in "GP Classification & Log-Gaussian Cox Processes".
"""

import numpy as np
from scipy.spatial.distance import pdist


# --------------------------------------------------------------------------- #
#  Ripley's K / L with translation edge correction (rectangular window)         #
# --------------------------------------------------------------------------- #

def ripley_L(pts, xr, yr, radii):
    """centred L-function L(r)-r for points in rectangle [xr]x[yr], translation edge correction.
    Returns an array aligned with `radii` (0 under CSR; >0 clustering, <0 inhibition)."""
    x = pts[:, 0]; y = pts[:, 1]; n = len(x); X = xr[1] - xr[0]; Y = yr[1] - yr[0]; A = X * Y
    d = pdist(pts); dx = pdist(x[:, None]); dy = pdist(y[:, None])
    w = A / ((X - dx) * (Y - dy))                                  # translation weights
    K = np.array([2.0 * np.sum(w[d <= r]) * A / (n ** 2) for r in radii])
    return np.sqrt(K / np.pi) - radii


def csr_envelope(n, xr, yr, radii, rng, nsim=39):
    """simulate CSR patterns of n points; return (lo, hi) pointwise envelope of L(r)-r."""
    sims = np.empty((nsim, len(radii)))
    for s in range(nsim):
        p = np.column_stack([rng.uniform(*xr, n), rng.uniform(*yr, n)])
        sims[s] = ripley_L(p, xr, yr, radii)
    return sims.min(0), sims.max(0)


# --------------------------------------------------------------------------- #
#  grid the pattern (count points + aggregate covariates per cell)              #
# --------------------------------------------------------------------------- #

def make_grid(pts, gx, gy, gcov, xr, yr, cell):
    """aggregate to a regular grid of side `cell`. gx,gy,gcov: fine covariate pixels (arrays, gcov
    is (npix, k)). Returns dict with counts, covariate means, centres, cell area, rook adjacency."""
    nx = int(round((xr[1] - xr[0]) / cell)); ny = int(round((yr[1] - yr[0]) / cell))
    ix = np.clip(((pts[:, 0] - xr[0]) / cell).astype(int), 0, nx - 1)
    iy = np.clip(((pts[:, 1] - yr[0]) / cell).astype(int), 0, ny - 1)
    counts = np.zeros((nx, ny))
    np.add.at(counts, (ix, iy), 1)
    px = np.clip(((gx - xr[0]) / cell).astype(int), 0, nx - 1)
    py = np.clip(((gy - yr[0]) / cell).astype(int), 0, ny - 1)
    k = gcov.shape[1]; csum = np.zeros((nx, ny, k)); ccnt = np.zeros((nx, ny))
    np.add.at(csum, (px, py), gcov); np.add.at(ccnt, (px, py), 1)
    cov = csum / np.maximum(ccnt[:, :, None], 1)
    cx = xr[0] + (np.arange(nx) + 0.5) * cell; cy = yr[0] + (np.arange(ny) + 0.5) * cell
    CX, CY = np.meshgrid(cx, cy, indexing="ij")
    counts = counts.ravel(); cov = cov.reshape(-1, k); centres = np.column_stack([CX.ravel(), CY.ravel()])
    m = nx * ny; W = np.zeros((m, m), np.int8)                     # rook adjacency
    idx = np.arange(m).reshape(nx, ny)
    for i in range(nx):
        for j in range(ny):
            if i + 1 < nx: W[idx[i, j], idx[i + 1, j]] = W[idx[i + 1, j], idx[i, j]] = 1
            if j + 1 < ny: W[idx[i, j], idx[i, j + 1]] = W[idx[i, j + 1], idx[i, j]] = 1
    return dict(counts=counts, cov=cov, centres=centres, area=cell * cell, W=W, nx=nx, ny=ny)


# --------------------------------------------------------------------------- #
#  inhomogeneous Poisson (grid Poisson GLM by IRLS) and LGCP                     #
# --------------------------------------------------------------------------- #

def inhom_poisson(counts, X, area):
    """Poisson GLM  n_c ~ Poisson(area * exp(X beta))  by IRLS. X includes an intercept column."""
    y = np.asarray(counts, float); off = np.log(area); beta = np.zeros(X.shape[1]); beta[0] = np.log(y.mean() + 1e-6) - off
    for _ in range(50):
        eta = X @ beta + off; mu = np.exp(eta)
        WX = X * mu[:, None]; H = X.T @ WX; g = X.T @ (y - mu)
        step = np.linalg.solve(H + 1e-8 * np.eye(X.shape[1]), g); beta += step
        if np.max(np.abs(step)) < 1e-8:
            break
    cov = np.linalg.inv(H); return beta, np.sqrt(np.diag(cov))


def lgcp_gibbs(counts, X, area, W, rng, draws=2000, burn=2000, a=1.0, b=0.01):
    """log-Gaussian Cox process on the grid: n_c ~ Poisson(area exp(X beta + phi)), phi ~ ICAR.
    Metropolis-within-Gibbs (as in the areal BYM). Returns beta, spatial SD, and the latent field."""
    y = np.asarray(counts, float); n, k = X.shape; off = np.log(area); nnb = W.sum(1)
    beta, _ = inhom_poisson(counts, X, area); phi = np.zeros(n); tau2 = 0.5
    sb = 0.05 * np.ones(k); sp = 0.4 * np.ones(n); accp = np.zeros(n); accb = np.zeros(k); nadapt = 0
    B = np.empty((draws, k)); TAU = np.empty(draws); PHI = np.empty((draws, n))
    for it in range(draws + burn):
        eta = X @ beta + phi + off; lam = np.exp(eta)
        for j in range(k):                                        # beta (RW-Metropolis)
            bp = beta.copy(); bp[j] += sb[j] * rng.standard_normal()
            lp = np.exp(X @ bp + phi + off)
            if np.log(rng.random()) < (y @ (X @ (bp - beta)) - (lp - lam).sum() - (bp[j] ** 2 - beta[j] ** 2) / (2 * 100)):
                beta = bp; lam = lp; accb[j] += 1
        for i in range(n):                                         # phi (ICAR-conditional Metropolis)
            nbm = (W[i] @ phi) / max(nnb[i], 1.0)                  # neighbours as they stand NOW
            dphi = sp[i] * rng.standard_normal()
            dll = y[i] * dphi - lam[i] * (np.exp(dphi) - 1)
            dpr = -(nnb[i] / (2 * tau2)) * ((phi[i] + dphi - nbm) ** 2 - (phi[i] - nbm) ** 2)
            if np.log(rng.random()) < dll + dpr:
                phi[i] += dphi; lam[i] *= np.exp(dphi); accp[i] += 1
        phi -= phi.mean()
        quad = 0.5 * (nnb * phi * phi).sum() - 0.5 * (phi * (W @ phi)).sum()
        tau2 = 1.0 / rng.gamma(a + (n - 1) / 2, 1.0 / (b + max(quad, 1e-6)))
        if it < burn and it % 100 == 99:                           # adapt every scale during burn-in
            nadapt += 1; g = 1.0 / np.sqrt(nadapt)
            sp *= np.exp((accp / 100 - 0.44) * g); sb *= np.exp((accb / 100 - 0.44) * g)
            accp[:] = 0; accb[:] = 0
        if it >= burn:
            t = it - burn; B[t] = beta; TAU[t] = np.sqrt(tau2); PHI[t] = phi
    return dict(beta=B, spatial_sd=TAU, phi=PHI)