"""Height Above Nearest Drainage (HAND) flood mapping in Python: the River Severn from Atcham to Buildwas.

Downloads the Environment Agency 1 m LIDAR composite DTM at 5 m, conditions it with pysheds,
derives a drainage network, computes HAND, and turns it into inundation extents for a range of
stage heights. Each extent is then scored against the Environment Agency Flood Map for Planning
(Flood Zones 2 and 3) with the hit rate, false-alarm ratio and critical success index (CSI).

Data: EA LIDAR Composite DTM (WCS), EA Flood Map for Planning Flood Zones 2 and 3 (OGC API Features),
EA flood-monitoring API (Buildwas gauge) and OpenStreetMap rivers (Overpass API). No accounts needed.
Run with:  python hand_flood_mapping.py
"""
import io
import json
import platform
import time
from importlib.metadata import version
from pathlib import Path

import numpy as np

# pysheds 0.5 still calls np.in1d, which NumPy 2.4 removed. np.isin is the drop-in replacement.
if not hasattr(np, "in1d"):
    np.in1d = np.isin
import geopandas as gpd
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import rasterio
import requests
from matplotlib.colors import LightSource, ListedColormap
from matplotlib.lines import Line2D
from matplotlib.patches import Patch
from pyproj import Transformer
from pysheds.grid import Grid
from pysheds.sview import Raster, ViewFinder
from rasterio import features
from rasterio.transform import from_origin, rowcol
from scipy import ndimage
from shapely.geometry import LineString, box, shape
from shapely.ops import linemerge

HERE = Path(__file__).parent
OUT = HERE / "output"
DATA = HERE / "data"
OUT.mkdir(exist_ok=True)
DATA.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
DIRMAP = (64, 128, 1, 2, 4, 8, 16, 32)


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


section("versions")
print(f"python {platform.python_version()} | pysheds {version('pysheds')} | numpy {np.__version__} | "
      f"rasterio {rasterio.__version__} | geopandas {gpd.__version__}")

# ------------------------------------------------------------------ 1. study area
RES = 5.0                                              # metres: the 1 m DTM resampled by the WCS
X0, Y0, X1, Y1 = 350_000, 300_000, 366_000, 314_000   # 16 km x 14 km, British National Grid
EX0, EY0, EX1, EY1 = 352_000, 301_000, 365_000, 311_000   # evaluation box, 1-2 km inside the edges
W, H = int((X1 - X0) / RES), int((Y1 - Y0) / RES)
TRANSFORM = from_origin(X0, Y1, RES, RES)
TO_LL = Transformer.from_crs(27700, 4326, always_xy=True)
lon0, lat0, lon1, lat1 = TO_LL.transform_bounds(X0, Y0, X1, Y1)
section("study area")
print(f"DTM box {X0}-{X1} E, {Y0}-{Y1} N ({(X1 - X0) / 1000:.0f} km x {(Y1 - Y0) / 1000:.0f} km) | "
      f"{W} x {H} cells at {RES:.0f} m | evaluation box {(EX1 - EX0) / 1000:.0f} km x {(EY1 - EY0) / 1000:.0f} km")

# ------------------------------------------------------------------ 2. EA LIDAR composite DTM
WCS = ("https://environment.data.gov.uk/spatialdata/lidar-composite-digital-terrain-model-dtm-1m/wcs?"
       "service=WCS&version=2.0.1&request=GetCoverage"
       "&CoverageId=13787b9a-26a4-4775-8523-806d13af58fc__Lidar_Composite_Elevation_DTM_1m"
       "&format=image/tiff&subset=E({},{})&subset=N({},{})&SCALEFACTOR={}")
