"""Voronoi (Thiessen) polygons in Python with shapely and GeoPandas, on real UK points.

Part 1: Thiessen-weighted catchment rainfall for the River Eden at Sheepmount (Carlisle) from
Environment Agency tipping-bucket rain gauges, compared with a plain gauge average and IDW.
Part 2: Voronoi "nearest fire station" areas for Cumbria compared with network-based (drive time)
service areas, and the population that a straight-line Voronoi assigns to the wrong station.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/voronoi-thiessen-polygons-python/

Data: Environment Agency real-time flood-monitoring API (OGL v3); NRFA catchment boundary for
station 76007 (UKCEH, NRFA terms, attribution required); OpenStreetMap via Geofabrik (ODbL);
ONS Open Geography boundaries and centroids and Census 2021 TS001 via Nomis (OGL v3).
"""
import json
import sys
import time
import warnings
import zipfile
from importlib.metadata import version
from pathlib import Path

import numpy as np
import pandas as pd
import geopandas as gpd
import networkx as nx
import osmium
import osmnx as ox
import requests
import shapely
from pyproj import Transformer
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import dijkstra
import matplotlib
matplotlib.use("Agg")   # render to files only; no GUI backend
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm
from matplotlib.lines import Line2D

warnings.filterwarnings("ignore")
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9,
                     "figure.facecolor": "white", "axes.facecolor": "white"})

DATA = Path("data"); DATA.mkdir(exist_ok=True)
OUT = Path("output"); OUT.mkdir(exist_ok=True)
CRS = "EPSG:27700"          # British National Grid, metres: Voronoi needs a projected CRS
UA = {"User-Agent": "TechnicalGIS-tutorial (https://www.technicalgis.com)"}

EA = "https://environment.data.gov.uk/flood-monitoring"
START, END = "2026-10-06", "2026-11-02"      # 28 whole days; enddate is inclusive
NRFA_ID = 76007                              # Eden at Sheepmount, Carlisle
NRFA_WFS = ("https://hydrologygeoserver.ceh.ac.uk/geoserver/ows?service=WFS&version=2.0.0&request=GetFeature"
            f"&typename=NRFA:NRFA_catchments&CQL_FILTER=ID={NRFA_ID}&outputformat=application/json"
            "&srsname=EPSG:27700")
GAUGE_BUFFER_M = 15_000     # gauges up to 15 km outside the catchment still shape the cells inside it

ONS = "https://services1.arcgis.com/ESMARspQHYMw9BZ9/arcgis/rest/services"
LAD_URL = f"{ONS}/Local_Authority_Districts_May_2026_Boundaries_UK_BGC/FeatureServer/0/query"
OA_PWC_URL = f"{ONS}/OA_December_2021_EW_PWC_V4/FeatureServer/0/query"
OA_URL = f"{ONS}/Output_Areas_2021_EW_BGC_V2/FeatureServer/0/query"
TS001 = "https://www.nomisweb.co.uk/output/census/2021/census2021-ts001.zip"
CUMBRIA = ["E06000063", "E06000064"]        # Cumberland, Westmorland and Furness
PBF = "https://download.geofabrik.de/europe/united-kingdom/england/cumbria-latest.osm.pbf"
SNAP_KMH = 20
NOT_DRIVE = {"abandoned", "bridleway", "bus_guideway", "construction", "corridor", "cycleway", "elevator",
             "escalator", "footway", "no", "path", "pedestrian", "planned", "platform", "proposed", "raceway",
             "razed", "rest_area", "service", "services", "steps", "track"}
NOT_SERVICE = {"alley", "driveway", "emergency_access", "parking", "parking_aisle", "private"}

RED, BLUE, LGREEN, GREEN, GREY = "#d53e4f", "#3288bd", "#abdda4", "#1a9850", "#4d4d4d"
CELLS = ["#fdae61", "#abdda4", "#66c2a5", "#9ecae1", "#fee08b", "#f4a6c0", "#c2a5cf"]
T0 = time.perf_counter()


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


def timer(label, t):
    print(f"[timer] {label}: {time.perf_counter() - t:.1f} s")


def get_json(url, path=None, **kw):
    if path is not None and path.exists():
        return json.loads(path.read_text())
    r = requests.get(url, timeout=300, headers=UA, **kw)
    r.raise_for_status()
    j = r.json()
    if path is not None:
        path.write_text(json.dumps(j))
    return j


def download(url, path):
    if not path.exists():
        r = requests.get(url, timeout=600, headers=UA)
        r.raise_for_status()
        path.write_bytes(r.content)
    return path


def arcgis(url, path, **params):
    """Query an ArcGIS FeatureServer layer page by page and cache the result as GeoJSON."""
    if not path.exists():
        feats, offset = [], 0
        while True:
            q = {"outFields": "*", "outSR": 27700, "f": "geojson", "resultOffset": offset,
                 "resultRecordCount": 2000, "where": "1=1", **params}
            r = requests.post(url, data=q, timeout=300, headers=UA)
            r.raise_for_status()
            page = r.json()["features"]
            feats += page
            if len(page) < 2000:
                break
            offset += len(page)
        path.write_text(json.dumps({"type": "FeatureCollection", "features": feats}))
    return gpd.read_file(path).set_crs(CRS, allow_override=True)


def envelope(x0, y0, x1, y1):
    return {"geometry": f"{x0},{y0},{x1},{y1}", "geometryType": "esriGeometryEnvelope", "inSR": 27700,
            "spatialRel": "esriSpatialRelIntersects"}


