"""Kernel density estimation in Python for GIS: bandwidth selection (Scott, Silverman, likelihood
cross-validation, a fixed planning distance), kernels, sample weights, a georeferenced GeoTIFF heatmap
and a case-control relative-risk surface, on London plane trees in the London Borough of Camden.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/kernel-density-estimation-python/
Part 2 of the "Point pattern analysis in Python" series. It downloads the same data as Part 1, so it runs on its own.
Data: Trees In Camden (London Borough of Camden, OGL v3), ONS local authority boundary (OGL v3),
basemap tiles (c) Esri (World Light Gray Canvas).
"""
import sys
import warnings
from importlib.metadata import version
from pathlib import Path

import numpy as np
import pandas as pd
import geopandas as gpd
import shapely
import requests
import rasterio
from rasterio.transform import from_origin
import contextily as cx
import matplotlib
matplotlib.use("Agg")   # render to files only; no GUI backend
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm
from scipy import ndimage
from scipy.stats import gaussian_kde
from sklearn.model_selection import GridSearchCV, KFold
from sklearn.neighbors import KernelDensity

warnings.filterwarnings("ignore")
np.seterr(all="ignore")
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})

DATA = Path("data"); DATA.mkdir(exist_ok=True)
OUT = Path("output"); OUT.mkdir(exist_ok=True)
SEED = 42
NSIM = 199
CRS = "EPSG:27700"
SPECIES = "London plane"
LAD = "E09000007"     # Camden
RES = 25              # output grid cell size (m)
PLANNING_RADIUS = 400  # a 5-minute walk: a fixed, explainable search radius
HA = 1e4              # m2 per hectare

TREES_URL = "https://opendata.camden.gov.uk/api/views/csqp-kdss/rows.csv?accessType=DOWNLOAD"
ONS = "https://services1.arcgis.com/ESMARspQHYMw9BZ9/arcgis/rest/services"
LAD_URL = f"{ONS}/Local_Authority_Districts_May_2026_Boundaries_UK_BFC/FeatureServer/0/query"


def download(url, path, **params):
    if not path.exists():
        r = requests.get(url, params=params or None, timeout=600, headers={"User-Agent": "TechnicalGIS-tutorial"})
        r.raise_for_status()
        path.write_bytes(r.content)
    return path


def section(title):
    print(f"\n=== {title} ===")


section("versions")
print(f"python {sys.version.split()[0]} | scikit-learn {version('scikit-learn')} | scipy {version('scipy')} | "
      f"rasterio {version('rasterio')} | contextily {version('contextily')} | geopandas {version('geopandas')}")

# ---------------------------------------------------------------- the same points and window as Part 1
section("data (same cleaning as Part 1)")
df = pd.read_csv(download(TREES_URL, DATA / "trees_in_camden.csv"), low_memory=False)
snapshot = pd.to_datetime(df["Last Uploaded"], format="%d/%m/%Y %I:%M:%S %p").max()
df = df[df.Easting > 0]
boundary = gpd.read_file(download(LAD_URL, DATA / "camden_lad.geojson", where=f"LAD26CD='{LAD}'",
                                  outFields="LAD26CD,LAD26NM", outSR=27700, f="geojson")).set_crs(CRS, allow_override=True)
W = boundary.geometry.iloc[0]
shapely.prepare(W)
trees = df[shapely.contains_xy(W, df.Easting, df.Northing)].drop_duplicates(["Easting", "Northing", "Common Name"])
is_plane = (trees["Common Name"] == SPECIES).to_numpy()
planes = trees[is_plane]
xy = planes[["Easting", "Northing"]].to_numpy(float)
n = len(xy)
print(f"snapshot uploaded {snapshot:%d %B %Y}: {len(trees):,} council trees in the borough, {n:,} {SPECIES}s")

# a 25 m grid of cell centres over the borough, top row first (the order a GeoTIFF expects)
x0, y0 = np.floor(np.array(W.bounds[:2]) / RES) * RES
x1, y1 = np.ceil(np.array(W.bounds[2:]) / RES) * RES
xs = np.arange(x0 + RES / 2, x1, RES)
ys = np.arange(y1 - RES / 2, y0, -RES)
XX, YY = np.meshgrid(xs, ys)
inside = shapely.contains_xy(W, XX, YY)
cells = np.column_stack([XX[inside], YY[inside]])
EXTENT = (x0, x1, y0, y1)
print(f"grid: {len(xs)} x {len(ys)} cells of {RES} m, {inside.sum():,} inside the borough")


