"""Extract building footprints and heights from classified LiDAR in Python.

Streams USGS 3DEP lidar (WI 4-County Statewide 2017, which classifies buildings as
ASPRS class 6) for 1 km x 1 km of downtown Madison, Wisconsin, builds a DTM and a
building nDSM, turns the building mask into regularised footprint polygons, splits
them into roof segments at height steps, computes height statistics per polygon,
compares them with OpenStreetMap buildings (IoU) and exports a GeoPackage and a 3D view.

Data: USGS 3DEP lidar as Cloud Optimised Point Clouds (COPC) from the Microsoft
Planetary Computer (no account needed), and OpenStreetMap buildings via Overpass.
Run with:  python building_footprints.py
"""
import platform
import time
from importlib.metadata import version
from pathlib import Path

import geopandas as gpd
import laspy
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import osmnx as ox
import pandas as pd
import planetary_computer
import pystac_client
import rasterio
import scipy
import shapely
import shapely.affinity
from matplotlib.colors import BoundaryNorm, LightSource, ListedColormap, Normalize
from matplotlib.patches import Patch
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from pyproj import Transformer
from rasterio import features
from rasterio.transform import from_origin
from scipy import ndimage
from scipy.interpolate import griddata
from shapely.geometry import Point, 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()} | laspy {laspy.__version__} | lazrs {version('lazrs')} | "
      f"numpy {np.__version__} | scipy {scipy.__version__} | rasterio {rasterio.__version__} | "
      f"shapely {shapely.__version__} | geopandas {gpd.__version__} | osmnx {ox.__version__}")

# ------------------------------------------------------------------ 1. find the tiles
LON, LAT = -89.38440, 43.07470   # Wisconsin State Capitol dome, Madison
EPSG = 26916                      # NAD83 / UTM zone 16N, the CRS of this lidar project
HALF = 500                        # 1 km x 1 km area of interest
RES = 1.0                         # grid cell size in metres

to_utm = Transformer.from_crs(4326, EPSG, always_xy=True)
to_ll = Transformer.from_crs(EPSG, 4326, always_xy=True)
cx, cy = (round(v) for v in to_utm.transform(LON, LAT))
xmin, xmax, ymin, ymax = cx - HALF, cx + HALF, cy - HALF, cy + HALF
lon0, lat0 = to_ll.transform(xmin, ymin)
lon1, lat1 = to_ll.transform(xmax, ymax)

catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1",
                                    modifier=planetary_computer.sign_inplace)
items = list(catalog.search(collections=["3dep-lidar-copc"], bbox=[lon0, lat0, lon1, lat1]).items())
items = [it for it in items if "WI_4CountyStatewide_2017" in it.id]
print("\n=== STAC search ===")
print(f"area of interest: {xmax - xmin} m x {ymax - ymin} m centred on the Wisconsin State Capitol")
for it in items:
    print(f"{it.id}  points={it.properties['pc:count']:,}")

# ------------------------------------------------- 2. stream only the points we need
laz_path = DATA / "madison_downtown.laz"
if not laz_path.exists():
    bounds = laspy.Bounds(mins=np.array([xmin, ymin]), maxs=np.array([xmax, ymax]))
    parts, header = [], None
    t = time.perf_counter()
    for it in items:
        with laspy.CopcReader.open(it.assets["data"].href) as reader:
            header = header or reader.header
            pts = reader.query(bounds)
            if len(pts):
                parts.append(pts)
    out_header = laspy.LasHeader(point_format=header.point_format.id, version=header.version)
    out_header.offsets = np.array([xmin, ymin, 0.0])
    out_header.scales = np.array([0.01, 0.01, 0.01])
    out_header.add_crs(header.parse_crs())
    merged = laspy.LasData(out_header)
    merged.points = laspy.ScaleAwarePointRecord.zeros(sum(len(p) for p in parts), header=out_header)
    for dim in header.point_format.dimension_names:
        if dim not in ("X", "Y", "Z"):
            merged[dim] = np.concatenate([np.asarray(p[dim]) for p in parts])
    for dim in ("x", "y", "z"):
        setattr(merged, dim, np.concatenate([np.asarray(getattr(p, dim)) for p in parts]))
    merged.write(laz_path)
    print(f"\nstreamed {len(merged.points):,} points in {time.perf_counter() - t:.1f} s")

las = laspy.read(laz_path)
print(f"\n=== {laz_path.name} ===")
print(f"CRS: {las.header.parse_crs().name}")
x, y, z = np.asarray(las.x), np.asarray(las.y), np.asarray(las.z)
cls = np.asarray(las.classification)
print(f"points: {len(x):,} | density {len(x) / ((xmax - xmin) * (ymax - ymin)):.1f} points/m²")
NAMES = {1: "Unclassified", 2: "Ground", 3: "Low vegetation", 4: "Medium vegetation",
         5: "High vegetation", 6: "Building", 7: "Low noise", 9: "Water", 10: "Rail", 17: "Bridge deck",
         18: "High noise"}