def thiessen(points, boundary):
    """Voronoi cells for `points` (GeoSeries), clipped to `boundary`, returned in the same order."""
    cells = shapely.voronoi_polygons(shapely.MultiPoint(list(points)), extend_to=boundary, ordered=True)
    out = gpd.GeoSeries(list(cells.geoms), index=points.index, crs=points.crs)
    return out.intersection(boundary)


def style(ax, geom, pad=3000):
    x0, y0, x1, y1 = geom.bounds
    ax.set_xlim(x0 - pad, x1 + pad); ax.set_ylim(y0 - pad, y1 + pad)
    ax.set_aspect("equal"); ax.set_xticks([]); ax.set_yticks([]); ax.set_xlabel(""); ax.set_ylabel("")
    for s in ax.spines.values():
        s.set_color("#bbbbbb")


def save(fig, path):
    for ax in fig.axes:          # GeoPandas labels map axes with the CRS axis names ("Easting [metre]")
        if "[metre]" in ax.get_xlabel() or ax.get_xlabel() == "x":
            ax.set_xlabel(""); ax.set_ylabel("")
    fig.savefig(path)


def scalebar(ax, km=10, loc=(0.05, 0.05)):
    x0, x1 = ax.get_xlim(); y0, y1 = ax.get_ylim()
    x = x0 + loc[0] * (x1 - x0); y = y0 + loc[1] * (y1 - y0)
    ax.plot([x, x + km * 1000], [y, y], color="k", lw=2.5, solid_capstyle="butt")
    ax.text(x + km * 500, y + 0.012 * (y1 - y0), f"{km} km", ha="center", va="bottom", fontsize=8)


section("versions")
print(f"python {sys.version.split()[0]} | shapely {version('shapely')} | geopandas {version('geopandas')} | "
      f"osmnx {version('osmnx')} | scipy {version('scipy')} | numpy {version('numpy')} | pandas {version('pandas')}")

# ================================================================ PART 1: Thiessen catchment rainfall
# ---------------------------------------------------------------- Step 1: catchment and gauges
section("Step 1: catchment boundary and rain gauges")
t = time.perf_counter()
cat = get_json(NRFA_WFS, DATA / f"nrfa_{NRFA_ID}.geojson")
catchment = gpd.GeoDataFrame.from_features(cat["features"], crs=CRS)
basin = shapely.make_valid(catchment.geometry.iloc[0])
props = cat["features"][0]["properties"]
print(f"{props['STATION_NAME']} (NRFA {NRFA_ID}): polygon area {basin.area / 1e6:,.1f} km2 "
      f"(NRFA catchment area {props['CATCHMENT_AREA']:,.1f} km2)")

lon, lat = Transformer.from_crs(CRS, 4326, always_xy=True).transform(*basin.centroid.coords[0])
stations = get_json(f"{EA}/id/stations?parameter=rainfall&lat={lat:.4f}&long={lon:.4f}&dist=80&_limit=1000",
                    DATA / "ea_rain_stations.json")["items"]
to_bng = Transformer.from_crs(4326, CRS, always_xy=True)
rows = []
for s in stations:
    la, lo = s.get("lat"), s.get("long")
    if isinstance(la, list) or la is None:       # a few records carry two positions; take the first
        la, lo = (la or [None])[0], (lo or [None])[0]
    if la is None:
        continue
    x, y = to_bng.transform(lo, la)
    rows.append({"station": s["notation"], "gridref": s.get("gridReference", ""), "x": x, "y": y})
gauges = gpd.GeoDataFrame(rows, geometry=gpd.points_from_xy([r["x"] for r in rows], [r["y"] for r in rows]),
                          crs=CRS).drop_duplicates("station").set_index("station")
gauges = gauges[gauges.distance(basin) <= GAUGE_BUFFER_M]
print(f"EA rainfall stations within 80 km of the catchment centroid: {len(stations)}; "
      f"inside or within {GAUGE_BUFFER_M / 1000:.0f} km of the catchment: {len(gauges)}")

# ---------------------------------------------------------------- Step 2: readings and totals
# The real-time API only keeps about four weeks of readings. If the dates above have aged out and there
# is no cached copy, fall back to the most recent 28 whole days (your numbers will then differ from ours).
if not (DATA / f"ea_rain_{START}_{END}.csv").exists():
    probe = get_json(f"{EA}/id/stations/{gauges.index[0]}/readings?startdate={START}&enddate={START}&_limit=10")
    if not probe["items"]:
        end = pd.Timestamp.now(tz="UTC").normalize() - pd.Timedelta(days=1)
        START, END = (end - pd.Timedelta(days=27)).strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d")
        print(f"note: readings for the original dates are no longer online; using {START} to {END}")
section(f"Step 2: 15-minute readings, {START} to {END}")
expected = (pd.Timestamp(END) - pd.Timestamp(START)).days * 96 + 96
cache = DATA / f"ea_rain_{START}_{END}.csv"
if not cache.exists():
    recs = []
    for sid in gauges.index:
        items = get_json(f"{EA}/id/stations/{sid}/readings?startdate={START}&enddate={END}&_sorted&_limit=10000")["items"]
        for it in items:
            if "15_min" in it["measure"] and isinstance(it.get("value"), (int, float)):
                recs.append((sid, it["dateTime"], it["value"]))
    pd.DataFrame(recs, columns=["station", "dateTime", "mm"]).to_csv(cache, index=False)