TILE = 2_000
dtm_path = DATA / "ea_dtm_severn_5m.tif"
section("EA LIDAR composite DTM")
if not dtm_path.exists():
    dtm = np.full((H, W), np.nan, dtype="float32")
    t0 = time.perf_counter()
    for ty in range(Y0, Y1, TILE):
        for tx in range(X0, X1, TILE):
            url = WCS.format(tx, tx + TILE, ty, ty + TILE, 1 / RES)
            for attempt in range(4):
                try:
                    r = requests.get(url, timeout=120)
                    r.raise_for_status()
                    break
                except requests.RequestException:
                    time.sleep(10 * (attempt + 1))
            with rasterio.MemoryFile(r.content) as mem, mem.open() as src:
                tile = src.read(1).astype("float32")
                tile[(tile < -100) | (tile > 2000)] = np.nan     # the WCS marks gaps with -3.4e38
                row, col = rowcol(TRANSFORM, src.bounds.left + RES / 2, src.bounds.top - RES / 2)
                dtm[row:row + tile.shape[0], col:col + tile.shape[1]] = tile[:H - row, :W - col]
    print(f"downloaded {((X1 - X0) // TILE) * ((Y1 - Y0) // TILE)} tiles in {time.perf_counter() - t0:.0f} s")
    with rasterio.open(dtm_path, "w", driver="GTiff", width=W, height=H, count=1, dtype="float32",
                       crs="EPSG:27700", transform=TRANSFORM, nodata=np.nan, compress="deflate") as dst:
        dst.write(dtm, 1)
with rasterio.open(dtm_path) as src:
    dtm = src.read(1).astype("float64")
gaps = np.isnan(dtm)
print(f"elevation {np.nanmin(dtm):.1f}-{np.nanmax(dtm):.1f} m | gaps {gaps.sum():,} cells ({gaps.mean():.3%})")
if gaps.any():   # fill small survey gaps with the nearest valid value so flow routing can cross them
    idx = ndimage.distance_transform_edt(gaps, return_distances=False, return_indices=True)
    dtm = dtm[tuple(idx)]
dem_path = DATA / "ea_dtm_severn_5m_filled.tif"
with rasterio.open(dem_path, "w", driver="GTiff", width=W, height=H, count=1, dtype="float32",
                   crs="EPSG:27700", transform=TRANSFORM, nodata=-9999, compress="deflate") as dst:
    dst.write(dtm.astype("float32"), 1)

# ------------------------------------------------------------------ 3. reference and context data
FZ_URL = ("https://environment.data.gov.uk/geoservices/datasets/04532375-a198-476e-985e-0579a0a11b47/ogc/features/v1/"
          f"collections/Flood_Zones_2_3_Rivers_and_Sea/items?f=application/json&bbox={lon0},{lat0},{lon1},{lat1}&limit=200")
OVERPASS = ["https://overpass-api.de/api/interpreter", "https://overpass.private.coffee/api/interpreter",
            "https://maps.mail.ru/osm/tools/overpass/api/interpreter"]


def get_json(url, cache, post=None):
    if cache.exists():
        return json.loads(cache.read_text())
    urls = OVERPASS if post else [url]
    for attempt in range(4):
        for u in urls:
            try:
                r = (requests.post(u, data={"data": post}, timeout=180, headers={"User-Agent": "technicalgis-tutorial"})
                     if post else requests.get(u, timeout=180))
                r.raise_for_status()
                out = r.json()
                cache.write_text(json.dumps(out))
                return out
            except (requests.RequestException, ValueError):
                continue
        time.sleep(15 * (attempt + 1))
    raise RuntimeError(f"download failed: {url}")


section("EA Flood Map for Planning: Flood Zones 2 and 3")
fz_cache = DATA / "ea_flood_zones_severn.geojson"
if not fz_cache.exists():
    feats, url = [], FZ_URL
    while url:                                   # the OGC API pages its results; follow the "next" links
        page = requests.get(url, timeout=180).json()
        feats += page["features"]
        url = next((l["href"] for l in page.get("links", []) if l["rel"] == "next"), None)
    fz_cache.write_text(json.dumps({"type": "FeatureCollection", "features": feats}))
