"""3D terrain visualisation in Python with PyVista: Helvellyn and Striding Edge in the Lake District.

Environment Agency LIDAR Composite DTM (5 m), a Sentinel-2 true-colour drape, hillshade lighting,
vertical exaggeration and a GPX route draped on the surface, rendered off-screen to PNG and MP4.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/3d-terrain-visualisation-pyvista/

Data: Environment Agency LIDAR Composite DTM 1 m via the Defra WCS (OGL v3); Copernicus Sentinel-2 L2A
from the Microsoft Planetary Computer (Copernicus open licence); OpenStreetMap paths from the Geofabrik
Cumbria extract (ODbL).
"""
import base64
import sys
import time
import warnings
from importlib.metadata import version
from pathlib import Path

import numpy as np
import requests
import rasterio
from rasterio.transform import from_origin
from rasterio.warp import reproject, Resampling
from rasterio.windows import from_bounds
from pyproj import Transformer
from scipy import ndimage
import networkx as nx
import osmium
import gpxpy
import gpxpy.gpx
import pystac_client
import planetary_computer
import pyvista as pv
import matplotlib
matplotlib.use("Agg")   # render to files only; no GUI backend
import matplotlib.pyplot as plt
from matplotlib.colors import LightSource

warnings.filterwarnings("ignore")
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9,
                     "figure.facecolor": "white", "axes.facecolor": "white"})
pv.OFF_SCREEN = True     # no window: works on Windows and macOS as is; a headless Linux server needs xvfb or EGL VTK

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)"}
X0, Y0, X1, Y1 = 331_000, 511_000, 340_000, 518_500   # 9 km x 7.5 km around Helvellyn, British National Grid
RES = 5.0                                              # metres: the WCS resamples the 1 m LIDAR to 5 m on the server
TILE = 1500                                            # metres per WCS request
WCS = ("https://environment.data.gov.uk/spatialdata/lidar-composite-digital-terrain-model-dtm-1m/wcs"
       "?service=WCS&version=2.0.1&request=GetCoverage&format=image/tiff"
       "&CoverageId=13787b9a-26a4-4775-8523-806d13af58fc__Lidar_Composite_Elevation_DTM_1m")
PBF = "https://download.geofabrik.de/europe/united-kingdom/england/cumbria-latest.osm.pbf"
WALK = {"path", "footway", "track", "bridleway", "steps", "unclassified", "residential", "service",
        "tertiary", "secondary", "living_street"}
GLENRIDDING = (338_560, 516_930)                       # the car park by the Glenridding Beck bridge
RED = "#d53e4f"                                        # Spectral red, used for the route


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


class Timer:
    def __init__(self, label):
        self.label = label

    def __enter__(self):
        self.t = time.perf_counter()

    def __exit__(self, *exc):
        print(f"[{self.label}: {time.perf_counter() - self.t:.1f} s]")


def get(url, path, **kw):
    if not path.exists():
        for attempt in range(4):
            try:
                r = requests.get(url, timeout=300, headers=UA, **kw)
                r.raise_for_status()
                break
            except requests.RequestException:
                if attempt == 3:
                    raise
                time.sleep(5)
        path.write_bytes(r.content)
    return path


section("versions")
print(f"python {sys.version.split()[0]} | pyvista {pv.__version__} | vtk {version('vtk')} | "
      f"rasterio {rasterio.__version__} | numpy {np.__version__} | gpxpy {version('gpxpy')} | "
      f"networkx {nx.__version__} | osmium {version('osmium')}")

# ------------------------------------------------------------------ Step 1: the DEM
section("Step 1: Environment Agency LIDAR DTM, resampled to 5 m by the WCS")
with Timer("DEM download and mosaic"):
    H, W = int((Y1 - Y0) / RES), int((X1 - X0) / RES)
    dem = np.full((H, W), np.nan)
    dem_tf = from_origin(X0, Y1, RES, RES)
    tiles = []
    for x in range(X0, X1, TILE):
        for y in range(Y0, Y1, TILE):
            tiles.append(get(f"{WCS}&subset=E({x},{x + TILE})&subset=N({y},{y + TILE})&SCALEFACTOR={1 / RES}",
                             DATA / f"ea_dtm_{x}_{y}_{RES:.0f}m.tif"))
            with rasterio.open(tiles[-1]) as src:          # tiles share the 5 m grid: paste each in place
                a = src.read(1).astype("float64")
                a[a < -1000] = np.nan                       # no-data is -3.4e38
                r, c = int((Y1 - src.bounds.top) / RES), int((src.bounds.left - X0) / RES)
                dem[r:r + a.shape[0], c:c + a.shape[1]] = a
