"""Deep learning for satellite imagery, part 3: run a trained U-Net over a large, unseen Sentinel-2 raster.

Streams a Sentinel-2 L2A summer 2021 composite of Leicester (UK), an area the model never saw, writes it as a Cloud
Optimised GeoTIFF, and runs the U-Net from part 2 over it with tiled sliding-window inference: overlapping 256 px
windows read straight from the GeoTIFF, blended with a linear-ramp weight so tile edges leave no seams. Writes the
class map and a confidence layer as COGs, vectorises the classes to simplified polygons in a GeoPackage, and
compares everything with the same reference labels used for training (Overture buildings over ESA WorldCover).
Run with:  python predict_large_raster.py      (it runs parts 1 and 2 first if the model is missing)
"""
import os
import platform
import subprocess
import sys
import time
import urllib.request
from importlib.metadata import version
from pathlib import Path

import dask
import geopandas as gpd
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import odc.stac
import pandas as pd
import planetary_computer
import pystac_client
import rasterio
import rasterio.features
import rasterio.shutil
import segmentation_models_pytorch as smp
import shapely
import torch
import xarray as xr
from matplotlib.colors import ListedColormap
from matplotlib.patches import Patch, Rectangle
from odc.geo.geobox import GeoBox
from overturemaps import core as overture
from pyproj import Transformer
from rasterio.transform import from_origin
from rasterio.windows import Window
from scipy import ndimage
from shapely.geometry import box, shape

HERE = Path(__file__).parent
OUT, DATA = HERE / "output", HERE / "data"
for d in (OUT, DATA):
    d.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 140)
dask.config.set(scheduler="threads", num_workers=8)
torch.set_num_threads(8)

print("=== versions ===")
print(f"python {platform.python_version()} | torch {torch.__version__} | segmentation-models-pytorch {smp.__version__} | "
      f"rasterio {rasterio.__version__} (GDAL {rasterio.__gdal_version__}) | geopandas {gpd.__version__} | "
      f"shapely {shapely.__version__} | odc-stac {odc.stac.__version__}")


def section(title: str) -> None:
    print(f"\n=== {title} ===")


# --- Step 1: load the model from part 2 ----------------------------------------------------------
PART2 = HERE.parent / "train-unet-satellite-segmentation-pytorch"
CKPT = Path(os.environ.get("MODEL_PATH", PART2 / "output" / "unet_resnet18_s2_nottingham.pt"))
if not CKPT.exists():
    script = PART2 / "train_unet.py"
    if not script.exists():                       # skipped part 2: fetch its script from the site
        PART2.mkdir(parents=True, exist_ok=True)
        url = os.environ.get("TGIS_CODE_URL", "https://www.technicalgis.com/code")
        urllib.request.urlretrieve(f"{url}/train-unet-satellite-segmentation-pytorch/train_unet.py", script)
    print(f"model not found: running {script.name} (part 2, which runs part 1 if needed) first")
    subprocess.run([sys.executable, str(script)], check=True, cwd=PART2)

ckpt = torch.load(CKPT, map_location="cpu", weights_only=False)
model = smp.Unet(ckpt["encoder"], encoder_weights=None, in_channels=ckpt["in_channels"], classes=len(ckpt["classes"]))
model.load_state_dict(ckpt["state_dict"])
model.eval()
CLASSES, COLOURS, BANDS = ckpt["classes"], ckpt["colours"], ckpt["bands"]
NC, IGNORE = len(CLASSES), 255
section("model")
print(f"{CKPT.name}: U-Net / {ckpt['encoder']}, {ckpt['in_channels']} input channels, classes {CLASSES}")


def to_input(x: np.ndarray) -> torch.Tensor:
    """Same preprocessing as in training: reflectance / 10000, with a zero channel where B10 would be."""
    x = torch.from_numpy(x.astype("float32") / ckpt["scale"])
    i = ckpt["b10_zero_index"]
    return torch.cat([x[..., :i, :, :], torch.zeros_like(x[..., :1, :, :]), x[..., i:, :, :]], dim=-3)