readings = pd.read_csv(cache, dtype={"station": str}).drop_duplicates(["station", "dateTime"])
print(f"{len(readings):,} readings from {readings.station.nunique()} gauges (cached in {cache.name}); "
      f"max single 15-minute value {readings.mm.max():.1f} mm, negative values {(readings.mm < 0).sum()}")
summary = readings.groupby("station").agg(n=("mm", "size"), total=("mm", "sum"))
summary["complete"] = summary.n / expected
gauges = gauges.join(summary, how="inner")
dropped = gauges[gauges.complete < 0.95]
for sid, r in dropped.iterrows():
    print(f"  dropped {sid} ({r.gridref}): {r.complete:.0%} of the {expected:,} expected readings")
gauges = gauges[gauges.complete >= 0.95].copy()
gauges["inside"] = gauges.within(basin)
print(f"gauges kept (>=95% complete): {len(gauges)} ({gauges.inside.sum()} inside the catchment)")
print(f"28-day totals: min {gauges.total.min():.1f} mm, median {gauges.total.median():.1f} mm, "
      f"max {gauges.total.max():.1f} mm")

# Quality control: compare each gauge with its Delaunay neighbours (the gauges whose cells touch its cell)
gauges["max15"] = readings.groupby("station").mm.max()
edges0 = shapely.delaunay_triangles(shapely.MultiPoint(list(gauges.geometry)), only_edges=True)
pos = {(round(p.x, 1), round(p.y, 1)): s for s, p in gauges.geometry.items()}
nbrs = {s: set() for s in gauges.index}
for ln in edges0.geoms:
    a_, b_ = (pos[(round(x, 1), round(y, 1))] for x, y in ln.coords)
    nbrs[a_].add(b_); nbrs[b_].add(a_)
gauges["nbr_median"] = [gauges.total[list(nbrs[s])].median() for s in gauges.index]
gauges["ratio"] = gauges.total / gauges.nbr_median
print("highest totals relative to their Delaunay neighbours:")
print(gauges.sort_values("ratio", ascending=False).head(5)[["gridref", "total", "nbr_median", "ratio", "max15"]]
      .round(1).to_string())
suspect = (gauges.ratio > 2) & (gauges.max15 > 25)
flagged = gauges[suspect].copy()
raw_mm = float((thiessen(gauges.geometry, basin).area / basin.area * gauges.total).sum())
for sid, r in flagged.iterrows():
    print(f"  flagged {sid} ({r.gridref}): {r.total:.0f} mm against a neighbour median of {r.nbr_median:.0f} mm, "
          f"with 15-minute readings up to {r.max15:.1f} mm")
gauges = gauges[~suspect].copy()
print(f"Thiessen mean with the flagged gauges left in: {raw_mm:.1f} mm; {len(gauges)} gauges go forward")
timer("gauges and readings", t)

# ---------------------------------------------------------------- Step 3: Voronoi construction
section("Step 3: Delaunay triangles -> Voronoi cells -> clip")
mp = shapely.MultiPoint(list(gauges.geometry))
tri = shapely.delaunay_triangles(mp)
print(f"Delaunay triangulation: {len(tri.geoms)} triangles from {len(gauges)} points")


def circumcentre(poly):
    (ax_, ay), (bx, by), (cx, cy) = list(poly.exterior.coords)[:3]
    d = 2 * (ax_ * (by - cy) + bx * (cy - ay) + cx * (ay - by))
    ux = ((ax_**2 + ay**2) * (by - cy) + (bx**2 + by**2) * (cy - ay) + (cx**2 + cy**2) * (ay - by)) / d
    uy = ((ax_**2 + ay**2) * (cx - bx) + (bx**2 + by**2) * (ax_ - cx) + (cx**2 + cy**2) * (bx - ax_)) / d
    return ux, uy


centres = np.array([circumcentre(p) for p in tri.geoms])
raw = shapely.voronoi_polygons(mp, extend_to=basin.buffer(30_000), ordered=True)
raw = gpd.GeoSeries(list(raw.geoms), index=gauges.index, crs=CRS)
gauges["cell"] = thiessen(gauges.geometry, basin)
gauges["area_km2"] = gauges.cell.area / 1e6
gauges["weight"] = gauges.cell.area / basin.area
has_cell = gauges.weight > 0
print(f"Voronoi vertices are the triangle circumcentres: {len(centres)} of them")
print(f"cells that reach the catchment: {has_cell.sum()} of {len(gauges)} "
      f"({(has_cell & ~gauges.inside).sum()} from gauges outside it)")
print(f"clipped cells cover {gauges.cell.area.sum() / basin.area:.4%} of the catchment")

# ---------------------------------------------------------------- Step 4: Thiessen weighting
section("Step 4: Thiessen-weighted catchment rainfall")
w = gauges[has_cell].sort_values("weight", ascending=False)
thiessen_mm = float((w.weight * w.total).sum())
print(w[["gridref", "total", "area_km2", "weight"]].round({"total": 1, "area_km2": 1, "weight": 3})
      .to_string())
print(f"Thiessen catchment rainfall: {thiessen_mm:.1f} mm (= sum of weight x gauge total)")
print(f"largest single weight {w.weight.max():.1%}, smallest {w.weight.min():.2%}")

