"""A 3D-printable city model: OpenStreetMap building footprints on Environment Agency LIDAR terrain.

Fetches the 1 m LIDAR DTM and first-return DSM for central Durham, takes building footprints from
OpenStreetMap, removes detail too thin to print, and builds two versions of the city as watertight solids
with manifold3d: flat-roofed blocks, one LIDAR height per building (90th percentile of the DSM inside its
footprint), and buildings whose roofs follow the DSM surface. It also writes the shaped buildings as a
separate part, and a 3MF with terrain and buildings, for a two-colour print. Every mesh is checked for
watertightness and orientation before it is written.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/3d-printed-city-model-osm-lidar/

Data: Environment Agency LIDAR Composite DTM and First Return DSM 1 m via the Defra WCS (Open
Government Licence v3); building footprints from OpenStreetMap (© OpenStreetMap contributors, ODbL).
"""
import json
import sys
import time
import zipfile
import xml.etree.ElementTree as ET
from importlib.metadata import version
from pathlib import Path

import numpy as np
import requests
import rasterio
import shapely
from rasterio import features
from rasterio.transform import from_origin
from pyproj import Transformer
from shapely import affinity
from shapely.geometry import LineString, MultiPolygon, Polygon, box
from shapely.geometry.polygon import orient
from shapely.ops import linemerge, polygonize, unary_union
import manifold3d as mf
import pyvista as pv
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.collections import PolyCollection
from matplotlib.colors import LightSource, Normalize

plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9,
                     "figure.facecolor": "white", "axes.facecolor": "white"})
pv.OFF_SCREEN = True

DATA = Path("data"); DATA.mkdir(exist_ok=True)
OUT = Path("output"); OUT.mkdir(exist_ok=True)
UA = {"User-Agent": "TechnicalGIS-tutorial (https://www.technicalgis.com)"}

AOI = (426_900, 541_700, 428_100, 542_900)              # 1.2 x 1.2 km of central Durham, British National Grid
TILE = 600                                              # metres per WCS request (600 x 600 cells of 1 m)
RES = 2.0                                               # metres per cell of the printed terrain
WCS = "https://environment.data.gov.uk/spatialdata/{}/wcs?service=WCS&version=2.0.1&request=GetCoverage" \
      "&format=image/tiff&CoverageId={}"
DTM = WCS.format("lidar-composite-digital-terrain-model-dtm-1m",
                 "13787b9a-26a4-4775-8523-806d13af58fc__Lidar_Composite_Elevation_DTM_1m")
DSM = WCS.format("lidar-composite-digital-surface-model-first-return-dsm-1m",
                 "df4e3ec3-315e-48aa-aaaf-b5ae74d7b2bb__Lidar_Composite_Elevation_FZ_DSM_1m")
OVERPASS = ["https://overpass-api.de/api/interpreter", "https://overpass.private.coffee/api/interpreter",
            "https://overpass.kumi.systems/api/interpreter"]

MODEL_MM = 160.0            # printed width of the square model
Z_EXAGGERATION = 1.0        # true scale: buildings and hills are not stretched
BASE_MM = 3.0
MIN_FEATURE_MM = 0.8        # two 0.4 mm extrusion lines: narrower parts of a block are removed
MIN_AREA_MM2 = 0.8          # smaller printable pieces are dropped
MIN_HEIGHT_M = 2.0          # below this, the LIDAR has no roof for the building
FALLBACK_M = 3.0            # roof height used for those buildings: one storey
SIMPLIFY_MM = 0.05

SAND, GREY = "#e8dcc0", "#9aa3ad"


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


def get(url, path, tries=6):
    """Download once and cache. The Defra WCS rate-limits bursts (HTTP 429), so back off and retry."""
    if not path.exists():
        for attempt in range(tries):
            try:
                r = requests.get(url, timeout=300, headers=UA)
                r.raise_for_status()
                if not r.headers.get("Content-Type", "").startswith("image/"):
                    raise requests.RequestException(f"not an image: {r.text[:200]}")
                break
            except requests.RequestException as err:
                if attempt == tries - 1:
                    raise
                wait = getattr(err.response, "headers", {}).get("Retry-After")
                time.sleep(float(wait) if wait and wait.isdigit() else 20 * (attempt + 1))
        path.write_bytes(r.content)
        time.sleep(1)
    return path