# --- Step 2: a Sentinel-2 composite of a new area, saved as a Cloud Optimised GeoTIFF ---------------
NAME = "leicester"
CENTRE = (-1.13, 52.635)          # Leicester, about 35 km south of the Nottingham training area
SIZE_PX, RES, CRS, YEAR = 3000, 10, "EPSG:32630", 2021
cx, cy = Transformer.from_crs(4326, CRS, always_xy=True).transform(*CENTRE)
size = SIZE_PX * RES
x0, y1 = round(cx - size / 2, -1), round(cy + size / 2, -1)
geobox = GeoBox.from_bbox((x0, y1 - size, x0 + size, y1), crs=CRS, resolution=RES)
TRANSFORM = from_origin(x0, y1, RES, RES)
H = W = SIZE_PX
bbox_ll = tuple(gpd.GeoSeries([box(x0, y1 - size, x0 + size, y1)], crs=CRS).to_crs(4326).total_bounds)
catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1",
                                    modifier=planetary_computer.sign_inplace)
section("new area")
print(f"{NAME}: {H} x {W} pixels of {RES} m ({size / 1000:.0f} km square); lon/lat bounds "
      f"{tuple(round(float(v), 3) for v in bbox_ll)}")

s2_cog = DATA / f"s2_{NAME}_{YEAR}_cog.tif"
if not s2_cog.exists():
    items = catalog.search(collections=["sentinel-2-l2a"], bbox=bbox_ll, datetime=f"{YEAR}-05-01/{YEAR}-08-31",
                           query={"eo:cloud_cover": {"lt": 40}}).item_collection()
    odc.stac.configure_rio(cloud_defaults=True)
    ds = odc.stac.load(items, bands=BANDS + ["SCL"], geobox=geobox, groupby="solar_day", dtype="uint16", nodata=0,
                       chunks={"time": 1, "x": 1500, "y": 1500}, resampling={"*": "bilinear", "SCL": "nearest"})
    base = {i.datetime.date(): float(i.properties["s2:processing_baseline"]) for i in items}
    offset = xr.DataArray([1000 if base[pd.Timestamp(t).date()] >= 4.0 else 0 for t in ds.time.values],
                          dims="time", coords={"time": ds.time})
    scl = ds["SCL"].compute()
    cloudy = ndimage.binary_dilation(np.isin(scl.values, [3, 8, 9, 10]), iterations=3, axes=(1, 2))
    clear = scl.isin([4, 5, 6, 7]) & ~xr.DataArray(cloudy, coords=scl.coords, dims=scl.dims)
    frac = clear.mean(("y", "x")).to_series()
    keep = frac.index[frac >= 0.2]
    print(f"{len(items)} items, {len(frac)} dates, {len(keep)} at least 20% clear: "
          f"{', '.join(f'{t:%d %b}' for t in keep)}")
    refl = (ds[BANDS].sel(time=keep).astype("float32") - offset.sel(time=keep))
    med = refl.where(clear.sel(time=keep)).chunk({"time": -1, "y": 500, "x": 500}).median("time").compute()
    raw = refl.where(ds["SCL"].sel(time=keep) > 0).chunk({"time": -1, "y": 500, "x": 500}).median("time").compute()
    med = med.fillna(raw)
    stack = np.stack([med[b].values for b in BANDS])
    stack = np.where(np.isnan(stack), 0, np.clip(stack, 1, 65535)).astype("uint16")
    with rasterio.open(s2_cog, "w", driver="COG", height=H, width=W, count=len(BANDS), dtype="uint16", crs=CRS,
                       transform=TRANSFORM, nodata=0, compress="deflate", blocksize=512) as dst:
        dst.write(stack)
        dst.descriptions = tuple(BANDS)
