"""Zonal statistics in Python: exactextract vs rasterstats on Cumbria's 2021 LSOAs.

Population (GHSL, Mollweide), elevation (Copernicus DEM, WGS84) and land cover (ESA WorldCover, WGS84)
summarised over ONS boundaries in British National Grid, with partial-pixel weighting, CRS alignment,
categorical fractions, a benchmark, GeoPackage export and a choropleth.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/zonal-statistics-python/
"""
import os
import sys
import time
import warnings
import zipfile
from importlib.metadata import version
from pathlib import Path

for var in ("PROJ_LIB", "PROJ_DATA", "GDAL_DATA"):   # stray installs (e.g. PostGIS) break the wheels' own copies
    os.environ.pop(var, None)

import numpy as np
import pandas as pd
import geopandas as gpd
import requests
import rasterio
import rasterio.merge
from rasterio.warp import transform_bounds, calculate_default_transform, reproject
from rasterio.enums import Resampling
from shapely.geometry import box, Point
import pyogrio
import pyproj
import pystac_client
import planetary_computer as pc
from exactextract import exact_extract
from rasterstats import zonal_stats
import matplotlib
matplotlib.use("Agg")   # render to files only
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

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)

ONS = ("https://services1.arcgis.com/ESMARspQHYMw9BZ9/arcgis/rest/services/"
       "Lower_layer_Super_Output_Areas_December_2021_Boundaries_EW_BFC_V10/FeatureServer/0/query")
# Cumbria's six former districts; LSOA 2021 names still start with them
DISTRICTS = ["Allerdale", "Barrow-in-Furness", "Carlisle", "Copeland", "Eden", "South Lakeland"]
NOMIS = "https://www.nomisweb.co.uk/api/v01/dataset/NM_2021_1.data.csv"   # Census 2021 TS001, usual residents
GHSL = ("https://jeodpp.jrc.ec.europa.eu/ftp/jrc-opendata/GHSL/GHS_POP_GLOBE_R2023A/"
        "GHS_POP_E2020_GLOBE_R2023A_54009_100/V1-0/tiles/GHS_POP_E2020_GLOBE_R2023A_54009_100_V1_0_R3_C18.zip")
STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
WORLDCOVER = {10: "Tree cover", 20: "Shrubland", 30: "Grassland", 40: "Cropland", 50: "Built-up",
              60: "Bare / sparse", 70: "Snow and ice", 80: "Water", 90: "Herbaceous wetland",
              95: "Mangroves", 100: "Moss and lichen"}


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


def timed(fn, repeats=3):
    """Best-of-n wall-clock time, plus the result of the last run."""
    best = np.inf
    for _ in range(repeats):
        t0 = time.perf_counter()
        res = fn()
        best = min(best, time.perf_counter() - t0)
    return best, res


section("versions")
print(f"python {sys.version.split()[0]} | exactextract {version('exactextract')} | rasterstats {version('rasterstats')}"
      f" | rasterio {version('rasterio')} (GDAL {rasterio.__gdal_version__}) | geopandas {version('geopandas')}")

# ---------------------------------------------------------------- Step 1: zones
section("Step 1: LSOA 2021 boundaries for Cumbria (ONS Open Geography Portal)")
lsoa_file = DATA / "cumbria_lsoa21.gpkg"
if not lsoa_file.exists():
    where = " OR ".join(f"LSOA21NM LIKE '{d} %'" for d in DISTRICTS)
    parts, offset = [], 0
    while True:
        r = requests.get(ONS, params={"where": where, "outFields": "LSOA21CD,LSOA21NM", "outSR": 27700,
                                      "f": "geojson", "resultOffset": offset, "resultRecordCount": 100}, timeout=120)
        r.raise_for_status()
        page = gpd.GeoDataFrame.from_features(r.json()["features"], crs=27700)
        parts.append(page)
        offset += len(page)
        if len(page) < 100:
            break
    pd.concat(parts, ignore_index=True).to_file(lsoa_file, layer="lsoa21")
lsoa = gpd.read_file(lsoa_file).sort_values("LSOA21CD").reset_index(drop=True)
lsoa["district"] = lsoa.LSOA21NM.str.rsplit(" ", n=1).str[0]
lsoa["area_km2"] = lsoa.area / 1e6
print(f"{len(lsoa)} LSOAs, CRS {lsoa.crs.to_string()}, total area {lsoa.area_km2.sum():,.0f} km2")
print(f"LSOA area: median {lsoa.area_km2.median():.2f} km2, smallest {lsoa.area_km2.min():.3f} km2, "
      f"largest {lsoa.area_km2.max():.0f} km2")