print(f"{len(tiles)} WCS tiles -> grid {W} x {H} cells of {RES:.0f} m ({W * RES / 1000:.1f} x {H * RES / 1000:.1f} km), "
      f"origin {dem_tf.c:.0f} E, {dem_tf.f:.0f} N")
gaps = np.isnan(dem)
print(f"no-data cells: {gaps.sum():,} ({gaps.mean() * 100:.2f}%), filled from the nearest valid cell")
idx = ndimage.distance_transform_edt(gaps, return_distances=False, return_indices=True)
dem = dem[tuple(idx)]
dy, dx = np.gradient(dem, RES)
slope = np.degrees(np.arctan(np.hypot(dx, dy)))
print(f"elevation {np.min(dem):.1f} to {np.max(dem):.1f} m (relief {np.ptp(dem):.0f} m) | slope median "
      f"{np.median(slope):.1f} deg, 95th percentile {np.percentile(slope, 95):.1f} deg, max {slope.max():.1f} deg")
ls = LightSource(azdeg=315, altdeg=45)
hs = ls.hillshade(dem, vert_exag=1, dx=RES, dy=RES)

# ------------------------------------------------------------------ Step 2: Sentinel-2
section("Step 2: a clear Sentinel-2 scene, reprojected onto the DEM grid")
to_ll = Transformer.from_crs(27700, 4326, always_xy=True)
bbox = to_ll.transform_bounds(X0, Y0, X1, Y1)
cat = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1",
                                modifier=planetary_computer.sign_inplace)
items = list(cat.search(collections=["sentinel-2-l2a"], bbox=bbox, datetime="2024-05-01/2026-09-30",
                        query={"eo:cloud_cover": {"lt": 20}}).items())
by_day = {}
for i in items:
    if 5 <= i.datetime.month <= 9:
        by_day.setdefault(i.datetime.date(), []).append(i)
print(f"{len(items)} scenes with under 20% tile cloud; {len(by_day)} summer days. The study area straddles "
      f"MGRS tiles {', '.join(sorted({i.id.split('_')[4] for i in items}))}, so each day is mosaicked")
H20, W20 = H // 4, W // 4                             # 20 m grid for the cloud check
best = None
for day in sorted(by_day, key=lambda d: max(i.properties["eo:cloud_cover"] for i in by_day[d]))[:10]:
    scl = np.zeros((H20, W20), "uint8")
    for item in cat.search(collections=["sentinel-2-l2a"], bbox=bbox, datetime=f"{day}").items():
        with rasterio.open(item.assets["SCL"].href) as src:
            part = np.zeros_like(scl)
            reproject(rasterio.band(src, 1), part, dst_transform=from_origin(X0, Y1, 20, 20),
                      dst_crs="EPSG:27700", resampling=Resampling.nearest, src_nodata=0, dst_nodata=0)
        scl = np.where(scl == 0, part, scl)
    bad = np.isin(scl, [0, 3, 8, 9, 10]).mean() * 100     # no data, cloud shadow, cloud, cirrus
    print(f"  {day}: cloud, shadow or no data over the study area {bad:5.2f}%")
    if best is None or bad < best[1]:
        best = (day, bad)
    if bad < 0.5:
        break
day = best[0]
day_items = list(cat.search(collections=["sentinel-2-l2a"], bbox=bbox, datetime=f"{day}").items())
print(f"chosen: {day:%d %B %Y}, tiles {', '.join(i.id for i in day_items)}; "
      f"{best[1]:.2f}% cloud or shadow over the study area")
utm = rasterio.crs.CRS.from_string(day_items[0].properties.get("proj:code")
                                   or f"EPSG:{day_items[0].properties['proj:epsg']}")