print("class     points   share   name")
for c, n in zip(*np.unique(cls, return_counts=True)):
    print(f"{c:>5}  {n:>9,}  {n / len(cls):6.1%}   {NAMES.get(c, '?')}")

# --------------------------------------------------------- 3. DTM and building nDSM
print("\n=== rasterise to 1 m ===")
W = H = int((xmax - xmin) / RES)
transform = from_origin(xmin, ymax, RES, RES)
col = np.clip(((x - xmin) / RES).astype(int), 0, W - 1)
row = np.clip(((ymax - y) / RES).astype(int), 0, H - 1)
cell = row * W + col


def grid_mean(mask):
    s = np.bincount(cell[mask], weights=z[mask], minlength=W * H)
    n = np.bincount(cell[mask], minlength=W * H)
    with np.errstate(invalid="ignore", divide="ignore"):
        return (s / n).reshape(H, W)


def grid_max(mask):
    g = np.full(W * H, -np.inf)
    np.maximum.at(g, cell[mask], z[mask])
    g[np.isinf(g)] = np.nan
    return g.reshape(H, W)


# Water returns are the "ground" of the lakes; buildings leave large holes to interpolate across.
ground = grid_mean(np.isin(cls, [2, 9]))
known = np.isfinite(ground)
print(f"DTM cells with a ground or water return: {known.mean():.1%}")
rr, cc = np.nonzero(known)
gr, gc = np.nonzero(~known)
filled = griddata((rr, cc), ground[known], (gr, gc), method="linear")
nearest = griddata((rr, cc), ground[known], (gr, gc), method="nearest")
dtm = ground.copy()
dtm[gr, gc] = np.where(np.isfinite(filled), filled, nearest)
dsm_all = grid_max(np.isin(cls, [1, 2, 3, 4, 5, 6, 9]))
ndsm_all = dsm_all - dtm
bld_top = grid_max(cls == 6)
ndsm_bld = bld_top - dtm
print(f"DTM {np.nanmin(dtm):.1f}..{np.nanmax(dtm):.1f} m | building nDSM max {np.nanmax(ndsm_bld):.1f} m")


# ------------------------------------------------------------ 4. building mask
print("\n=== building mask ===")
MIN_HEIGHT = 2.5     # metres above ground: lower class 6 cells are walls, steps or misclassified points
MIN_AREA = 20        # m²: drop garden sheds and stray clusters
raw = np.isfinite(ndsm_bld) & (ndsm_bld > MIN_HEIGHT)
print(f"cells with a class 6 return > {MIN_HEIGHT} m above ground: {raw.sum():,} ({raw.mean():.1%})")
# At ~5 points/m² some 1 m cells inside a roof get no return: close those gaps, fill holes < 9 m².
mask = ndimage.binary_closing(raw, structure=np.ones((3, 3)), iterations=1)
holes = ndimage.binary_fill_holes(mask) & ~mask
hole_lab, _ = ndimage.label(holes)
small = np.nonzero(np.bincount(hole_lab.ravel()) < 9)[0]
mask |= np.isin(hole_lab, small[small > 0])
mask = ndimage.binary_opening(mask, structure=np.ones((2, 2)))
labels, n = ndimage.label(mask)
keep = np.nonzero(np.bincount(labels.ravel()) >= MIN_AREA / RES ** 2)[0]
labels, n = ndimage.label(np.isin(labels, keep[keep > 0]))
mask = labels > 0
print(f"after closing, hole filling and opening: {mask.sum():,} cells, {n} connected footprints "
      f"of {MIN_AREA} m² or more")

# ------------------------------------- 5. split joined buildings at roof-height steps
print("\n=== split at roof-height steps ===")
STEP = 3.0           # a jump of more than 3 m within 1 m is a wall between two roofs, not a roof slope
# Fill the few roof cells with no class 6 return from their nearest neighbour.
idx = ndimage.distance_transform_edt(~np.isfinite(ndsm_bld), return_distances=False, return_indices=True)
roof = ndsm_bld[tuple(idx)]
step = ndimage.maximum_filter(roof, size=3) - ndimage.minimum_filter(roof, size=3)
core = mask & (step <= STEP)
seg, _ = ndimage.label(core)
seg_keep = np.nonzero(np.bincount(seg.ravel()) >= MIN_AREA / RES ** 2)[0]
seg = np.where(np.isin(seg, seg_keep[seg_keep > 0]), seg, 0)
# Give the wall cells (and dropped slivers) the label of the nearest roof segment...
idx = ndimage.distance_transform_edt(seg == 0, return_distances=False, return_indices=True)
seg = np.where(mask, seg[tuple(idx)], 0)
# ...but never across a gap between two separate footprints.
pair = labels.astype("int64") * 1_000_000 + seg
_, segments = np.unique(pair, return_inverse=True)
segments = segments.reshape(H, W).astype("int32")
segments[~mask] = 0
n_seg = int(segments.max())
n_split = sum(len(np.unique(segments[labels == i])) > 1 for i in range(1, n + 1))
print(f"{n_seg} roof segments; {n_split} of {n} footprints were split at height steps > {STEP} m")