fz = gpd.read_file(fz_cache).set_crs(4326, allow_override=True).to_crs(27700)
fz = fz.explode(index_parts=False)
fz = fz[fz.geom_type.isin(["Polygon", "MultiPolygon"])]          # a few features arrive as mixed collections
fz["geometry"] = fz.geometry.buffer(0).intersection(box(X0, Y0, X1, Y1))
fz = fz[~fz.geometry.is_empty]
print(f"{len(fz)} polygons | by zone and source: "
      + ", ".join(f"{z}/{s} {n}" for (z, s), n in fz.groupby(["flood_zone", "flood_source"]).size().items()))
fz3_geom = fz[fz.flood_zone == "FZ3"].geometry.union_all()
fz2_geom = fz[fz.flood_zone == "FZ2"].geometry.union_all()
print(f"FZ3 {fz3_geom.area / 1e6:.2f} km2 | FZ2 polygons {fz2_geom.area / 1e6:.2f} km2, "
      f"of which {fz2_geom.intersection(fz3_geom).area / 1e6:.2f} km2 overlaps FZ3")
FZ3 = features.rasterize([(fz3_geom, 1)], out_shape=(H, W), transform=TRANSFORM, dtype="uint8").astype(bool)
FZ2 = features.rasterize([(fz2_geom.union(fz3_geom), 1)], out_shape=(H, W), transform=TRANSFORM,
                         dtype="uint8").astype(bool)          # Flood Zone 2 extent = zone 2 plus zone 3

section("OpenStreetMap rivers and the Buildwas gauge")
osm = get_json(None, DATA / "osm_rivers_severn.json", post=f"""[out:json][timeout:120];
way["waterway"="river"]({lat0},{lon0},{lat1},{lon1});(._;>;);out;""")["elements"]
TO_BNG = Transformer.from_crs(4326, 27700, always_xy=True)
nodes = {e["id"]: TO_BNG.transform(e["lon"], e["lat"]) for e in osm if e["type"] == "node"}
rivers = {}
for e in osm:
    if e["type"] == "way":
        rivers.setdefault(e.get("tags", {}).get("name", "unnamed"), []).append(
            LineString([nodes[n] for n in e["nodes"]]))
rivers = {k: linemerge(v).intersection(box(X0, Y0, X1, Y1)) for k, v in rivers.items()}
print(", ".join(f"{k} {v.length / 1000:.1f} km" for k, v in sorted(rivers.items(), key=lambda kv: -kv[1].length)))
gauge = get_json("https://environment.data.gov.uk/flood-monitoring/id/stations/2134", DATA / "buildwas_station.json")
item = gauge["items"] if isinstance(gauge["items"], dict) else gauge["items"][0]
scale = item["stageScale"]
GX, GY = item["easting"], item["northing"]
print(f"{item['label']} gauge on the {item['riverName']} at {GX} E, {GY} N | gauge datum {scale['datum']} m AOD | "
      f"typical high {scale['typicalRangeHigh']} m | highest recent {scale['highestRecent']['value']} m "
      f"({scale['highestRecent']['dateTime'][:10]}) | maximum on record {scale['maxOnRecord']['value']} m "
      f"({scale['maxOnRecord']['dateTime'][:10]})")

# ------------------------------------------------------------------ 4. condition the DTM and route flow
section("conditioning and flow routing (pysheds)")
t0 = time.perf_counter()
grid = Grid.from_raster(str(dem_path))
dem = grid.read_raster(str(dem_path))
# Burn the mapped rivers 5 m into a copy of the DTM so that D8 routing follows the real channels across the
# flat floodplain. The burned copy is used for routing only; HAND is measured on the unburned DTM.
BURN_M = 5.0
river_px = features.rasterize([(g, 1) for g in rivers.values()], out_shape=(H, W), transform=TRANSFORM,
                              dtype="uint8").astype(bool)
# Round to 1 cm first: pysheds' resolve_flats can lift a flat above a neighbour that was only a fraction of a
# millimetre higher, which leaves new pits that HAND cannot drain through. LIDAR is only good to about 0.1 m anyway.
burned = Raster(np.round(np.where(river_px, np.asarray(dem) - BURN_M, np.asarray(dem)).astype("float64"), 2),
                viewfinder=grid.viewfinder)