ux0, uy0, ux1, uy1 = Transformer.from_crs(27700, utm, always_xy=True).transform_bounds(X0, Y0, X1, Y1)
ux0, uy0, ux1, uy1 = np.floor(ux0 / 10) * 10, np.floor(uy0 / 10) * 10, np.ceil(ux1 / 10) * 10, np.ceil(uy1 / 10) * 10
utm_tf = from_origin(ux0, uy1, 10, 10)
with Timer("Sentinel-2 read and reproject"):
    rgb_bng = np.zeros((3, H, W), "float32")
    native = np.zeros((3, int((uy1 - uy0) / 10), int((ux1 - ux0) / 10)), "float32")
    for b, band in enumerate(["B04", "B03", "B02"]):
        for item in day_items:
            with rasterio.open(item.assets[band].href) as src:
                part = np.zeros((H, W), "float32")
                reproject(rasterio.band(src, 1), part, dst_transform=dem_tf, dst_crs="EPSG:27700",
                          resampling=Resampling.bilinear, src_nodata=0, dst_nodata=0)
                rgb_bng[b] = np.where(rgb_bng[b] == 0, part, rgb_bng[b])
                part = np.zeros(native.shape[1:], "float32")      # the same pixels in UTM, for Step 5
                reproject(rasterio.band(src, 1), part, dst_transform=utm_tf, dst_crs=utm,
                          resampling=Resampling.nearest, src_nodata=0, dst_nodata=0)
                native[b] = np.where(native[b] == 0, part, native[b])
print(f"reprojected to the DEM grid: {W} x {H} cells of {RES:.0f} m, empty cells {(rgb_bng[0] == 0).mean() * 100:.2f}% | "
      f"the same area in {utm} (UTM 30N): {native.shape[2]} x {native.shape[1]} pixels of 10 m")

def stretch(a, lo=None, hi=None):
    lo = np.percentile(a, 1, axis=(1, 2), keepdims=True) if lo is None else lo
    hi = np.percentile(a, 99, axis=(1, 2), keepdims=True) if hi is None else hi
    return (np.clip((a - lo) / (hi - lo), 0, 1) ** 0.8 * 255).astype("uint8").transpose(1, 2, 0), lo, hi


rgb, lo, hi = stretch(rgb_bng)
rgb_native, _, _ = stretch(native, lo, hi)
shaded = (rgb * (0.35 + 0.65 * hs[..., None])).astype("uint8")   # imagery with the hillshade baked in

# How far out is "just resize the image to the DEM shape"? Compare true and naive positions per cell.
cx, cy = np.meshgrid(X0 + RES * (np.arange(W) + 0.5), Y1 - RES * (np.arange(H) + 0.5))
ux, uy = Transformer.from_crs(27700, utm, always_xy=True).transform(cx, cy)
true_col, true_row = ~utm_tf * (ux, uy)
naive_col = (np.arange(W) + 0.5) / W * native.shape[2]
naive_row = (np.arange(H) + 0.5) / H * native.shape[1]
err = np.hypot(true_col - naive_col[None, :], true_row - naive_row[:, None]) * 10
print(f"naive resize vs proper reprojection: mean misplacement {err.mean():.0f} m, max {err.max():.0f} m "
      f"(UTM 30N grid north is rotated against BNG grid north)")

# ------------------------------------------------------------------ Step 3: the route
section("Step 3: Glenridding - Striding Edge - Helvellyn - Swirral Edge from OpenStreetMap, as GPX")
pbf = get(PBF, DATA / "cumbria-latest.osm.pbf")
to_bng = Transformer.from_crs(4326, 27700, always_xy=True)
G, named = nx.Graph(), {}
with Timer("read OSM paths"):
    for w in osmium.FileProcessor(str(pbf), osmium.osm.NODE | osmium.osm.WAY).with_locations():
        if not w.is_way() or w.tags.get("highway") not in WALK:
            continue
        try:
            ll = np.array([(n.lon, n.lat) for n in w.nodes])
        except osmium.InvalidLocationError:
            continue
        ex, ny = to_bng.transform(ll[:, 0], ll[:, 1])
        if not ((ex > X0) & (ex < X1) & (ny > Y0) & (ny < Y1)).any():
            continue
        ids = [n.ref for n in w.nodes]
        for i, (a, b) in enumerate(zip(ids[:-1], ids[1:])):
            G.add_edge(a, b, length=float(np.hypot(ex[i + 1] - ex[i], ny[i + 1] - ny[i])))
            G.nodes[a]["xy"], G.nodes[b]["xy"] = (ex[i], ny[i]), (ex[i + 1], ny[i + 1])
        if w.tags.get("name") in ("Striding Edge", "Swirral Edge"):
            named[w.tags["name"]] = ids