# --------------------------------------------- 6. polygonise, simplify and regularise
print("\n=== polygonise and regularise ===")
RECT_MIN = 0.85      # if a footprint fills >= 85% of its minimum rotated rectangle, use the rectangle
DP_TOL = 1.0         # Douglas-Peucker tolerance (one cell)
ANGLE_TOL = 20       # edges within 20° of the dominant direction (or its perpendicular) are snapped to it
MIN_EDGE = 2.5       # jogs shorter than 2.5 m between parallel edges are removed


def polygonise(lab):
    """One (multi)polygon per label, with the 1 m staircase edges of the raster."""
    polys = {}
    for geom, val in features.shapes(lab, mask=lab > 0, transform=transform, connectivity=4):
        polys.setdefault(int(val), []).append(shape(geom))
    return gpd.GeoDataFrame({"bid": list(polys)}, geometry=[shapely.union_all(p) for p in polys.values()],
                            crs=EPSG).sort_values("bid", ignore_index=True)


def dominant_angle(poly):
    """Direction of the longest side of the minimum rotated rectangle, in degrees."""
    c = np.asarray(poly.minimum_rotated_rectangle.exterior.coords)
    e = np.diff(c, axis=0)[:2]
    i = np.argmax(np.hypot(e[:, 0], e[:, 1]))
    return np.degrees(np.arctan2(e[i, 1], e[i, 0]))


def orthogonal_ring(coords):
    """Snap near-horizontal and near-vertical edges of a ring (already rotated to the dominant
    direction) to exact horizontals and verticals, then rebuild the corners as line intersections."""
    pts = np.asarray(coords)[:-1]
    edges = []                                            # [kind, offset, p, q, length]
    for p, q in zip(pts, np.roll(pts, -1, axis=0)):
        d = q - p
        length = np.hypot(*d)
        if length < 1e-6:
            continue
        ang = np.degrees(np.arctan2(d[1], d[0])) % 180
        kind = "h" if min(ang, 180 - ang) <= ANGLE_TOL else "v" if abs(ang - 90) <= ANGLE_TOL else "d"
        off = (p[1] + q[1]) / 2 if kind == "h" else (p[0] + q[0]) / 2
        if edges and kind != "d" and edges[-1][0] == kind:   # merge runs of parallel edges (length-weighted)
            e = edges[-1]
            e[1] = (e[1] * e[4] + off * length) / (e[4] + length)
            e[3], e[4] = q, e[4] + length
        else:
            edges.append([kind, off, p, q, length])
    if len(edges) > 1 and edges[0][0] == edges[-1][0] != "d":
        first, last = edges[0], edges.pop()
        first[1] = (first[1] * first[4] + last[1] * last[4]) / (first[4] + last[4])
        first[2], first[4] = last[2], first[4] + last[4]
    # Remove short jogs: an edge under MIN_EDGE between two parallel neighbours is a staircase step,
    # so drop it and merge the neighbours into one straight edge.
    while len(edges) > 4:
        n_e = len(edges)
        cand = [i for i in range(n_e) if edges[i][0] != "d" and edges[i][4] < MIN_EDGE
                and edges[i - 1][0] == edges[(i + 1) % n_e][0] not in ("d", edges[i][0])]
        if not cand:
            break
        i = min(cand, key=lambda j: edges[j][4])
        a, c = edges[i - 1], edges[(i + 1) % n_e]
        a[1] = (a[1] * a[4] + c[1] * c[4]) / (a[4] + c[4])
        a[3], a[4] = c[3], a[4] + c[4] + edges[i][4]
        for j in sorted({i, (i + 1) % n_e}, reverse=True):
            edges.pop(j)
    if len(edges) < 3:
        return None

    def line(e):
        kind, off, p, q, _ = e
        return ((0.0, off), (1.0, 0.0)) if kind == "h" else ((off, 0.0), (0.0, 1.0)) if kind == "v" else (p, q - p)

    verts = []
    for prev, cur in zip(np.roll(np.arange(len(edges)), 1), range(len(edges))):
        (p1, d1), (p2, d2) = line(edges[prev]), line(edges[cur])
        den = d1[0] * d2[1] - d1[1] * d2[0]
        if abs(den) < 1e-9:                                   # parallel neighbours: keep the original vertex
            verts.append(tuple(edges[cur][2]))
            continue
        t = ((p2[0] - p1[0]) * d2[1] - (p2[1] - p1[1]) * d2[0]) / den
        verts.append((p1[0] + t * d1[0], p1[1] + t * d1[1]))
    return verts