pit_filled = grid.fill_pits(burned)
flooded = grid.fill_depressions(pit_filled)
inflated = grid.resolve_flats(flooded)
fdir = grid.flowdir(inflated, dirmap=DIRMAP)
acc = grid.accumulation(fdir, dirmap=DIRMAP)
fill_depth = np.asarray(flooded) - np.asarray(burned)
print(f"burned {river_px.sum():,} river cells by {BURN_M:g} m | {int((np.asarray(fdir)[1:-1, 1:-1] <= 0).sum())} unresolved "
      f"interior pits | fill, resolve flats, D8 and accumulation in "
      f"{time.perf_counter() - t0:.1f} s | {(fill_depth > 0.01).mean():.1%} of cells raised by depression filling, "
      f"deepest fill {fill_depth.max():.1f} m")

# A 16 km DTM cannot see the 2,000 km2 the Severn drains before it enters the box, so accumulation alone
# would only call it a stream several kilometres downstream. The drainage network is therefore the mapped
# rivers plus every derived channel draining at least STREAM_KM2 inside the box.
acc_km2 = np.asarray(acc) * RES ** 2 / 1e6
BOOL_VF = ViewFinder(affine=grid.affine, shape=grid.shape, crs=grid.crs, nodata=False)
STREAM_KM2 = 3.0          # the Flood Map for Planning covers watercourses draining roughly 3 km2 or more


def hand_for(threshold_km2):
    streams = (acc_km2 >= threshold_km2) | river_px
    h = grid.compute_hand(fdir, dem, Raster(streams, viewfinder=BOOL_VF), dirmap=DIRMAP)
    return streams, np.asarray(h, dtype="float64")


t0 = time.perf_counter()
streams, hand_raw = hand_for(STREAM_KM2)
print(f"streams draining >= {STREAM_KM2:g} km2: {streams.sum():,} cells, about {streams.sum() * RES / 1000:.0f} km "
      f"of channel | HAND computed in {time.perf_counter() - t0:.1f} s")

# ------------------------------------------------------------------ 5. HAND
section("height above nearest drainage")
neg = hand_raw < 0
hand = np.where(np.isfinite(hand_raw), np.clip(hand_raw, 0, None), np.nan)
ev = np.zeros((H, W), bool)
er0, ec0 = rowcol(TRANSFORM, EX0 + RES / 2, EY1 - RES / 2)
er1, ec1 = rowcol(TRANSFORM, EX1 - RES / 2, EY0 + RES / 2)
ev[er0:er1 + 1, ec0:ec1 + 1] = True
valid = ev & np.isfinite(hand)
print(f"evaluation box: {ev.sum():,} cells, {(ev & ~valid).sum():,} without a HAND value (they drain off the grid "
      f"before meeting a channel and count as dry)")
q = np.percentile(hand[valid], [10, 25, 50, 75, 90])
print("HAND percentiles 10/25/50/75/90: " + " / ".join(f"{v:.1f}" for v in q) + " m"
      f" | {neg[ev].mean():.2%} of cells sit below their drainage cell (set to 0)")
for label, zone in (("Flood Zone 3", FZ3), ("Flood Zone 2", FZ2), ("outside Flood Zone 2", ~FZ2)):
    m = valid & zone
    print(f"  {label:<21} {m.sum() * RES ** 2 / 1e6:6.2f} km2 | median HAND {np.median(hand[m]):5.1f} m | "
          f"90th percentile {np.percentile(hand[m], 90):5.1f} m")