print(f"walkable OSM network: {G.number_of_nodes():,} nodes, {G.number_of_edges():,} segments; "
      f"named ways found: {', '.join(named)}")
xy = np.array([G.nodes[n]["xy"] for n in G.nodes])
nodes = list(G.nodes)


def nearest(p):
    return nodes[int(np.argmin(np.hypot(xy[:, 0] - p[0], xy[:, 1] - p[1])))]


start = nearest(GLENRIDDING)
summit = nearest((334_243, 515_107))                  # Helvellyn summit (OSM peak node, 949.8 m)
se, sw = named["Striding Edge"], named["Swirral Edge"]
se = se if np.hypot(*np.subtract(G.nodes[se[0]]["xy"], GLENRIDDING)) < np.hypot(*np.subtract(G.nodes[se[-1]]["xy"], GLENRIDDING)) else se[::-1]
sw = sw if np.hypot(*np.subtract(G.nodes[sw[0]]["xy"], G.nodes[summit]["xy"])) < np.hypot(*np.subtract(G.nodes[sw[-1]]["xy"], G.nodes[summit]["xy"])) else sw[::-1]
legs = [(start, se[0]), (se[0], se[-1]), (se[-1], summit), (summit, sw[0]), (sw[0], sw[-1]), (sw[-1], start)]
path = [start]
for a, b in legs:
    path += nx.shortest_path(G, a, b, weight="length")[1:]
route_xy = np.array([G.nodes[n]["xy"] for n in path])

gpx = gpxpy.gpx.GPX()
trk = gpxpy.gpx.GPXTrack(name="Helvellyn by Striding Edge and Swirral Edge")
seg = gpxpy.gpx.GPXTrackSegment()
lon, lat = to_ll.transform(route_xy[:, 0], route_xy[:, 1])
seg.points = [gpxpy.gpx.GPXTrackPoint(la, lo_) for lo_, la in zip(lon, lat)]
trk.segments.append(seg); gpx.tracks.append(trk)
(OUT / "helvellyn_route.gpx").write_text(gpx.to_xml())

# From here on we work from the GPX file, exactly as you would with a track from a watch or an app.
track = gpxpy.parse((OUT / "helvellyn_route.gpx").read_text())
pts = np.array([(p.longitude, p.latitude) for p in track.tracks[0].segments[0].points])
rx, ry = to_bng.transform(pts[:, 0], pts[:, 1])
d = np.r_[0, np.cumsum(np.hypot(np.diff(rx), np.diff(ry)))]
step = np.arange(0, d[-1], RES)                        # densify to one point per DEM cell
rx, ry = np.interp(step, d, rx), np.interp(step, d, ry)


def sample(z, x, y):
    """Bilinear DEM height at BNG coordinates."""
    col, row = (x - X0) / RES - 0.5, (Y1 - y) / RES - 0.5
    return ndimage.map_coordinates(z, [row, col], order=1, mode="nearest")


rz = sample(dem, rx, ry)
climb = np.diff(ndimage.uniform_filter1d(rz, 5))
print(f"GPX: {len(pts)} track points written and read back; densified to {len(rx):,} points every {RES:.0f} m")
print(f"route length {step[-1] / 1000:.2f} km | lowest {rz.min():.0f} m, highest {rz.max():.0f} m | "
      f"total ascent {climb[climb > 0].sum():.0f} m (5-point smoothed LIDAR profile)")