with rasterio.open(s2_cog) as src:
    print(f"{s2_cog.name}: {src.count} bands, {src.profile['blockxsize']} px internal tiles, "
          f"overviews {src.overviews(1)}, {s2_cog.stat().st_size / 1e6:.0f} MB on disk")

# --- Step 3: tiled sliding-window inference with overlap and blending -------------------------------
TILE, OVERLAP = 256, 64
STRIDE = TILE - OVERLAP


def starts(n: int) -> list[int]:
    """Window origins every STRIDE pixels, with the last window flush with the raster edge."""
    s = list(range(0, n - TILE, STRIDE))
    return s + [n - TILE]


# A weight that ramps linearly from near 0 at a window's edge to 1 at OVERLAP pixels in, so each pixel's
# prediction comes mostly from the windows in which it sits far from the edge.
ramp = np.minimum(np.arange(TILE) + 1, OVERLAP) / OVERLAP
ramp = np.minimum(ramp, ramp[::-1])
WEIGHT = np.outer(ramp, ramp).astype("float32")


def predict(path: Path, overlap: bool) -> tuple[np.ndarray, np.ndarray, float]:
    """Return (class map, max softmax probability, seconds). Reads one window at a time from the GeoTIFF."""
    t0 = time.time()
    rows = starts(H) if overlap else list(range(0, H, TILE))
    cols = starts(W) if overlap else list(range(0, W, TILE))
    probs = np.zeros((NC, H, W), "float32")
    wsum = np.zeros((H, W), "float32")
    with rasterio.open(path) as src, torch.no_grad():
        for r in rows:
            for c in cols:
                x = src.read(window=Window(c, r, TILE, TILE), boundless=True, fill_value=0)   # pads past the edge
                h, w = min(TILE, H - r), min(TILE, W - c)
                p = torch.softmax(model(to_input(x[None])), 1)[0].numpy()[:, :h, :w]
                wt = WEIGHT[:h, :w] if overlap else 1.0
                probs[:, r:r + h, c:c + w] += p * wt
                wsum[r:r + h, c:c + w] += wt
    probs /= wsum
    return probs.argmax(0).astype("uint8"), probs.max(0), time.time() - t0


pred, conf, secs = predict(s2_cog, overlap=True)
naive, _, secs_naive = predict(s2_cog, overlap=False)
n_win = len(starts(H)) * len(starts(W))
section("sliding-window inference")
print(f"blended: {n_win} windows of {TILE} px with {OVERLAP} px overlap in {secs:.0f} s "
      f"({H * W / 1e6 / secs:.2f} Mpx/s, {size ** 2 / 1e6 / secs:.0f} km2/s)")
print(f"naive:   {int(np.ceil(H / TILE)) ** 2} windows, no overlap, in {secs_naive:.0f} s")
edge = np.zeros((H, W), bool)
for k in range(TILE, H, TILE):
    edge[k - 8:k + 8, :] = True
    edge[:, k - 8:k + 8] = True
diff = naive != pred
print(f"naive and blended maps differ on {diff.mean():.1%} of pixels: {diff[edge].mean():.1%} within 8 px of a "
      f"naive tile edge, {diff[~edge].mean():.1%} elsewhere")

# --- Step 4: write the class map and confidence as Cloud Optimised GeoTIFFs --------------------------
cog = dict(driver="COG", height=H, width=W, count=1, crs=CRS, transform=TRANSFORM, compress="deflate",
           blocksize=512)
class_cog, conf_cog = OUT / f"unet_classes_{NAME}_{YEAR}.tif", OUT / f"unet_confidence_{NAME}_{YEAR}.tif"
# The COG driver builds overviews and cannot take a colour map afterwards, so write a GTiff first and translate.
tmp = DATA / "classes_tmp.tif"
with rasterio.open(tmp, "w", driver="GTiff", height=H, width=W, count=1, dtype="uint8", crs=CRS,
                   transform=TRANSFORM, nodata=IGNORE) as dst:
    dst.write(pred, 1)
    dst.write_colormap(1, {i: tuple(int(c[k:k + 2], 16) for k in (1, 3, 5)) + (255,) for i, c in enumerate(COLOURS)})
    dst.update_tags(1, **{f"class_{i}": n for i, n in enumerate(CLASSES)})
