"""Watershed (catchment) delineation and stream networks from a DEM in Python with pysheds.

Delineates the catchment of the River Derwent above Yorkshire Bridge (NRFA gauging
station 28001, just below Ladybower Reservoir in the Peak District) from the
Copernicus GLO-30 DEM: fill pits and depressions, D8 flow direction, flow
accumulation, snap the pour point to the stream, catchment polygon and area,
stream extraction and Strahler order. The result is checked against the official
NRFA catchment boundary and the Environment Agency (EA) Water Framework Directive
(WFD) catchments.

Data: Copernicus DEM GLO-30 from the Microsoft Planetary Computer (no account needed),
NRFA station metadata and catchment boundary (UKCEH), EA Catchment Data Explorer.
Run with:  python watershed_delineation.py
"""
import json
import platform
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 pandas as pd
import planetary_computer
import pystac_client
import rasterio
import requests
import shapely
from matplotlib.colors import BoundaryNorm, LightSource, ListedColormap, LogNorm
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
from rasterio.warp import Resampling, reproject
from shapely.geometry import LineString, shape

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})
SPECTRAL = plt.get_cmap("Spectral_r")

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

# ------------------------------------------------------ 1. the gauge and the official answer
STATION = 28001                       # River Derwent at Yorkshire Bridge, below Ladybower dam
NRFA = "https://nrfaapps.ceh.ac.uk/nrfa/ws/station-info"
info_path = DATA / f"nrfa_{STATION}.json"
if not info_path.exists():
    r = requests.get(NRFA, params={"station": STATION, "format": "json-object", "fields": "all"}, timeout=60)
    r.raise_for_status()
    info_path.write_text(json.dumps(r.json()["data"][0], indent=1))
info = json.loads(info_path.read_text())
gx, gy = info["easting"], info["northing"]
print(f"\n=== NRFA station {STATION}: {info['name']} ===")
print(f"gauge at British National Grid {gx:.0f} E, {gy:.0f} N | catchment area {info['catchment-area']} km² "
      f"(IHDTM {info['ihdtm-catchment-area']} km²) | altitude {info['minimum-altitude']}–"
      f"{info['maximum-altitude']} m | drainage density {info['draindens']} km/km²")

bnd_path = DATA / f"nrfa_{STATION}_catchment.geojson"
if not bnd_path.exists():
    r = requests.get("https://hydrologygeoserver.ceh.ac.uk/geoserver/ows",
                     params={"service": "WFS", "version": "2.0.0", "request": "GetFeature",
                             "typename": "NRFA:NRFA_catchments", "CQL_FILTER": f"ID={STATION}",
                             "outputformat": "application/json", "srsname": "EPSG:27700"}, timeout=60)
    r.raise_for_status()
    bnd_path.write_text(r.text)
nrfa = gpd.read_file(bnd_path).set_crs(27700, allow_override=True)
print(f"NRFA boundary polygon: {nrfa.area.sum() / 1e6:.2f} km²")

# ------------------------------------------------------------------- 2. the DEM
RES = 30.0                                       # metres, close to GLO-30's native 1 arc-second
X0, Y0, X1, Y1 = 402_000, 381_000, 427_000, 402_000   # 25 km x 21 km of the Dark Peak in BNG
W, H = int((X1 - X0) / RES), int((Y1 - Y0) / RES)
dem_path = DATA / "glo30_derwent_bng.tif"
if not dem_path.exists():
    lon0, lat0, lon1, lat1 = Transformer.from_crs(27700, 4326, always_xy=True).transform_bounds(X0, Y0, X1, Y1)
    cat = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1",
                                    modifier=planetary_computer.sign_inplace)
    items = list(cat.search(collections=["cop-dem-glo-30"], bbox=[lon0, lat0, lon1, lat1]).items())
    print(f"\n=== Copernicus GLO-30 tiles: {', '.join(i.id for i in items)} ===")
    dem = np.full((H, W), np.nan, dtype="float32")
    for item in items:
        with rasterio.open(item.assets["data"].href) as src:
            part = np.full((H, W), np.nan, dtype="float32")
            reproject(rasterio.band(src, 1), part, dst_transform=from_origin(X0, Y1, RES, RES),
                      dst_crs="EPSG:27700", dst_nodata=np.nan, resampling=Resampling.bilinear)
            dem = np.where(np.isnan(dem), part, dem)
    with rasterio.open(dem_path, "w", driver="GTiff", width=W, height=H, count=1, dtype="float32",
                       crs="EPSG:27700", transform=from_origin(X0, Y1, RES, RES), nodata=-9999,
                       compress="deflate") as dst:
        dst.write(np.nan_to_num(dem, nan=-9999), 1)