print(lsoa.groupby("district").size().to_string())

# Census 2021 usual residents for the same LSOAs, an independent check on the population raster
census_file = DATA / "census2021_ts001_lsoa.csv"
if not census_file.exists():
    r = requests.get(NOMIS, params={"date": "latest", "geography": "TYPE151", "measures": 20100,
                                    "c2021_restype_3": 0, "select": "geography_code,obs_value"}, timeout=300)
    r.raise_for_status()
    census_file.write_bytes(r.content)
census = pd.read_csv(census_file).rename(columns={"GEOGRAPHY_CODE": "LSOA21CD", "OBS_VALUE": "census_2021"})
lsoa = lsoa.merge(census, on="LSOA21CD", how="left")
print(f"Census 2021 usual residents in these LSOAs: {lsoa.census_2021.sum():,.0f}")

# ---------------------------------------------------------------- Step 2: rasters
section("Step 2: three rasters in three different CRSs")
pop_tif = DATA / "ghs_pop_e2020_cumbria_54009.tif"
if not pop_tif.exists():
    zpath = DATA / Path(GHSL).name
    if not zpath.exists():
        zpath.write_bytes(requests.get(GHSL, timeout=600).content)
    tif_name = [n for n in zipfile.ZipFile(zpath).namelist() if n.endswith(".tif")][0]
    with rasterio.open(f"/vsizip/{zpath.as_posix()}/{tif_name}") as src:
        b = transform_bounds(lsoa.crs, src.crs, *lsoa.total_bounds)
        win = rasterio.windows.from_bounds(b[0] - 2000, b[1] - 2000, b[2] + 2000, b[3] + 2000, src.transform)
        win = win.round_offsets().round_lengths()
        prof = src.profile | {"width": win.width, "height": win.height, "transform": src.window_transform(win),
                              "compress": "deflate", "tiled": True, "blockxsize": 256, "blockysize": 256}
        with rasterio.open(pop_tif, "w", **prof) as dst:
            dst.write(src.read(1, window=win), 1)

catalog = pystac_client.Client.open(STAC_URL, modifier=pc.sign_inplace)
bbox_ll = transform_bounds(lsoa.crs, "EPSG:4326", *lsoa.total_bounds)
bbox_ll = (bbox_ll[0] - 0.01, bbox_ll[1] - 0.01, bbox_ll[2] + 0.01, bbox_ll[3] + 0.01)


def mosaic_from_stac(collection, asset, out_path, **search):
    """Mosaic every STAC item covering Cumbria into one local, tiled GeoTIFF."""
    if out_path.exists():
        return
    items = catalog.search(collections=[collection], bbox=bbox_ll, **search).item_collection()
    hrefs = [it.assets[asset].href for it in items]
    print(f"  {collection}: mosaicking {len(hrefs)} tiles")
    arr, transform = rasterio.merge.merge(hrefs, bounds=bbox_ll)
    with rasterio.open(hrefs[0]) as src:
        prof = src.profile
    prof.update(driver="GTiff", width=arr.shape[2], height=arr.shape[1], transform=transform, compress="deflate",
                tiled=True, blockxsize=512, blockysize=512, BIGTIFF="IF_SAFER")
    with rasterio.open(out_path, "w", **prof) as dst:
        dst.write(arr)


dem_tif = DATA / "cop_dem_glo30_cumbria.tif"
lc_tif = DATA / "esa_worldcover_2021_cumbria.tif"
mosaic_from_stac("cop-dem-glo-30", "data", dem_tif)
mosaic_from_stac("esa-worldcover", "map", lc_tif, datetime="2021")   # the 2021 map (v200)

rasters = {"population": pop_tif, "elevation": dem_tif, "land cover": lc_tif}
for name, path in rasters.items():
    with rasterio.open(path) as src:
        print(f"{name:10s} {path.name:36s} {src.crs.to_string() if src.crs.to_epsg() else 'ESRI:54009 (Mollweide)':24s}"
              f" pixel {abs(src.res[0]):.6g} {'m' if src.crs.is_projected else 'deg'}  "
              f"{src.width} x {src.height}  {src.dtypes[0]}")
with rasterio.open(pop_tif) as src:
    POP_CRS = src.crs
with rasterio.open(dem_tif) as src:
    LL_CRS = src.crs