def mosaic(url, name):
    """Fetch the AOI in TILE x TILE m requests at 1 m and place them in one array (north up)."""
    X0, Y0, X1, Y1 = AOI
    out = np.full((Y1 - Y0, X1 - X0), np.nan)
    for x in range(X0, X1, TILE):
        for y in range(Y0, Y1, TILE):
            f = get(f"{url}&subset=E({x},{x + TILE})&subset=N({y},{y + TILE})", DATA / f"{name}_{x}_{y}_{TILE}m.tif")
            with rasterio.open(f) as src:
                a = src.read(1).astype("float64")
                a[a < -1000] = np.nan
                r, c = int(round(Y1 - src.bounds.top)), int(round(src.bounds.left - X0))
                out[r:r + a.shape[0], c:c + a.shape[1]] = a
    return out


# ------------------------------------------------------------------ mesh helpers (as in parts 1 to 3)
def grid_triangles(r, c):
    idx = np.arange(r * c).reshape(r, c)
    a, b = idx[:-1, :-1].ravel(), idx[:-1, 1:].ravel()
    d, e = idx[1:, :-1].ravel(), idx[1:, 1:].ravel()
    return np.concatenate([np.column_stack([a, d, b]), np.column_stack([b, d, e])])


def heightfield_solid(Z, cell_mm):
    """Closed solid: the surface Z (mm), a flat base at z = 0 and four side walls."""
    rows, cols = Z.shape
    X, Y = np.meshgrid(np.arange(cols) * cell_mm, (rows - 1 - np.arange(rows)) * cell_mm)
    top = np.column_stack([X.ravel(), Y.ravel(), Z.ravel()])
    bot = top.copy()
    bot[:, 2] = 0.0
    n = rows * cols
    tri_top = grid_triangles(rows, cols)
    idx = np.arange(n).reshape(rows, cols)
    ring = np.concatenate([idx[0, :], idx[1:, -1], idx[-1, -2::-1], idx[-2:0:-1, 0]])
    p, q = ring, np.roll(ring, -1)
    walls = np.concatenate([np.column_stack([p, q, p + n]), np.column_stack([q, q + n, p + n])])
    return np.vstack([top, bot]), np.vstack([tri_top, tri_top[:, ::-1] + n, walls])


def to_manifold(verts, tris):
    m = mf.Manifold(mf.Mesh(vert_properties=np.ascontiguousarray(verts, dtype=np.float32),
                            tri_verts=np.ascontiguousarray(tris, dtype=np.uint32)))
    assert m.status() == mf.Error.NoError, m.status()
    return m


def from_manifold(m):
    mesh = m.to_mesh()
    return np.asarray(mesh.vert_properties)[:, :3].astype("float64"), np.asarray(mesh.tri_verts).astype("int64")


def check_mesh(name, verts, tris):
    """Watertight: every undirected edge used twice. Oriented: every directed edge used once."""
    e = np.concatenate([tris[:, [0, 1]], tris[:, [1, 2]], tris[:, [2, 0]]])
    n = len(verts)
    und = np.unique(np.sort(e, axis=1) @ np.array([n, 1]), return_counts=True)[1]
    drc = np.unique(e @ np.array([n, 1]), return_counts=True)[1]
    t = verts[tris]
    vol = np.einsum("ij,ij->i", t[:, 0], np.cross(t[:, 1], t[:, 2])).sum() / 6
    ok = bool((und == 2).all() and (drc == 1).all() and vol > 0)
    print(f"{name:<10s} {len(tris):>9,} triangles | watertight {bool((und == 2).all())} | "
          f"oriented {bool((drc == 1).all())} | volume {vol / 1000:7.2f} cm3")
    return ok, vol