def regularise_part(poly):
    mrr = poly.minimum_rotated_rectangle
    if poly.area / mrr.area >= RECT_MIN:
        return mrr, "rectangle"
    # A 0.5 m closing removes the 1 m staircase, then Douglas-Peucker drops redundant vertices.
    simp = poly.buffer(0.5, join_style="mitre").buffer(-0.5, join_style="mitre").simplify(DP_TOL)
    if simp.geom_type != "Polygon" or simp.is_empty:
        return poly, "staircase"
    # Rotate the building so its dominant direction is horizontal and re-grid it at 0.5 m there:
    # the outline becomes horizontal and vertical edges plus short jogs, which orthogonal_ring removes.
    theta, origin = dominant_angle(poly), poly.centroid
    rot = shapely.affinity.rotate(poly.buffer(1, join_style="round").buffer(-1, join_style="round"),
                                  -theta, origin=origin)
    x0, y0, x1, y1 = rot.bounds
    tr = from_origin(x0 - 1, y1 + 1, 0.5, 0.5)
    grid = features.rasterize([(rot, 1)], out_shape=(int((y1 - y0) / 0.5) + 4, int((x1 - x0) / 0.5) + 4),
                              transform=tr, dtype="uint8")
    parts = [shape(g) for g, v in features.shapes(grid, mask=grid > 0, transform=tr, connectivity=4)]
    if parts:
        big = max(parts, key=lambda p: p.area)
        rings = [orthogonal_ring(r.coords) for r in [big.exterior, *big.interiors]]
        if rings[0] is not None:
            ortho = shapely.Polygon(rings[0], [r for r in rings[1:] if r is not None and len(r) >= 4])
            ortho = shapely.affinity.rotate(shapely.make_valid(ortho), theta, origin=origin)
            if ortho.geom_type == "Polygon" and ortho.intersection(poly).area / ortho.union(poly).area >= 0.8:
                return ortho, "orthogonal"
    return simp, "simplified"


def regularise(geom):
    """Oriented rectangle for simple blocks, edges snapped to the dominant direction for the rest."""
    if geom.geom_type == "Polygon":
        return regularise_part(geom)
    parts = [regularise_part(p) for p in geom.geoms]
    big = max(range(len(parts)), key=lambda i: geom.geoms[i].area)
    return shapely.union_all([g for g, _ in parts]), parts[big][1]


def regularised(raw_gdf):
    out = raw_gdf.copy()
    reg = raw_gdf.geometry.apply(regularise)
    out["geometry"] = [g for g, _ in reg]
    out["method"] = [m for _, m in reg]
    out["iou_raw"] = (out.geometry.intersection(raw_gdf.geometry).area /
                      out.geometry.union(raw_gdf.geometry).area)
    return out


raw_fp, raw_seg = polygonise(labels.astype("int32")), polygonise(segments)
fp, sg = regularised(raw_fp), regularised(raw_seg)
for name, r_, g_ in [("footprints", raw_fp, fp), ("roof segments", raw_seg, sg)]:
    print(f"{name}: {len(g_)} polygons | vertices {r_.geometry.count_coordinates().mean():.0f} -> "
          f"{g_.geometry.count_coordinates().mean():.0f} on average | "
          f"{g_['method'].value_counts().to_dict()} | IoU with staircase median {g_['iou_raw'].median():.3f}, "
          f"min {g_['iou_raw'].min():.3f} | area {r_.area.sum():,.0f} -> {g_.area.sum():,.0f} m²")

# ---------------------------------------------------- 7. height statistics per building
print("\n=== heights per polygon ===")
# Use the points themselves, not the raster: every class 6 return gets a height above the DTM.
b = cls == 6
xb, yb = x[b], y[b]
hb = z[b] - dtm[row[b], col[b]]


def add_heights(gdf, lab):
    lab_b = lab[row[b], col[b]]
    pts = pd.DataFrame({"bid": lab_b, "h": hb})[lab_b > 0]
    stats = pts.groupby("bid")["h"].agg(n_points="size", h_median="median",
                                        h_p95=lambda s: s.quantile(0.95), h_max="max")
    ground = pd.Series(ndimage.median(dtm, lab, index=gdf["bid"]), index=gdf["bid"], name="ground_m")
    gdf = gdf.join(stats, on="bid").join(ground, on="bid")
    gdf["area_m2"] = gdf.area
    gdf["volume_m3"] = ndimage.sum(np.clip(roof, 0, None), lab, index=gdf["bid"])
    gdf["floors_est"] = np.maximum(1, np.round(gdf["h_median"] / 3.5)).astype(int)
    return gdf.round({"h_median": 1, "h_p95": 1, "h_max": 1, "ground_m": 1, "area_m2": 0, "volume_m3": 0,
                      "iou_raw": 3})


fp, sg = add_heights(fp, labels), add_heights(sg, segments)
print(f"{len(fp)} footprints contain {int(fp['n_points'].sum()):,} of {b.sum():,} class 6 points")
print("footprints:")
print(fp[["area_m2", "h_median", "h_p95", "h_max", "n_points"]].describe(
    percentiles=[0.5, 0.9]).round(1).loc[["mean", "50%", "90%", "max"]].to_string())