# ---------------------------------------------------------------- Step 3: CRS alignment
section("Step 3: CRS alignment")
# rasterstats never looks at a CRS: it reads the zones' BNG bounds (metres) as if they were degrees
try:
    rs_wrong = zonal_stats(lsoa.head(5), str(dem_tif), stats=["count", "mean"], nodata=-32767)
    print("rasterstats, BNG zones on the WGS84 DEM:", [s["count"] for s in rs_wrong])
except MemoryError as e:
    print("rasterstats, BNG zones on the WGS84 DEM:", type(e).__name__, str(e).split(" for an array")[0])
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    ee_wrong = exact_extract(str(dem_tif), lsoa.head(5), ["count", "mean"], output="pandas")
print("exactextract, BNG zones on the WGS84 DEM:", ee_wrong["count"].round(1).tolist(),
      "| warning:", caught[0].message if caught else "none")

# The fix: reproject the zones (cheap, lossless for statistics) and never the raster
zones = {"population": lsoa.to_crs(POP_CRS), "elevation": lsoa.to_crs(LL_CRS), "land cover": lsoa.to_crs(LL_CRS)}
for name, z in zones.items():
    with rasterio.open(rasters[name]) as src:
        assert z.crs == pyproj.CRS.from_user_input(src.crs), name
print("zones reprojected to each raster's CRS: OK")

# Why not reproject the raster instead? Resampling a count raster does not conserve the total.
with rasterio.open(pop_tif) as src:
    pop = src.read(1, masked=True).filled(0).astype("float64")
    t, w, h = calculate_default_transform(src.crs, "EPSG:27700", src.width, src.height, *src.bounds, resolution=100)
    for rs in [Resampling.nearest, Resampling.bilinear, Resampling.sum]:
        dst = np.zeros((h, w), "float64")
        reproject(pop, dst, src_transform=src.transform, src_crs=src.crs, dst_transform=t, dst_crs="EPSG:27700",
                  resampling=rs, src_nodata=0, dst_nodata=0)
        print(f"GHSL total in the window: {pop.sum():,.0f}; after warping to BNG with {rs.name:8s}: {dst.sum():,.0f} "
              f"({(dst.sum() / pop.sum() - 1) * 100:+.2f}%)")

# ---------------------------------------------------------------- Step 4: partial pixels vs centroids
section("Step 4: population per LSOA, coverage-weighted vs centroid")
zp = zones["population"]
ee_pop = exact_extract(str(pop_tif), zp, ["sum", "count"], include_cols=["LSOA21CD"], output="pandas")
rs_pop = zonal_stats(zp, str(pop_tif), stats=["sum", "count"], nodata=-200)
rs_pop_at = zonal_stats(zp, str(pop_tif), stats=["sum", "count"], nodata=-200, all_touched=True)
lsoa["pop_exact"] = ee_pop["sum"].values
lsoa["pop_centroid"] = [s["sum"] or 0 for s in rs_pop]
lsoa["pop_all_touched"] = [s["sum"] or 0 for s in rs_pop_at]
lsoa["cells_exact"] = ee_pop["count"].values
lsoa["cells_centroid"] = [s["count"] for s in rs_pop]
assert (ee_pop.LSOA21CD.values == lsoa.LSOA21CD.values).all()

print(f"{'method':22s} {'total':>9s} {'vs census':>10s} {'median |err| vs census':>23s}")
for col, label in [("pop_exact", "exactextract (weighted)"), ("pop_centroid", "rasterstats (centroid)"),
                   ("pop_all_touched", "rasterstats all_touched")]:
    err = (lsoa[col] / lsoa.census_2021 - 1) * 100
    print(f"{label:22s} {lsoa[col].sum():9,.0f} {(lsoa[col].sum() / lsoa.census_2021.sum() - 1) * 100:+9.1f}% "
          f"{err.abs().median():22.1f}%")
print(f"{'Census 2021':22s} {lsoa.census_2021.sum():9,.0f}")

diff = (lsoa.pop_centroid / lsoa.pop_exact - 1) * 100
lsoa["size_class"] = pd.cut(lsoa.area_km2, [0, 0.5, 2, 10, 1e4], labels=["< 0.5 km2", "0.5-2 km2", "2-10 km2", "> 10 km2"])
tab = lsoa.assign(d=diff.abs()).groupby("size_class", observed=True).agg(
    lsoas=("d", "size"), median_cells=("cells_exact", "median"), median_abs_diff_pct=("d", "median"),
    max_abs_diff_pct=("d", "max"))