def write_stl(path, verts, tris, header):
    t = verts[tris]
    nrm = np.cross(t[:, 1] - t[:, 0], t[:, 2] - t[:, 0])
    ln = np.linalg.norm(nrm, axis=1, keepdims=True)
    rec = np.zeros(len(tris), dtype=[("normal", "<f4", 3), ("v", "<f4", (3, 3)), ("attr", "<u2")])
    rec["normal"] = np.divide(nrm, ln, out=np.zeros_like(nrm), where=ln > 0)
    rec["v"] = t
    with open(path, "wb") as f:
        f.write(header.encode("ascii")[:80].ljust(80, b" "))
        f.write(np.uint32(len(rec)).tobytes())
        f.write(rec.tobytes())


def write_3mf(path, parts):
    """Minimal 3MF: one mesh object per part, each with its own display colour."""
    ns = "http://schemas.microsoft.com/3dmanufacturing/core/2015/02"
    xml = [f'<?xml version="1.0" encoding="UTF-8"?>\n<model unit="millimeter" xml:lang="en-GB" xmlns="{ns}">',
           "<resources>", '<basematerials id="1">']
    xml += [f'<base name="{name}" displaycolor="{colour}FF"/>' for name, _, _, colour in parts]
    xml.append("</basematerials>")
    for i, (name, v, t, _) in enumerate(parts):
        xml.append(f'<object id="{i + 2}" name="{name}" type="model" pid="1" pindex="{i}"><mesh><vertices>')
        xml.append("".join(f'<vertex x="{x:.4f}" y="{y:.4f}" z="{z:.4f}"/>' for x, y, z in v))
        xml.append("</vertices><triangles>")
        xml.append("".join(f'<triangle v1="{a}" v2="{b}" v3="{c}"/>' for a, b, c in t))
        xml.append("</triangles></mesh></object>")
    xml += ["</resources>", "<build>"] + [f'<item objectid="{i + 2}"/>' for i in range(len(parts))]
    xml += ["</build>", "</model>"]
    with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as z:
        z.writestr("[Content_Types].xml",
                   '<?xml version="1.0" encoding="UTF-8"?><Types xmlns="http://schemas.openxmlformats.org/'
                   'package/2006/content-types"><Default Extension="rels" ContentType="application/'
                   'vnd.openxmlformats-package.relationships+xml"/><Default Extension="model" '
                   'ContentType="application/vnd.ms-package.3dmanufacturing-3dmodel+xml"/></Types>')
        z.writestr("_rels/.rels",
                   '<?xml version="1.0" encoding="UTF-8"?><Relationships xmlns="http://schemas.openxmlformats.org/'
                   'package/2006/relationships"><Relationship Target="/3D/3dmodel.model" Id="rel0" '
                   'Type="http://schemas.microsoft.com/3dmanufacturing/2013/01/3dmodel"/></Relationships>')
        z.writestr("3D/3dmodel.model", "\n".join(xml))


def polygons(g):
    return list(g.geoms) if isinstance(g, MultiPolygon) else [g] if isinstance(g, Polygon) and not g.is_empty else \
        [p for p in getattr(g, "geoms", []) if isinstance(p, Polygon)]


section("versions")
print(f"python {sys.version.split()[0]} | numpy {np.__version__} | rasterio {rasterio.__version__} | "
      f"shapely {version('shapely')} | manifold3d {version('manifold3d')} | pyvista {pv.__version__}")
X0, Y0, X1, Y1 = AOI

# ------------------------------------------------------------------ Step 1: LIDAR terrain and surface
section("Step 1: Environment Agency LIDAR DTM and DSM at 1 m")
dtm = mosaic(DTM, "ea_dtm")
dsm = mosaic(DSM, "ea_dsm")
print(f"DTM {dtm.shape[1]} x {dtm.shape[0]} cells | no-data {np.isnan(dtm).sum():,} | "
      f"{np.nanmin(dtm):.1f} - {np.nanmax(dtm):.1f} m")