rasterio.shutil.copy(tmp, class_cog, driver="COG", compress="deflate", blocksize=512, overview_resampling="mode")
tmp.unlink()
with rasterio.open(conf_cog, "w", dtype="uint8", nodata=0, overview_resampling="average", **cog) as dst:
    dst.write(np.clip(np.round(conf * 100), 1, 100).astype("uint8"), 1)
section("Cloud Optimised GeoTIFFs")
for f in (class_cog, conf_cog):
    with rasterio.open(f) as src:
        print(f"{f.name}: {src.profile['driver']}, {src.dtypes[0]}, tiles {src.block_shapes[0]}, "
              f"overviews {src.overviews(1)}, layout {src.tags(ns='IMAGE_STRUCTURE').get('LAYOUT')}, "
              f"{f.stat().st_size / 1e6:.1f} MB")
print(f"mean confidence {conf.mean():.2f}; pixels below 0.6: {(conf < 0.6).mean():.1%}")

# --- Step 5: vectorise to simplified polygons and save a GeoPackage --------------------------------
t0 = time.time()
MIN_PX = 4                                     # drop specks smaller than 4 pixels (400 m2)
sieved = rasterio.features.sieve(pred, size=MIN_PX, connectivity=4)
polys = [(shape(g), int(v)) for g, v in rasterio.features.shapes(sieved, transform=TRANSFORM, connectivity=4)]
gdf = gpd.GeoDataFrame({"class_id": [v for _, v in polys]}, geometry=[g for g, _ in polys], crs=CRS)
gdf["class"] = gdf["class_id"].map(dict(enumerate(CLASSES)))
n_raw = int(shapely.get_num_coordinates(gdf.geometry.values).sum())
gdf["geometry"] = gdf.simplify(RES / 2, preserve_topology=True)   # 5 m: removes the pixel staircase
n_simple = int(shapely.get_num_coordinates(gdf.geometry.values).sum())
gdf["area_m2"] = gdf.area.round(0)
gpkg = OUT / f"unet_landcover_{NAME}_{YEAR}.gpkg"
gdf.to_file(gpkg, layer="landcover", driver="GPKG")
section("vectorise")
print(f"sieved specks < {MIN_PX} px: {(sieved != pred).mean():.1%} of pixels changed")
print(f"{len(gdf):,} polygons in {time.time() - t0:.0f} s; vertices {n_raw:,} -> {n_simple:,} after simplify "
      f"({1 - n_simple / n_raw:.0%} fewer)")
print(gdf.groupby("class").agg(polygons=("area_m2", "size"), area_km2=("area_m2", lambda a: a.sum() / 1e6),
                               median_m2=("area_m2", "median")).reindex(CLASSES).round(2).to_string())
print(f"saved {gpkg.name} ({gpkg.stat().st_size / 1e6:.1f} MB)")

# --- Step 6: reference labels for the new area and an honest score ---------------------------------
wc_items = catalog.search(collections=["esa-worldcover"], bbox=bbox_ll,
                          query={"esa_worldcover:product_version": {"eq": "2.0.0"}}).item_collection()
wc = odc.stac.load(wc_items, bands=["map"], geobox=geobox, resampling="nearest")["map"].isel(time=0).values
bld_file = DATA / f"overture_buildings_{NAME}.parquet"
if not bld_file.exists():
    b = overture.geodataframe("building", bbox=bbox_ll, stac=True).set_crs(4326, allow_override=True)
    b[["id", "geometry"]].to_parquet(bld_file)
