"""Canopy height mapping in Python: read the Meta/WRI global 1 m canopy height map from AWS open data,
compare it with the ETH Global Canopy Height 2020 (10 m) map, and validate both against a canopy
height model (CHM) made from Environment Agency LiDAR (first-return DSM minus DTM) in the Forest
of Dean, Gloucestershire.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/canopy-height-map-python/
Data: Meta and World Resources Institute High Resolution Canopy Height Maps (CC BY 4.0, AWS open
data); ETH Global Canopy Height 2020 (Lang et al. 2023, CC BY 4.0); Environment Agency LIDAR
Composite DSM and DTM 1 m (Open Government Licence v3.0); Natural Earth (public domain).
"""
import json
import sys
import urllib.parse
import warnings
from importlib.metadata import version
from pathlib import Path

import numpy as np
import pandas as pd
import requests
import rasterio
from rasterio.crs import CRS
from rasterio.enums import Resampling
from rasterio.transform import from_origin
from rasterio.warp import reproject, transform_bounds
from rasterio.windows import from_bounds
from scipy import ndimage
from shapely.geometry import box, shape
import matplotlib
matplotlib.use("Agg")   # render to files only; no GUI backend
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm, ListedColormap, TwoSlopeNorm
from matplotlib.patches import Rectangle

warnings.filterwarnings("ignore")
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
DATA = Path("data"); DATA.mkdir(exist_ok=True)
OUT = Path("output"); OUT.mkdir(exist_ok=True)

# Study area: a 4 km x 4 km block of the Forest of Dean around the Cannop Valley, in British National Grid
AOI_BNG = (360000, 209000, 364000, 213000)          # xmin, ymin, xmax, ymax in EPSG:27700
BNG = CRS.from_epsg(27700)
META_ROOT = "https://dataforgood-fb-data.s3.amazonaws.com/forests/v1/alsgedi_global_v6_float"
ETH_URL = ("https://libdrive.ethz.ch/index.php/s/cO8or7iOe5dT2Rt/download?path=%2F3deg_cogs"
           "&files=ETH_GlobalCanopyHeight_10m_2020_N51W003_Map.tif")
EA_WCS = {
    "dtm": ("lidar-composite-digital-terrain-model-dtm-1m",
            "13787b9a-26a4-4775-8523-806d13af58fc__Lidar_Composite_Elevation_DTM_1m"),
    "dsm": ("lidar-composite-digital-surface-model-first-return-dsm-1m",
            "df4e3ec3-315e-48aa-aaaf-b5ae74d7b2bb__Lidar_Composite_Elevation_FZ_DSM_1m"),
}
CLASS_EDGES = [0, 2, 5, 10, 20, 30, 99]
CLASS_NAMES = ["<2 m open", "2-5 m shrub", "5-10 m", "10-20 m", "20-30 m", ">=30 m"]
CHM_CMAP = plt.get_cmap("Spectral_r")
CLASS_CMAP = ListedColormap(plt.get_cmap("Spectral_r")(np.linspace(0.05, 0.95, 6)))


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


def write_tif(path, arr, transform, crs, nodata=None):
    profile = dict(driver="GTiff", height=arr.shape[0], width=arr.shape[1], count=1, dtype=arr.dtype,
                   crs=crs, transform=transform, nodata=nodata, tiled=True, compress="deflate")
    with rasterio.open(path, "w", **profile) as dst:
        dst.write(arr, 1)


section("versions")
print(f"python {sys.version.split()[0]} | rasterio {version('rasterio')} (GDAL {rasterio.__gdal_version__}) | "
      f"numpy {version('numpy')} | scipy {version('scipy')} | pandas {version('pandas')} | "
      f"matplotlib {version('matplotlib')}")

# ---------------------------------------------------------------- Step 1: find and read the Meta tile
section("Meta canopy height tiles")


def quadkey(lon, lat, zoom):
    """Bing/Web Mercator QuadKey of the tile that contains a point."""
    n = 2 ** zoom
    x = int((lon + 180) / 360 * n)
    lat_r = np.radians(lat)
    y = int((1 - np.log(np.tan(lat_r) + 1 / np.cos(lat_r)) / np.pi) / 2 * n)
    digits = []
    for i in range(zoom, 0, -1):
        mask = 1 << (i - 1)
        digits.append(str((1 if x & mask else 0) + (2 if y & mask else 0)))
    return "".join(digits)


