"""kriging.py

Variogram estimation, ordinary kriging and regression kriging,
written against numpy and scipy only.

Every parameter a commercial tool hides behind a default appears here as a
named argument: the lag distance, the bin count, the model form, the nugget,
the sill, the range, the anisotropy angle and ratio, and the search
neighborhood. Exercises 0-B and 0-C move them one at a time and score the
result by leave one out cross validation.

Distances are metric because the caller supplies projected coordinates. The
AZMet stations arrive as longitude and latitude and get projected to Albers
Equal Area Conic, EPSG:5070, by `project_albers` below.
"""
from __future__ import annotations

import numpy as np
from scipy.optimize import least_squares

MODELS = ("spherical", "exponential", "gaussian")


# --------------------------------------------------------------------------
# Projection. Albers Equal Area Conic, USGS parameters, implemented directly
# so the student stack needs no projection library.
# --------------------------------------------------------------------------
def project_albers(lon, lat, lat1=29.5, lat2=45.5, lon0=-96.0, lat0=23.0):
    """Longitude and latitude in degrees to EPSG:5070 meters, GRS80 sphere."""
    R = 6378137.0
    lon, lat = np.radians(np.asarray(lon, float)), np.radians(np.asarray(lat, float))
    p1, p2, p0, l0 = map(np.radians, (lat1, lat2, lat0, lon0))
    n = 0.5 * (np.sin(p1) + np.sin(p2))
    C = np.cos(p1) ** 2 + 2 * n * np.sin(p1)
    rho0 = R * np.sqrt(C - 2 * n * np.sin(p0)) / n
    rho = R * np.sqrt(C - 2 * n * np.sin(lat)) / n
    theta = n * (lon - l0)
    return rho * np.sin(theta), rho0 - rho * np.cos(theta)


def pairwise_distance(x, y):
    dx = x[:, None] - x[None, :]
    dy = y[:, None] - y[None, :]
    return np.sqrt(dx ** 2 + dy ** 2)


# --------------------------------------------------------------------------
# Empirical variogram.
# --------------------------------------------------------------------------
def empirical_variogram(x, y, z, *, lag=None, n_bins=12, max_dist_frac=0.5,
                        angle=None, tolerance_deg=22.5):
    """Binned semivariance.

    `angle` in degrees clockwise from north restricts pairs to a direction,
    which is how anisotropy gets detected. Leaving it None pools all
    directions, which assumes the phenomenon behaves the same way in every
    orientation.
    """
    x, y, z = map(lambda a: np.asarray(a, float), (x, y, z))
    d = pairwise_distance(x, y)
    iu = np.triu_indices(len(z), k=1)
    dist = d[iu]
    semi = 0.5 * (z[iu[0]] - z[iu[1]]) ** 2

    if angle is not None:
        dx = (x[iu[1]] - x[iu[0]])
        dy = (y[iu[1]] - y[iu[0]])
        bearing = (np.degrees(np.arctan2(dx, dy))) % 180.0
        target = angle % 180.0
        sep = np.abs(bearing - target)
        sep = np.minimum(sep, 180.0 - sep)
        keep = sep <= tolerance_deg
        dist, semi = dist[keep], semi[keep]

    cutoff = dist.max() * max_dist_frac
    if lag is None:
        lag = cutoff / n_bins
    edges = np.arange(0, cutoff + lag, lag)
    centers, values, counts = [], [], []
    for i in range(len(edges) - 1):
        m = (dist > edges[i]) & (dist <= edges[i + 1])
        if m.sum() >= 5:                       # too few pairs means no estimate
            centers.append(0.5 * (edges[i] + edges[i + 1]))
            values.append(semi[m].mean())
            counts.append(int(m.sum()))
    return (np.array(centers), np.array(values), np.array(counts),
            float(lag), int(len(dist)))


def model_function(name):
    def spherical(h, nugget, sill, rng):
        h = np.asarray(h, float)
        out = np.where(h >= rng, sill,
                       nugget + (sill - nugget) *
                       (1.5 * h / rng - 0.5 * (h / rng) ** 3))
        return np.where(h == 0, 0.0, out)

    def exponential(h, nugget, sill, rng):
        h = np.asarray(h, float)
        return np.where(h == 0, 0.0,
                        nugget + (sill - nugget) * (1 - np.exp(-3 * h / rng)))

    def gaussian(h, nugget, sill, rng):
        h = np.asarray(h, float)
        return np.where(h == 0, 0.0,
                        nugget + (sill - nugget) *
                        (1 - np.exp(-3 * (h / rng) ** 2)))
    return {"spherical": spherical, "exponential": exponential,
            "gaussian": gaussian}[name]