print("\ntallest five roof segments by 95th-percentile height:")
top = sg.sort_values("h_p95", ascending=False).head(5).copy()
c = top.geometry.centroid.to_crs(4326)
top["lon"], top["lat"] = c.x.round(5), c.y.round(5)
print(top[["bid", "area_m2", "h_median", "h_p95", "h_max", "floors_est", "lon", "lat"]].to_string(index=False))

# Wisconsin State Capitol: 284 ft 5 in (86.7 m) from the ground floor to the top of the "Wisconsin" statue.
cap = fp[fp.contains(Point(cx, cy))].iloc[0]
lab_cap = labels[row[b], col[b]] == int(cap["bid"])
k = np.argmax(np.where(lab_cap, hb, -np.inf))
cen = cap.geometry.centroid
print(f"\nCapitol footprint: {cap['area_m2']:,.0f} m², median roof {cap['h_median']} m, highest return "
      f"{cap['h_max']} m above ground, {np.hypot(xb[k] - cen.x, yb[k] - cen.y):.1f} m from the footprint centroid")
print(f"official height 86.7 m: difference {cap['h_max'] - 86.7:+.1f} m")
print(f"the Capitol footprint was split into {(sg.within(cap.geometry.buffer(2))).sum()} roof segments")

# ---------------------------------------------------------------- 8. OpenStreetMap
print("\n=== OpenStreetMap buildings ===")
osm_path = DATA / "osm_buildings.gpkg"
if not osm_path.exists():
    ox.settings.cache_folder = DATA / "osmnx"
    osm = None
    for attempt in range(5):   # public Overpass servers are often busy: retry with backoff and a mirror
        for url in ("https://overpass-api.de/api", "https://overpass.kumi.systems/api",
                    "https://maps.mail.ru/osm/tools/overpass/api"):
            ox.settings.overpass_url = url
            try:
                osm = ox.features_from_bbox((lon0, lat0, lon1, lat1), {"building": True})
                break
            except Exception as exc:
                print(f"{url} failed: {str(exc).splitlines()[0][:70]}")
        if osm is not None:
            break
        time.sleep(30 * (attempt + 1))
    osm = osm[osm.geometry.type.isin(["Polygon", "MultiPolygon"])].reset_index()
    cols = [c for c in ["id", "building", "name", "height", "building:levels", "start_date"] if c in osm]
    osm[cols + ["geometry"]].to_file(osm_path)
osm = gpd.read_file(osm_path).to_crs(EPSG)
aoi = shapely.box(xmin, ymin, xmax, ymax)
osm = osm[osm.intersects(aoi)].copy()
osm["geometry"] = osm.intersection(aoi)
osm = osm[osm.area > 1].reset_index(drop=True)
osm["oid"] = osm.index
print(f"{len(osm)} OSM building polygons in the area, {osm.area.sum():,.0f} m²; "
      f"{pd.to_numeric(osm.get('height'), errors='coerce').notna().sum()} have a numeric height tag")


def rasterise(gdf):
    return features.rasterize(((g, 1) for g in gdf.geometry), out_shape=(H, W), transform=transform,
                              dtype="uint8").astype(bool)


osm_r, lid_r = rasterise(osm), rasterise(fp)
# NAD83 and WGS84 differ by about a metre in the US: check for a systematic shift before scoring.
scores = {(dx, dy): (np.roll(osm_r, (dy, dx), axis=(0, 1)) & lid_r).sum() /
          (np.roll(osm_r, (dy, dx), axis=(0, 1)) | lid_r).sum()
          for dx in range(-3, 4) for dy in range(-3, 4)}
(dx, dy), best = max(scores.items(), key=lambda kv: kv[1])
print(f"pixel IoU as downloaded: {scores[(0, 0)]:.3f}; best whole-metre shift {dx:+d} m east, "
      f"{-dy:+d} m north gives {best:.3f}")
if best - scores[(0, 0)] > 0.005:
    osm["geometry"] = osm.translate(dx * RES, -dy * RES)
    osm_r = rasterise(osm)
    print("applied the shift to the OSM layer")
tp, fp_px, fn_px = (lid_r & osm_r).sum(), (lid_r & ~osm_r).sum(), (~lid_r & osm_r).sum()
pix_iou = tp / (tp + fp_px + fn_px)
print(f"pixel level: IoU {pix_iou:.3f} | precision {tp / (tp + fp_px):.3f} | recall {tp / (tp + fn_px):.3f}")