print("\ncentroid vs coverage-weighted sum, by LSOA size:")
print(tab.round(1).to_string())
print(f"LSOAs where the two differ by more than 10%: {np.sum(diff.abs() > 10)}")

# ---------------------------------------------------------------- Step 5: elevation
section("Step 5: mean and maximum elevation (Copernicus DEM GLO-30)")
zd = zones["elevation"]
ee_dem = exact_extract(str(dem_tif), zd, ["mean", "max", "stdev"], output="pandas")
rs_dem = zonal_stats(zd, str(dem_tif), stats=["mean", "max", "std"], nodata=-32767)
lsoa["elev_mean"], lsoa["elev_max"], lsoa["elev_sd"] = ee_dem["mean"].values, ee_dem["max"].values, ee_dem["stdev"].values
d_mean = ee_dem["mean"].values - np.array([s["mean"] for s in rs_dem])
print(f"mean elevation, exactextract minus rasterstats: median |diff| {np.median(np.abs(d_mean)):.2f} m, "
      f"max |diff| {np.max(np.abs(d_mean)):.2f} m")
top = lsoa.loc[lsoa.elev_max.idxmax()]
print(f"highest point: {top.elev_max:.0f} m in {top.LSOA21NM} ({top.LSOA21CD}); "
      f"mean elevation range {lsoa.elev_mean.min():.0f}-{lsoa.elev_mean.max():.0f} m")

# ---------------------------------------------------------------- Step 6: categorical fractions
section("Step 6: land-cover fractions (ESA WorldCover 2021)")
zl = zones["land cover"]
ee_lc = exact_extract(str(lc_tif), zl, ["unique", "frac"], output="pandas")
rs_lc = zonal_stats(zl, str(lc_tif), categorical=True, nodata=0)

classes = sorted(WORLDCOVER)
ee_frac = pd.DataFrame([dict(zip(u.astype(int), f)) for u, f in zip(ee_lc["unique"], ee_lc["frac"])],
                       columns=classes).fillna(0)
rs_frac = pd.DataFrame([{k: v / sum(s.values()) for k, v in s.items()} for s in rs_lc], columns=classes).fillna(0)
present = [c for c in classes if ee_frac[c].max() > 0]
summary = pd.DataFrame({
    "class": [WORLDCOVER[c] for c in present],
    "area_weighted_%": [(ee_frac[c] * lsoa.area_km2).sum() / lsoa.area_km2.sum() * 100 for c in present],
    "median_abs_diff_pp": [(ee_frac[c] - rs_frac[c]).abs().median() * 100 for c in present],
    "max_abs_diff_pp": [(ee_frac[c] - rs_frac[c]).abs().max() * 100 for c in present]})
print("share of Cumbria and exactextract-vs-rasterstats difference per LSOA (percentage points):")
print(summary.round(2).to_string(index=False))
for c, col in [(10, "tree_pct"), (30, "grass_pct"), (50, "built_pct"), (80, "water_pct")]:
    lsoa[col] = ee_frac[c].values * 100
print(f"fractions sum to 1 for every LSOA: {np.allclose(ee_frac.sum(axis=1), 1)}")

# ---------------------------------------------------------------- Step 7: benchmark
section("Step 7: benchmark (best of 3 runs, all 319 LSOAs)")
bench = []
jobs = {
    "population sum (100 m)": (pop_tif, zp, ["sum"], dict(stats=["sum"], nodata=-200)),
    "elevation mean (30 m)": (dem_tif, zd, ["mean"], dict(stats=["mean"], nodata=-32767)),
    "land-cover fractions (10 m)": (lc_tif, zl, ["unique", "frac"], dict(categorical=True, nodata=0)),
}
for job, (path, z, ee_ops, rs_kw) in jobs.items():
    with rasterio.open(path) as src:
        cells = src.width * src.height
    t_ee, _ = timed(lambda: exact_extract(str(path), z, ee_ops, output="pandas"))
    t_ee_rs, _ = timed(lambda: exact_extract(str(path), z, ee_ops, output="pandas", strategy="raster-sequential"))
    t_rs, _ = timed(lambda: zonal_stats(z, str(path), **rs_kw))
    bench.append({"job": job, "raster_Mcells": cells / 1e6, "exactextract_s": t_ee,
                  "exactextract_raster_seq_s": t_ee_rs, "rasterstats_s": t_rs, "speed-up": t_rs / t_ee})
bench = pd.DataFrame(bench)
print(bench.round(2).to_string(index=False))