# ------------------------------------------------------------------ Figure 1: the inputs
ext = (X0 / 1000, X1 / 1000, Y0 / 1000, Y1 / 1000)
fig = plt.figure(figsize=(11, 7.4))
gs = fig.add_gridspec(2, 2, height_ratios=[3, 1.25], hspace=0.28, wspace=0.08)
ax = fig.add_subplot(gs[0, 0])
ax.imshow(dem, extent=ext, cmap="gist_earth", vmin=100, vmax=1000)
ax.imshow(hs, extent=ext, cmap="gray", alpha=0.45)
cs = ax.contour(cx / 1000, cy / 1000, dem, levels=range(200, 1000, 200), colors="k", linewidths=0.3, alpha=0.6)
ax.plot(rx / 1000, ry / 1000, color=RED, lw=1.6)
ax.set_title(f"a) EA LIDAR DTM ({RES:.0f} m), hillshade, 200 m contours", loc="left")
ax.set_xlabel("easting (km, BNG)"); ax.set_ylabel("northing (km, BNG)")
ax2 = fig.add_subplot(gs[0, 1], sharey=ax)
ax2.imshow(rgb, extent=ext)
ax2.plot(rx / 1000, ry / 1000, color=RED, lw=1.6)
for name, (px, py) in {"Helvellyn": (334243, 515107), "Glenridding": GLENRIDDING, "Red Tarn": (334850, 515270),
                       "Ullswater": (338700, 517900)}.items():
    ax2.annotate(name, (px / 1000, py / 1000), xytext=(-4, 6), textcoords="offset points", fontsize=8, ha="right",
                 color="white", fontweight="bold")