bld = gpd.read_parquet(bld_file).to_crs(CRS)
SUB = 5
fine = rasterio.features.rasterize(((g, 1) for g in bld.geometry), out_shape=(H * SUB, W * SUB),
                                   transform=from_origin(x0, y1, RES / SUB, RES / SUB), dtype="uint8")
cover = fine.reshape(H, SUB, W, SUB).mean((1, 3))
del fine
ref = np.full((H, W), IGNORE, "uint8")
ref[np.isin(wc, [20, 30, 40, 90, 100])] = 3
ref[wc == 10] = 2
ref[np.isin(wc, [50, 60])] = 1
ref[wc == 80] = 4
ref[cover >= 0.5] = 0
ok = ref != IGNORE
cm = np.bincount(ref[ok].astype(int) * NC + pred[ok], minlength=NC * NC).reshape(NC, NC)
tp = np.diag(cm)
fp, fn = cm.sum(0) - tp, cm.sum(1) - tp
score = pd.DataFrame({"IoU": tp / (tp + fp + fn), "F1": 2 * tp / (2 * tp + fp + fn), "precision": tp / (tp + fp),
                      "recall": tp / (tp + fn), "ref share": cm.sum(1) / cm.sum(),
                      "pred share": cm.sum(0) / cm.sum()}, index=CLASSES)
notts = pd.DataFrame(ckpt["test_scores"])
section(f"accuracy against the reference labels ({len(bld):,} Overture buildings, ESA WorldCover 2021)")
print(score.round(3).to_string())
print(f"overall accuracy {tp.sum() / cm.sum():.3f}, mean IoU {score['IoU'].mean():.3f} "
      f"(Nottingham test blocks: mean IoU {notts['IoU'].mean():.3f})")
print("confusion matrix (rows: reference, columns: prediction, % of each reference class):")
print(pd.DataFrame(100 * cm / cm.sum(1, keepdims=True), index=CLASSES, columns=CLASSES).round(1).to_string())

# Object view: which Overture buildings does the predicted building layer find?
pred_bld = gdf[gdf["class"] == "Building"].union_all()
bld["area_m2"] = bld.area
bld["found"] = bld.intersects(pred_bld)
bins = [0, 50, 100, 200, 500, 2000, np.inf]
labels = ["<50", "50-100", "100-200", "200-500", "500-2000", ">2000"]
by_size = bld.groupby(pd.cut(bld["area_m2"], bins, labels=labels), observed=False)["found"].agg(["size", "mean"])
by_size.columns = ["buildings", "share found"]
print("\nOverture buildings touched by a predicted building polygon, by footprint area (m2):")
print(by_size.round(3).to_string())
print(f"all buildings: {bld['found'].mean():.1%}; area-weighted: "
      f"{(bld['area_m2'] * bld['found']).sum() / bld['area_m2'].sum():.1%}")
bpoly = gdf[gdf["class"] == "Building"]
hit = bpoly.intersects(bld.union_all())
print(f"predicted building polygons that touch no Overture footprint: {(~hit).sum()} of {len(bpoly)} "
      f"({bpoly.loc[~hit, 'area_m2'].sum() / bpoly['area_m2'].sum():.1%} of predicted building area)")

# --- Figures ------------------------------------------------------------------------------------
cmap = ListedColormap(COLOURS)
legend = [Patch(color=c, label=n) for c, n in zip(COLOURS, CLASSES)]
with rasterio.open(s2_cog) as src:
    s2 = src.read()
RGB = np.clip(s2[[3, 2, 1]].transpose(1, 2, 0) / 2500, 0, 1) ** 0.8
extent = (x0 / 1000, (x0 + size) / 1000, (y1 - size) / 1000, y1 / 1000)


def frame(ax, title, km=True):
    ax.set_title(title, loc="left", fontsize=9)
    if km:
        ax.set_xlabel("UTM 30N easting (km)")
        ax.set_ylabel("northing (km)")