def fit_variogram(centers, values, counts=None, model="spherical",
                  nugget=None, sill=None, rng=None):
    """Weighted least squares fit of a model to the binned means.

    Fitting runs on the binned means and never on the raw pair cloud, because
    the cloud is a chi-square scatter whose individual points carry almost no
    information; the bin mean is the estimate the theory describes. Pair
    counts weight the fit, so a bin resting on nine pairs pulls less than one
    resting on four hundred.

    Any of nugget, sill or range supplied by the caller is held fixed, which
    is how the exercise overrides one parameter at a time.
    """
    f = model_function(model)
    fixed = {"nugget": nugget, "sill": sill, "rng": rng}
    free = [k for k, v in fixed.items() if v is None]
    start = {"nugget": max(values.min(), 1e-9),
             "sill": values.max(),
             "rng": centers.max() * 0.5}
    w = np.sqrt(counts) if counts is not None else np.ones_like(values)

    def resid(p):
        vals = dict(fixed)
        for k, v in zip(free, p):
            vals[k] = v
        pred = f(centers, vals["nugget"], vals["sill"], vals["rng"])
        return (pred - values) * w

    p0 = [start[k] for k in free]
    lo = [0.0 for _ in free]
    # The range is bounded at the largest separation actually observed. A
    # looser bound lets a sparse, noisy set of bins drive the fit to a range
    # far beyond the data, which produces a surface that is one flat plane
    # and a cross validated error two orders of magnitude too large.
    hi = [values.max() * 5 if k != "rng" else centers.max() * 1.5
          for k in free]
    sol = least_squares(resid, p0, bounds=(lo, hi))
    out = dict(fixed)
    for k, v in zip(free, sol.x):
        out[k] = float(v)
    out["model"] = model
    if out["sill"] <= out["nugget"]:
        out["sill"] = out["nugget"] * 1.0000001
    return out


# --------------------------------------------------------------------------
# Ordinary kriging.
# --------------------------------------------------------------------------
def _cov(h, p):
    f = model_function(p["model"])
    return p["sill"] - f(h, p["nugget"], p["sill"], p["rng"])


def ordinary_kriging(x, y, z, xi, yi, params, *, secondary=None,
                     secondary_i=None):
    """Predict at (xi, yi). Supplying `secondary` runs regression kriging.

    Regression kriging removes a linear trend on the secondary variable,
    krige the residual, then adds the trend back at the prediction sites.
    That requires the secondary variable to be known at the prediction sites,
    which is the condition Exercise 0-C exists to test.
    """
    x, y, z = map(lambda a: np.asarray(a, float), (x, y, z))
    xi, yi = np.asarray(xi, float), np.asarray(yi, float)
    trend_i = 0.0
    if secondary is not None and secondary_i is not None:
        s = np.asarray(secondary, float)
        b, a = np.polyfit(s, z, 1)
        z = z - (a + b * s)
        trend_i = a + b * np.asarray(secondary_i, float)

    n = len(z)

    # A pure nugget model carries no spatial structure, so every off-diagonal
    # covariance is zero and the kriging system is singular. The correct
    # answer in that case is the global mean everywhere with the sill as the
    # variance, which is precisely what the exercise wants a student to see.
    # Solving it numerically instead returns values near 1e20.
    if params["sill"] - params["nugget"] <= params["sill"] * 1e-6:
        return (np.full(len(xi), z.mean()) + trend_i,
                np.full(len(xi), float(params["sill"])))

    d = pairwise_distance(x, y)
    C = _cov(d, params)
    A = np.ones((n + 1, n + 1))
    A[:n, :n] = C
    A[n, n] = 0.0
    # Ridge on the diagonal keeps the solve stable where two stations sit
    # almost on top of one another.
    A[:n, :n] += np.eye(n) * params["sill"] * 1e-10

    dx = xi[:, None] - x[None, :]
    dy = yi[:, None] - y[None, :]
    d0 = np.sqrt(dx ** 2 + dy ** 2)
    B = np.ones((n + 1, len(xi)))
    B[:n, :] = _cov(d0, params).T

    try:
        W = np.linalg.solve(A, B)
    except np.linalg.LinAlgError:
        W = np.linalg.lstsq(A, B, rcond=None)[0]
    pred = (W[:n, :] * z[:, None]).sum(0) + trend_i
    var = params["sill"] - (W[:n, :] * B[:n, :]).sum(0) - W[n, :]
    return pred, np.maximum(var, 0.0)


def idw(x, y, z, xi, yi, power=2.0):
    x, y, z = map(lambda a: np.asarray(a, float), (x, y, z))
    dx = np.asarray(xi, float)[:, None] - x[None, :]
    dy = np.asarray(yi, float)[:, None] - y[None, :]
    d = np.sqrt(dx ** 2 + dy ** 2)
    d = np.where(d < 1e-9, 1e-9, d)
    w = 1.0 / d ** power
    return (w * z[None, :]).sum(1) / w.sum(1)


def loocv(x, y, z, params=None, *, method="kriging", secondary=None,
          power=2.0):
    """Leave one out cross validation. The referee for every parameter set."""
    x, y, z = map(lambda a: np.asarray(a, float), (x, y, z))
    n = len(z)
    pred = np.empty(n)
    for i in range(n):
        m = np.ones(n, bool)
        m[i] = False
        if method == "idw":
            pred[i] = idw(x[m], y[m], z[m], x[i:i + 1], y[i:i + 1], power)[0]
        else:
            sec = None if secondary is None else np.asarray(secondary)[m]
            sec_i = None if secondary is None else np.asarray(secondary)[i:i + 1]
            pred[i] = ordinary_kriging(x[m], y[m], z[m], x[i:i + 1],
                                       y[i:i + 1], params,
                                       secondary=sec, secondary_i=sec_i)[0][0]
    err = pred - z
    return {"rmse": float(np.sqrt((err ** 2).mean())),
            "mae": float(np.abs(err).mean()),
            "me": float(err.mean()), "pred": pred, "error": err}


def grid(x, y, nx=140, ny=140, pad=0.05):
    x, y = np.asarray(x, float), np.asarray(y, float)
    px, py = (x.max() - x.min()) * pad, (y.max() - y.min()) * pad
    gx = np.linspace(x.min() - px, x.max() + px, nx)
    gy = np.linspace(y.min() - py, y.max() + py, ny)
    XX, YY = np.meshgrid(gx, gy)
    return XX, YY, gx, gy