ax2.set_title(f"b) Sentinel-2, {day:%d %b %Y}, on the same grid", loc="left")
ax2.set_xlabel("easting (km, BNG)"); ax2.tick_params(labelleft=False)
ax3 = fig.add_subplot(gs[1, :])
ax3.fill_between(step / 1000, rz, rz.min() - 30, color="#abdda4")
ax3.plot(step / 1000, rz, color="#3288bd", lw=1)
for name, ids in (("Striding Edge", se), ("Swirral Edge", sw)):
    ex0 = np.array([G.nodes[n]["xy"] for n in ids])
    on = np.hypot(rx[:, None] - ex0[None, :, 0], ry[:, None] - ex0[None, :, 1]).min(axis=1) < 6
    ax3.plot(step[on] / 1000, rz[on], ".", color=RED, ms=2)
    k = np.flatnonzero(on)[len(np.flatnonzero(on)) // 2]
    ax3.annotate(name, (step[k] / 1000, rz[k]), xytext=(0, 8), textcoords="offset points", ha="center", fontsize=8)
ax3.set_xlim(0, step[-1] / 1000); ax3.set_ylim(rz.min() - 30, rz.max() + 80)
ax3.set_xlabel("distance along the GPX route (km)"); ax3.set_ylabel("elevation (m)")
ax3.set_title("c) The route's profile sampled from the DTM (bilinear); the two edges in red", loc="left")
ax3.grid(alpha=0.3)
fig.savefig(OUT / "01-inputs.png"); plt.close(fig)

# ------------------------------------------------------------------ Step 4: DEM to a PyVista grid
section("Step 4: the DEM as a PyVista StructuredGrid")
EXAG = 1.5


# The WCS resamples 1 m to 5 m by interpolation, not a 5 x 5 block average, so 1 m micro-relief leaks through.
# Lit at screen resolution that noise turns into moire stripes, so the 3D meshes use a 1-cell Gaussian.
dem_vis = ndimage.gaussian_filter(dem, sigma=1.0)
print(f"display DEM: Gaussian sigma 1 cell, mean absolute change {np.abs(dem_vis - dem).mean():.2f} m, "
      f"95th percentile {np.percentile(np.abs(dem_vis - dem), 95):.2f} m")


def terrain(z, exag, stride=1, r0=0, c0=0):
    """StructuredGrid in local metres (origin at the SW corner) with texture coordinates.
    r0, c0: row and column offset when z is a window of the full DEM."""
    zz = z[::stride, ::stride]
    rows, cols = slice(r0, r0 + z.shape[0], stride), slice(c0, c0 + z.shape[1], stride)
    xs, ys = cx[rows, cols] - X0, cy[rows, cols] - Y0
    g = pv.StructuredGrid(xs, ys, zz * exag)
    g.texture_map_to_plane(origin=(0, 0, 0), point_u=(X1 - X0, 0, 0), point_v=(0, Y1 - Y0, 0), inplace=True)
    g["elevation"] = zz.ravel(order="F")
    return g


with Timer("build grid"):
    grid = terrain(dem_vis, EXAG)
print(f"StructuredGrid: dimensions {grid.dimensions}, {grid.n_points:,} points, {grid.n_cells:,} quad cells, "
      f"{grid.actual_memory_size / 1024:.0f} MiB")
print(f"texture coordinates: u {grid.active_texture_coordinates[:, 0].min():.3f}-{grid.active_texture_coordinates[:, 0].max():.3f}, "
      f"v {grid.active_texture_coordinates[:, 1].min():.3f}-{grid.active_texture_coordinates[:, 1].max():.3f}")
surf = grid.extract_surface()
print(f"as triangles: {surf.triangulate().n_cells:,} triangles")
tex = pv.Texture(rgb)                    # row 0 of the image is north, which is v = 1
tex_shaded = pv.Texture(shaded)
for tx in (tex, tex_shaded):                 # mipmaps stop the 10 m texture shimmering on distant slopes
    tx.mipmap = True
    tx.interpolate = True

CENTRE = np.array([(X1 - X0) / 2, (Y1 - Y0) / 2, 0])
SUN = np.array([-0.5, 0.5, np.sqrt(0.5)])             # azimuth 315 (from the north-west), altitude 45 degrees


def plotter(size=(1200, 760)):
    pl = pv.Plotter(off_screen=True, window_size=size)
    pl.set_background("white")
    return pl


def sun_light(pl):
    pl.remove_all_lights()
    pl.add_light(pv.Light(position=tuple(CENTRE + SUN * 1e5), focal_point=tuple(CENTRE),
                          light_type="scene light", positional=False, intensity=1.0))


def camera(pl, exag, az=100.0, el=28.0, dist=10_500, zoom=1.0):
    """Look at the centre of the scene from compass bearing az, el degrees above the horizon."""
    a, e = np.radians(az), np.radians(el)
    focus = CENTRE + [0, 0, 600 * exag]
    pl.camera.focal_point = tuple(focus)
    pl.camera.position = tuple(focus + dist * np.array([np.sin(a) * np.cos(e), np.cos(a) * np.cos(e), np.sin(e)]))
    pl.camera.up = (0, 0, 1)
    pl.camera.view_angle = 30
    pl.camera.zoom(zoom)


# Close-up of the grid around the summit: every vertex is a DEM cell centre
r0, c0, n = int((Y1 - 515_400) / RES), int((334_100 - X0) / RES), 100
sub = terrain(dem_vis[r0:r0 + n, c0:c0 + n], 1.0, r0=r0, c0=c0)
pl = plotter((900, 700))
pl.add_mesh(sub, scalars="elevation", cmap="gist_earth", clim=(600, 1000), show_edges=True, edge_color="#222222",
            line_width=0.3, show_scalar_bar=False, ambient=0.3, diffuse=0.8)
sun_light(pl)
bx, by = sub.center[0], sub.center[1]
zc = np.mean(dem[r0:r0 + n, c0:c0 + n])
pl.camera.focal_point = (bx, by, zc)
pl.camera.position = (bx + 520, by - 420, zc + 380)
pl.camera.up = (0, 0, 1)
close = pl.screenshot(return_img=True); pl.close()

fig, axs = plt.subplots(1, 2, figsize=(11, 4.6), gridspec_kw={"width_ratios": [1, 1.25]})
k = 12
zz = dem[r0:r0 + k, c0:c0 + k]
axs[0].imshow(zz, cmap="gist_earth", vmin=600, vmax=1000, extent=(0, k * RES, 0, k * RES))
gx, gy = np.meshgrid(np.arange(k) * RES + RES / 2, k * RES - np.arange(k) * RES - RES / 2)
axs[0].plot(gx, gy, color="k", lw=0.4); axs[0].plot(gx.T, gy.T, color="k", lw=0.4)
axs[0].plot(gx, gy, "o", color="k", ms=2.5)
for i in range(0, k, 4):
    for j in range(0, k, 4):
        axs[0].annotate(f"{zz[i, j]:.0f}", (gx[i, j], gy[i, j]), xytext=(2, 2), textcoords="offset points",
                        fontsize=6.5, color="white")
axs[0].set_title(f"a) {k} x {k} DEM cells: one vertex per cell centre,\nz = elevation x exaggeration", loc="left")
axs[0].set_xlabel("metres"); axs[0].set_ylabel("metres")
axs[1].imshow(close); axs[1].axis("off")
axs[1].set_title(f"b) The {n} x {n} cell StructuredGrid by Helvellyn summit,\nedges shown, coloured by elevation, no exaggeration", loc="left")
fig.savefig(OUT / "02-dem-to-grid.png"); plt.close(fig)

def trim(img, pad=6):
    """Crop a screenshot to the rendered terrain (drop the white margins)."""
    rows, cols = np.where((img[..., :3] < 245).any(axis=2))
    return img[max(rows.min() - pad, 0):rows.max() + pad, max(cols.min() - pad, 0):cols.max() + pad]


# ------------------------------------------------------------------ Step 5: texture alignment
section("Step 5: texture alignment")
fig, axs = plt.subplots(1, 3, figsize=(12, 3.9), gridspec_kw={"width_ratios": [1, 1, 0.9]})
lev = [150, 300, 500, 700, 900]
naive_img = rgb_native[(naive_row.astype(int))[:, None], (naive_col.astype(int))[None, :]]
for ax, img, title in ((axs[0], naive_img, "a) Naive resize to the DEM shape"),
                       (axs[1], rgb, "b) Reprojected with rasterio")):
    ax.imshow(img, extent=ext)
    ax.contour(cx / 1000, cy / 1000, dem, levels=lev, colors="white", linewidths=0.5)
    ax.contour(cx / 1000, cy / 1000, dem, levels=[146.5], colors=RED, linewidths=0.9)
    ax.set_xlim(336.5, 340); ax.set_ylim(515.0, 518.5)
    ax.set_title(title, loc="left"); ax.set_xlabel("easting (km)")
axs[0].set_ylabel("northing (km)")
im = axs[2].imshow(err, extent=ext, cmap="Spectral_r")
plt.colorbar(im, ax=axs[2], shrink=0.85, label="misplacement (m)")
axs[2].set_title("c) Naive misplacement", loc="left"); axs[2].set_xlabel("easting (km)")
fig.savefig(OUT / "03-texture-alignment.png"); plt.close(fig)

# ------------------------------------------------------------------ Step 6: vertical exaggeration
section("Step 6: vertical exaggeration")
shots = []
for ex in (1.0, 1.5, 3.0):
    g = terrain(dem_vis, ex)
    pl = plotter((900, 560))
    pl.add_mesh(g, texture=tex, ambient=0.25, diffuse=0.85, specular=0.0, smooth_shading=True)
    sun_light(pl); camera(pl, ex, az=100, el=24, dist=12_500)
    shots.append(pl.screenshot(return_img=True)); pl.close()
    zr = np.ptp(g.points[:, 2])
    print(f"exaggeration {ex:.1f}x: z range in the scene {zr:,.0f} m against {X1 - X0:,} m across; "
          f"95th percentile slope as drawn {np.degrees(np.arctan(np.tan(np.radians(np.percentile(slope, 95))) * ex)):.0f} deg")
fig, axs = plt.subplots(1, 3, figsize=(13, 3.0))
for ax, img, ex in zip(axs, shots, (1.0, 1.5, 3.0)):
    ax.imshow(trim(img)); ax.axis("off"); ax.set_title(f"{ex:g}x vertical exaggeration", loc="left")
fig.savefig(OUT / "04-exaggeration.png"); plt.close(fig)

# ------------------------------------------------------------------ Step 7: lighting
section("Step 7: lighting")
variants = [("a) Texture only, lighting off", dict(texture=tex, lighting=False), False),
            ("b) Texture with a directional sun (NW, 45 deg)", dict(texture=tex, ambient=0.25, diffuse=0.85), True),
            ("c) Hillshade baked into the texture, plus the sun", dict(texture=tex_shaded, ambient=0.35, diffuse=0.7), True),
            ("d) Elevation colour map with the sun", dict(scalars="elevation", cmap="gist_earth", clim=(100, 1000),
                                                       show_scalar_bar=False, ambient=0.25, diffuse=0.85), True)]
shots = []
for title, kw, lit in variants:
    pl = plotter((900, 560))
    pl.add_mesh(grid, smooth_shading=True, **kw)
    if lit:
        sun_light(pl)
    camera(pl, EXAG, az=100, el=24, dist=12_500)
    shots.append(pl.screenshot(return_img=True)); pl.close()
fig, axs = plt.subplots(2, 2, figsize=(11, 5.6))
for ax, img, (title, _, _) in zip(axs.ravel(), shots, variants):
    ax.imshow(trim(img)); ax.axis("off"); ax.set_title(title, loc="left")
fig.subplots_adjust(wspace=0.02, hspace=0.12)
fig.savefig(OUT / "05-lighting.png"); plt.close(fig)

# ------------------------------------------------------------------ Step 8: drape the route
section("Step 8: drape the GPX route and render the final scene")
LIFT = 4.0                                    # metres above the ground, so the line is not hidden in the mesh
line = pv.lines_from_points(np.c_[rx - X0, ry - Y0, (rz + LIFT) * EXAG])
tube = line.tube(radius=9)
print(f"route polyline: {line.n_points:,} points -> tube of {tube.n_points:,} points")
labels = {"Helvellyn 950 m": (334243, 515107), "Striding Edge": (335750, 514900), "Swirral Edge": (334650, 515650),
          "Red Tarn": (334850, 515270), "Catstycam": (334811, 515824)}
lab_pts = np.array([(x - X0, y - Y0, (sample(dem, np.array([x]), np.array([y]))[0] + 60) * EXAG)
                    for x, y in labels.values()])


def scene(pl, mesh, texture):
    pl.add_mesh(mesh, texture=texture, ambient=0.35, diffuse=0.7, specular=0.0, smooth_shading=True)
    pl.add_mesh(tube, color=RED, ambient=0.5, diffuse=0.6)
    sun_light(pl)


pl = plotter((1600, 950))
scene(pl, grid, tex_shaded)
pl.add_point_labels(lab_pts, list(labels), font_size=15, point_size=1, shape_opacity=0.65, shape_color="white",
                    text_color="black", always_visible=True, show_points=False)
camera(pl, EXAG, az=95, el=27, dist=8_200)
with Timer("final render"):
    img = pl.screenshot(return_img=True, window_size=(1600, 950))
plt.imsave(OUT / "06-helvellyn-3d.png", img[150:])          # trim the empty sky above the skyline
pl.close()

# ------------------------------------------------------------------ Step 9: orbit animation
section("Step 9: a rotating MP4, embedded in a small HTML page")
small = terrain(dem_vis, EXAG, stride=2)          # 10 m mesh: a quarter of the points, the same look at video size
print(f"animation mesh: {small.n_points:,} points (every 2nd DEM cell)")
pl = plotter((768, 432))
scene(pl, small, tex_shaded)
mp4 = OUT / "helvellyn_orbit.mp4"
pl.open_movie(str(mp4), framerate=15, quality=5)
FRAMES = 120
with Timer(f"render {FRAMES} frames"):
    for i in range(FRAMES):
        camera(pl, EXAG, az=95 + 360 * i / FRAMES, el=30, dist=11_000)
        pl.write_frame()
pl.close()
size = mp4.stat().st_size
print(f"{mp4.name}: {FRAMES} frames at 768 x 432, {FRAMES / 15:.0f} s, {size / 1e6:.2f} MB")
b64 = base64.b64encode(mp4.read_bytes()).decode()
(OUT / "helvellyn-orbit.html").write_text(f"""<!doctype html>
<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1">
<title>Helvellyn in 3D: a PyVista orbit</title>
<style>html,body{{margin:0;height:100%;background:#ffffff}}video{{width:100%;height:100%;object-fit:contain;display:block}}</style>
</head><body>
<video autoplay loop muted playsinline controls aria-label="A 3D model of Helvellyn rotating through 360 degrees, with the Striding Edge route in red">
<source type="video/mp4" src="data:video/mp4;base64,{b64}"></video>
</body></html>""", encoding="utf-8")
print(f"helvellyn-orbit.html: {(OUT / 'helvellyn-orbit.html').stat().st_size / 1e6:.2f} MB")