grid = Grid.from_raster(str(dem_path))
dem = grid.read_raster(str(dem_path))
z = np.where(dem == -9999, np.nan, dem.astype("float64"))
print(f"\n=== DEM ===\n{W} x {H} cells of {RES:.0f} m (EPSG:27700) | elevation {np.nanmin(z):.1f}–{np.nanmax(z):.1f} m "
      f"| missing cells {np.isnan(z).sum()}")

# ------------------------------------------------------- 3. condition the DEM
print("\n=== condition the DEM ===")
pit_filled = grid.fill_pits(dem)
pits = np.asarray(pit_filled) > np.asarray(dem)          # copy now: the next step can modify in place
flooded = grid.fill_depressions(pit_filled)
inflated = grid.resolve_flats(flooded)
fill = np.asarray(flooded) - np.asarray(dem)
dep = (fill > 0) & ~pits
print(f"single-cell pits raised: {pits.sum():,}")
print(f"cells raised by depression filling (beyond the pits): {dep.sum():,} ({dep.mean():.1%}) | "
      f"max fill {fill.max():.1f} m | volume {fill.sum() * RES ** 2 / 1e6:.2f} million m³")
r_, c_ = np.unravel_index(np.argmax(fill), fill.shape)
print(f"deepest fill at ({X0 + (c_ + 0.5) * RES:.0f}, {Y1 - (r_ + 0.5) * RES:.0f}), "
      f"DEM {z[r_, c_]:.1f} m raised to {np.asarray(flooded)[r_, c_]:.1f} m")
flats = grid.detect_flats(flooded)
print(f"flat cells after filling: {np.asarray(flats).sum():,}; resolve_flats adds a tiny gradient so they drain")

# ------------------------------------------------- 4. flow direction and accumulation
print("\n=== D8 flow direction and accumulation ===")
DIRMAP = (64, 128, 1, 2, 4, 8, 16, 32)           # N, NE, E, SE, S, SW, W, NW (ESRI convention)
fdir = grid.flowdir(inflated, dirmap=DIRMAP)
acc = grid.accumulation(fdir, dirmap=DIRMAP)
acc_km2 = np.asarray(acc) * RES ** 2 / 1e6
print(f"largest upstream area in the grid: {acc_km2.max():.1f} km²")
codes, counts = np.unique(np.asarray(fdir), return_counts=True)
names = dict(zip(DIRMAP, ["N", "NE", "E", "SE", "S", "SW", "W", "NW"]))
print("flow directions: " + ", ".join(f"{names[c]} {n / counts.sum():.1%}" for c, n in zip(codes, counts) if c in names))