# 1. the input and the sliding windows
fig, ax = plt.subplots(figsize=(6.4, 6.2), constrained_layout=True)
ax.imshow(RGB, extent=extent)
for r in starts(H):
    for c in starts(W):
        ax.add_patch(Rectangle(((x0 + c * RES) / 1000, (y1 - (r + TILE) * RES) / 1000), TILE * RES / 1000,
                               TILE * RES / 1000, fill=False, edgecolor="white", lw=0.4, alpha=0.7))
r, c = starts(H)[6], starts(W)[6]
ax.add_patch(Rectangle(((x0 + c * RES) / 1000, (y1 - (r + TILE) * RES) / 1000), TILE * RES / 1000, TILE * RES / 1000,
                       fill=False, edgecolor="#ff5c74", lw=1.6))
frame(ax, f"Leicester, Sentinel-2 composite 2021: {n_win} overlapping {TILE} px windows")
fig.savefig(OUT / "01-input-and-windows.png")
plt.close(fig)

# 2. seams: naive tiling against blended windows, zoomed on one tile corner
zr, zc = 3 * TILE - 160, 3 * TILE - 160
z = np.s_[zr:zr + 320, zc:zc + 320]
fig, axes = plt.subplots(1, 3, figsize=(11, 4), constrained_layout=True)
for ax, data, title in [(axes[0], naive[z], "Naive tiles, no overlap"),
                        (axes[1], pred[z], f"Overlapping windows, blended ({OVERLAP} px)"),
                        (axes[2], None, "Pixels that differ")]:
    if data is None:
        ax.imshow(RGB[z])
        d = diff[z]
        ax.imshow(np.ma.masked_where(~d, d), cmap=ListedColormap(["#ff00c8"]), interpolation="nearest")
    else:
        ax.imshow(data, cmap=cmap, vmin=-0.5, vmax=4.5, interpolation="nearest")
    for k in (160,):
        ax.axhline(k - 0.5, color="black", lw=0.6, ls=":")
        ax.axvline(k - 0.5, color="black", lw=0.6, ls=":")
    ax.set_title(title, loc="left", fontsize=9)
    ax.set_xticks([])
    ax.set_yticks([])
fig.legend(handles=legend, loc="lower center", ncol=5, fontsize=8, frameon=False, bbox_to_anchor=(0.5, -0.07))
fig.savefig(OUT / "02-seams-naive-vs-blended.png")
plt.close(fig)

# 3. the final map against the reference
fig, axes = plt.subplots(1, 2, figsize=(11, 5.3), constrained_layout=True)
axes[0].imshow(pred, cmap=cmap, vmin=-0.5, vmax=4.5, extent=extent, interpolation="nearest")
frame(axes[0], f"U-Net prediction (mean IoU {score['IoU'].mean():.2f} against the reference)")
axes[1].imshow(np.ma.masked_equal(ref, IGNORE), cmap=cmap, vmin=-0.5, vmax=4.5, extent=extent,
               interpolation="nearest")
frame(axes[1], "Reference: Overture buildings over ESA WorldCover 2021")
axes[1].legend(handles=legend, loc="upper left", bbox_to_anchor=(1.01, 1), fontsize=8, frameon=False)
fig.savefig(OUT / "03-prediction-vs-reference.png")
plt.close(fig)