# Relate HAND to the gauge: HAND is measured from the LIDAR water surface, not from the gauge datum.
gr, gc = rowcol(TRANSFORM, GX, GY)
near = streams[gr - 6:gr + 7, gc - 6:gc + 7]
wl_lidar = np.asarray(dem)[gr - 6:gr + 7, gc - 6:gc + 7][near].min()
stage_lidar = wl_lidar - scale["datum"]
GAUGE_LEVELS = {"typical high": scale["typicalRangeHigh"], "highest recent (2014)": scale["highestRecent"]["value"],
                "maximum on record (2000)": scale["maxOnRecord"]["value"]}
print(f"LIDAR water surface at Buildwas {wl_lidar:.2f} m AOD = stage {stage_lidar:.2f} m on the gauge")
for k, v in GAUGE_LEVELS.items():
    print(f"  {k:<26} stage {v:.2f} m -> {v - stage_lidar:.2f} m above the LIDAR water surface")

# ------------------------------------------------------------------ 6. inundation and agreement
section("inundation extents scored against the Flood Map for Planning")


def scores(pred, ref, mask):
    a = (pred & ref & mask).sum(); b = (pred & ~ref & mask).sum(); c = (~pred & ref & mask).sum()
    return {"area_km2": (pred & mask).sum() * RES ** 2 / 1e6, "hit": a / (a + c), "far": b / (a + b),
            "csi": a / (a + b + c), "bias": (a + b) / (a + c)}


STAGES = np.round(np.arange(0.5, 12.01, 0.25), 2)
table = {z: [scores(hand <= h, ref, ev) for h in STAGES] for z, ref in (("FZ3", FZ3), ("FZ2", FZ2))}
print(f"{'stage':>6} | {'area km2':>8} | {'FZ3 hit':>7} {'FAR':>5} {'CSI':>5} | {'FZ2 hit':>7} {'FAR':>5} {'CSI':>5}")
for i, h in enumerate(STAGES):
    if h in (1, 2, 3, 4, 5, 6, 7, 8, 10, 12):
        a, b = table["FZ3"][i], table["FZ2"][i]
        print(f"{h:5.1f}m | {a['area_km2']:8.2f} | {a['hit']:7.2f} {a['far']:5.2f} {a['csi']:5.2f} | "
              f"{b['hit']:7.2f} {b['far']:5.2f} {b['csi']:5.2f}")
best = {}
for z in table:
    i = int(np.argmax([s["csi"] for s in table[z]]))
    best[z] = STAGES[i]
    s = table[z][i]
    print(f"best match to {z}: HAND <= {STAGES[i]:.2f} m | CSI {s['csi']:.2f} | hit rate {s['hit']:.2f} | "
          f"false-alarm ratio {s['far']:.2f} | bias {s['bias']:.2f} | {s['area_km2']:.2f} km2 "
          f"(reference {(FZ3 if z == 'FZ3' else FZ2)[ev].sum() * RES ** 2 / 1e6:.2f} km2)")

section("sensitivity to the stream threshold")
sens = {}
for thr in (0.5, 1.0, 3.0, 10.0):
    s_t, h_t = hand_for(thr) if thr != STREAM_KM2 else (streams, hand_raw)
    h_t = np.where(np.isfinite(h_t), np.clip(h_t, 0, None), np.nan)
    res = [scores(h_t <= h, FZ3, ev) for h in STAGES]
    sens[thr] = [r["csi"] for r in res]
    i = int(np.argmax(sens[thr]))
    print(f"streams >= {thr:4.1f} km2: {s_t.sum() * RES / 1000:5.0f} km of channel | best FZ3 CSI {res[i]['csi']:.2f} "
          f"at HAND <= {STAGES[i]:.2f} m | false-alarm ratio {res[i]['far']:.2f}")

# ------------------------------------------------------------------ 7. figures
section("figures")
dem_a = np.asarray(dem)
ls = LightSource(azdeg=315, altdeg=45)
shade = ls.hillshade(dem_a, vert_exag=3, dx=RES, dy=RES)
FULL = [X0, X1, Y0, Y1]
EVAL = [EX0, EX1, EY0, EY1]
SPECTRAL = plt.get_cmap("Spectral_r")


def crop(a):
    return a[er0:er1 + 1, ec0:ec1 + 1]