# ------------------------------------------------------------ 5. snap the pour point
print("\n=== pour point ===")
col_g, row_g = int((gx - X0) // RES), int((Y1 - gy) // RES)
print(f"gauge cell drains {acc_km2[row_g, col_g]:.2f} km² before snapping")
SNAP_KM2 = 50                                       # snap to a cell that drains at least 50 km²
sx, sy = grid.snap_to_mask(acc > SNAP_KM2 * 1e6 / RES ** 2, (gx, gy))
col_s, row_s = int((sx - X0) // RES), int((sy - Y1) // -RES)
print(f"snapped to ({sx:.0f}, {sy:.0f}), {np.hypot(sx - gx, sy - gy):.0f} m from the gauge; "
      f"that cell drains {acc_km2[row_s, col_s]:.1f} km²")

# --------------------------------------------------------------- 6. the catchment
print("\n=== catchment ===")
catch = grid.catchment(x=sx, y=sy, fdir=fdir, dirmap=DIRMAP, xytype="coordinate")
cmask = np.asarray(catch).astype(bool)
area = cmask.sum() * RES ** 2 / 1e6
polys = [shape(g) for g, v in features.shapes(cmask.astype("uint8"), mask=cmask,
                                                transform=grid.affine, connectivity=8)]
ours = gpd.GeoDataFrame({"name": ["Derwent at Yorkshire Bridge (GLO-30)"], "area_km2": [round(area, 2)]},
                        geometry=[shapely.union_all(polys)], crs=27700)
ref = nrfa.geometry.iloc[0]
inter, union = ours.geometry.iloc[0].intersection(ref).area, ours.geometry.iloc[0].union(ref).area
print(f"delineated area {area:.2f} km² vs NRFA {info['catchment-area']} km² "
      f"({(area / info['catchment-area'] - 1):+.1%}); vs NRFA boundary polygon {ref.area / 1e6:.2f} km²")
print(f"IoU with the NRFA boundary {inter / union:.3f} | only in ours {(ours.area.iloc[0] - inter) / 1e6:.2f} km² "
      f"| only in NRFA {(ref.area - inter) / 1e6:.2f} km²")
zc = z[cmask]
qs = np.percentile(zc, [0, 10, 50, 90, 100])
nrfa_q = [info["minimum-altitude"], info["10-percentile-altitude"], info["50-percentile-altitude"],
          info["90-percentile-altitude"], info["maximum-altitude"]]
print(pd.DataFrame({"GLO-30 (m)": qs.round(1), "NRFA (m)": nrfa_q},
                   index=["min", "10%", "median", "90%", "max"]).T.to_string())
print(f"mean elevation {zc.mean():.0f} m (NRFA ALTBAR {info['altbar']} m)")
cell_len = np.where(np.isin(np.asarray(fdir), [128, 2, 8, 32]), RES * np.sqrt(2), RES)
# distance_to_outlet counts cells unless you weight each step by its length (diagonals are 42.4 m)
dist_m = grid.distance_to_outlet(x=sx, y=sy, fdir=fdir, dirmap=DIRMAP, xytype="coordinate",
                                 weights=Raster(cell_len, viewfinder=grid.viewfinder))
dist_m = np.where(cmask, np.asarray(dist_m), np.nan)
print(f"longest flow path {np.nanmax(dist_m) / 1000:.1f} km; mean flow path {np.nanmean(dist_m) / 1000:.2f} km "
      f"(NRFA DPLBAR {info['dplbar']} km)")

# ---------------------------------------------------- 7. streams and Strahler order
print("\n=== stream network ===")
STREAM_KM2 = 0.5                                    # a channel starts where 0.5 km² drains to a cell
BOOL_VF = ViewFinder(affine=grid.affine, shape=grid.shape, crs=grid.crs, nodata=False)   # for mask rasters
streams = Raster((acc_km2 > STREAM_KM2) & cmask, viewfinder=BOOL_VF)
# pysheds also labels the cell just past the outlet (order + 1), so keep only stream cells
order = np.where(np.asarray(streams), np.asarray(grid.stream_order(fdir, streams, dirmap=DIRMAP)), 0)
branches = grid.extract_river_network(fdir, streams, dirmap=DIRMAP)
lines, orders = [], []
for feat in branches["features"]:
    coords = feat["geometry"]["coordinates"]
    if len(coords) < 2:
        continue
    xy = np.asarray(coords[:-1])            # the last vertex is the junction with the next segment
    o = order[((Y1 - xy[:, 1]) // RES).astype(int), ((xy[:, 0] - X0) // RES).astype(int)]
    lines.append(LineString(coords))
    orders.append(int(np.bincount(o.astype(int)).argmax()))
net = gpd.GeoDataFrame({"strahler": orders}, geometry=lines, crs=27700)
net["length_km"] = net.length / 1000
by_order = net.groupby("strahler").agg(segments=("length_km", "size"), length_km=("length_km", "sum"))
by_order["length_km"] = by_order["length_km"].round(1)
print(f"channel heads at {STREAM_KM2} km²: {len(net)} segments, {net['length_km'].sum():.1f} km, "
      f"highest Strahler order {net['strahler'].max()}")
print(by_order.to_string())
print(f"drainage density {net['length_km'].sum() / area:.3f} km/km² (NRFA DRAINDENS {info['draindens']})")
rows = []
for t in [0.1, 0.25, 0.5, 1, 2, 5]:
    s = (acc_km2 > t) & cmask
    s_order = grid.stream_order(fdir, Raster(s, viewfinder=BOOL_VF), dirmap=DIRMAP)
    # length of the stream cells, counting diagonal steps as RES * sqrt(2)
    rows.append((t, (s * cell_len).sum() / 1000 / area, int(np.asarray(s_order)[s].max())))
sens = pd.DataFrame(rows, columns=["threshold_km2", "drainage_density", "max_strahler"])
print("\nthreshold sensitivity:")
print(sens.round(3).to_string(index=False))

# ------------------------------------------------------- 8. EA WFD catchments
print("\n=== EA Water Framework Directive catchments ===")
WFD = {"GB104028057960": "river", "GB104028057950": "river", "GB104028057930": "river",
       "GB104028057940": "river", "GB104028057910": "river", "GB104028057900": "river",
       "GB104028057880": "river", "GB30432299": "lake", "GB30432359": "lake", "GB30432459": "lake"}
wfd_path = DATA / "ea_wfd_waterbodies.gpkg"
if not wfd_path.exists():
    parts = []
    for wb in WFD:
        r = requests.get(f"https://environment.data.gov.uk/catchment-planning/WaterBody/{wb}.geojson", timeout=60)
        r.raise_for_status()
        g = gpd.GeoDataFrame.from_features(r.json()["features"], crs=4326)
        g = g[g.geometry.type.isin(["Polygon", "MultiPolygon"])]
        parts.append(gpd.GeoDataFrame({"id": [wb], "name": [g["name"].iloc[0]]}, geometry=[g.union_all()], crs=4326))
    pd.concat(parts).to_crs(27700).to_file(wfd_path)
wfd = gpd.read_file(wfd_path)
wfd["kind"] = wfd["id"].map(WFD)
wfd["area_km2"] = (wfd.area / 1e6).round(2)
wfd["share_inside"] = (wfd.intersection(ours.geometry.iloc[0]).area / wfd.area).round(3)
print(wfd[["id", "name", "kind", "area_km2", "share_inside"]].to_string(index=False))
up = wfd[(wfd["kind"] == "river") & (wfd["share_inside"] > 0.5)]
up_u = up.union_all()
print(f"river water bodies more than half inside: {len(up)}, {up_u.area / 1e6:.2f} km² together; "
      f"IoU with our catchment {up_u.intersection(ours.geometry.iloc[0]).area / up_u.union(ours.geometry.iloc[0]).area:.3f}")

# ------------------------------------------------------------------- 9. export
gpkg = OUT / "derwent_catchment.gpkg"
gpkg.unlink(missing_ok=True)
ours.to_file(gpkg, layer="catchment")
net.to_file(gpkg, layer="streams")
nrfa.drop(columns=["id"], errors="ignore").to_file(gpkg, layer="nrfa_catchment")
wfd.to_file(gpkg, layer="ea_wfd_waterbodies")
gpd.GeoDataFrame({"name": ["gauge", "snapped pour point"]},
                 geometry=gpd.points_from_xy([gx, sx], [gy, sy]), crs=27700).to_file(gpkg, layer="pour_point")
with rasterio.open(OUT / "flow_accumulation_km2.tif", "w", driver="GTiff", width=W, height=H, count=1,
                   dtype="float32", crs="EPSG:27700", transform=grid.affine, compress="deflate") as dst:
    dst.write(acc_km2.astype("float32"), 1)
print(f"\n=== export ===\nwrote {gpkg.name} (layers: {', '.join(gpd.list_layers(gpkg)['name'])}) "
      f"and flow_accumulation_km2.tif")

# ------------------------------------------------------------------- 10. figures
ext = (X0, X1, Y0, Y1)
hs = LightSource(azdeg=315, altdeg=40).hillshade(np.nan_to_num(z, nan=np.nanmin(z)), vert_exag=2, dx=RES, dy=RES)
lakes = wfd[wfd["kind"] == "lake"]


def frame(ax, title, extent=ext):
    ax.set_title(title, loc="left")
    ax.set_xticks([]), ax.set_yticks([])
    ax.set_xlim(extent[0], extent[1]), ax.set_ylim(extent[2], extent[3])


def scalebar(ax, x0, y0, length=5000):
    ax.plot([x0, x0 + length], [y0, y0], color="black", lw=3, solid_capstyle="butt")
    ax.text(x0 + length / 2, y0 + 250, f"{length // 1000} km", ha="center", va="bottom", fontsize=8,
            bbox=dict(boxstyle="round,pad=0.15", fc="white", ec="none", alpha=0.8))


# 01: DEM, reservoirs and gauge
fig, ax = plt.subplots(figsize=(8.5, 6.6), layout="constrained")
ax.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
im = ax.imshow(z, cmap="terrain", extent=ext, alpha=0.55, vmin=100, vmax=650)
lakes.plot(ax=ax, color="#2b83ba", edgecolor="none")
nrfa.boundary.plot(ax=ax, color="black", lw=1.2, ls="--")
ax.plot(gx, gy, marker="v", ms=9, mfc="#d7191c", mec="white")
ax.annotate("NRFA 28001\nYorkshire Bridge", (gx, gy), (gx + 1200, gy - 1300), fontsize=8,
            bbox=dict(boxstyle="round,pad=0.25", fc="white", ec="none", alpha=0.85))
for _, r in lakes.iterrows():
    p = r.geometry.representative_point()
    ax.annotate(r["name"].replace(" Reservoir", ""), (p.x, p.y), (p.x + 1500, p.y), fontsize=7.5, color="#08306b",
                arrowprops=dict(arrowstyle="-", color="#08306b", lw=0.6),
                bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none", alpha=0.8))
ax.legend(handles=[Line2D([], [], color="black", ls="--", label="NRFA catchment boundary"),
                   Patch(color="#2b83ba", label="reservoirs (EA WFD lakes)")], loc="upper right", fontsize=8)
fig.colorbar(im, ax=ax, shrink=0.75, label="elevation (m, Copernicus GLO-30)")
frame(ax, "Upper Derwent valley, Peak District: GLO-30 DEM at 30 m")
scalebar(ax, X0 + 1000, Y0 + 1000)
fig.savefig(OUT / "01-dem-pour-point.png")
plt.close(fig)

# 02: what depression filling changed
fig, axes = plt.subplots(1, 2, figsize=(11, 4.6), gridspec_kw={"width_ratios": [1.3, 1]}, layout="constrained")
ax = axes[0]
ax.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
im = ax.imshow(np.ma.masked_less_equal(fill, 0), cmap="magma_r", norm=LogNorm(0.001, max(fill.max(), 1)),
               extent=ext, interpolation="nearest")
lakes.boundary.plot(ax=ax, color="#2b83ba", lw=0.8)
fig.colorbar(im, ax=ax, shrink=0.8, label="raised by filling (m, log scale)")
frame(ax, "Depression filling: cells raised so every cell can drain")
ax = axes[1]
ax.hist(fill[fill > 0], bins=np.logspace(-3, np.log10(max(fill.max(), 1)), 30), color="#f46d43", edgecolor="white")
ax.set_xscale("log")
ax.set_xlabel("fill depth (m)"), ax.set_ylabel("cells")
ax.set_title(f"{(fill > 0).sum():,} cells raised, median {np.median(fill[fill > 0]):.2f} m", loc="left")
fig.savefig(OUT / "02-depression-filling.png")
plt.close(fig)

# 03: flow direction and accumulation, zoomed on the gauge, plus the whole area
zoom = (sx - 2400, sx + 2400, sy - 1600, sy + 2000)
fig, axes = plt.subplots(1, 2, figsize=(11.5, 5), gridspec_kw={"width_ratios": [1, 1.25]}, layout="constrained")
ax = axes[0]
r0, r1 = int((Y1 - zoom[3]) / RES), int((Y1 - zoom[2]) / RES)
c0, c1 = int((zoom[0] - X0) / RES), int((zoom[1] - X0) / RES)
fd = np.asarray(fdir)[r0:r1, c0:c1]
idx = np.searchsorted(np.array(sorted(DIRMAP)), fd).clip(0, 7)
cmap8 = ListedColormap(plt.get_cmap("Spectral")(np.linspace(0, 1, 8)))
ax.imshow(idx, cmap=cmap8, extent=zoom, vmin=-0.5, vmax=7.5, interpolation="nearest")
step = 4
yy, xx = np.mgrid[r0:r1:step, c0:c1:step]
vec = {64: (0, 1), 128: (1, 1), 1: (1, 0), 2: (1, -1), 4: (0, -1), 8: (-1, -1), 16: (-1, 0), 32: (-1, 1)}
u = np.vectorize(lambda c: vec.get(c, (0, 0))[0])(np.asarray(fdir)[yy, xx])
v = np.vectorize(lambda c: vec.get(c, (0, 0))[1])(np.asarray(fdir)[yy, xx])
ax.quiver(X0 + (xx + 0.5) * RES, Y1 - (yy + 0.5) * RES, u, v, scale=40, width=0.003, color="black")
ax.plot(gx, gy, marker="v", ms=9, mfc="white", mec="black")
ax.plot(sx, sy, marker="o", ms=7, mfc="none", mec="black", mew=1.5)
ax.legend(handles=[Patch(color=cmap8(i), label=names[d]) for i, d in enumerate(sorted(DIRMAP))],
          loc="lower left", ncol=4, fontsize=7, title="D8 direction", title_fontsize=7, framealpha=0.9)
frame(ax, "D8 flow direction near the gauge", zoom)
ax = axes[1]
ax.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
im = ax.imshow(np.ma.masked_less(acc_km2, 0.05), cmap="viridis_r", norm=LogNorm(0.05, acc_km2.max()), extent=ext,
               interpolation="nearest")
ax.add_patch(plt.Rectangle((zoom[0], zoom[2]), zoom[1] - zoom[0], zoom[3] - zoom[2], fill=False, ec="#d7191c", lw=1.2))
fig.colorbar(im, ax=ax, shrink=0.8, label="upstream area (km², log scale)")
frame(ax, "Flow accumulation: area draining through each cell")
fig.savefig(OUT / "03-flow-direction-accumulation.png")
plt.close(fig)

# 04: our catchment vs the NRFA boundary and the EA WFD catchments
ref_r = features.rasterize([(ref, 1)], out_shape=(H, W), transform=grid.affine, dtype="uint8").astype(bool)
comp = np.zeros((H, W), dtype="uint8")
comp[cmask & ref_r], comp[cmask & ~ref_r], comp[~cmask & ref_r] = 1, 2, 3
cb = ours.total_bounds
cext = (cb[0] - 1500, cb[2] + 1500, cb[1] - 1500, cb[3] + 1500)
fig, axes = plt.subplots(1, 2, figsize=(11.5, 5.2), layout="constrained")
ax = axes[0]
ax.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
ax.imshow(comp, cmap=ListedColormap(["#00000000", "#abdda4", "#d7191c", "#2b83ba"]), extent=ext,
          vmin=-0.5, vmax=3.5, interpolation="nearest", alpha=0.85)
ax.plot(sx, sy, marker="o", ms=6, mfc="white", mec="black")
ax.legend(handles=[Patch(color="#abdda4", label="both"), Patch(color="#d7191c", label="GLO-30 only"),
                   Patch(color="#2b83ba", label="NRFA only")], loc="upper right", fontsize=8, framealpha=0.9)
frame(ax, f"GLO-30 catchment {area:.1f} km² vs NRFA {info['catchment-area']} km² (IoU {inter / union:.3f})", cext)
scalebar(ax, cext[0] + 800, cext[2] + 800)
ax = axes[1]
ax.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
rivers = wfd[wfd["kind"] == "river"].reset_index(drop=True)
rivers.plot(ax=ax, column="name", cmap="Spectral", alpha=0.45, edgecolor="#333333", lw=0.6)
lakes.plot(ax=ax, color="#2b83ba")
ours.boundary.plot(ax=ax, color="black", lw=1.6)
for _, r in rivers.iterrows():
    p = r.geometry.representative_point()
    ax.annotate(f"{r['name'].split(' (')[0]}\n{r['share_inside']:.0%} inside", (p.x, p.y), fontsize=6.5,
                ha="center", va="center", bbox=dict(boxstyle="round,pad=0.15", fc="white", ec="none", alpha=0.8))
frame(ax, "EA WFD river water-body catchments and our boundary", cext)
fig.savefig(OUT / "04-catchment-vs-nrfa-wfd.png")
plt.close(fig)

# 05: Strahler order
fig, ax = plt.subplots(figsize=(8, 6.6), layout="constrained")
ax.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
ours.boundary.plot(ax=ax, color="black", lw=1)
lakes.plot(ax=ax, color="#9ecae1")
omax = int(net["strahler"].max())
norm_o = BoundaryNorm(np.arange(0.5, omax + 1), SPECTRAL.N)
for o in range(1, omax + 1):
    net[net["strahler"] == o].plot(ax=ax, color=SPECTRAL(norm_o(o)), lw=0.5 + 0.6 * o)
ax.legend(handles=[Line2D([], [], color=SPECTRAL(norm_o(o)), lw=0.5 + 0.6 * o,
                          label=f"order {o}: {by_order.loc[o, 'length_km']} km") for o in range(1, omax + 1)],
          loc="lower left", fontsize=7.5, title=f"Strahler order (channels from {STREAM_KM2} km²)",
          title_fontsize=7.5, framealpha=0.9)
frame(ax, f"Stream network: {net['length_km'].sum():.0f} km of channel, Strahler order up to {omax}", cext)
scalebar(ax, cext[1] - 6000, cext[2] + 800)
fig.savefig(OUT / "05-strahler-order.png")
plt.close(fig)

# 06: threshold sensitivity and hypsometry
fig, axes = plt.subplots(1, 2, figsize=(11, 4), layout="constrained")
ax = axes[0]
ax.plot(sens["threshold_km2"], sens["drainage_density"], marker="o", color="#3288bd")
ax.axhline(info["draindens"], color="#d53e4f", ls="--", lw=1)
ax.text(sens["threshold_km2"].max(), info["draindens"] * 1.06, f"NRFA DRAINDENS {info['draindens']}", ha="right",
        color="#d53e4f", fontsize=8)
for _, r in sens.iterrows():
    ax.annotate(f"order {int(r['max_strahler'])}", (r["threshold_km2"], r["drainage_density"]),
                (5, 5), textcoords="offset points", fontsize=7)
ax.set_xscale("log")
ax.set_xlabel("channel initiation threshold (km² upstream)"), ax.set_ylabel("drainage density (km/km²)")
ax.set_title("Your stream network is a choice of threshold", loc="left")
ax = axes[1]
zs = np.sort(zc)[::-1]
ax.plot(np.arange(1, len(zs) + 1) / len(zs) * 100, zs, color="#3288bd", label="GLO-30 catchment")
ax.scatter([100, 90, 50, 10, 0], nrfa_q, color="#d53e4f", zorder=3, label="NRFA altitude percentiles")
ax.set_xlabel("% of catchment above this elevation"), ax.set_ylabel("elevation (m)")
ax.set_title("Hypsometric curve", loc="left")
ax.legend(fontsize=8)
fig.savefig(OUT / "06-threshold-hypsometry.png")
plt.close(fig)
print("figures written")