def match(gdf):
    """Best one-to-one IoU of every lidar polygon and every OSM building, plus a many-to-one IoU."""
    pairs = gpd.overlay(gdf[["bid", "geometry"]], osm[["oid", "geometry"]], how="intersection",
                        keep_geom_type=True)
    pairs["inter"] = pairs.area
    a_l, a_o = gdf.set_index("bid").area, osm.set_index("oid").area
    pairs["iou"] = pairs["inter"] / (pairs["bid"].map(a_l) + pairs["oid"].map(a_o) - pairs["inter"])
    gdf["iou_osm"] = pairs.groupby("bid")["iou"].max().reindex(gdf["bid"]).fillna(0).round(3).values
    best_o = pairs.groupby("oid")["iou"].max().reindex(osm["oid"]).fillna(0)
    # Many-to-one: each lidar polygon against the union of the OSM buildings it mostly covers.
    grp = pairs[pairs["inter"] / pairs["oid"].map(a_o) > 0.5].groupby("bid")["oid"].apply(list)
    merged = {}
    for bid_, oids in grp.items():
        u = shapely.union_all(osm.loc[oids, "geometry"].values)
        g = gdf.loc[gdf["bid"] == bid_, "geometry"].iloc[0]
        merged[bid_] = g.intersection(u).area / g.union(u).area
    gdf["iou_osm_merged"] = gdf["bid"].map(merged).fillna(0).round(3)
    gdf["n_osm"] = gdf["bid"].map(grp.apply(len)).fillna(0).astype(int)
    return {"lidar polygons": len(gdf),
            "lidar matched 1:1": f"{(gdf['iou_osm'] >= 0.5).mean():.1%}",
            "OSM matched 1:1": f"{(best_o >= 0.5).mean():.1%}",
            "lidar matched many:1": f"{(gdf['iou_osm_merged'] >= 0.5).mean():.1%}",
            "median IoU 1:1": round(gdf["iou_osm"].median(), 3),
            "covering 2+ OSM": int((gdf["n_osm"] >= 2).sum())}, best_o


m_fp, _ = match(fp)
m_sg, best_o = match(sg)
print(f"\nobject level (a match is IoU >= 0.5), {len(osm)} OSM buildings:")
print(pd.DataFrame({"footprints": m_fp, "roof segments": m_sg}).to_string())
osm["iou_lidar"] = best_o.values.round(3)
lid_union, osm_union = shapely.union_all(fp.geometry.values), shapely.union_all(osm.geometry.values)
osm["lidar_cover"] = (osm.intersection(lid_union).area / osm.area).round(3)
fp["osm_cover"] = (fp.intersection(osm_union).area / fp.area).round(3)
missed = osm[osm["lidar_cover"] < 0.1]
print(f"\nOSM buildings less than 10% covered by lidar footprints: {len(missed)}, median area "
      f"{missed.area.median():.0f} m², total {missed.area.sum():,.0f} m²; tags {missed['building'].value_counts().head(3).to_dict()}")
extra = fp[fp["osm_cover"] < 0.1]
print(f"lidar footprints less than 10% covered by OSM: {len(extra)}, median area {extra['area_m2'].median():.0f} m², "
      f"median p95 height {extra['h_p95'].median():.1f} m")

# ------------------------------------------------------------------- 9. export
gpkg = OUT / "madison_buildings.gpkg"
gpkg.unlink(missing_ok=True)
cols = ["bid", "method", "area_m2", "ground_m", "h_median", "h_p95", "h_max", "floors_est",
        "volume_m3", "n_points", "iou_raw", "iou_osm", "iou_osm_merged", "n_osm"]
fp[cols + ["osm_cover", "geometry"]].to_file(gpkg, layer="footprints")
sg[cols + ["geometry"]].to_file(gpkg, layer="roof_segments")
raw_fp.to_file(gpkg, layer="footprints_staircase")
osm.drop(columns=["oid"]).to_file(gpkg, layer="osm_buildings")
with rasterio.open(OUT / "ndsm_buildings.tif", "w", driver="GTiff", width=W, height=H, count=1,
                   dtype="float32", crs=f"EPSG:{EPSG}", transform=transform, nodata=-9999,
                   compress="deflate") as dst:
    dst.write(np.where(mask, roof, -9999).astype("float32"), 1)
sg[cols].to_csv(OUT / "roof_segment_heights.csv", index=False)
print(f"\n=== export ===\nwrote {gpkg.name} (layers: {', '.join(gpd.list_layers(gpkg)['name'])}), "
      f"ndsm_buildings.tif and roof_segment_heights.csv")

# ------------------------------------------------------------------- 10. figures
ext = (xmin, xmax, ymin, ymax)
surface = dtm + np.nan_to_num(np.clip(ndsm_all, 0, None))
hs = LightSource(azdeg=315, altdeg=45).hillshade(surface, vert_exag=1, dx=RES, dy=RES)


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


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


# 01: classified points and nDSM
rng = np.random.default_rng(0)
sel = rng.choice(len(x), size=min(len(x), 1_500_000), replace=False)
cls_style = {2: ("Ground", "#c9b28f"), 5: ("High vegetation", "#3a8f3a"), 6: ("Building", "#d7301f"),
             9: ("Water", "#2b83ba"), 1: ("Unclassified", "#9e9e9e")}