tiles_path = DATA / "tiles.geojson"
if not tiles_path.exists():
    tiles_path.write_bytes(requests.get(f"{META_ROOT}/tiles.geojson", timeout=300).content)
tiles = json.loads(tiles_path.read_text())["features"]
aoi_ll = transform_bounds(BNG, "EPSG:4326", *AOI_BNG)
aoi_poly = box(*aoi_ll)
hits = [f["properties"]["tile"] for f in tiles if shape(f["geometry"]).intersects(aoi_poly)]
centre = ((aoi_ll[0] + aoi_ll[2]) / 2, (aoi_ll[1] + aoi_ll[3]) / 2)
print(f"{len(tiles):,} tiles in the global index (zoom-9 QuadKeys)")
print(f"AOI lon/lat bounds: {aoi_ll[0]:.4f}, {aoi_ll[1]:.4f}, {aoi_ll[2]:.4f}, {aoi_ll[3]:.4f}")
print(f"tiles intersecting the AOI: {hits}; QuadKey computed from the AOI centre: {quadkey(*centre, 9)}")
TILE = hits[0]
meta_url = f"{META_ROOT}/chm/{TILE}.tif"

with rasterio.open("/vsicurl/" + meta_url) as src:
    print(f"chm/{TILE}.tif: {src.width} x {src.height} px, {src.dtypes[0]}, pixel {src.res[0]:.4f} m, "
          f"tiled={src.profile.get('tiled')}, block {src.block_shapes[0]}, nodata={src.nodata}")
    meta_crs = src.crs
    print(f"CRS: {meta_crs} (Web Mercator); origin x {src.transform.c:.1f}, y {src.transform.f:.1f}")
    b3857 = transform_bounds(BNG, meta_crs, *AOI_BNG)
    pad = 60
    win = from_bounds(b3857[0] - pad, b3857[1] - pad, b3857[2] + pad, b3857[3] + pad,
                      src.transform).round_offsets().round_lengths()
    meta_path = DATA / f"meta_chm_{TILE}_forest_of_dean_3857.tif"
    if not meta_path.exists():
        write_tif(meta_path, src.read(1, window=win), src.window_transform(win), meta_crs)
with rasterio.open(meta_path) as src:
    meta_3857 = src.read(1)
    meta_3857_tf = src.transform
ground = meta_3857_tf.a * np.cos(np.radians(centre[1]))
print(f"read window {win.width} x {win.height} px; heights {meta_3857.min()}..{meta_3857.max()} m")
print(f"Web Mercator pixel {meta_3857_tf.a:.3f} m on the map = {ground:.3f} m on the ground at {centre[1]:.2f} N")

dates = json.loads(requests.get(f"{META_ROOT}/metadata/{TILE}.geojson", timeout=120).text)["features"]
acq = sorted({f["properties"]["acq_date"] for f in dates if shape(f["geometry"]).intersects(aoi_poly)})
all_acq = sorted(f["properties"]["acq_date"] for f in dates)
print(f"imagery dates in tile {TILE}: {len(all_acq)} footprints, {all_acq[0]} .. {all_acq[-1]}")
print(f"imagery date(s) under the AOI: {acq}")