# 4. where it fails: a 3 x 3 km zoom on the city centre
cr, cc = H // 2 - 150, W // 2 - 150
z = np.s_[cr:cr + 300, cc:cc + 300]
err = np.where(ref[z] == IGNORE, 0, (pred[z] != ref[z]).astype(int))
fig, axes = plt.subplots(1, 4, figsize=(13, 3.6), constrained_layout=True)
axes[0].imshow(RGB[z])
axes[1].imshow(np.ma.masked_equal(ref[z], IGNORE), cmap=cmap, vmin=-0.5, vmax=4.5, interpolation="nearest")
axes[2].imshow(pred[z], cmap=cmap, vmin=-0.5, vmax=4.5, interpolation="nearest")
axes[3].imshow(conf[z], cmap="magma", vmin=0.2, vmax=1)
axes[3].contour(err, levels=[0.5], colors="#00e5ff", linewidths=0.4)
for ax, t in zip(axes, ["Sentinel-2 true colour, central Leicester", "Reference labels", "U-Net prediction",
                        f"Confidence, with errors outlined ({err.mean():.0%} wrong)"]):
    ax.set_title(t, loc="left", fontsize=9)
    ax.set_xticks([])
    ax.set_yticks([])
fig.colorbar(axes[3].images[0], ax=axes[3], shrink=0.8, label="max softmax probability")
fig.legend(handles=legend, loc="lower center", ncol=5, fontsize=8, frameon=False, bbox_to_anchor=(0.45, -0.08))
fig.savefig(OUT / "04-city-centre-errors.png")
plt.close(fig)

# 5. vector polygons over the image, with Overture buildings for comparison
zx0, zy1 = x0 + (cc + 75) * RES, y1 - (cr + 75) * RES
zb = box(zx0, zy1 - 1500, zx0 + 1500, zy1)
fig, axes = plt.subplots(1, 2, figsize=(10.5, 5.2), constrained_layout=True)
for ax, title in zip(axes, ["Simplified U-Net polygons (GeoPackage)", "Overture building footprints (reference)"]):
    ax.imshow(RGB, extent=(x0, x0 + size, y1 - size, y1))
    ax.set_xlim(zb.bounds[0], zb.bounds[2])
    ax.set_ylim(zb.bounds[1], zb.bounds[3])
    ax.set_title(title, loc="left", fontsize=9)
    ax.set_xticks([])
    ax.set_yticks([])
sub = gdf[gdf.intersects(zb) & gdf["class"].isin(["Building", "Water", "Trees"])].clip(zb)
for k, col in [("Trees", "#7fd17f"), ("Water", "#4fa3ff"), ("Building", "#ff3b5c")]:
    sub[sub["class"] == k].boundary.plot(ax=axes[0], color=col, lw=0.9)
bld[bld.intersects(zb)].clip(zb).boundary.plot(ax=axes[1], color="#ffd400", lw=0.4)
axes[0].legend(handles=[Patch(edgecolor=c, facecolor="none", label=k) for k, c in
                        [("Building", "#ff3b5c"), ("Trees", "#7fd17f"), ("Water", "#4fa3ff")]],
               loc="lower left", fontsize=8, framealpha=0.85)
fig.savefig(OUT / "05-vector-polygons.png")
plt.close(fig)

# 6. per-class IoU: Nottingham test blocks against Leicester
fig, ax = plt.subplots(figsize=(8.5, 3.4), constrained_layout=True)
xs = np.arange(NC)
for k, (vals, lab, col) in enumerate([(notts["IoU"].reindex(CLASSES), "Nottingham test blocks", "#4a5553"),
                                      (score["IoU"], "Leicester (new area)", "#c21f3a")]):
    ax.bar(xs + (k - 0.5) * 0.36, vals, 0.36, color=col, label=lab)
    for x, v in zip(xs + (k - 0.5) * 0.36, vals):
        ax.text(x, v + 0.01, f"{v:.2f}", ha="center", va="bottom", fontsize=7)
ax.set_xticks(xs, CLASSES)
ax.set_ylim(0, 1)
ax.set_ylabel("IoU")
ax.set_title("Per-class IoU: held-out blocks in the training city against a new city", loc="left", fontsize=9)
ax.legend(fontsize=8, frameon=False)
ax.grid(axis="y", alpha=0.3)
fig.savefig(OUT / "06-iou-nottingham-vs-leicester.png")
plt.close(fig)
print(f"\nfigures written to {OUT.name}/")