# ---------------------------------------------------------------- Step 8: GeoPackage
section("Step 8: export to GeoPackage")
gpkg = OUT / "cumbria_lsoa_zonal_stats.gpkg"
cols = ["LSOA21CD", "LSOA21NM", "district", "area_km2", "census_2021", "pop_exact", "pop_centroid",
        "elev_mean", "elev_max", "elev_sd", "tree_pct", "grass_pct", "built_pct", "water_pct", "geometry"]
lsoa["pop_density"] = lsoa.census_2021 / lsoa.area_km2
out = lsoa[cols[:-1] + ["pop_density", "geometry"]].copy()
out.to_file(gpkg, layer="lsoa21_zonal_stats", driver="GPKG")
summary.to_csv(OUT / "landcover_summary.csv", index=False)
print(pyogrio.list_layers(gpkg))
info = pyogrio.read_info(gpkg, layer="lsoa21_zonal_stats")
print(f"{info['features']} features, {len(info['fields'])} fields, CRS {info['crs']}, {gpkg.stat().st_size / 1e6:.1f} MB")

# ---------------------------------------------------------------- figures
section("figures")
districts = lsoa.dissolve("district")

# 1. Partial-pixel illustration for one small urban LSOA
small = lsoa[(lsoa.area_km2 > 0.15) & (lsoa.area_km2 < 0.3) & (lsoa.district == "Carlisle")].iloc[0]
one = zp[zp.LSOA21CD == small.LSOA21CD]
px = exact_extract(str(pop_tif), one, ["coverage", "center_x", "center_y", "values"], output="pandas").iloc[0]
geom = one.geometry.iloc[0]
fig, ax = plt.subplots(figsize=(6.2, 5.6))
cmap = plt.get_cmap("Spectral_r")
for cov, cx, cy in zip(px["coverage"], px["center_x"], px["center_y"]):
    ax.add_patch(Rectangle((cx - 50, cy - 50), 100, 100, fc=cmap(cov), ec="#555555", lw=0.4))
    if geom.contains(Point(cx, cy)):
        ax.plot(cx, cy, "o", ms=4, mfc="black", mec="white", mew=0.6)
gpd.GeoSeries([geom], crs=zp.crs).boundary.plot(ax=ax, color="black", lw=1.6)
ax.set_aspect("equal"); ax.set_xticks([]); ax.set_yticks([]); ax.set_xlabel(""); ax.set_ylabel("")
sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(0, 1))
fig.colorbar(sm, ax=ax, shrink=0.8, label="fraction of the 100 m pixel inside the LSOA")
ax.set_title(f"{small.LSOA21NM} ({small.area_km2:.2f} km²): {len(px['coverage'])} pixels touch it,\n"
             f"{int(small.cells_centroid)} pixel centres (dots) fall inside")
fig.savefig(OUT / "01-partial-pixels.png")
plt.close(fig)
print(f"figure 1: {small.LSOA21NM} {small.area_km2:.2f} km2, weighted {small.pop_exact:.0f}, "
      f"centroid {small.pop_centroid:.0f}, census {small.census_2021:.0f}")

# 2. Population against the census
fig, axes = plt.subplots(1, 2, figsize=(11, 4.6), layout="constrained")
for ax, col, label, c in [(axes[0], "pop_exact", "exactextract (coverage-weighted)", "#3288bd"),
                          (axes[0], "pop_centroid", "rasterstats (centroid)", "#f46d43")]:
    ax.scatter(lsoa.census_2021, lsoa[col], s=10, alpha=0.7, color=c, label=label)
lim = [0, lsoa[["census_2021", "pop_exact", "pop_centroid"]].max().max() * 1.05]
axes[0].plot(lim, lim, color="#888888", lw=0.8, ls="--")
axes[0].set_xlim(lim); axes[0].set_ylim(lim)
axes[0].set_xlabel("Census 2021 usual residents"); axes[0].set_ylabel("GHSL 2020 population summed per LSOA")
axes[0].legend(frameon=False); axes[0].set_title("Population per LSOA")
order = ["< 0.5 km2", "0.5-2 km2", "2-10 km2", "> 10 km2"]
data = [diff[lsoa.size_class == k].values for k in order]
axes[1].boxplot(data, tick_labels=[k.replace("km2", "km²") for k in order], showfliers=True,
                flierprops={"ms": 3}, medianprops={"color": "#d53e4f"})
axes[1].axhline(0, color="#888888", lw=0.8)
axes[1].set_ylabel("centroid sum vs coverage-weighted sum (%)"); axes[1].set_xlabel("LSOA area")
axes[1].set_title("Centroid error grows as zones shrink")
fig.savefig(OUT / "02-centroid-vs-weighted.png")
plt.close(fig)