def outline(ax, geom, **kw):
    for g in getattr(geom, "geoms", [geom]):
        ax.plot(*g.exterior.xy, **kw)


def rivers_on(ax, **kw):
    for g in rivers.values():
        for gg in getattr(g, "geoms", [g]):
            ax.plot(*gg.xy, **kw)


def tidy(ax, ext):
    ax.set_xlim(ext[0], ext[1]); ax.set_ylim(ext[2], ext[3]); ax.set_xticks([]); ax.set_yticks([])


# Figure 1: study area
fig, ax = plt.subplots(figsize=(9, 8))
ax.imshow(shade, extent=FULL, cmap="gray", vmin=0, vmax=1)
im = ax.imshow(dem_a, extent=FULL, cmap=SPECTRAL, alpha=0.55, vmin=35, vmax=150)
fz3_poly = [g for g in getattr(fz3_geom, "geoms", [fz3_geom]) if g.area > 2_000]
for g in fz3_poly:
    ax.fill(*g.exterior.xy, color="#2b83ba", alpha=0.35, lw=0)
outline(ax, gpd.GeoSeries([fz2_geom.union(fz3_geom)]).simplify(5).iloc[0], color="#2b83ba", lw=0.5)
rivers_on(ax, color="#08306b", lw=1.1)
ax.plot([EX0, EX1, EX1, EX0, EX0], [EY0, EY0, EY1, EY1, EY0], color="black", ls="--", lw=1)
ax.plot(GX, GY, "^", ms=9, mfc="white", mec="black")
ax.annotate("Buildwas gauge", (GX, GY), xytext=(-8, 10), textcoords="offset points", ha="right", fontsize=8,
            bbox=dict(boxstyle="square,pad=0.15", fc="white", ec="none", alpha=0.85))
for name, frac in (("River Severn", 0.62), ("River Tern", 0.35)):
    g = linemerge(rivers[name]) if rivers[name].geom_type == "MultiLineString" else rivers[name]
    g = max(getattr(g, "geoms", [g]), key=lambda x: x.length)
    pt = g.interpolate(frac, normalized=True)
    ax.annotate(name, (pt.x, pt.y), xytext=(6, 6), textcoords="offset points", fontsize=8, style="italic",
                bbox=dict(boxstyle="square,pad=0.15", fc="white", ec="none", alpha=0.85))
tidy(ax, FULL)
fig.colorbar(im, ax=ax, shrink=0.7, label="elevation (m AOD, EA LIDAR composite DTM)")
ax.legend(handles=[Line2D([], [], color="#08306b", lw=1.1, label="OSM rivers"),
                   Patch(fc="#2b83ba", alpha=0.35, label="Flood Zone 3 (EA Flood Map for Planning)"),
                   Line2D([], [], color="#2b83ba", lw=0.6, label="Flood Zone 2 outline"),
                   Line2D([], [], color="black", ls="--", lw=1, label=f"evaluation box ({(EX1 - EX0) // 1000} km x {(EY1 - EY0) // 1000} km)")],
          loc="lower left", fontsize=8, framealpha=0.9)
ax.set_title(f"River Severn, Shrewsbury to Buildwas: {(X1 - X0) // 1000} km x {(Y1 - Y0) // 1000} km at {RES:.0f} m")
fig.savefig(OUT / "01-study-area.png"); plt.close(fig)