print(f"DSM {dsm.shape[1]} x {dsm.shape[0]} cells | no-data {np.isnan(dsm).sum():,} | "
      f"{np.nanmin(dsm):.1f} - {np.nanmax(dsm):.1f} m")
ndsm = dsm - dtm                                       # height above the ground: buildings, trees, bridges
print(f"DSM - DTM: median {np.nanmedian(ndsm):.2f} m | 99th percentile {np.nanpercentile(ndsm, 99):.1f} m | "
      f"below -0.5 m: {(ndsm < -0.5).mean() * 100:.2f} % of cells")

# ------------------------------------------------------------------ Step 2: OpenStreetMap buildings
section("Step 2: OpenStreetMap building footprints")
to_wgs = Transformer.from_crs("EPSG:27700", "EPSG:4326", always_xy=True)
lon, lat = to_wgs.transform([X0, X1, X0, X1], [Y0, Y0, Y1, Y1])
bbox = f"{min(lat)},{min(lon)},{max(lat)},{max(lon)}"
cache = DATA / "osm_buildings.json"
if not cache.exists():
    query = f'[out:json][timeout:120];(way["building"]({bbox});relation["building"]({bbox}););out geom;'
    for attempt in range(12):                          # public Overpass servers are busy: rotate and back off
        endpoint = OVERPASS[attempt % len(OVERPASS)]
        try:
            r = requests.post(endpoint, data={"data": query}, headers=UA, timeout=200)
            r.raise_for_status()
            json.loads(r.text)["elements"]
            cache.write_text(r.text, encoding="utf-8")
            break
        except (requests.RequestException, ValueError, KeyError) as err:
            print(f"{endpoint}: {str(err)[:80]} - retrying")
            time.sleep(10 * (attempt // len(OVERPASS) + 1))
osm = json.loads(cache.read_text(encoding="utf-8"))["elements"]
to_bng = Transformer.from_crs("EPSG:4326", "EPSG:27700", always_xy=True).transform


def ring(geom):
    return [(p["lon"], p["lat"]) for p in geom]


buildings = []
for e in osm:
    if e["type"] == "way" and len(e.get("geometry", [])) >= 4 and e["geometry"][0] == e["geometry"][-1]:
        g = Polygon(ring(e["geometry"]))
    elif e["type"] == "relation":
        parts = {r: [LineString(ring(m["geometry"])) for m in e["members"]
                     if m["type"] == "way" and m.get("role") == r and len(m.get("geometry", [])) >= 2]
                 for r in ("outer", "inner")}
        outer = unary_union(list(polygonize(linemerge(parts["outer"])))) if parts["outer"] else Polygon()
        inner = unary_union(list(polygonize(linemerge(parts["inner"])))) if parts["inner"] else Polygon()
        g = outer.difference(inner)
    else:
        continue
    g = shapely.transform(g, lambda xy: np.column_stack(to_bng(xy[:, 0], xy[:, 1]))).buffer(0)
    if not g.is_empty:
        buildings.append({"id": f"{e['type']}/{e['id']}", "geom": g, "tags": e.get("tags", {})})
inside = box(X0 + 1, Y0 + 1, X1 - 1, Y1 - 1)
for b in buildings:
    b["geom"] = b["geom"].intersection(inside)
buildings = [b for b in buildings if b["geom"].area > 0]
print(f"Overpass: {len(osm):,} elements -> {len(buildings):,} footprints inside the model area "
      f"({sum(b['id'].startswith('relation') for b in buildings)} multipolygons)")
print(f"tagged: height {sum('height' in b['tags'] for b in buildings)} | "
      f"building:levels {sum('building:levels' in b['tags'] for b in buildings)}")
print(f"footprint area: total {sum(b['geom'].area for b in buildings) / 1e4:.1f} ha | "
      f"median {np.median([b['geom'].area for b in buildings]):.0f} m2")

# ------------------------------------------------------------------ Step 3: roof heights from LIDAR
section("Step 3: a roof height for every building from the LIDAR")
labels = features.rasterize([(b["geom"], i + 1) for i, b in enumerate(buildings)], out_shape=dtm.shape,
                            transform=from_origin(X0, Y1, 1.0, 1.0), fill=0, dtype="int32")
lab = labels.ravel()
order = np.argsort(lab, kind="stable")
starts = np.searchsorted(lab[order], np.arange(1, len(buildings) + 2))
n_fallback = n_nocell = 0
for i, b in enumerate(buildings):
    cells = order[starts[i]:starts[i + 1]]
    surface, ground = dsm.ravel()[cells], dtm.ravel()[cells]
    ok = ~np.isnan(surface) & ~np.isnan(ground)
    if ok.sum() == 0:
        c = b["geom"].representative_point()
        r, k = int(Y1 - c.y), int(c.x - X0)
        b["ground"], b["height"], b["source"] = dtm[r, k], FALLBACK_M, "no cell"
        n_nocell += 1
        continue
    b["ground"] = float(np.median(ground[ok]))
    h = float(np.percentile(surface[ok], 90)) - b["ground"]   # 90th percentile: the roof, not the eaves
    if h < MIN_HEIGHT_M:
        h, b["source"] = FALLBACK_M, "fallback"
        n_fallback += 1
    else:
        b["source"] = "lidar"
    b["height"] = h
hts = np.array([b["height"] for b in buildings if b["source"] == "lidar"])
print(f"LIDAR roof heights for {len(hts):,} buildings: median {np.median(hts):.1f} m | "
      f"90th percentile {np.percentile(hts, 90):.1f} m | max {hts.max():.1f} m")
print(f"too small to hold a 1 m cell: {n_nocell} | no roof in the LIDAR (< {MIN_HEIGHT_M:.0f} m): {n_fallback} "
      f"-> {FALLBACK_M:.0f} m")
tallest = max(buildings, key=lambda b: b["height"])
c = tallest["geom"].centroid
print(f"tallest: {tallest['tags'].get('name', tallest['id'])} {tallest['height']:.1f} m above ground "
      f"({c.x:,.0f} E {c.y:,.0f} N)")
lv = [(float(b["tags"]["building:levels"]), b["height"]) for b in buildings
      if b["source"] == "lidar" and b["tags"].get("building:levels", "").replace(".", "", 1).isdigit()]
lv = np.array(lv)
per_storey = lv[:, 1] / lv[:, 0]
print(f"buildings with building:levels and a LIDAR height: {len(lv)} | LIDAR height per storey: median "
      f"{np.median(per_storey):.2f} m (interquartile {np.percentile(per_storey, 25):.2f} - "
      f"{np.percentile(per_storey, 75):.2f} m)")

# ------------------------------------------------------------------ Step 4: scale and printable footprints
section("Step 4: scale and printable footprints")
k = int(RES)
rows, cols = dtm.shape[0] // k, dtm.shape[1] // k
dem = np.nanmean(dtm.reshape(rows, k, cols, k), axis=(1, 3))
zmin = float(np.nanmin(dem))
scale = MODEL_MM / ((cols - 1) * RES)                  # mm per metre
cell_mm = RES * scale
Z = BASE_MM + (dem - zmin) * scale * Z_EXAGGERATION
W_mm = D_mm = (cols - 1) * cell_mm
print(f"terrain {cols} x {rows} cells of {RES:.0f} m | scale 1:{1000 / scale:,.0f} | cell {cell_mm:.3f} mm")
print(f"model {W_mm:.1f} x {D_mm:.1f} mm | terrain {Z.min():.1f} - {Z.max():.1f} mm high | "
      f"{MIN_FEATURE_MM} mm on the print = {MIN_FEATURE_MM / scale:.1f} m on the ground")


def to_mm(g):
    return affinity.affine_transform(g, [scale, 0, 0, scale, -(X0 + RES / 2) * scale, -(Y0 + RES / 2) * scale])


footprints = [to_mm(b["geom"]).simplify(SIMPLIFY_MM) for b in buildings]
blocks = unary_union(footprints)                       # terraces and joined buildings become one block
r = MIN_FEATURE_MM / 2
printable = blocks.buffer(-r, join_style="mitre").buffer(r, join_style="mitre").intersection(
    box(0.5, 0.5, W_mm - 0.5, D_mm - 0.5))            # morphological opening: removes anything thinner than 0.8 mm
kept = []
for b, fp in zip(buildings, footprints):
    parts = [p for p in polygons(fp.intersection(printable).buffer(0)) if p.area >= MIN_AREA_MM2]
    if parts:
        b["mm"] = MultiPolygon(parts)
        kept.append(b)
area_in, area_out = blocks.area, sum(b["mm"].area for b in kept)
print(f"blocks after merging touching footprints: {len(polygons(blocks)):,}")
print(f"printable: {len(kept):,} of {len(buildings):,} buildings kept | footprint area {area_out:,.0f} of "
      f"{area_in:,.0f} mm2 ({area_out / area_in * 100:.1f} %)")

# ------------------------------------------------------------------ Step 5: flat-roofed blocks
section("Step 5: flat-roofed blocks on the terrain")


def prism(poly, top):
    """Extrude a footprint (mm) from z = 1 mm, inside the base, up to `top`: it is always sunk into the terrain."""
    p = orient(poly, 1.0)
    cs = mf.CrossSection([np.array(p.exterior.coords)[:-1]] + [np.array(h.coords)[:-1] for h in p.interiors])
    return mf.Manifold.extrude(cs, top - 1.0).translate((0, 0, 1.0))


t0 = time.perf_counter()
terrain = to_manifold(*heightfield_solid(Z, cell_mm))
flat = mf.Manifold.batch_boolean([prism(p, BASE_MM + (b["ground"] + b["height"] - zmin) * scale * Z_EXAGGERATION)
                                  for b in kept for p in b["mm"].geoms], mf.OpType.Add)
city_flat = terrain + flat
meshes = {}
for name, m in {"terrain": terrain, "flat city": city_flat}.items():
    v, t = from_manifold(m)
    ok, vol = check_mesh(name, v, t)
    assert ok, name
    meshes[name] = (v, t, vol)
print(f"[{sum(len(b['mm'].geoms) for b in kept):,} prisms, booleans and checks: {time.perf_counter() - t0:.1f} s]")

# ------------------------------------------------------------------ Step 6: roofs shaped by the LIDAR
section("Step 6: roofs shaped by the LIDAR surface")
t0 = time.perf_counter()
dsm2 = dsm.reshape(rows, k, cols, k).max(axis=(1, 3))  # block maximum keeps roof edges at roof height
roof_m = np.maximum(dsm2, dem + FALLBACK_M)            # every building at least one storey high
Zr = BASE_MM + (roof_m - zmin) * scale * Z_EXAGGERATION
surface = to_manifold(*heightfield_solid(Zr, cell_mm))
tall = mf.Manifold.batch_boolean([prism(p, Zr.max() + 1) for b in kept for p in b["mm"].geoms], mf.OpType.Add)
shaped = tall ^ surface                                # vertical walls from OSM, roof shape from the DSM
city = terrain + shaped
roofs = shaped - terrain                               # the part of each building above the ground
for name, m in {"city": city, "buildings": roofs}.items():
    v, t = from_manifold(m)
    ok, vol = check_mesh(name, v, t)
    assert ok, name
    meshes[name] = (v, t, vol)
print(f"[booleans and checks: {time.perf_counter() - t0:.1f} s]")
overlap = (terrain ^ roofs).volume()
print(f"terrain + buildings = {(meshes['terrain'][2] + meshes['buildings'][2]) / 1000:.3f} cm3 | "
      f"city = {meshes['city'][2] / 1000:.3f} cm3 | overlap between the two parts {overlap:.4f} mm3")
pieces = roofs.decompose()
vols = np.array([p.volume() for p in pieces])
print(f"buildings part: {len(pieces):,} separate pieces | smallest {vols.min():.3f} mm3 | "
      f"largest {vols.max():.1f} mm3 | top of the model {city.bounding_box()[5]:.1f} mm")
cells = order[starts[buildings.index(tallest)]:starts[buildings.index(tallest) + 1]]
print(f"{tallest['tags'].get('name', 'tallest building')}: flat roof {tallest['height']:.1f} m | LIDAR roof up to "
      f"{np.nanmax(dsm.ravel()[cells]) - tallest['ground']:.1f} m above its median ground")

# ------------------------------------------------------------------ Step 7: files
section("Step 7: write STL and 3MF")
files = {"flat city": "durham_city_flat_roofs.stl", "city": "durham_city.stl", "terrain": "durham_terrain.stl",
         "buildings": "durham_buildings.stl"}
for name, fn in files.items():
    v, t, _ = meshes[name]
    write_stl(OUT / fn, v, t, f"TechnicalGIS Durham {name} - EA LIDAR, (c) OpenStreetMap contributors")
    size = (OUT / fn).stat().st_size
    with open(OUT / fn, "rb") as fh:
        fh.seek(80)
        n = int(np.frombuffer(fh.read(4), "<u4")[0])
    print(f"{fn:<27s} {size / 1e6:5.1f} MB | {n:,} facets | 84 + 50 x facets = file size: {size == 84 + 50 * n}")
write_3mf(OUT / "durham_two_colour.3mf", [("terrain", *meshes["terrain"][:2], SAND),
                                          ("buildings", *meshes["buildings"][:2], GREY)])
with zipfile.ZipFile(OUT / "durham_two_colour.3mf") as z:
    model = ET.fromstring(z.read("3D/3dmodel.model"))
objs = model.findall(".//{http://schemas.microsoft.com/3dmanufacturing/core/2015/02}object")
print(f"durham_two_colour.3mf       {(OUT / 'durham_two_colour.3mf').stat().st_size / 1e6:5.1f} MB | "
      f"parts: {', '.join(o.get('name') for o in objs)}")

# ------------------------------------------------------------------ Figures
section("figures")
ext = (X0 / 1000, X1 / 1000, Y0 / 1000, Y1 / 1000)
hs = LightSource(azdeg=315, altdeg=45).hillshade(np.nan_to_num(dtm, nan=zmin), vert_exag=2, dx=1, dy=1)
norm = Normalize(0, np.percentile(hts, 99))
fig, ax = plt.subplots(figsize=(7.2, 6.4))
ax.imshow(hs, cmap="gray", extent=ext, alpha=0.6)
for src_, ec in [("lidar", "none"), ("fallback", "#d53e4f"), ("no cell", "#d53e4f")]:
    sel = [b for b in buildings if b["source"] == src_]
    polys = [np.array(p.exterior.coords) / 1000 for b in sel for p in polygons(b["geom"])]
    cols_ = [plt.cm.viridis(norm(b["height"])) for b in sel for _ in polygons(b["geom"])]
    ax.add_collection(PolyCollection(polys, facecolors=cols_ if src_ == "lidar" else "white", edgecolors=ec,
                                     linewidths=0.6 if ec != "none" else 0))
plt.colorbar(plt.cm.ScalarMappable(norm=norm, cmap="viridis"), ax=ax, shrink=0.7, label="Roof height above ground (m)")
ax.plot([], [], "s", mfc="white", mec="#d53e4f", label=f"no roof in the LIDAR: {FALLBACK_M:.0f} m used")
ax.legend(loc="lower left", fontsize=8)
ax.set_xlim(ext[:2]); ax.set_ylim(ext[2:])
ax.set_xlabel("Easting (km, British National Grid)"); ax.set_ylabel("Northing (km)")
ax.set_title("OpenStreetMap buildings coloured by LIDAR roof height, central Durham")
fig.savefig(OUT / "01-building-heights.png"); plt.close(fig)

fig, ax = plt.subplots(figsize=(5.6, 4.2))
jit = np.random.default_rng(0).uniform(-0.12, 0.12, len(lv))
ax.scatter(lv[:, 0] + jit, lv[:, 1], s=14, alpha=0.6, color="#3288bd", label="OSM buildings")
xs = np.array([0.5, lv[:, 0].max() + 0.5])
ax.plot(xs, xs * np.median(per_storey), color="#d53e4f", lw=1.2,
        label=f"median {np.median(per_storey):.2f} m per storey")
ax.set_xlabel("building:levels tag in OpenStreetMap"); ax.set_ylabel("LIDAR roof height above ground (m)")
ax.legend(fontsize=8); ax.grid(alpha=0.3)
ax.set_title(f"LIDAR height against tagged storeys ({len(lv)} buildings)")
fig.savefig(OUT / "02-lidar-vs-levels.png"); plt.close(fig)

cx, cy = to_mm(tallest["geom"].centroid).coords[0]
win = box(cx - 14, cy - 10, cx + 14, cy + 10)
fig, ax = plt.subplots(figsize=(7.4, 5.4))
for g, kw in [(blocks, dict(fc="#f4a6a6", ec="#d53e4f", lw=0.5)), (printable, dict(fc=GREY, ec="#444444", lw=0.5))]:
    for p in polygons(g.intersection(win)):
        ax.fill(*p.exterior.xy, **kw)
        for h in p.interiors:
            ax.fill(*h.xy, fc="white", ec=kw["ec"], lw=0.5)
ax.fill([], [], fc=GREY, ec="#444444", label="printed")
ax.fill([], [], fc="#f4a6a6", ec="#d53e4f", label=f"removed: narrower than {MIN_FEATURE_MM} mm")
ax.set_xlim(cx - 14, cx + 14); ax.set_ylim(cy - 10, cy + 10); ax.set_aspect("equal")
ax.set_xlabel("x on the print (mm)"); ax.set_ylabel("y on the print (mm)")
ax.legend(loc="lower left", fontsize=8)
ax.set_title("Footprints at print scale around the tallest building")
fig.savefig(OUT / "03-printable-footprints.png"); plt.close(fig)


def render(parts, png, zoom=1.0, focal=None, elev_deg=35, azim_deg=-60, size=(1600, 1150)):
    p = pv.Plotter(off_screen=True, window_size=size)
    p.set_background("white")
    for fn, colour in parts:
        p.add_mesh(pv.read(OUT / fn), color=colour, smooth_shading=False, ambient=0.25, diffuse=0.75, specular=0.1)
    p.enable_lightkit()
    c = np.array(focal if focal is not None else (W_mm / 2, D_mm / 2, 0))
    dist = 1.25 * max(W_mm, D_mm) / zoom
    el, az = np.radians(elev_deg), np.radians(azim_deg)
    p.camera_position = [tuple(c + dist * np.array([np.cos(el) * np.cos(az), np.cos(el) * np.sin(az), np.sin(el)])),
                         tuple(c), (0, 0, 1)]
    p.camera.view_angle = 30
    p.screenshot(OUT / png)
    p.close()


render([(files["terrain"], SAND), (files["buildings"], GREY)], "04-stl-preview-city.png")
for fn, png in [(files["flat city"], "_flat.png"), (files["city"], "_lidar.png")]:
    render([(fn, SAND)], png, zoom=4.0, focal=(cx, cy, BASE_MM + 4), elev_deg=30, azim_deg=-110, size=(1100, 900))
fig, axes = plt.subplots(1, 2, figsize=(11, 4.6))
for ax, png, title in [(axes[0], "_flat.png", "Flat roofs: one LIDAR height per building"),
                       (axes[1], "_lidar.png", "Roofs shaped by the LIDAR surface")]:
    ax.imshow(plt.imread(OUT / png)); ax.set_axis_off(); ax.set_title(title)
    (OUT / png).unlink()
fig.suptitle("Durham Cathedral and Palace Green, rendered from the two STL files")
fig.savefig(OUT / "05-stl-preview-roofs.png"); plt.close(fig)
print("wrote 01-building-heights, 02-lidar-vs-levels, 03-printable-footprints, 04-stl-preview-city, "
      "05-stl-preview-roofs")