"""spatial_stats.py

Spatial dependence measures written against numpy and scipy only.

Everything here is deliberately visible. Commercial tools present Moran's I
and Getis-Ord Gi* as a menu item with the weights matrix chosen for you and
the multiple-comparison problem unmentioned. This module makes both of those
decisions explicit parameters, because both change the answer.

Contiguity is derived from shared polygon vertices, which works because the
shipped GeoJSON rounds every coordinate to a whole meter in EPSG:5070.
Neighboring polygons therefore share exact vertices and a dictionary finds
them without any topology library.
"""
from __future__ import annotations

import json
from collections import defaultdict
from pathlib import Path

import numpy as np
from scipy import stats

WEIGHT_SCHEMES = ("rook", "queen", "knn4", "knn8", "distance")


# --------------------------------------------------------------------------
# Geometry, read straight from the shipped GeoJSON.
# --------------------------------------------------------------------------
def read_polygons(path):
    obj = json.loads(Path(path).read_text())
    ids, rings = [], []
    for f in obj["features"]:
        ids.append(f["properties"]["id"])
        rings.append([np.asarray(r[0], dtype=float)
                      for r in f["geometry"]["coordinates"]])
    return ids, rings


def _ring_area_centroid(p):
    """Shoelace area and centroid of one closed ring."""
    x, y = p[:, 0], p[:, 1]
    cross = x[:-1] * y[1:] - x[1:] * y[:-1]
    a = cross.sum() / 2.0
    if abs(a) < 1e-9:
        return 0.0, p[:, 0].mean(), p[:, 1].mean()
    cx = ((x[:-1] + x[1:]) * cross).sum() / (6 * a)
    cy = ((y[:-1] + y[1:]) * cross).sum() / (6 * a)
    return a, cx, cy


def centroids(rings):
    out = np.zeros((len(rings), 2))
    for i, polys in enumerate(rings):
        tot = 0.0
        acc = np.zeros(2)
        for p in polys:
            a, cx, cy = _ring_area_centroid(p)
            tot += abs(a)
            acc += abs(a) * np.array([cx, cy])
        out[i] = acc / tot if tot else np.array([p[:, 0].mean(),
                                                 p[:, 1].mean()])
    return out


# --------------------------------------------------------------------------
# Weights. The choice among these is the analyst's and it changes the answer.
# --------------------------------------------------------------------------
def build_weights(rings, scheme="queen", *, band_m=40000, row_standardize=True):
    n = len(rings)
    W = np.zeros((n, n))
    if scheme in ("rook", "queen"):
        shared = defaultdict(set)
        for i, polys in enumerate(rings):
            for p in polys:
                for xy in map(tuple, p):
                    shared[xy].add(i)
        counts = defaultdict(int)
        for owners in shared.values():
            owners = sorted(owners)
            for a in range(len(owners)):
                for b in range(a + 1, len(owners)):
                    counts[(owners[a], owners[b])] += 1
        need = 1 if scheme == "queen" else 2
        for (a, b), c in counts.items():
            if c >= need:
                W[a, b] = W[b, a] = 1.0
    else:
        c = centroids(rings)
        d = np.sqrt(((c[:, None, :] - c[None, :, :]) ** 2).sum(-1))
        np.fill_diagonal(d, np.inf)
        if scheme.startswith("knn"):
            k = int(scheme[3:])
            for i in range(n):
                W[i, np.argsort(d[i])[:k]] = 1.0
            W = np.maximum(W, W.T)          # force symmetry
        else:
            W[d <= band_m] = 1.0
    if row_standardize:
        s = W.sum(1, keepdims=True)
        W = np.divide(W, s, out=np.zeros_like(W), where=s > 0)
    return W


# --------------------------------------------------------------------------
# Global and local statistics.
# --------------------------------------------------------------------------
def morans_i(x, W, permutations=999, seed=20260824):
    x = np.asarray(x, dtype=float)
    n = len(x)
    z = x - x.mean()
    S0 = W.sum()
    I = (n / S0) * (z @ W @ z) / (z @ z)
    EI = -1.0 / (n - 1)
    rng = np.random.default_rng(seed)
    sim = np.empty(permutations)
    for p in range(permutations):
        zp = rng.permutation(z)
        sim[p] = (n / S0) * (zp @ W @ zp) / (zp @ zp)
    # Pseudo p value, the inference a menu-driven tool usually hides.
    ge = (np.abs(sim - EI) >= abs(I - EI)).sum()
    return {"I": float(I), "EI": float(EI),
            "z_sim": float((I - EI) / sim.std()),
            "p_sim": float((ge + 1) / (permutations + 1))}


def local_morans(x, W, permutations=999, seed=20260824):
    x = np.asarray(x, dtype=float)
    n = len(x)
    z = (x - x.mean()) / x.std()
    lag = W @ z
    Ii = z * lag
    rng = np.random.default_rng(seed)
    sim = np.empty((permutations, n))
    for p in range(permutations):
        zp = rng.permutation(z)
        sim[p] = z * (W @ zp)
    p_sim = (np.abs(sim) >= np.abs(Ii)).sum(0)
    p_sim = (p_sim + 1) / (permutations + 1)
    quad = np.where(z > 0, np.where(lag > 0, "HH", "HL"),
                    np.where(lag > 0, "LH", "LL"))
    return {"Ii": Ii, "z": z, "lag": lag, "p": p_sim, "quadrant": quad}


def getis_ord(x, W):
    """Gi* z scores. W is row standardized on input, so rebuild binary."""
    x = np.asarray(x, dtype=float)
    n = len(x)
    B = (W > 0).astype(float) + np.eye(n)      # Gi* includes the focal unit
    xbar = x.mean()
    S = np.sqrt((x ** 2).mean() - xbar ** 2)
    num = B @ x - xbar * B.sum(1)
    den = S * np.sqrt((n * (B ** 2).sum(1) - B.sum(1) ** 2) / (n - 1))
    g = num / den
    return {"z": g, "p": 2 * (1 - stats.norm.cdf(np.abs(g)))}


def fdr_mask(p, q=0.05):
    """Benjamini-Hochberg. Returns a boolean mask of survivors."""
    p = np.asarray(p, dtype=float)
    m = len(p)
    order = np.argsort(p)
    thresh = q * np.arange(1, m + 1) / m
    passed = np.where(p[order] <= thresh)[0]
    mask = np.zeros(m, dtype=bool)
    if len(passed):
        mask[order[:passed[-1] + 1]] = True
    return mask