# Figure 2: conditioning and drainage
fig, axs = plt.subplots(1, 3, figsize=(15, 5.2))
ax = axs[0]
ax.imshow(shade, extent=FULL, cmap="gray", vmin=0, vmax=1)
fd = np.where(fill_depth > 0.01, fill_depth, np.nan)
im = ax.imshow(fd, extent=FULL, cmap="magma_r", vmin=0, vmax=3)
fig.colorbar(im, ax=ax, shrink=0.75, orientation="horizontal", pad=0.03, label="depression fill depth (m, capped at 3)")
tidy(ax, FULL); ax.set_title("1. Burn mapped rivers, fill depressions")
ax = axs[1]
ax.imshow(shade, extent=FULL, cmap="gray", vmin=0, vmax=1)
la = ndimage.maximum_filter(np.log10(acc_km2 + 1e-6), size=7)   # thickened so 5 m channels survive downsampling
im = ax.imshow(np.where(la >= -1, la, np.nan), extent=FULL, cmap="viridis_r", vmin=-1, vmax=1.5)
fig.colorbar(im, ax=ax, shrink=0.75, orientation="horizontal", pad=0.03,
             label="log10 drainage area (km2), cells draining >= 0.1 km2")
tidy(ax, FULL); ax.set_title("2. Flow accumulation")
ax = axs[2]
ax.imshow(shade, extent=FULL, cmap="gray", vmin=0, vmax=1)
derived = ndimage.binary_dilation(streams & ~river_px, iterations=2)
mapped = ndimage.binary_dilation(river_px, iterations=3)
ax.imshow(np.where(derived, 1.0, np.nan), extent=FULL, cmap=ListedColormap(["#f46d43"]), vmin=0, vmax=1)
ax.imshow(np.where(mapped, 1.0, np.nan), extent=FULL, cmap=ListedColormap(["#08306b"]), vmin=0, vmax=1)
tidy(ax, FULL); ax.set_title(f"3. Drainage network ({streams.sum() * RES / 1000:.0f} km)")
spacer = fig.colorbar(im, ax=ax, shrink=0.75, orientation="horizontal", pad=0.03)
spacer.ax.set_visible(False)   # keeps the third map the same size as the other two
ax.legend(handles=[Patch(fc="#08306b", label="mapped rivers (OSM)"),
                   Patch(fc="#f46d43", label=f"derived channels draining >= {STREAM_KM2:g} km2")],
          loc="upper center", bbox_to_anchor=(0.5, -0.04), ncol=1, fontsize=9, frameon=False)
fig.savefig(OUT / "02-conditioning-and-drainage.png"); plt.close(fig)

# Figure 3: HAND
fig, ax = plt.subplots(figsize=(10, 7.4))
ax.imshow(crop(shade), extent=EVAL, cmap="gray", vmin=0, vmax=1)
im = ax.imshow(crop(hand), extent=EVAL, cmap=SPECTRAL, vmin=0, vmax=20, alpha=0.8)
ax.imshow(np.where(crop(ndimage.binary_dilation(streams, iterations=1)), 1.0, np.nan), extent=EVAL,
          cmap=ListedColormap(["#08306b"]), vmin=0, vmax=1)
tidy(ax, EVAL)
fig.colorbar(im, ax=ax, shrink=0.75, extend="max", label="height above nearest drainage (m)")
ax.set_title("HAND: height of every cell above the channel it drains to")
fig.savefig(OUT / "03-hand.png"); plt.close(fig)

# Figure 4: inundation at three stage heights
SHOW = (2.0, float(best["FZ3"]), 6.0)
fig, axs = plt.subplots(1, 3, figsize=(15, 4.6))
for ax, h in zip(axs, SHOW):
    i = int(np.argmin(np.abs(STAGES - h)))
    s = table["FZ3"][i]
    ax.imshow(crop(shade), extent=EVAL, cmap="gray", vmin=0, vmax=1)
    ax.imshow(np.where(crop(hand <= h), 1.0, np.nan), extent=EVAL, cmap=ListedColormap(["#3288bd"]),
              vmin=0, vmax=1, alpha=0.75)
    for g in fz3_poly:
        ax.plot(*g.exterior.xy, color="black", lw=0.6)
    tidy(ax, EVAL)
    ax.set_title(f"HAND <= {h:g} m: {s['area_km2']:.1f} km2 | CSI {s['csi']:.2f} against FZ3")
axs[0].legend(handles=[Patch(fc="#3288bd", alpha=0.75, label="HAND inundation extent"),
                       Line2D([], [], color="black", lw=0.8, label="Flood Zone 3 outline")],
              loc="lower left", fontsize=8, framealpha=0.9)