# 3. Land-cover fractions
lc_colours = {10: "#006400", 20: "#ffbb22", 30: "#ffff4c", 40: "#f096ff", 50: "#fa0000", 60: "#b4b4b4",
              80: "#0064c8", 90: "#0096a0", 100: "#fae6a0"}
mix = ee_frac.assign(district=lsoa.district.values, area=lsoa.area_km2.values)
by_d = mix.groupby("district").apply(lambda g: (g[present].multiply(g.area, axis=0)).sum() / g.area.sum() * 100,
                                     include_groups=False)
fig, axes = plt.subplots(1, 2, figsize=(12, 4.4), layout="constrained", gridspec_kw={"width_ratios": [1.2, 1]})
left = np.zeros(len(by_d))
for c in [c for c in present if by_d[c].max() >= 0.05]:   # skip classes that would be invisible slivers
    axes[0].barh(by_d.index, by_d[c], left=left, color=lc_colours.get(c, "#999999"), label=WORLDCOVER[c],
                 edgecolor="white", lw=0.3)
    left += by_d[c].values
axes[0].invert_yaxis(); axes[0].set_xlim(0, 100); axes[0].set_xlabel("% of district area")
axes[0].legend(frameon=False, fontsize=7.5, ncol=4, loc="upper center", bbox_to_anchor=(0.5, -0.14))
axes[0].set_title("WorldCover 2021 mix, area-weighted from LSOA fractions")
dd = [(ee_frac[c] - rs_frac[c]).values * 100 for c in [10, 30, 50, 80]]
axes[1].boxplot(dd, tick_labels=["Tree cover", "Grassland", "Built-up", "Water"], flierprops={"ms": 3},
                medianprops={"color": "#d53e4f"})
axes[1].axhline(0, color="#888888", lw=0.8)
axes[1].set_ylabel("exactextract minus rasterstats (percentage points)")
axes[1].set_title("Categorical fractions: weighted vs centroid")
fig.savefig(OUT / "03-landcover-fractions.png")
plt.close(fig)

# 4. Benchmark
fig, ax = plt.subplots(figsize=(8, 3.6))
y = np.arange(len(bench)); hgt = 0.27
ax.barh(y - hgt, bench.exactextract_s, hgt, color="#3288bd", label="exactextract (feature-sequential)")
ax.barh(y, bench.exactextract_raster_seq_s, hgt, color="#66c2a5", label="exactextract (raster-sequential)")
ax.barh(y + hgt, bench.rasterstats_s, hgt, color="#f46d43", label="rasterstats")
for i, r in bench.iterrows():
    for off, v in [(-hgt, r.exactextract_s), (0, r.exactextract_raster_seq_s), (hgt, r.rasterstats_s)]:
        ax.text(v, i + off, f" {v:.2f} s", va="center", fontsize=7.5)
ax.set_yticks(y, bench.job); ax.invert_yaxis(); ax.set_xlabel("seconds (best of 3), 319 LSOAs")
ax.set_xlim(0, bench[["exactextract_s", "exactextract_raster_seq_s", "rasterstats_s"]].max().max() * 1.2)
ax.legend(frameon=False, fontsize=8, ncol=3, loc="upper center", bbox_to_anchor=(0.4, -0.16))
ax.set_title("Zonal statistics run time")
fig.savefig(OUT / "04-benchmark.png")
plt.close(fig)

# 5. Choropleths
fig, axes = plt.subplots(1, 3, figsize=(15, 6.2), layout="constrained")
specs = [("elev_mean", "Mean elevation (m)", "Spectral_r", None),
         ("tree_pct", "Tree cover (% of LSOA)", "Greens", None),
         ("pop_density", "Residents per km² (Census 2021, log scale)", "Spectral_r", "log")]
for ax, (col, title, cm, scale) in zip(axes, specs):
    kw = {"norm": matplotlib.colors.LogNorm(vmin=max(lsoa[col].min(), 1), vmax=lsoa[col].max())} if scale else {}
    lsoa.plot(col, ax=ax, cmap=cm, edgecolor="white", linewidth=0.15, legend=True,
              legend_kwds={"shrink": 0.6}, **kw)
    districts.boundary.plot(ax=ax, color="#333333", lw=0.6)
    ax.set_title(title); ax.set_axis_off()
fig.savefig(OUT / "05-choropleth.png")
plt.close(fig)
print("wrote 5 figures and", gpkg)