# ---------------------------------------------------------------- Step 5: compare methods
section("Step 5: Thiessen vs arithmetic mean vs IDW, and sensitivity")
arith = gauges.loc[gauges.inside, "total"].mean()
xs = np.arange(basin.bounds[0], basin.bounds[2], 500) + 250
ys = np.arange(basin.bounds[1], basin.bounds[3], 500) + 250
gx, gy = np.meshgrid(xs, ys)
inside = shapely.contains_xy(basin, gx, gy)
px, py, pv = gauges.geometry.x.values, gauges.geometry.y.values, gauges.total.values
d = np.hypot(gx[inside][:, None] - px[None], gy[inside][:, None] - py[None])
wts = 1 / np.maximum(d, 1) ** 2
idw_vals = (wts * pv).sum(1) / wts.sum(1)
idw_mm = idw_vals.mean()
idw_grid = np.full(gx.shape, np.nan); idw_grid[inside] = idw_vals
print(f"arithmetic mean of the {gauges.inside.sum()} gauges inside: {arith:.1f} mm")
print(f"IDW (power 2, 500 m grid, {inside.sum():,} cells, all {len(gauges)} gauges): {idw_mm:.1f} mm")
print(f"Thiessen: {thiessen_mm:.1f} mm")

loo = {}
for sid in w.index:
    rest = gauges.drop(sid)
    cells = thiessen(rest.geometry, basin)
    loo[sid] = float((cells.area / basin.area * rest.total).sum())
loo = pd.Series(loo)
print(f"leave-one-out Thiessen (drop each contributing gauge in turn): {loo.min():.1f} to {loo.max():.1f} mm")
print(f"  biggest change from dropping {loo.sub(thiessen_mm).abs().idxmax()} "
      f"({loo[loo.sub(thiessen_mm).abs().idxmax()] - thiessen_mm:+.1f} mm)")