fig, axes = plt.subplots(1, 2, figsize=(11, 5.4), layout="constrained")
ax = axes[0]
for c_, (name, colr) in cls_style.items():
    m = cls[sel] == c_
    ax.scatter(x[sel][m], y[sel][m], s=0.08, c=colr, linewidths=0, rasterized=True)
ax.set_aspect("equal")
frame(ax, "Lidar points coloured by ASPRS class")
ax.legend(handles=[Patch(color=v[1], label=f"{k}  {v[0]}") for k, v in cls_style.items()],
          loc="lower left", fontsize=7.5, framealpha=0.9)
ax.plot(cx, cy, marker="*", ms=12, mec="black", mfc="white")
ax = axes[1]
ax.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
im = ax.imshow(np.ma.masked_invalid(np.where(ndsm_all > 1, ndsm_all, np.nan)), cmap=SPECTRAL, extent=ext,
               vmin=0, vmax=40, alpha=0.9)
frame(ax, "Normalised DSM: height above the DTM")
scalebar(ax, xmin + 30, ymin + 30)
fig.colorbar(im, ax=ax, shrink=0.8, label="metres above ground (clipped at 40 m)")
fig.savefig(OUT / "01-classified-points-ndsm.png")
plt.close(fig)

# 02: mask -> staircase polygons -> regularised footprints -> roof segments, zoomed on a mixed block
zx0, zy0 = cx + 90, cy - 300
zext = (zx0, zx0 + 260, zy0, zy0 + 190)
r0, r1 = int((ymax - zext[3]) / RES), int((ymax - zext[2]) / RES)
c0, c1 = int((zext[0] - xmin) / RES), int((zext[1] - xmin) / RES)
fig, axes = plt.subplots(2, 2, figsize=(10, 7.6), layout="constrained")
axes = axes.ravel()
for ax in axes:
    ax.imshow(hs[r0:r1, c0:c1], cmap="gray", extent=zext, vmin=0, vmax=1)
axes[0].imshow(np.ma.masked_where(~raw[r0:r1, c0:c1], raw[r0:r1, c0:c1]), cmap=ListedColormap(["#d7301f"]),
               extent=zext, alpha=0.75, interpolation="nearest")
raw_fp.boundary.plot(ax=axes[1], color="#d7301f", lw=0.9)
for meth, colr in [("rectangle", "#2b83ba"), ("orthogonal", "#d53e4f"), ("simplified", "#fdae61")]:
    fp[fp["method"] == meth].boundary.plot(ax=axes[2], color=colr, lw=1.5)
axes[2].legend(handles=[Patch(color="#2b83ba", label="oriented rectangle"),
                        Patch(color="#d53e4f", label="snapped to dominant direction"),
                        Patch(color="#fdae61", label="simplified only")], loc="lower right",
               fontsize=7, framealpha=0.9)
sg.plot(ax=axes[3], column="h_p95", cmap=SPECTRAL, norm=Normalize(0, 60), edgecolor="black", lw=0.6, alpha=0.85)
frame(axes[0], "1. Class 6 cells > 2.5 m above ground", zext)
frame(axes[1], "2. Cleaned mask, polygonised (1 m staircase)", zext)
frame(axes[2], "3. Regularised footprints", zext)
frame(axes[3], "4. Split at roof-height steps > 3 m", zext)
scalebar(axes[0], zext[0] + 10, zext[2] + 10, 50)
fig.savefig(OUT / "02-mask-to-footprints.png")
plt.close(fig)

# 03: roof segments coloured by height
bounds_h = [2.5, 5, 10, 15, 20, 30, 45, 60, 90]
norm_h = BoundaryNorm(bounds_h, SPECTRAL.N)
fig, ax = plt.subplots(figsize=(8, 7), layout="constrained")
ax.imshow(hs, cmap="gray", extent=ext, vmin=-0.4, vmax=1.3)
sg.plot(ax=ax, column="h_p95", cmap=SPECTRAL, norm=norm_h, edgecolor="black", lw=0.3)
fig.colorbar(plt.cm.ScalarMappable(norm=norm_h, cmap=SPECTRAL), ax=ax, shrink=0.75,
             label="95th-percentile roof height above ground (m)")