def to_grid(values):
    out = np.full(XX.shape, np.nan)
    out[inside] = values
    return out


def kde_sklearn(points, bandwidth, kernel="gaussian", weights=None):
    """Density per m2 (integrates to 1 over the plane) at every grid cell inside the borough."""
    kd = KernelDensity(bandwidth=bandwidth, kernel=kernel, rtol=1e-4).fit(points, sample_weight=weights)
    return to_grid(np.exp(kd.score_samples(cells)))


# ---------------------------------------------------------------- Step 1: bandwidth selection
section("bandwidth selection")
# Rules of thumb. In two dimensions Scott's factor n^(-1/(d+4)) and Silverman's (n(d+2)/4)^(-1/(d+4)) are identical.
scott = gaussian_kde(xy.T, bw_method="scott")
silver = gaussian_kde(xy.T, bw_method="silverman")
sd_xy = np.sqrt(np.diag(scott.covariance))
print(f"scipy gaussian_kde: Scott factor {scott.factor:.4f}, Silverman factor {silver.factor:.4f}")
print(f"  kernel standard deviation (from the data covariance): {sd_xy[0]:.0f} m east-west, "
      f"{sd_xy[1]:.0f} m north-south")
kd_scott = KernelDensity(bandwidth="scott").fit(xy)
print(f"scikit-learn KernelDensity(bandwidth='scott').bandwidth_ = {kd_scott.bandwidth_:.3f} "
      f"-> read as metres on raw coordinates: standardise first or pass a number")

# Likelihood cross-validation: choose h to maximise the log-likelihood of held-out trees
grid_h = np.geomspace(10, 1000, 25)
cv = GridSearchCV(KernelDensity(kernel="gaussian"), {"bandwidth": grid_h},
                  cv=KFold(5, shuffle=True, random_state=SEED))
cv.fit(xy)
h_cv = cv.best_params_["bandwidth"]
cv_ll = cv.cv_results_["mean_test_score"] / (n / 5)   # mean held-out log-likelihood per tree
print(f"likelihood cross-validation (5-fold, {len(grid_h)} bandwidths from 10 to 1,000 m): best h = {h_cv:.0f} m")
for h in (10, 21, 46, 100, 215, 464, 1000):
    i = np.argmin(np.abs(grid_h - h))
    print(f"  h = {grid_h[i]:6.0f} m: mean held-out log-likelihood {cv_ll[i]:.3f} per tree")

# A fixed planning distance: an Epanechnikov kernel that reaches exactly 400 m and is zero beyond
h_epan = PLANNING_RADIUS
print(f"fixed planning distance: Epanechnikov kernel, radius {h_epan} m "
      f"(same spread as a Gaussian with sd {h_epan / np.sqrt(6):.0f} m)")

dens = {
    f"Scott/Silverman (Gaussian, sd {sd_xy[0]:.0f} × {sd_xy[1]:.0f} m)": to_grid(scott(cells.T)),
    f"likelihood CV (Gaussian, sd {h_cv:.0f} m)": kde_sklearn(xy, h_cv),
    f"planning distance (Epanechnikov, radius {h_epan} m)": kde_sklearn(xy, h_epan, "epanechnikov"),
}
for name, d in dens.items():
    t = d * n * HA
    print(f"{name}: peak {np.nanmax(t):.1f} planes/ha, 99th percentile {np.nanpercentile(t, 99):.1f}, "
          f"share of the kernel mass inside the borough {np.nansum(d) * RES * RES:.3f}")

fig, axes = plt.subplots(1, 3, figsize=(13, 4.6))
for ax, (name, d) in zip(axes, dens.items()):
    t = d * n * HA
    im = ax.imshow(t, extent=EXTENT, cmap="viridis", vmin=0, vmax=np.nanpercentile(t, 99.5))
    fig.colorbar(im, ax=ax, shrink=0.7, orientation="horizontal", pad=0.03, label="planes per hectare", extend="max")
    boundary.boundary.plot(ax=ax, color="black", linewidth=0.6)
    ax.set_title(name, fontsize=9); ax.set_axis_off()
fig.suptitle("Same 3,292 trees, three bandwidths (each map has its own colour scale, capped at the 99.5th percentile)",
             y=0.98)