fig.savefig(OUT / "04-inundation-stages.png"); plt.close(fig)

# Figure 5: agreement scores against stage height
fig, axs = plt.subplots(1, 3, figsize=(15, 4.4))
for ax, z in zip(axs[:2], ("FZ3", "FZ2")):
    for key, col, lab in (("csi", "#d53e4f", "critical success index"), ("hit", "#3288bd", "hit rate"),
                          ("far", "#fdae61", "false-alarm ratio")):
        ax.plot(STAGES, [s[key] for s in table[z]], color=col, lw=2, label=lab)
    for (k, v), st in zip(GAUGE_LEVELS.items(), (":", "--", "-.")):
        ax.axvline(v - stage_lidar, color="grey", ls=st, lw=1, label=f"Buildwas {k}")
    ax.set_xlabel("stage height: HAND threshold (m)"); ax.set_ylim(0, 1.02); ax.grid(alpha=0.3)
    ax.set_title(f"Against Flood Zone {z[-1]} | best CSI {max(s['csi'] for s in table[z]):.2f} at {best[z]:g} m")
axs[0].legend(fontsize=7, loc="lower right")
ax = axs[2]
for (thr, csi), col in zip(sens.items(), ("#fdae61", "#f46d43", "#d53e4f", "#5e4fa2")):
    ax.plot(STAGES, csi, color=col, lw=2, label=f"channels draining >= {thr:g} km2")
ax.set_xlabel("stage height: HAND threshold (m)"); ax.set_ylabel("CSI against Flood Zone 3")
ax.set_ylim(0, 1.02); ax.grid(alpha=0.3); ax.legend(fontsize=7); ax.set_title("Sensitivity to the stream threshold")
fig.savefig(OUT / "05-agreement-scores.png"); plt.close(fig)

# Figure 6: where the best extent agrees and disagrees
hb = hand <= best["FZ3"]
cat = np.full((H, W), np.nan)
cat[hb & FZ3] = 0; cat[~hb & FZ3] = 1; cat[hb & ~FZ3] = 2
fig = plt.figure(figsize=(13, 5.6))
ax = fig.add_axes([0.0, 0.0, 0.58, 0.94])
ax.imshow(crop(shade), extent=EVAL, cmap="gray", vmin=0, vmax=1)
ax.imshow(crop(cat), extent=EVAL, cmap=ListedColormap(["#3288bd", "#d53e4f", "#fdae61"]), vmin=-0.5, vmax=2.5,
          interpolation="nearest")
tidy(ax, EVAL)
ax.legend(handles=[Patch(fc="#3288bd", label="hit: HAND extent and FZ3"),
                   Patch(fc="#d53e4f", label="miss: FZ3 only"),
                   Patch(fc="#fdae61", label="false alarm: HAND extent only")],
          loc="lower left", fontsize=8, framealpha=0.9)
ax.set_title(f"HAND <= {best['FZ3']:g} m against Flood Zone 3")
ax = fig.add_axes([0.66, 0.1, 0.33, 0.78])
bins = np.arange(0, 20.25, 0.25)
for label, zone, col in (("inside Flood Zone 3", FZ3, "#3288bd"), ("Flood Zone 2 only", FZ2 & ~FZ3, "#fdae61"),
                         ("outside Flood Zone 2", ~FZ2, "#bdbdbd")):
    ax.hist(hand[valid & zone], bins=bins, color=col, alpha=0.75, density=True, label=label)
ax.axvline(best["FZ3"], color="#d53e4f", lw=1.5, label=f"best threshold ({best['FZ3']:g} m)")
ax.set_xlabel("HAND (m)"); ax.set_ylabel("share of cells (density)"); ax.legend(fontsize=8)
ax.set_title("HAND inside and outside the flood zones")
fig.savefig(OUT / "06-agreement-map.png"); plt.close(fig)
print(f"saved 6 figures to {OUT}")