for _, r in sg.sort_values("h_p95", ascending=False).head(5).iterrows():
    p = r.geometry.representative_point()
    ax.annotate(f"{r['h_p95']:.0f} m", (p.x, p.y), fontsize=7.5, ha="center", va="center",
                bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none", alpha=0.85))
frame(ax, f"{len(sg)} roof segments in {len(fp)} footprints, downtown Madison (USGS 3DEP, 2017)")
scalebar(ax, xmin + 30, ymin + 30)
fig.savefig(OUT / "03-footprint-heights.png")
plt.close(fig)

# 04: comparison with OSM
comp = np.zeros((H, W), dtype="uint8")
comp[lid_r & osm_r], comp[lid_r & ~osm_r], comp[~lid_r & osm_r] = 1, 2, 3
fig, axes = plt.subplots(1, 2, figsize=(11.5, 5.2), gridspec_kw={"width_ratios": [1.25, 1]},
                         layout="constrained")
ax = axes[0]
ax.imshow(hs, cmap="gray", extent=ext, vmin=-0.4, vmax=1.3)
ax.imshow(comp, cmap=ListedColormap(["#00000000", "#abdda4", "#d7191c", "#2b83ba"]), extent=ext,
          vmin=-0.5, vmax=3.5, interpolation="nearest")
ax.legend(handles=[Patch(color="#abdda4", label="both"), Patch(color="#d7191c", label="lidar only"),
                   Patch(color="#2b83ba", label="OSM only")], loc="lower left", fontsize=8, framealpha=0.9)
frame(ax, f"Lidar footprints vs OpenStreetMap, pixel IoU {pix_iou:.2f}")
ax = axes[1]
bins = np.linspace(0, 1, 21)
ax.hist(fp["iou_osm"], bins=bins, color="#fdae61", edgecolor="white", label="footprint vs best single OSM building")
ax.hist(sg["iou_osm"], bins=bins, histtype="step", color="#d53e4f", lw=2,
        label="roof segment vs best single OSM building")
ax.hist(fp["iou_osm_merged"], bins=bins, histtype="step", color="#3288bd", lw=2,
        label="footprint vs union of OSM buildings it covers")
ax.axvline(0.5, color="black", ls="--", lw=0.8)
ax.set_xlabel("IoU with OpenStreetMap")
ax.set_ylabel("lidar polygons")
ax.legend(fontsize=7.5, loc="upper left")
ax.set_title("Per-building agreement", loc="left")
fig.savefig(OUT / "04-osm-comparison.png")
plt.close(fig)

# 05: 3D extrusion with matplotlib
fig = plt.figure(figsize=(10, 6.6))
ax = fig.add_subplot(projection="3d")
norm3 = Normalize(0, 60)
for _, r in sg.iterrows():
    geoms = r.geometry.geoms if r.geometry.geom_type == "MultiPolygon" else [r.geometry]
    h = r["h_p95"]
    for g in geoms:
        xy = np.asarray(g.exterior.coords)[:, :2] - [cx, cy]
        walls = [[(*xy[i], 0), (*xy[i + 1], 0), (*xy[i + 1], h), (*xy[i], h)] for i in range(len(xy) - 1)]
        roof_poly = [[(px, py, h) for px, py in xy]]
        ax.add_collection3d(Poly3DCollection(walls + roof_poly, facecolor=SPECTRAL(norm3(h)),
                                             edgecolor="#333333", lw=0.1))
ax.set_xlim(-HALF, HALF), ax.set_ylim(-HALF, HALF), ax.set_zlim(0, 100)
ax.set_box_aspect((1, 1, 0.22))
ax.view_init(elev=30, azim=-60)
ax.set_xlabel("east (m)"), ax.set_ylabel("north (m)"), ax.set_zlabel("height (m)", labelpad=8)
ax.set_title("Roof segments extruded to their 95th-percentile lidar height (true scale)", loc="left")
fig.colorbar(plt.cm.ScalarMappable(norm=norm3, cmap=SPECTRAL), ax=ax, shrink=0.5, pad=0.14,
             label="height above ground (m)")
fig.savefig(OUT / "05-3d-extrusion.png")
plt.close(fig)

# 06: profile through the Capitol and distribution of heights
fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), gridspec_kw={"width_ratios": [1.6, 1]}, layout="constrained")
ax = axes[0]
strip = (np.abs(y - cen.y) < 3) & (np.abs(x - cen.x) < 180)
for c_, (name, colr) in cls_style.items():
    m = strip & (cls == c_)
    ax.scatter(x[m] - cen.x, z[m], s=2, c=colr, linewidths=0, label=name)
ax.plot(np.arange(xmin, xmax) - cen.x + 0.5, dtm[int((ymax - cen.y) / RES)], color="black", lw=0.8, label="DTM")
ax.set_xlim(-180, 180)
ax.set_xlabel("distance east of the dome centre (m)"), ax.set_ylabel("elevation (m, NAVD88)")
ax.set_title("6 m-wide east–west strip through the Capitol", loc="left")
ax.legend(fontsize=7, markerscale=4, loc="upper right")
ax = axes[1]
ax.hist(sg["h_p95"], bins=np.arange(0, 95, 5), weights=sg["area_m2"] / 1e4, color="#66c2a5", edgecolor="white")
ax.set_xlabel("95th-percentile roof height (m)"), ax.set_ylabel("roof area (ha)")
ax.set_title("How much roof is at each height", loc="left")
fig.savefig(OUT / "06-capitol-profile-heights.png")
plt.close(fig)
print("figures written")