fig.savefig(OUT / "01-bandwidth-comparison.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(7, 3.8))
ax.plot(grid_h, cv_ll, "o-", color="#0072B2", ms=4)
ax.set_xscale("log")
ax.axvline(h_cv, color="#D55E00", lw=2, label=f"likelihood CV maximum: {h_cv:.0f} m")
ax.axvline(sd_xy.mean(), color="black", ls="--", lw=1.5, label=f"Scott/Silverman: about {sd_xy.mean():.0f} m")
ax.axvline(h_epan / np.sqrt(6), color="#E69F00", ls=":", lw=2,
           label=f"400 m Epanechnikov ≈ Gaussian sd {h_epan / np.sqrt(6):.0f} m")
ax.set_xlabel("Gaussian bandwidth h (standard deviation, m)")
ax.set_ylabel("held-out log-likelihood per tree")
ax.set_title("Likelihood cross-validation picks a small bandwidth for clustered street trees")
ax.legend(frameon=False, fontsize=8); ax.grid(alpha=0.3, which="both")
fig.tight_layout()
fig.savefig(OUT / "02-cv-bandwidth.png")
plt.close(fig)

# ---------------------------------------------------------------- Step 2: the kernel matters much less than h
section("kernel comparison at the same spread (sd 100 m per axis)")
SD = 100
kernels = {  # scikit-learn h giving each 2-D kernel a 100 m standard deviation along each axis
    "gaussian": SD, "epanechnikov": SD * np.sqrt(6), "tophat": SD * 2, "exponential": SD / np.sqrt(3)}
kd_maps = {k: kde_sklearn(xy, h, k) for k, h in kernels.items()}
ref = kd_maps["gaussian"][inside]
for k, h in kernels.items():
    m = kd_maps[k][inside]
    print(f"{k:>12}: h = {h:5.0f} m, peak {m.max() * n * HA:5.1f} planes/ha, "
          f"correlation with Gaussian r = {np.corrcoef(ref, m)[0, 1]:.3f}")
zoom = (528300, 530800, 181000, 183300)   # Bloomsbury, King's Cross and Somers Town
zx = (xy[:, 0] > zoom[0]) & (xy[:, 0] < zoom[1]) & (xy[:, 1] > zoom[2]) & (xy[:, 1] < zoom[3])
vmax_k = max(np.nanmax(m) for m in kd_maps.values()) * n * HA
fig, axes = plt.subplots(1, 4, figsize=(13, 3.9), layout="constrained")
for ax, (k, m) in zip(axes, kd_maps.items()):
    im = ax.imshow(m * n * HA, extent=EXTENT, cmap="viridis", vmin=0, vmax=vmax_k)
    ax.scatter(xy[zx, 0], xy[zx, 1], s=0.15, color="white", alpha=0.5, linewidths=0)
    ax.set_xlim(zoom[:2]); ax.set_ylim(zoom[2:])
    ax.set_title(f"{k} (h = {kernels[k]:.0f} m)", fontsize=9); ax.set_xticks([]); ax.set_yticks([])
fig.colorbar(im, ax=axes, shrink=0.9, label="London planes per hectare")
fig.suptitle("Four kernels with the same 100 m spread, south Camden (white dots: planes; white area: outside Camden)")
fig.savefig(OUT / "03-kernel-comparison.png")
plt.close(fig)

# ---------------------------------------------------------------- Step 3: sample weights
section("sample weights: crown area")
spread = pd.to_numeric(planes["Spread In Metres"], errors="coerce").to_numpy()
ok = spread > 0
crown = np.pi * (spread[ok] / 2) ** 2
print(f"{ok.sum():,} of {n:,} planes have a crown spread > 0 (median {np.median(spread[ok]):.0f} m); "
      f"total crown area {crown.sum() / HA:.1f} ha")
H_MAP = 100
unweighted = kde_sklearn(xy, H_MAP)
weighted = kde_sklearn(xy[ok], H_MAP, weights=crown)
u_map = unweighted * n * HA               # planes per hectare
w_map = weighted * crown.sum()            # m2 of plane crown per m2 of ground
print(f"Gaussian h = {H_MAP} m: peak {np.nanmax(u_map):.1f} planes/ha; peak crown density "
      f"{np.nanmax(w_map) * 100:.1f}% of the ground under plane crowns")
print(f"correlation between count and crown-area surfaces: r = {np.corrcoef(u_map[inside], w_map[inside])[0, 1]:.3f}")
iu, iw = np.nanargmax(u_map), np.nanargmax(w_map)
print(f"count peak at E {XX.flat[iu]:.0f} N {YY.flat[iu]:.0f}; crown-area peak at E {XX.flat[iw]:.0f} N {YY.flat[iw]:.0f} "
      f"({np.hypot(XX.flat[iu] - XX.flat[iw], YY.flat[iu] - YY.flat[iw]):,.0f} m apart)")

fig, axes = plt.subplots(1, 2, figsize=(11, 4.8))
im0 = axes[0].imshow(u_map, extent=EXTENT, cmap="viridis", vmin=0)
fig.colorbar(im0, ax=axes[0], shrink=0.75, label="London planes per hectare")
axes[0].set_title(f"Unweighted: every tree counts once (h = {H_MAP} m)")
im1 = axes[1].imshow(w_map * 100, extent=EXTENT, cmap="viridis", vmin=0)
fig.colorbar(im1, ax=axes[1], shrink=0.75, label="plane crown area, % of ground")
axes[1].set_title("Weighted by crown area: big trees count more")
for ax in axes:
    boundary.boundary.plot(ax=ax, color="black", linewidth=0.6); ax.set_axis_off()
fig.tight_layout()
fig.savefig(OUT / "04-weighted-kde.png")
plt.close(fig)

# ---------------------------------------------------------------- Step 4: edge correction, GeoTIFF export and a basemap
section("GeoTIFF export")
# Edge correction (Diggle 1985): divide by the share of each cell's kernel that falls inside the borough.
edge = ndimage.gaussian_filter(inside.astype(float), sigma=H_MAP / RES, mode="constant", cval=0.0)
corrected = np.where(inside, u_map / edge, np.nan)
print(f"edge correction factor 1/c(x): median {np.nanmedian(1 / edge[inside]):.2f}, "
      f"max {np.nanmax(1 / edge[inside]):.2f} (cells on the boundary)")
tif = OUT / "camden_london_plane_density.tif"
profile = dict(driver="GTiff", width=len(xs), height=len(ys), count=1, dtype="float32", crs=CRS,
               transform=from_origin(x0, y1, RES, RES), nodata=-9999.0, compress="deflate", tiled=True,
               blockxsize=256, blockysize=256)
with rasterio.open(tif, "w", **profile) as dst:
    dst.write(np.where(inside, corrected, -9999.0).astype("float32"), 1)
    dst.update_tags(units="trees per hectare", kernel="gaussian", bandwidth_m=str(H_MAP),
                    edge_correction="Diggle: divided by the kernel mass inside the borough",
                    source=f"Trees In Camden, London Borough of Camden, snapshot {snapshot:%Y-%m-%d}, OGL v3")
    dst.set_band_description(1, "London plane density (trees/ha)")
with rasterio.open(tif) as src:
    arr = src.read(1, masked=True)
    bounds = src.bounds
    print(f"wrote {tif}: {src.width} x {src.height} cells, {src.crs.to_string()}, {src.res[0]:.0f} m, "
          f"{src.dtypes[0]}, nodata {src.nodata:g}, {tif.stat().st_size / 1024:.0f} KB")
    print(f"  read back: max {arr.max():.1f} trees/ha, mean {arr.mean():.2f} trees/ha "
          f"(the borough average is {n / (W.area / HA):.2f})")

fig, ax = plt.subplots(figsize=(8, 7.5))
show = np.ma.masked_less(arr, 2)   # leave low densities transparent so the streets show through
im = ax.imshow(show, extent=(bounds.left, bounds.right, bounds.bottom, bounds.top), cmap="magma_r", alpha=0.75,
               vmin=2, vmax=float(np.percentile(arr.compressed(), 99.5)), zorder=2)
boundary.boundary.plot(ax=ax, color="black", linewidth=0.8, zorder=3)
cx.add_basemap(ax, crs=CRS, source=cx.providers.Esri.WorldGrayCanvas, zoom=14, attribution_size=6)
fig.colorbar(im, ax=ax, shrink=0.7, label="London planes per hectare (below 2: transparent)", extend="max")
ax.set_axis_off()
ax.set_title(f"camden_london_plane_density.tif on an Esri basemap (Gaussian h = {H_MAP} m, edge-corrected)")
fig.tight_layout()
fig.savefig(OUT / "05-geotiff-basemap.png")
plt.close(fig)

# ---------------------------------------------------------------- Step 5: relative risk (cases vs controls)
section("relative risk: London planes vs all other council trees")
H_RR = 150
all_xy = trees[["Easting", "Northing"]].to_numpy(float)
xedges = np.arange(x0, x1 + RES / 2, RES)
yedges = np.arange(y0, y1 + RES / 2, RES)


def binned_kde(points):
    """Fast binned Gaussian KDE: counts per 25 m cell smoothed with a Gaussian filter (trees per cell)."""
    h, _, _ = np.histogram2d(points[:, 1], points[:, 0], bins=[yedges, xedges])
    return ndimage.gaussian_filter(np.flipud(h), sigma=H_RR / RES, mode="constant", cval=0.0)


case = binned_kde(all_xy[is_plane])
control = binned_kde(all_xy[~is_plane])
n_case, n_ctrl = int(is_plane.sum()), int((~is_plane).sum())
# relative risk = ratio of the two densities (each scaled to integrate to 1), so the overall plane share cancels
valid = inside & (control / (RES * RES) * HA >= 5)   # at least 5 other trees per hectare for a stable ratio
log_rr = np.where(valid, np.log2((case / n_case) / (control / n_ctrl)), np.nan)
r_raw = np.corrcoef(case[inside], (case + control)[inside])[0, 1]
print(f"cases {n_case:,} planes, controls {n_ctrl:,} other trees; common Gaussian h = {H_RR} m")
print(f"correlation of the raw plane density with the density of all council trees: r = {r_raw:.2f}")
print(f"{valid.sum() / inside.sum():.0%} of the borough has >= 5 other trees per ha (elsewhere the ratio is masked)")
lo, hi = np.nanpercentile(log_rr, [1, 99])
print(f"log2 relative risk, 1st to 99th percentile: {lo:.2f} to {hi:.2f} "
      f"(planes {2 ** lo:.2f}x to {2 ** hi:.1f}x their borough-wide share)")

# Monte Carlo tolerance contours: shuffle the plane/other labels over the same tree locations
rng = np.random.default_rng(SEED)
total = case + control
target = np.nan_to_num(log_rr, nan=np.inf)
exceed = np.zeros_like(case)
for _ in range(NSIM):
    lab = np.zeros(len(all_xy), bool)
    lab[rng.choice(len(all_xy), n_case, replace=False)] = True
    c = binned_kde(all_xy[lab])
    exceed += np.log2((c / n_case) / ((total - c) / n_ctrl)) >= target
p_upper = (exceed + 1) / (NSIM + 1)
sig = valid & (p_upper <= 0.05)
print(f"pointwise p <= 0.05 (planes over-represented), {NSIM} random labellings: {sig.sum() / valid.sum():.1%} of "
      f"valid cells ({sig.sum() * RES * RES / 1e6:.2f} km2); about 5% would be flagged by chance alone")

fig, axes = plt.subplots(1, 3, figsize=(14, 4.9))
for ax, d, nm, lab in ((axes[0], case, f"{n_case:,} London planes", "planes per hectare"),
                       (axes[1], control, f"{n_ctrl:,} other council trees", "other trees per hectare")):
    im = ax.imshow(np.where(inside, d / (RES * RES) * HA, np.nan), extent=EXTENT, cmap="viridis", vmin=0)
    fig.colorbar(im, ax=ax, shrink=0.7, orientation="horizontal", pad=0.03, label=lab)
    ax.set_title(f"{nm} (h = {H_RR} m)", fontsize=9)
lim = float(np.nanmax(np.abs([lo, hi])))
gpd.GeoSeries([W], crs=CRS).plot(ax=axes[2], facecolor="#d0d0d0", edgecolor="none", zorder=0)
im = axes[2].imshow(np.clip(log_rr, -lim, lim), extent=EXTENT, cmap="PuOr_r", norm=TwoSlopeNorm(0, -lim, lim),
                    zorder=1)
axes[2].contour(XX, YY, sig.astype(float), levels=[0.5], colors="black", linewidths=0.6, zorder=2)
fig.colorbar(im, ax=axes[2], shrink=0.7, orientation="horizontal", pad=0.03, extend="both",
             label="log₂ relative risk (0 = planes at their borough-wide share)")
axes[2].set_title("Planes vs other trees; black outline: pointwise p ≤ 0.05\n"
                  "grey: fewer than 5 other trees per ha (masked)", fontsize=9)
for ax in axes:
    boundary.boundary.plot(ax=ax, color="black", linewidth=0.6, zorder=3); ax.set_axis_off()
fig.savefig(OUT / "06-relative-risk.png")
plt.close(fig)

# a coarse sanity check readers can reproduce: plane share by contract area
section("plane share by land owner")
share = trees.groupby("Contract Area")["Common Name"].apply(lambda s: (s == SPECIES).mean())
print("London plane share of council trees by contract area: "
      + ", ".join(f"{k} {v:.0%}" for k, v in share.sort_values(ascending=False).items()))
print(f"borough-wide share: {n_case / len(trees):.1%}")