rng = np.random.default_rng(42)
half = []
for _ in range(200):
    keep = rng.choice(gauges.index, size=len(gauges) // 2, replace=False)
    sub = gauges.loc[keep]
    cells = thiessen(sub.geometry, basin)
    half.append(float((cells.area / basin.area * sub.total).sum()))
half = np.array(half)
print(f"half the gauges (200 random draws of {len(gauges) // 2}): 5th-95th percentile "
      f"{np.percentile(half, 5):.1f} to {np.percentile(half, 95):.1f} mm")
volume = thiessen_mm / 1000 * basin.area
print(f"rain volume over the catchment: {volume / 1e6:,.0f} million m3 in 28 days")

# ---------------------------------------------------------------- Figures 1-4
section("figures, part 1")
inner = gauges[gauges.inside]
fig, ax = plt.subplots(figsize=(7.2, 7.6))
gpd.GeoSeries([basin]).plot(ax=ax, color="#f2f7f2", edgecolor=GREEN, lw=1.4)
gpd.GeoSeries([basin.buffer(GAUGE_BUFFER_M)]).boundary.plot(ax=ax, color="#999999", lw=0.8, ls="--")
sc = ax.scatter(gauges.geometry.x, gauges.geometry.y, c=gauges.total, cmap="Blues", vmin=0,
                vmax=gauges.total.max(), s=55, edgecolor="k", lw=0.6, zorder=3)
for sid, r in gauges.iterrows():
    ax.annotate(f"{r.total:.0f}", (r.geometry.x, r.geometry.y), xytext=(5, 4), textcoords="offset points",
                fontsize=7, color="#222222")
ax.scatter(flagged.geometry.x, flagged.geometry.y, marker="X", s=90, color=RED, edgecolor="k", lw=0.6, zorder=4)
for sid, r in flagged.iterrows():
    ax.annotate(f"{r.total:.0f}?", (r.geometry.x, r.geometry.y), xytext=(6, 4), textcoords="offset points",
                fontsize=8, color=RED, fontweight="bold")
ax.plot(props["EASTING"], props["NORTHING"], marker="v", color=RED, ms=9, mec="k", zorder=4)
ax.annotate("Sheepmount gauging station", (props["EASTING"], props["NORTHING"]), xytext=(-9, 9), ha="right",
            textcoords="offset points", fontsize=8, color=RED, fontweight="bold")
style(ax, basin.buffer(GAUGE_BUFFER_M), pad=1000); scalebar(ax)
cb = fig.colorbar(sc, ax=ax, shrink=0.6, pad=0.01); cb.set_label(f"Rain, {START} to {END} (mm)")
ax.legend(handles=[Line2D([], [], color=GREEN, lw=1.4, label=f"Eden catchment to Sheepmount ({basin.area / 1e6:,.0f} km²)"),
                   Line2D([], [], color="#999999", ls="--", lw=0.8, label="15 km search buffer"),
                   Line2D([], [], marker="o", ls="", mfc=BLUE, mec="k", label=f"EA rain gauge ({len(gauges)})"),
                   Line2D([], [], marker="X", ls="", mfc=RED, mec="k", ms=8,
                          label=f"failed the neighbour check ({len(flagged)})")],
          loc="upper left", fontsize=8, framealpha=0.95)
ax.set_title("EA tipping-bucket rain gauges and 28-day totals (mm)", fontsize=10)
save(fig, OUT / "01-gauges-and-catchment.png"); plt.close(fig)

fig, axs = plt.subplots(2, 2, figsize=(11, 11.4))
view = basin.buffer(GAUGE_BUFFER_M)
for ax in axs.flat:
    gpd.GeoSeries([basin]).boundary.plot(ax=ax, color=GREEN, lw=1.3, zorder=2)
    style(ax, view, pad=1000)
a, b, c, dd = axs.flat
a.scatter(gauges.geometry.x, gauges.geometry.y, s=22, color=BLUE, edgecolor="k", lw=0.5, zorder=3)
a.set_title(f"a) {len(gauges)} gauges and the catchment boundary", fontsize=10)
gpd.GeoSeries(list(tri.geoms)).boundary.plot(ax=b, color="#888888", lw=0.6)
b.scatter(gauges.geometry.x, gauges.geometry.y, s=22, color=BLUE, edgecolor="k", lw=0.5, zorder=3)
b.set_title(f"b) Delaunay triangulation ({len(tri.geoms)} triangles)", fontsize=10)
gpd.GeoSeries(list(tri.geoms)).boundary.plot(ax=c, color="#cccccc", lw=0.5)
raw.boundary.plot(ax=c, color=RED, lw=0.9)
c.scatter(centres[:, 0], centres[:, 1], s=8, color=RED, zorder=3)
c.scatter(gauges.geometry.x, gauges.geometry.y, s=22, color=BLUE, edgecolor="k", lw=0.5, zorder=3)
c.set_title("c) Circumcentres join up into Voronoi edges", fontsize=10)
cmap_rain = plt.get_cmap("Blues")
gpd.GeoDataFrame(gauges[has_cell], geometry="cell", crs=CRS).plot(
    ax=dd, column="total", cmap=cmap_rain, vmin=0, vmax=gauges.total.max(), edgecolor="white", lw=0.8)
raw.boundary.plot(ax=dd, color="#cccccc", lw=0.5, zorder=1)
dd.scatter(gauges.geometry.x, gauges.geometry.y, s=22, color=BLUE, edgecolor="k", lw=0.5, zorder=3)
dd.set_title(f"d) Clipped to the catchment: {has_cell.sum()} cells, shaded by gauge total", fontsize=10)
fig.tight_layout(); save(fig, OUT / "02-voronoi-construction.png"); plt.close(fig)

fig, (ax, bx) = plt.subplots(1, 2, figsize=(12, 6.6), gridspec_kw={"width_ratios": [1.15, 1]})
cg = gpd.GeoDataFrame(w, geometry="cell", crs=CRS)
cg.plot(ax=ax, column="total", cmap="Blues", vmin=0, vmax=gauges.total.max(), edgecolor="white", lw=1)
gpd.GeoSeries([basin]).boundary.plot(ax=ax, color=GREEN, lw=1.3)
for sid, r in cg.iterrows():
    if r.weight >= 0.02:
        p = r.cell.representative_point()
        ax.text(p.x, p.y, f"{r.weight:.0%}", ha="center", va="center", fontsize=7.5,
                color="white" if r.total > 0.6 * gauges.total.max() else "k", fontweight="bold")
ax.scatter(gauges.geometry.x, gauges.geometry.y, s=10, color="k", zorder=3)
style(ax, basin, pad=2000); scalebar(ax)
ax.set_title("a) Thiessen cells with their area weights", fontsize=10)
sm = plt.cm.ScalarMappable(cmap="Blues", norm=plt.Normalize(0, gauges.total.max()))
fig.colorbar(sm, ax=ax, shrink=0.6, pad=0.01).set_label("Gauge total (mm)")
top = w.head(15)[::-1]
bx.barh(top.gridref, top.weight * 100, color=[BLUE if i else LGREEN for i in top.inside], edgecolor="k", lw=0.4)
for i, (sid, r) in enumerate(top.iterrows()):
    bx.text(r.weight * 100 + 0.3, i, f"{r.total:.0f} mm", va="center", fontsize=7.5)
bx.set_xlabel("Share of catchment area (%)"); bx.set_xlim(0, top.weight.max() * 100 * 1.25)
bx.legend(handles=[plt.Rectangle((0, 0), 1, 1, fc=BLUE, ec="k", lw=0.4, label="gauge inside the catchment"),
                   plt.Rectangle((0, 0), 1, 1, fc=LGREEN, ec="k", lw=0.4, label="gauge outside it")],
          loc="lower right", fontsize=8)
bx.set_title(f"b) The 15 largest weights (Thiessen mean {thiessen_mm:.0f} mm)", fontsize=10)
for s in ("top", "right"):
    bx.spines[s].set_visible(False)
fig.tight_layout(); save(fig, OUT / "03-thiessen-weights.png"); plt.close(fig)

fig, (ax, bx) = plt.subplots(1, 2, figsize=(12.5, 6.2), gridspec_kw={"width_ratios": [1, 1.05]})
im = ax.imshow(idw_grid, origin="lower", extent=(xs[0] - 250, xs[-1] + 250, ys[0] - 250, ys[-1] + 250),
               cmap="Blues", vmin=0, vmax=gauges.total.max())
cg.boundary.plot(ax=ax, color=GREY, lw=0.5)
gpd.GeoSeries([basin]).boundary.plot(ax=ax, color=GREEN, lw=1.3)
ax.scatter(gauges.geometry.x, gauges.geometry.y, s=10, color="k", zorder=3)
style(ax, basin, pad=2000); scalebar(ax)
fig.colorbar(im, ax=ax, shrink=0.6, pad=0.01).set_label("IDW rainfall (mm)")
ax.set_title("a) IDW surface (power 2) with the Thiessen cells over it", fontsize=10)
labels = ["Thiessen before\nquality control", "Arithmetic mean\n(gauges inside)", "IDW\n(500 m grid)", "Thiessen"]
vals = [raw_mm, arith, idw_mm, thiessen_mm]
bx.bar(labels, vals, color=["#f4a582", LGREEN, "#9ecae1", BLUE], edgecolor="k", lw=0.5, width=0.55)
for i, v in enumerate(vals):
    bx.text(i, v + 2, f"{v:.0f} mm", ha="center", fontsize=9)
bx.errorbar(3.42, thiessen_mm, yerr=[[thiessen_mm - loo.min()], [loo.max() - thiessen_mm]], fmt="none",
            ecolor=RED, capsize=4, lw=1.5)
bx.text(3.42, loo.max() + 2, "drop any\none gauge", color=RED, fontsize=7.5, va="bottom", ha="center")
bx.errorbar(3.75, thiessen_mm, yerr=[[thiessen_mm - np.percentile(half, 5)],
                                     [np.percentile(half, 95) - thiessen_mm]], fmt="none", ecolor=GREY, capsize=4, lw=1.5)
bx.text(3.75, np.percentile(half, 5) - 2, "half the\ngauges\n(5-95%)", color=GREY, fontsize=7.5, va="top",
        ha="center")
bx.set_ylabel(f"Catchment rainfall, {START} to {END} (mm)"); bx.set_xlim(-0.5, 4.2)
bx.set_ylim(0, max(vals + [np.percentile(half, 95), loo.max()]) * 1.15)
for s in ("top", "right"):
    bx.spines[s].set_visible(False)
bx.set_title("b) Catchment averages from the same gauges", fontsize=10)
fig.tight_layout(); save(fig, OUT / "04-method-comparison.png"); plt.close(fig)
gauges.drop(columns=["cell"]).to_csv(OUT / "eden_gauge_weights.csv")

# ================================================================ PART 2: fire stations, Voronoi vs network
section("Step 6: Cumbria fire stations from OpenStreetMap")
t = time.perf_counter()
lads = arcgis(LAD_URL, DATA / "lads_cumbria.geojson", **envelope(290000, 460000, 400000, 590000))
county = shapely.make_valid(lads[lads.LAD26CD.isin(CUMBRIA)].union_all())
print(f"Cumbria (Cumberland + Westmorland and Furness): {county.area / 1e6:,.0f} km2")
pbf = download(PBF, DATA / "cumbria-latest.osm.pbf")
found = []
for o in osmium.FileProcessor(str(pbf), osmium.osm.NODE | osmium.osm.WAY).with_locations().with_filter(
        osmium.filter.TagFilter(("amenity", "fire_station"))):
    if o.is_node():
        lo, la = o.location.lon, o.location.lat
    else:
        xy = np.array([(n.lon, n.lat) for n in o.nodes if n.location.valid()])
        lo, la = xy.mean(0)
    found.append({"osm": f"{'n' if o.is_node() else 'w'}{o.id}", "name": o.tags.get("name", ""),
                  "operator": o.tags.get("operator", ""), "fstype": o.tags.get("fire_station:type", ""),
                  "lon": lo, "lat": la})
fs = gpd.GeoDataFrame(found, geometry=gpd.points_from_xy([f["lon"] for f in found], [f["lat"] for f in found]),
                      crs=4326).to_crs(CRS)
print(f"amenity=fire_station objects in the Geofabrik extract: {len(fs)}")
fs = fs[fs.within(county)]
# OSM is a volunteer map, so check it against the fire service's own list of 38 stations before trusting it
EXCLUDE = {"n986915522": "former Ulverston station (The Ellers), replaced by the Blue Light Hub in 2021"}
NAMES = {"n4362774327": "Workington", "w55730638": "Cockermouth", "w94399820": "Broughton",
         "w303483840": "Whitehaven", "w314846127": "Millom", "w325535002": "Millom", "w320689383": "Egremont",
         "w324483620": "Aspatria", "w346290116": "Shap", "w360994949": "Milnthorpe", "w1199436772": "Seascale",
         "w1355628909": "Maryport", "w1450009652": "Brampton"}   # unnamed in OSM

for osm_id, why in EXCLUDE.items():
    print(f"  removed {osm_id}: {why}")
airport = fs.fstype == "airport"
print(f"  removed {airport.sum()} airport fire service (fire_station:type=airport, Walney airfield)")
fs = fs[~fs.osm.isin(EXCLUDE) & ~airport].copy()
fs["name"] = fs.osm.map(NAMES).fillna(fs.name)
print(f"names added for {fs.osm.isin(NAMES).sum()} features mapped without one")
# a station is often mapped twice (a node and the building outline): merge features within 150 m
merged = []
for g in (gpd.GeoSeries([shapely.unary_union(shapely.buffer(list(fs.geometry), 75))]).explode(index_parts=False)):
    members = fs[fs.within(g)]
    best = members.sort_values("name", key=lambda s: s.str.len(), ascending=False).iloc[0]
    merged.append({**best.drop("geometry").to_dict(), "geometry": members.union_all().centroid})
stations = gpd.GeoDataFrame(merged, crs=CRS).reset_index(drop=True)
stations["name"] = stations.name.str.replace(r"(?i)\s*(community )?fire( and rescue)? station$", "", regex=True)
stations["name"] = stations.name.replace({"Barrow & South Area Office": "Barrow",
                                          "Furness Peninsula Blue Light Hub": "Ulverston"})
print(f"inside Cumbria after merging duplicates within 150 m: {len(stations)} stations")
print("names: " + ", ".join(sorted(stations.name)))
timer("stations", t)

section("Step 7: Voronoi 'nearest station' areas")
stations["cell"] = thiessen(stations.geometry, county)
areas = stations.cell.area / 1e6
print(f"Voronoi cell area: median {areas.median():,.0f} km2, smallest {areas.min():,.0f} km2 "
      f"({stations.name[areas.idxmin()]}), largest {areas.max():,.0f} km2 ({stations.name[areas.idxmax()]})")

section("Step 8: drive network and network-nearest station")
t = time.perf_counter()
osm_file = DATA / "drive_network.osm"
if not osm_file.exists():
    poly = gpd.GeoSeries([county.buffer(5000)], crs=CRS).to_crs(4326).iloc[0]
    keep_nodes, ways = set(), {}
    for wy in osmium.FileProcessor(str(pbf), osmium.osm.NODE | osmium.osm.WAY).with_locations():
        if not wy.is_way():
            continue
        tags = wy.tags
        hw = tags.get("highway")
        if (hw is None or hw in NOT_DRIVE or tags.get("area") == "yes" or tags.get("motor_vehicle") == "no"
                or tags.get("motorcar") == "no" or tags.get("access") == "private"
                or tags.get("service") in NOT_SERVICE):
            continue
        xy = np.array([(n.lon, n.lat) for n in wy.nodes])
        if shapely.contains_xy(poly, xy[:, 0], xy[:, 1]).any():
            ways[wy.id] = ([n.ref for n in wy.nodes], dict(tags))
            keep_nodes.update(n.ref for n in wy.nodes)
    with osmium.SimpleWriter(str(osm_file), overwrite=True) as writer:
        for n in osmium.FileProcessor(str(pbf), osmium.osm.NODE):
            if n.id in keep_nodes:
                writer.add_node(n)
        for wid, (refs, tags) in sorted(ways.items()):
            writer.add_way(osmium.osm.mutable.Way(id=wid, version=1, nodes=refs, tags=tags))
G = ox.graph_from_xml(osm_file, simplify=True)
G = ox.routing.add_edge_speeds(G)
G = ox.routing.add_edge_travel_times(G)
G = ox.truncate.largest_component(G, strongly=True)
G = ox.project_graph(G, to_crs=CRS)
nodes, edges = ox.graph_to_gdfs(G)
idx = {n: i for i, n in enumerate(nodes.index)}
e = edges.reset_index()[["u", "v", "travel_time"]].groupby(["u", "v"], as_index=False).min()
M = csr_matrix((e.travel_time.values, (e.u.map(idx).values, e.v.map(idx).values)), shape=(len(nodes),) * 2)
print(f"drive graph: {len(nodes):,} nodes, {len(edges):,} edges")
timer("drive graph", t)

t = time.perf_counter()
snode = ox.distance.nearest_nodes(G, stations.geometry.x, stations.geometry.y)
s_snap = np.hypot(nodes.loc[snode].geometry.x.values - stations.geometry.x.values,
                  nodes.loc[snode].geometry.y.values - stations.geometry.y.values) / (SNAP_KMH / 3.6)
# station -> every node, along one-way streets in the direction an appliance would drive
T = dijkstra(M, directed=True, indices=[idx[n] for n in snode]) + s_snap[:, None]
nearest_net = T.argmin(0)
print(f"travel times from {len(stations)} stations to {T.shape[1]:,} nodes: {T.size:,} values")
timer("dijkstra", t)

# ---------------------------------------------------------------- Step 9: population
section("Step 9: output areas and population")
pop = pd.read_csv(zipfile.ZipFile(download(TS001, DATA / "census2021-ts001.zip")).open(
    "census2021-ts001-oa.csv"), encoding="utf-8-sig")
pop = pop.rename(columns={"geography code": "OA21CD", "Residence type: Total; measures: Value": "pop"})
pwc = arcgis(OA_PWC_URL, DATA / "oa_pwc_cumbria.geojson", **envelope(*county.bounds))
pwc = pwc[pwc.within(county)].merge(pop[["OA21CD", "pop"]], on="OA21CD")
oas = arcgis(OA_URL, DATA / "oa_cumbria.geojson", **envelope(*county.bounds))
oas = oas[oas.OA21CD.isin(pwc.OA21CD)].merge(pwc[["OA21CD", "pop"]], on="OA21CD")
print(f"{len(pwc):,} output areas in Cumbria, population {pwc['pop'].sum():,}")

onode = ox.distance.nearest_nodes(G, pwc.geometry.x, pwc.geometry.y)
o_snap = np.hypot(nodes.loc[onode].geometry.x.values - pwc.geometry.x.values,
                  nodes.loc[onode].geometry.y.values - pwc.geometry.y.values) / (SNAP_KMH / 3.6)
Toa = T[:, [idx[n] for n in onode]] + o_snap[None]
pwc["net_station"] = Toa.argmin(0)
pwc["net_min"] = Toa.min(0) / 60
cells = gpd.GeoDataFrame({"vor_station": stations.index}, geometry=stations.cell, crs=CRS)
pwc = gpd.sjoin(pwc, cells, predicate="within", how="left").drop(columns="index_right")
pwc["vor_min"] = Toa[pwc.vor_station.values, np.arange(len(pwc))] / 60
pwc["extra_min"] = pwc.vor_min - pwc.net_min
mis = pwc[pwc.net_station != pwc.vor_station]
print(f"output areas whose Voronoi station is not their quickest station by road: {len(mis):,} "
      f"({len(mis) / len(pwc):.1%}), {mis['pop'].sum():,} people ({mis['pop'].sum() / pwc['pop'].sum():.1%})")
print(f"extra drive time if sent to the Voronoi station instead: median {mis.extra_min.median():.1f} min, "
      f"95th percentile {mis.extra_min.quantile(0.95):.1f} min, max {mis.extra_min.max():.1f} min")
print(f"people with more than 5 extra minutes: {mis.loc[mis.extra_min > 5, 'pop'].sum():,}")
print(f"drive time to the quickest station, population-weighted: median "
      f"{np.repeat(pwc.net_min, pwc['pop']).median():.1f} min; share within 10 min "
      f"{pwc.loc[pwc.net_min <= 10, 'pop'].sum() / pwc['pop'].sum():.1%}, within 20 min "
      f"{pwc.loc[pwc.net_min <= 20, 'pop'].sum() / pwc['pop'].sum():.1%}")
worst = mis.groupby("vor_station").agg(pop=("pop", "sum")).sort_values("pop", ascending=False).head(3)
for s_, r in worst.iterrows():
    print(f"  Voronoi cell of {stations.name[s_]}: {r['pop']:,} people are quicker to reach from another station")
oas = oas.merge(pwc[["OA21CD", "net_station", "vor_station", "extra_min"]], on="OA21CD")
net_areas = oas.dissolve("net_station")
shared = sum(stations.cell[i].intersection(net_areas.geometry.get(i, shapely.Polygon())).area
             for i in stations.index) / county.area
print(f"area where the Voronoi cell and the network service area agree: {shared:.1%} of the county")

# ---------------------------------------------------------------- Figures 5-6
section("figures, part 2")
adj = nx.Graph()
adj.add_nodes_from(stations.index)
for i in stations.index:
    for j in stations.index[stations.cell.intersects(stations.cell[i].buffer(50))]:
        if i != j:
            adj.add_edge(i, j)
colours = nx.coloring.greedy_color(adj, strategy="largest_first")
print(f"colours needed so that no two neighbouring cells match: {max(colours.values()) + 1}")
col = stations.index.map(lambda i: CELLS[colours[i] % len(CELLS)])

fig, axs = plt.subplots(1, 2, figsize=(12, 7.4))
for ax in axs:
    style(ax, county, pad=2000)
a, b = axs
gpd.GeoSeries(stations.cell, crs=CRS).plot(ax=a, color=list(col), edgecolor="white", lw=1.0)
a.set_title(f"a) Voronoi: nearest station in a straight line ({len(stations)} cells)", fontsize=10)
net_areas = net_areas.reindex(stations.index)
gpd.GeoSeries(net_areas.geometry, crs=CRS).plot(ax=b, color=list(col), edgecolor="none")
edges.plot(ax=b, color="white", lw=0.15, alpha=0.6)
gpd.GeoSeries(stations.cell, crs=CRS).boundary.plot(ax=b, color="k", lw=0.6)
b.set_title("b) Quickest station by road (output areas), Voronoi edges in black", fontsize=10)
for ax in axs:
    gpd.GeoSeries([county]).boundary.plot(ax=ax, color=GREY, lw=0.8)
    ax.scatter(stations.geometry.x, stations.geometry.y, s=26, marker="^", color=RED, edgecolor="k",
               lw=0.5, zorder=4)
scalebar(a, km=20); scalebar(b, km=20)
OFFSET = {"Whitehaven": (-4, -9, "right"), "Windermere": (3, -9, "left"), "Walney": (-4, -9, "right"),
          "Grange": (3, -9, "left"), "Carlisle West": (-4, 4, "right"), "Arnside": (4, -9, "left"),
          "Millom": (4, -9, "left")}
for _, r in stations.iterrows():
    dx, dy, ha = OFFSET.get(r["name"], (4, 3, "left"))
    a.annotate(r["name"], (r.geometry.x, r.geometry.y), xytext=(dx, dy), textcoords="offset points", fontsize=6,
               ha=ha)
fig.tight_layout(); save(fig, OUT / "05-voronoi-vs-network.png"); plt.close(fig)

fig, ax = plt.subplots(figsize=(8.4, 8.4))
gpd.GeoSeries(stations.cell, crs=CRS).plot(ax=ax, color="#f4f4f4", edgecolor="#bdbdbd", lw=0.6)
bins = [0, 2, 5, 10, 60]
cm = ListedColormap(["#fee08b", "#fdae61", "#f46d43", "#a50026"])
mo = oas[oas.net_station != oas.vor_station]
mo.plot(ax=ax, column="extra_min", cmap=cm, norm=BoundaryNorm(bins, cm.N), edgecolor="none")
gpd.GeoSeries([county]).boundary.plot(ax=ax, color=GREY, lw=0.8)
ax.scatter(stations.geometry.x, stations.geometry.y, s=26, marker="^", color=RED, edgecolor="k", lw=0.5, zorder=4)
style(ax, county, pad=2000); scalebar(ax, km=20)
ax.legend(handles=[plt.Rectangle((0, 0), 1, 1, fc=cm(i), label=f"{bins[i]}–{bins[i + 1]} min" if i < 3 else
                                 f"over {bins[i]} min") for i in range(4)] +
          [Line2D([], [], marker="^", ls="", mfc=RED, mec="k", label="fire station (OSM)")],
          title="Extra drive time to the\nVoronoi station", loc="upper right", fontsize=8, title_fontsize=8)
ax.set_title(f"Output areas the Voronoi gives to the wrong station: {mis['pop'].sum():,} people "
             f"({mis['pop'].sum() / pwc['pop'].sum():.1%})", fontsize=10)
save(fig, OUT / "06-mismatch.png"); plt.close(fig)
stations.drop(columns="cell").to_csv(OUT / "cumbria_fire_stations.csv", index=False)
pwc.drop(columns="geometry").to_csv(OUT / "oa_station_assignment.csv", index=False)
timer("total", T0)