# ---------------------------------------------------------------- Step 2: ETH Global Canopy Height 2020
section("ETH Global Canopy Height 2020")
eth_path = DATA / "eth_chm_2020_forest_of_dean_4326.tif"
eth_vsi = "/vsicurl?use_head=no&max_retry=3&url=" + urllib.parse.quote(ETH_URL, safe="")
with rasterio.Env(GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR"):
    with rasterio.open(eth_vsi) as src:
        print(f"N51W003 tile: {src.width} x {src.height} px of {src.res[0]:.8f} deg, {src.crs}, "
              f"nodata={src.nodata}, overviews {src.overviews(1)}")
        if not eth_path.exists():
            d = 0.002
            w = from_bounds(aoi_ll[0] - d, aoi_ll[1] - d, aoi_ll[2] + d, aoi_ll[3] + d,
                            src.transform).round_offsets().round_lengths()
            write_tif(eth_path, src.read(1, window=w), src.window_transform(w), src.crs, src.nodata)
with rasterio.open(eth_path) as src:
    eth_ll = src.read(1)
    eth_tf, eth_crs, eth_nd = src.transform, src.crs, src.nodata
dx_m = eth_tf.a * 111320 * np.cos(np.radians(centre[1]))
print(f"ETH pixel on the ground: {dx_m:.1f} m east-west x {abs(eth_tf.e) * 111260:.1f} m north-south")

# ---------------------------------------------------------------- Step 3: CHM from Environment Agency LiDAR
section("Environment Agency LiDAR: DSM - DTM")


def ea_wcs(product, x0, y0, size=1000):
    """One 1 km x 1 km GeoTIFF from the Environment Agency WCS (no key needed)."""
    service, cov = EA_WCS[product]
    url = (f"https://environment.data.gov.uk/spatialdata/{service}/wcs?service=WCS&version=2.0.1"
           f"&request=GetCoverage&CoverageId={cov}&format=image/tiff"
           f"&subset=E({x0},{x0 + size})&subset=N({y0},{y0 + size})")
    for attempt in range(4):
        try:
            r = requests.get(url, timeout=180)
            r.raise_for_status()
            return r.content
        except requests.RequestException:
            if attempt == 3:
                raise


def ea_mosaic(product):
    """Download 1 km tiles and paste them into one 1 m array on the AOI grid."""
    path = DATA / f"ea_{product}_1m.tif"
    if not path.exists():
        tf = from_origin(AOI_BNG[0], AOI_BNG[3], 1, 1)
        a = np.full((AOI_BNG[3] - AOI_BNG[1], AOI_BNG[2] - AOI_BNG[0]), np.nan, "float32")
        for x0 in range(AOI_BNG[0], AOI_BNG[2], 1000):
            for y0 in range(AOI_BNG[1], AOI_BNG[3], 1000):
                tile = DATA / f"ea_{product}_{x0}_{y0}.tif"
                if not tile.exists():
                    tile.write_bytes(ea_wcs(product, x0, y0))
                with rasterio.open(tile) as src:
                    t = src.read(1).astype("float32")
                    t[(t == src.nodata) | (t < -1000)] = np.nan
                    c, r = (int(round(v)) for v in ~tf * (src.bounds.left, src.bounds.top))
                    a[r:r + t.shape[0], c:c + t.shape[1]] = t
        write_tif(path, a, tf, BNG, np.nan)
    with rasterio.open(path) as src:
        return src.read(1), src.transform


dtm, grid_tf = ea_mosaic("dtm")
dsm, _ = ea_mosaic("dsm")
with rasterio.open(DATA / f"ea_dtm_{AOI_BNG[0]}_{AOI_BNG[1]}.tif") as probe:
    print(f"one WCS tile: {probe.width} x {probe.height} px, {probe.dtypes[0]}, {probe.crs}, nodata {probe.nodata:.3g}")
chm = dsm - dtm
print(f"grid: {dtm.shape[1]} x {dtm.shape[0]} px of 1 m, origin E {grid_tf.c:.0f} N {grid_tf.f:.0f}")
print(f"DTM {np.nanmin(dtm):.1f}..{np.nanmax(dtm):.1f} m; DSM {np.nanmin(dsm):.1f}..{np.nanmax(dsm):.1f} m; "
      f"no data: {np.isnan(chm).mean() * 100:.2f}%")
print(f"raw DSM - DTM: min {np.nanmin(chm):.2f} m, max {np.nanmax(chm):.2f} m, "
      f"{(chm < -0.5).mean() * 100:.2f}% below -0.5 m, {(chm > 50).sum()} px above 50 m")
chm[chm < 0] = 0
chm[chm > 50] = np.nan          # isolated spikes (birds, power lines) are not trees
chm_raw = chm.copy()
# Leaf-off first returns pass through gaps in broadleaf crowns and leave "pits"; a 3 x 3 grey closing fills them
chm = ndimage.grey_closing(np.nan_to_num(chm_raw), size=(3, 3)).astype("float32")
chm[np.isnan(chm_raw)] = np.nan
h, w = chm.shape[0] // 10, chm.shape[1] // 10
tall_cells = np.nanpercentile(chm_raw.reshape(h, 10, w, 10).swapaxes(1, 2).reshape(h, w, 100), 95, axis=2) >= 15
tall_px = np.kron(tall_cells, np.ones((10, 10), bool))
print(f"pits: in 10 m cells whose p95 is >= 15 m, {np.mean(chm_raw[tall_px] < 2) * 100:.1f}% of 1 m pixels "
      f"are < 2 m before pit filling and {np.mean(chm[tall_px] < 2) * 100:.1f}% after")
canopy = chm >= 2
print(f"pit-filled LiDAR CHM: mean {np.nanmean(chm):.2f} m, canopy >= 2 m on {np.mean(canopy) * 100:.1f}% of the AOI, "
      f"p95 {np.nanpercentile(chm, 95):.1f} m, p99 {np.nanpercentile(chm, 99):.1f} m, "
      f"max {np.nanmax(chm):.1f} m")
write_tif(OUT / "lidar_chm_1m_bng.tif", chm.astype("float32"), grid_tf, BNG, np.nan)

# ---------------------------------------------------------------- Step 4: align everything on the BNG grid
section("CRS alignment onto the British National Grid")
meta = np.full(chm.shape, np.nan, "float32")
reproject(meta_3857.astype("float32"), meta, src_transform=meta_3857_tf, src_crs=meta_crs,
          dst_transform=grid_tf, dst_crs=BNG, dst_nodata=np.nan, resampling=Resampling.average)
write_tif(OUT / "meta_chm_1m_bng.tif", meta, grid_tf, BNG, np.nan)
grid10_tf = from_origin(AOI_BNG[0], AOI_BNG[3], 10, 10)
n10 = (AOI_BNG[3] - AOI_BNG[1]) // 10
eth = np.full((n10, n10), np.nan, "float32")
eth_src = eth_ll.astype("float32")
eth_src[eth_ll == eth_nd] = np.nan
reproject(eth_src, eth, src_transform=eth_tf, src_crs=eth_crs, src_nodata=np.nan,
          dst_transform=grid10_tf, dst_crs=BNG, dst_nodata=np.nan, resampling=Resampling.nearest)
print(f"Meta resampled to 1 m BNG (average): {meta.shape}, NaN {np.isnan(meta).mean() * 100:.2f}%")
print(f"ETH resampled to 10 m BNG (nearest): {eth.shape}, NaN {np.isnan(eth).mean() * 100:.2f}%")


def block(a, k=10):
    """Reshape a 1 m array into k x k blocks so each 10 m cell can be summarised."""
    h, w = a.shape[0] // k, a.shape[1] // k
    return a.reshape(h, k, w, k).swapaxes(1, 2).reshape(h, w, k * k)


lid10_mean = np.nanmean(block(chm), axis=2)
lid10_p95 = np.nanpercentile(block(chm), 95, axis=2)
meta10_mean = np.nanmean(block(meta), axis=2)
meta10_p95 = np.nanpercentile(block(meta), 95, axis=2)

# co-registration check: shift Meta against the LiDAR CHM and find the lowest RMSE
inner = (slice(20, -20), slice(20, -20))
ref = chm[inner][::2, ::2]
shift_rmse = {}
for dy in range(-10, 11):
    for dx in range(-10, 11):
        shifted = np.roll(meta, (dy, dx), axis=(0, 1))[inner][::2, ::2]
        ok = ~np.isnan(ref) & ~np.isnan(shifted)
        shift_rmse[(dx, dy)] = float(np.sqrt(np.mean((shifted[ok] - ref[ok]) ** 2)))
best = min(shift_rmse, key=shift_rmse.get)
print(f"co-registration: RMSE {shift_rmse[(0, 0)]:.2f} m as aligned; lowest {shift_rmse[best]:.2f} m "
      f"when Meta is moved {best[0]:+d} m east and {-best[1]:+d} m north (positive = east/north)")

# ---------------------------------------------------------------- Step 5: pixel-wise comparison
section("pixel-wise comparison with the LiDAR CHM")


def stats(pred, obs, label, subset=None):
    ok = ~np.isnan(pred) & ~np.isnan(obs)
    if subset is not None:
        ok &= subset
    p, o = pred[ok].astype("float64"), obs[ok].astype("float64")
    e = p - o
    r = np.corrcoef(p, o)[0, 1]
    return {"comparison": label, "n": int(ok.sum()), "bias_m": e.mean(), "mae_m": np.abs(e).mean(),
            "rmse_m": np.sqrt((e ** 2).mean()), "r2": r ** 2,
            "mean_lidar_m": o.mean(), "mean_product_m": p.mean()}


rows = [
    stats(meta, chm, "Meta 1 m vs LiDAR 1 m"),
    stats(meta, chm_raw, "Meta 1 m vs LiDAR 1 m, no pit filling"),
    stats(meta, chm, "Meta 1 m vs LiDAR 1 m, LiDAR >= 5 m", chm >= 5),
    stats(meta10_mean, lid10_mean, "Meta 10 m mean vs LiDAR 10 m mean"),
    stats(meta10_p95, lid10_p95, "Meta 10 m p95 vs LiDAR 10 m p95"),
    stats(eth, lid10_p95, "ETH 10 m vs LiDAR 10 m p95"),
    stats(meta10_p95, lid10_p95, "Meta 10 m p95 vs LiDAR p95, LiDAR >= 5 m", lid10_p95 >= 5),
    stats(eth, lid10_p95, "ETH 10 m vs LiDAR p95, LiDAR >= 5 m", lid10_p95 >= 5),
]
metrics = pd.DataFrame(rows)
metrics.to_csv(OUT / "comparison_metrics.csv", index=False, float_format="%.3f")
with pd.option_context("display.width", 160, "display.max_columns", 20):
    print(metrics.to_string(index=False, float_format=lambda v: f"{v:.2f}"))

# ---------------------------------------------------------------- Step 6: saturation and edge effects
section("error by LiDAR height (saturation of tall canopy)")
bins = [0, 2, 5, 10, 15, 20, 25, 30, 35, 50]
labels = [f"{a}-{b}" for a, b in zip(bins[:-1], bins[1:])]
sat = []
for name, pred, obs in [("Meta 1 m", meta, chm), ("Meta 10 m p95", meta10_p95, lid10_p95),
                        ("ETH 10 m", eth, lid10_p95)]:
    ok = ~np.isnan(pred) & ~np.isnan(obs)
    p, o = pred[ok], obs[ok]
    idx = np.digitize(o, bins) - 1
    for i, lab in enumerate(labels):
        sel = idx == i
        if sel.sum() < 50:
            continue
        sat.append({"product": name, "lidar_bin_m": lab, "n": int(sel.sum()),
                    "lidar_mean_m": o[sel].mean(), "product_mean_m": p[sel].mean(),
                    "product_p25_m": np.percentile(p[sel], 25), "product_p75_m": np.percentile(p[sel], 75),
                    "bias_m": (p[sel] - o[sel]).mean()})
sat = pd.DataFrame(sat)
sat.to_csv(OUT / "bias_by_height.csv", index=False, float_format="%.3f")
print("mean product - LiDAR (m) by LiDAR height bin:")
print(sat.pivot(index="lidar_bin_m", columns="product", values="bias_m").reindex(labels)
      .to_string(float_format=lambda v: f"{v:+.2f}"))
print(f"LiDAR pixels >= 30 m: {np.nanmean(chm >= 30) * 100:.2f}% of the AOI; LiDAR p99.9 "
      f"{np.nanpercentile(chm, 99.9):.1f} m; Meta max {np.nanmax(meta_3857)} m; ETH max {np.nanmax(eth):.0f} m")

section("error by distance to the canopy edge")
inside = ndimage.distance_transform_edt(canopy)
outside = ndimage.distance_transform_edt(~canopy)
signed = np.where(canopy, inside, -outside)
ebins = [-1000, -10, -5, -2, 0, 2, 5, 10, 20, 1000]
elabels = ["< -10", "-10..-5", "-5..-2", "-2..0", "0..2", "2..5", "5..10", "10..20", "> 20"]
ok = ~np.isnan(meta) & ~np.isnan(chm)
eidx = np.digitize(signed[ok], ebins) - 1
err = (meta - chm)[ok]
lid_ok, meta_ok = chm[ok], meta[ok]
edge = []
for i, lab in enumerate(elabels):
    sel = eidx == i
    edge.append({"signed_distance_m": lab, "n": int(sel.sum()), "bias_m": err[sel].mean(),
                 "mae_m": np.abs(err[sel]).mean(), "lidar_mean_m": lid_ok[sel].mean(),
                 "meta_mean_m": meta_ok[sel].mean()})
edge = pd.DataFrame(edge)
edge.to_csv(OUT / "error_by_edge_distance.csv", index=False, float_format="%.3f")
print("signed distance: negative = metres outside the LiDAR canopy (>= 2 m), positive = inside it")
print(edge.to_string(index=False, float_format=lambda v: f"{v:.2f}"))

# ---------------------------------------------------------------- Step 7: height-class maps
section("height classes")


def classify(a):
    c = np.digitize(a, CLASS_EDGES[1:-1]).astype("uint8")
    c[np.isnan(a)] = 255
    return c


cls_lid, cls_meta, cls_eth = classify(chm), classify(meta), classify(eth)
write_tif(OUT / "meta_height_classes_1m_bng.tif", cls_meta, grid_tf, BNG, 255)
area = []
for name, c, px_ha in [("LiDAR 1 m", cls_lid, 1e-4), ("Meta 1 m", cls_meta, 1e-4), ("ETH 10 m", cls_eth, 1e-2)]:
    for k, lab in enumerate(CLASS_NAMES):
        area.append({"product": name, "class": lab, "area_ha": (c == k).sum() * px_ha})
area = pd.DataFrame(area)
area.to_csv(OUT / "height_class_area.csv", index=False, float_format="%.2f")
print("area (ha) per height class:")
print(area.pivot(index="class", columns="product", values="area_ha").reindex(CLASS_NAMES)
      [["LiDAR 1 m", "Meta 1 m", "ETH 10 m"]].to_string(float_format=lambda v: f"{v:8.1f}"))
valid = (cls_lid != 255) & (cls_meta != 255)
cls_lid10, cls_meta10 = classify(lid10_p95), classify(meta10_p95)
valid10 = (cls_lid10 != 255) & (cls_eth != 255)
print(f"same class as LiDAR: Meta 1 m {(cls_lid == cls_meta)[valid].mean() * 100:.1f}% of pixels; "
      f"10 m p95 classes: Meta {(cls_meta10 == cls_lid10)[valid10].mean() * 100:.1f}%, "
      f"ETH {(cls_eth == cls_lid10)[valid10].mean() * 100:.1f}%")
within1 = (np.abs(cls_lid.astype(int) - cls_meta.astype(int)) <= 1)[valid].mean()
print(f"Meta 1 m class within one class of LiDAR: {within1 * 100:.1f}%")

# ---------------------------------------------------------------- figures
section("figures")
ext = (AOI_BNG[0] / 1000, AOI_BNG[2] / 1000, AOI_BNG[1] / 1000, AOI_BNG[3] / 1000)

# 01: the Meta tile grid over Britain, with our tile and the study area
ne_path = DATA / "ne_50m_admin_0_countries.geojson"
if not ne_path.exists():
    ne_path.write_bytes(requests.get("https://raw.githubusercontent.com/nvkelso/natural-earth-vector/master/"
                                     "geojson/ne_50m_admin_0_countries.geojson", timeout=120).content)
ne = [f for f in json.loads(ne_path.read_text(encoding="utf-8"))["features"]
      if f["properties"]["ADM0_A3"] in ("GBR", "IRL", "FRA", "IMN")]
fig, ax = plt.subplots(figsize=(5.2, 6))
view = box(-8.5, 49.6, 2.2, 59.2)
for f in ne:
    geoms = shape(f["geometry"])
    for g in getattr(geoms, "geoms", [geoms]):
        x, y = g.exterior.xy
        ax.fill(x, y, color="#e3e7e4", ec="#7a8582", lw=0.5, zorder=1)
for f in tiles:
    g = shape(f["geometry"])
    if not g.intersects(view):
        continue
    x, y = g.exterior.xy
    hit = f["properties"]["tile"] == TILE
    ax.plot(x, y, color="#c21f3a" if hit else "#3288bd", lw=1.8 if hit else 0.5, zorder=3 if hit else 2)
tb = shape(next(f for f in tiles if f["properties"]["tile"] == TILE)["geometry"]).bounds
ax.annotate(f"tile {TILE}", (tb[2], tb[3]), xytext=(0.0, 53.3), color="#c21f3a", fontsize=9,
            arrowprops=dict(arrowstyle="-", color="#c21f3a", lw=0.8))
ax.plot(*aoi_poly.exterior.xy, color="black", lw=1.5, zorder=4)
ax.annotate("Forest of Dean\nstudy area", (aoi_ll[0], aoi_ll[1]), xytext=(-7.9, 50.2), fontsize=9,
            arrowprops=dict(arrowstyle="-", color="black", lw=0.8))
ax.set_xlim(-8.5, 2.2); ax.set_ylim(49.6, 59.2); ax.set_aspect(1 / np.cos(np.radians(54)))
ax.set_xlabel("Longitude"); ax.set_ylabel("Latitude")
ax.set_title("Meta/WRI canopy height: zoom-9 QuadKey tiles", fontsize=10)
fig.savefig(OUT / "01-meta-tile-index.png", facecolor="white"); plt.close(fig)

# 02: the three canopy height maps, full study area and a 400 m zoom
zx, zy, zs = 361300, 211300, 400             # lower-left corner and size of the zoom box (BNG, m)
r0, c0 = int(AOI_BNG[3] - zy - zs), int(zx - AOI_BNG[0])
fig, axes = plt.subplots(2, 3, figsize=(11, 7.6), constrained_layout=True)
panels = [("LiDAR CHM (EA, pit-filled), 1 m", chm), ("Meta/WRI canopy height, 1 m", meta),
          ("ETH Global Canopy Height 2020, 10 m", eth)]
for ax, (t, a) in zip(axes[0], panels):
    im = ax.imshow(a, extent=ext, cmap=CHM_CMAP, vmin=0, vmax=35, interpolation="nearest")
    ax.add_patch(Rectangle((zx / 1000, zy / 1000), zs / 1000, zs / 1000, fill=False, ec="black", lw=1.2))
    ax.set_title(t, fontsize=9); ax.set_xlabel("Easting (km, BNG)")
axes[0, 0].set_ylabel("Northing (km, BNG)")
zext = (zx / 1000, (zx + zs) / 1000, zy / 1000, (zy + zs) / 1000)
zooms = [chm[r0:r0 + zs, c0:c0 + zs], meta[r0:r0 + zs, c0:c0 + zs],
         eth[r0 // 10:(r0 + zs) // 10, c0 // 10:(c0 + zs) // 10]]
for ax, a, (t, _) in zip(axes[1], zooms, panels):
    ax.imshow(a, extent=zext, cmap=CHM_CMAP, vmin=0, vmax=35, interpolation="nearest")
    ax.set_title("400 m zoom: " + t.split(" (")[0].split(",")[0], fontsize=9); ax.set_xlabel("Easting (km, BNG)")
axes[1, 0].set_ylabel("Northing (km, BNG)")
fig.colorbar(im, ax=axes, shrink=0.6, label="Canopy height (m)")
fig.savefig(OUT / "02-canopy-height-maps.png", facecolor="white"); plt.close(fig)

# 03: hexbin scatter plots
fig, axes = plt.subplots(1, 3, figsize=(12, 4.2), constrained_layout=True)
rng = np.random.default_rng(0)
for ax, (t, p, o, m) in zip(axes, [("Meta 1 m vs LiDAR 1 m", meta, chm, rows[0]),
                                   ("Meta vs LiDAR, 10 m p95", meta10_p95, lid10_p95, rows[4]),
                                   ("ETH 2020 vs LiDAR, 10 m p95", eth, lid10_p95, rows[5])]):
    ok = ~np.isnan(p) & ~np.isnan(o)
    pp, oo = p[ok], o[ok]
    if pp.size > 2_000_000:
        sel = rng.choice(pp.size, 2_000_000, replace=False)
        pp, oo = pp[sel], oo[sel]
    hb = ax.hexbin(oo, pp, gridsize=70, extent=(0, 45, 0, 45), bins="log", cmap="Spectral_r", mincnt=1)
    ax.plot([0, 45], [0, 45], color="black", lw=0.8, ls="--")
    ax.text(1.5, 43.5, f"bias {m['bias_m']:+.2f} m\nRMSE {m['rmse_m']:.2f} m\nr² {m['r2']:.2f}\nn = {m['n']:,}",
            va="top", fontsize=8, bbox=dict(fc="white", ec="#999999", lw=0.5))
    ax.set_xlim(0, 45); ax.set_ylim(0, 45); ax.set_aspect("equal")
    ax.set_xlabel("LiDAR canopy height (m)"); ax.set_ylabel("Product canopy height (m)"); ax.set_title(t, fontsize=9)
    fig.colorbar(hb, ax=ax, shrink=0.8, label="pixels per hexagon (log)")
fig.savefig(OUT / "03-hexbin-comparison.png", facecolor="white"); plt.close(fig)

# 04: difference maps at 10 m
fig, axes = plt.subplots(1, 2, figsize=(10, 4.6), constrained_layout=True)
norm = TwoSlopeNorm(vmin=-20, vcenter=0, vmax=20)
for ax, (t, d) in zip(axes, [("Meta p95 - LiDAR p95 (10 m)", meta10_p95 - lid10_p95),
                             ("ETH 2020 - LiDAR p95 (10 m)", eth - lid10_p95)]):
    im = ax.imshow(d, extent=ext, cmap="RdBu_r", norm=norm, interpolation="nearest")
    ax.set_title(t, fontsize=9); ax.set_xlabel("Easting (km, BNG)")
axes[0].set_ylabel("Northing (km, BNG)")
fig.colorbar(im, ax=axes, shrink=0.8, label="Difference (m): red = product taller, blue = shorter")
fig.savefig(OUT / "04-difference-maps.png", facecolor="white"); plt.close(fig)

# 05: saturation and edge effects
fig, axes = plt.subplots(1, 2, figsize=(11, 4.2), constrained_layout=True)
colours = {"Meta 1 m": "#3288bd", "Meta 10 m p95": "#66c2a5", "ETH 10 m": "#f46d43"}
for name, g in sat.groupby("product", sort=False):
    axes[0].plot(g["lidar_mean_m"], g["product_mean_m"], "o-", color=colours[name], label=name)
    axes[0].fill_between(g["lidar_mean_m"], g["product_p25_m"], g["product_p75_m"], color=colours[name], alpha=0.15)
axes[0].plot([0, 40], [0, 40], color="black", lw=0.8, ls="--", label="1:1")
axes[0].set_xlabel("LiDAR canopy height, bin mean (m)"); axes[0].set_ylabel("Product height, mean and IQR (m)")
axes[0].set_title("Product height against LiDAR height", fontsize=9); axes[0].legend(fontsize=8)
x = np.arange(len(edge))
axes[1].bar(x, edge["bias_m"], color=["#3288bd" if v < 0 else "#d53e4f" for v in edge["bias_m"]])
axes[1].axhline(0, color="black", lw=0.8)
axes[1].axvline(3.5, color="#555555", lw=0.8, ls=":")
axes[1].set_xticks(x, edge["signed_distance_m"], rotation=30)
axes[1].set_xlabel("Distance from the LiDAR canopy edge (m): outside < 0 < inside")
axes[1].set_ylabel("Meta - LiDAR, mean (m)")
axes[1].set_title("Meta 1 m bias by distance to the canopy edge", fontsize=9)
fig.savefig(OUT / "05-saturation-and-edges.png", facecolor="white"); plt.close(fig)

# 06: height-class maps and areas
fig = plt.figure(figsize=(12, 4.4), constrained_layout=True)
gs = fig.add_gridspec(1, 3, width_ratios=[1, 1, 1.05])
cnorm = BoundaryNorm(np.arange(-0.5, 6.5), 6)
for i, (t, c) in enumerate([("LiDAR CHM classes, 1 m", cls_lid), ("Meta/WRI classes, 1 m", cls_meta)]):
    ax = fig.add_subplot(gs[0, i])
    ax.imshow(np.ma.masked_equal(c, 255), extent=ext, cmap=CLASS_CMAP, norm=cnorm, interpolation="nearest")
    ax.set_title(t, fontsize=9); ax.set_xlabel("Easting (km, BNG)")
    if i == 0:
        ax.set_ylabel("Northing (km, BNG)")
ax = fig.add_subplot(gs[0, 2])
piv = area.pivot(index="class", columns="product", values="area_ha").reindex(CLASS_NAMES)
w = 0.27
for j, (name, hatch) in enumerate([("LiDAR 1 m", ""), ("Meta 1 m", "//"), ("ETH 10 m", "..")]):
    ax.bar(np.arange(6) + (j - 1) * w, piv[name], w, color=CLASS_CMAP.colors, edgecolor="#333333", lw=0.5,
           hatch=hatch, label=name)
ax.set_xticks(np.arange(6), CLASS_NAMES, rotation=30, ha="right")
ax.set_ylabel("Area (ha) of the 1,600 ha study area"); ax.legend(fontsize=8)
ax.set_title("Area per height class (colours as in the maps)", fontsize=9)
fig.savefig(OUT / "06-height-classes.png", facecolor="white"); plt.close(fig)
print("wrote", ", ".join(sorted(p.name for p in OUT.glob("0*.png"))))