"""Least-cost path analysis in Python: Pen-y-Pass to the summit of Yr Wyddfa (Snowdon).

Builds a walking cost surface from real slope (Copernicus GLO-30 DEM) and land cover
(ESA WorldCover 2021), computes the accumulated cost (travel time) surface from Pen-y-Pass
with scikit-image's MCP_Geometric, traces the least-cost path to the summit, and compares it
with the straight line and with the footpaths mapped in OpenStreetMap. It then changes the
cost weights to show how each choice moves the route.

Data: Copernicus DEM GLO-30 and ESA WorldCover from the Microsoft Planetary Computer (no
account needed) and OpenStreetMap footpaths from the Overpass API.
Run with:  python least_cost_path.py
"""
import json
import platform
import time
from importlib.metadata import version
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import planetary_computer
import pystac_client
import rasterio
import requests
import skimage
from matplotlib.colors import BoundaryNorm, LightSource, ListedColormap, LogNorm
from matplotlib.lines import Line2D
from pyproj import Transformer
from rasterio import features
from rasterio.transform import from_origin, rowcol
from rasterio.warp import Resampling, reproject
from scipy import ndimage, sparse
from scipy.sparse import csgraph
from shapely.geometry import LineString
from skimage.graph import MCP_Geometric

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")


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


def tobler_kmh(s):
    """Tobler's hiking function. Pass a signed gradient for direction, or abs(gradient) for the isotropic form."""
    return 6.0 * np.exp(-3.5 * np.abs(s + 0.05))


section("versions")
print(f"python {platform.python_version()} | scikit-image {skimage.__version__} | numpy {np.__version__} | "
      f"rasterio {rasterio.__version__} | networkx {nx.__version__} | scipy {version('scipy')}")

# ------------------------------------------------------------------ 1. study area and end points
RES = 10.0                                         # metres: WorldCover's native cell size
X0, Y0, X1, Y1 = 258_000, 351_000, 267_000, 359_000   # 9 km x 8 km around Yr Wyddfa, British National Grid
W, H = int((X1 - X0) / RES), int((Y1 - Y0) / RES)
TRANSFORM = from_origin(X0, Y1, RES, RES)
TO_BNG = Transformer.from_crs(4326, 27700, always_xy=True)
TO_LL = Transformer.from_crs(27700, 4326, always_xy=True)
START = ("Pen-y-Pass", 53.08066, -4.01934)        # car park and youth hostel at the top of the Llanberis Pass
END = ("Yr Wyddfa summit", 53.06850, -4.07620)     # summit trig point, 1,085 m
sx, sy = TO_BNG.transform(START[2], START[1])
ex, ey = TO_BNG.transform(END[2], END[1])
section("end points")
print(f"start {START[0]}: {sx:.0f} E, {sy:.0f} N | end {END[0]}: {ex:.0f} E, {ey:.0f} N | "
      f"straight-line distance {np.hypot(ex - sx, ey - sy):.0f} m")
lon0, lat0, lon1, lat1 = TO_LL.transform_bounds(X0, Y0, X1, Y1)
catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1",
                                    modifier=planetary_computer.sign_inplace)


def mosaic(collection, asset, path, resampling, dtype, query=None):
    """Reproject every STAC item that touches the box onto the 10 m BNG grid."""
    if not path.exists():
        items = [i for i in catalog.search(collections=[collection], bbox=[lon0, lat0, lon1, lat1]).items()
                 if query is None or query in i.id]
        print(f"{collection}: {', '.join(i.id for i in items)}")
        nodata = np.nan if dtype == "float32" else 0
        out = np.full((H, W), nodata, dtype=dtype)
        for item in items:
            with rasterio.open(item.assets[asset].href) as src:
                part = np.full((H, W), nodata, dtype=dtype)
                reproject(rasterio.band(src, 1), part, dst_transform=TRANSFORM, dst_crs="EPSG:27700",
                          dst_nodata=nodata, resampling=resampling)
                empty = np.isnan(out) if dtype == "float32" else out == 0
                out = np.where(empty, part, out)
        with rasterio.open(path, "w", driver="GTiff", width=W, height=H, count=1, dtype=dtype, crs="EPSG:27700",
                           transform=TRANSFORM, compress="deflate") as dst:
            dst.write(out, 1)
    with rasterio.open(path) as src:
        return src.read(1)


section("rasters on a 10 m British National Grid")
dem = mosaic("cop-dem-glo-30", "data", DATA / "glo30_snowdon_10m.tif", Resampling.bilinear, "float32").astype("float64")
lc = mosaic("esa-worldcover", "map", DATA / "worldcover_snowdon_10m.tif", Resampling.nearest, "uint8", query="2021")
print(f"grid {W} x {H} cells | DEM {np.nanmin(dem):.0f}-{np.nanmax(dem):.0f} m, missing {np.isnan(dem).sum()} | "
      f"WorldCover classes {sorted(np.unique(lc).tolist())}")

# ------------------------------------------------------------------ 2. OpenStreetMap footpaths
OVERPASS = ["https://overpass-api.de/api/interpreter", "https://overpass.private.coffee/api/interpreter",
            "https://maps.mail.ru/osm/tools/overpass/api/interpreter"]
QUERY = f"""[out:json][timeout:90];
way["highway"~"^(path|footway|track|bridleway|steps)$"]({lat0},{lon0},{lat1},{lon1});
(._;>;);
out;"""


def overpass(query, cache):
    if cache.exists():
        return json.loads(cache.read_text())
    for attempt in range(4):
        for url in OVERPASS:
            try:
                r = requests.post(url, data={"data": query}, timeout=120,
                                  headers={"User-Agent": "technicalgis-tutorial"})
                r.raise_for_status()
                els = r.json()["elements"]
                cache.write_text(json.dumps(els))
                return els
            except (requests.RequestException, ValueError):
                continue
        time.sleep(15 * (attempt + 1))
    raise RuntimeError("all Overpass mirrors failed")


section("OpenStreetMap footpaths")
els = overpass(QUERY, DATA / "osm_paths_snowdon.json")
nodes = {e["id"]: TO_BNG.transform(e["lon"], e["lat"]) for e in els if e["type"] == "node"}
ways = [e for e in els if e["type"] == "way"]
SCRAMBLES = {"demanding_mountain_hiking", "alpine_hiking", "demanding_alpine_hiking", "difficult_alpine_hiking"}


def z_at(x, y):
    """Bilinear DEM height at British National Grid coordinates."""
    col = (np.asarray(x) - X0) / RES - 0.5
    row = (Y1 - np.asarray(y)) / RES - 0.5
    return ndimage.map_coordinates(dem, [np.atleast_1d(row), np.atleast_1d(col)], order=1, mode="nearest")


G = nx.Graph()
for w in ways:
    nd, tags = w["nodes"], w.get("tags", {})
    for a, b in zip(nd[:-1], nd[1:]):
        (xa, ya), (xb, yb) = nodes[a], nodes[b]
        length = float(np.hypot(xb - xa, yb - ya))
        grade = abs(float(z_at(xb, yb)[0] - z_at(xa, ya)[0])) / max(length, 0.1)
        G.add_edge(a, b, length=length, minutes=length * 3.6 / tobler_kmh(grade) / 60,
                   name=tags.get("name", ""), sac=tags.get("sac_scale", ""))
paths_len = sum(d["length"] for *_, d in G.edges(data=True))
print(f"{len(ways)} ways, {G.number_of_nodes()} nodes, {paths_len / 1000:.1f} km of path in the box")
named = sorted({w["tags"]["name"] for w in ways if "name" in w.get("tags", {})})
print("named paths:", "; ".join(named))
sac = {}
for *_, d in G.edges(data=True):
    sac[d["sac"] or "(no sac_scale)"] = sac.get(d["sac"] or "(no sac_scale)", 0) + d["length"] / 1000
print("sac_scale:", ", ".join(f"{k} {v:.1f} km" for k, v in sorted(sac.items(), key=lambda kv: -kv[1])))
comps = sorted(nx.connected_components(G), key=len, reverse=True)
print(f"{len(comps)} connected pieces; the largest holds {len(comps[0]) / G.number_of_nodes() * 100:.0f}% of the nodes")
node_ids = np.array(list(comps[0]))                 # snap only to the main network, not to stray fragments
node_xy = np.array([nodes[n] for n in node_ids])


def snap(x, y):
    return node_ids[np.argmin(np.hypot(node_xy[:, 0] - x, node_xy[:, 1] - y))]


n_start, n_end = snap(sx, sy), snap(ex, ey)
print(f"start snapped {np.hypot(*(np.array(nodes[n_start]) - (sx, sy))):.0f} m, "
      f"end snapped {np.hypot(*(np.array(nodes[n_end]) - (ex, ey))):.0f} m to the path network")
(sx, sy), (ex, ey) = nodes[n_start], nodes[n_end]   # every route below uses the same two snapped points


def osm_route(graph, weight, label):
    ids = nx.shortest_path(graph, n_start, n_end, weight=weight)
    used = []
    for a, b in zip(ids[:-1], ids[1:]):
        nm = graph[a][b]["name"] or "(unnamed)"
        if not used or used[-1][0] != nm:
            used.append([nm, 0.0])
        used[-1][1] += graph[a][b]["length"]
    line = LineString([nodes[n] for n in ids])
    mins = sum(graph[a][b]["minutes"] for a, b in zip(ids[:-1], ids[1:]))
    hard = sum(graph[a][b]["length"] for a, b in zip(ids[:-1], ids[1:]) if graph[a][b]["sac"] in SCRAMBLES)
    print(f"{label}: {line.length:.0f} m, {mins:.0f} min, {hard:.0f} m graded as scrambling | via " +
          ", ".join(f"{n} {l:.0f} m" for n, l in used if l > 100))
    return line


shortest_line = osm_route(G, "length", "shortest OSM route")
walkers = G.edge_subgraph([(a, b) for a, b, d in G.edges(data=True) if d["sac"] not in SCRAMBLES])
osm_line = osm_route(walkers, "minutes", "fastest OSM route without scrambles")

# ------------------------------------------------------------------ 3. cost surface
section("cost surface")
gy, gx = np.gradient(dem, RES)
slope = np.hypot(gx, gy)                            # rise over run
slope_deg = np.degrees(np.arctan(slope))
speed = tobler_kmh(slope)                           # km/h; slope >= 0, so up and down cost the same
slope_cost = 3.6 / speed                            # seconds per metre
print(f"slope {np.percentile(slope_deg, 50):.1f} deg median, {np.percentile(slope_deg, 99):.1f} deg 99th percentile, "
      f"{slope_deg.max():.1f} deg max")
print(f"Tobler speed: {tobler_kmh(0):.2f} km/h on the flat, {tobler_kmh(np.tan(np.radians(20))):.2f} km/h at 20 deg, "
      f"{tobler_kmh(np.tan(np.radians(40))):.2f} km/h at 40 deg")

WC = {10: ("Tree cover", "#2b8a3e", 1.5), 20: ("Shrubland", "#a9c25d", 1.3), 30: ("Grassland", "#f3e48c", 1.0),
      40: ("Cropland", "#f0a73a", 1.2), 50: ("Built-up", "#d7301f", 1.0), 60: ("Bare / sparse", "#bdbdbd", 1.25),
      70: ("Snow and ice", "#f7f7f7", 1.5), 80: ("Permanent water", "#3182bd", np.inf),
      90: ("Herbaceous wetland", "#5ab4ac", 2.0), 100: ("Moss and lichen", "#e0c9a6", 1.1)}
lc_mult = np.ones_like(dem)
for code, (_, _, m) in WC.items():
    lc_mult[lc == code] = m
for code in np.unique(lc):
    if code in WC:
        print(f"  {code:>3} {WC[code][0]:<20} {(lc == code).mean() * 100:5.1f}% of cells  multiplier {WC[code][2]}")
cost = slope_cost * lc_mult
print(f"combined cost: {np.nanmin(cost):.2f}-{np.nanmax(cost[np.isfinite(cost)]):.2f} s/m, "
      f"{(~np.isfinite(cost)).sum()} impassable water cells")

# ------------------------------------------------------------------ 4. accumulated cost and least-cost path
section("accumulated cost and least-cost path")
r0, c0 = rowcol(TRANSFORM, sx, sy)
r1, c1 = rowcol(TRANSFORM, ex, ey)
xs = X0 + (np.arange(W) + 0.5) * RES
ys = Y1 - (np.arange(H) + 0.5) * RES
osm_mask = features.rasterize([(osm_line, 1)], out_shape=(H, W), transform=TRANSFORM, all_touched=True)
all_paths = [LineString([nodes[a], nodes[b]]) for a, b in G.edges]
paths_mask = features.rasterize([(g, 1) for g in all_paths], out_shape=(H, W), transform=TRANSFORM, all_touched=True)
dist_to_osm = ndimage.distance_transform_edt(osm_mask == 0) * RES
dist_to_any = ndimage.distance_transform_edt(paths_mask == 0) * RES


def least_cost(costs):
    t = time.perf_counter()
    mcp = MCP_Geometric(costs, fully_connected=True, sampling=(RES, RES))
    acc, _ = mcp.find_costs([(r0, c0)])
    rc = np.array(mcp.traceback((r1, c1)))
    return acc, rc, time.perf_counter() - t


def describe(rc):
    line = LineString(np.c_[xs[rc[:, 1]], ys[rc[:, 0]]])
    z = dem[rc[:, 0], rc[:, 1]]
    d_osm = dist_to_osm[rc[:, 0], rc[:, 1]]
    d_any = dist_to_any[rc[:, 0], rc[:, 1]]
    stats = dict(length=line.length, ascent=float(np.clip(np.diff(z), 0, None).sum()),
                 within50=float((d_osm <= 50).mean() * 100), within50_any=float((d_any <= 50).mean() * 100),
                 mean_d=float(d_osm.mean()))
    return line, stats


def walk(line, step=10.0):
    """Walk a line uphill from Pen-y-Pass: Tobler speed from the gradient along the line, sampled every 10 m."""
    d = np.arange(0, line.length + step, step).clip(max=line.length)
    p = np.array([line.interpolate(v).coords[0] for v in d])
    z = z_at(p[:, 0], p[:, 1])
    ds = np.diff(d)
    ok = ds > 0
    grade = np.diff(z)[ok] / ds[ok]
    rr, cc = rowcol(TRANSFORM, p[:, 0], p[:, 1])
    terrain = slope_deg[np.array(rr), np.array(cc)]
    return dict(minutes=float(np.sum(ds[ok] * 3.6 / tobler_kmh(grade)) / 60),
                ascent=float(np.clip(np.diff(z), 0, None).sum()),
                water=float((lc[np.array(rr), np.array(cc)] == 80).mean() * line.length),
                steepest_terrain=float(terrain.max()), over30=float((terrain > 30).mean() * 100))


acc, rc, dt = least_cost(cost)
lcp, st = describe(rc)
print(f"MCP_Geometric on {W * H:,} cells: {dt:.2f} s | accumulated cost at the summit {acc[r1, c1] / 60:.0f} min")
print(f"least-cost path: {st['length']:.0f} m, {len(rc)} cells | {st['within50']:.0f}% within 50 m of the OSM route, "
      f"{st['within50_any']:.0f}% within 50 m of any mapped path, mean distance from the OSM route {st['mean_d']:.0f} m")

straight = LineString([(sx, sy), (ex, ey)])
naismith = lambda length, w: length / 5000 * 60 + w["ascent"] / 100 * 10
print(f"{'route':<22} {'length':>8} {'Tobler':>8} {'Naismith':>9} {'ascent':>7} {'water':>6} {'>30 deg':>8} {'steepest':>9}")
for name, line in (("straight line", straight), ("OSM walkers' route", osm_line),
                   ("OSM shortest route", shortest_line), ("least-cost path", lcp)):
    w = walk(line)
    print(f"{name:<22} {line.length:7.0f}m {w['minutes']:6.0f}mn {naismith(line.length, w):7.0f}mn "
          f"{w['ascent']:6.0f}m {w['water']:5.0f}m {w['over30']:7.0f}% {w['steepest_terrain']:7.0f}deg")

# ------------------------------------------------------------------ 5. how the weights change the route
section("cost weights and the route")
steep = np.where(slope_deg > 30, 4.0, 1.0)           # extra penalty for scrambling ground
scenarios = {
    "Distance only": np.where(lc == 80, np.inf, 1.0),
    "Slope only (Tobler)": slope_cost,
    "Slope + land cover": cost,
    "Slope + land cover + 4x above 30 deg": cost * steep,
    "As above + 0.5x on mapped paths": cost * steep * np.where(paths_mask == 1, 0.5, 1.0),
}
results = {}
for name, cst in scenarios.items():
    _, rcs, dts = least_cost(cst)
    line, s = describe(rcs)
    s.update(walk(line))                                           # every route walked with the same Tobler rule
    results[name] = (line, s)
    print(f"{name:<37} {s['length']:5.0f} m | {s['minutes']:4.0f} min | ascent {s['ascent']:4.0f} m | water "
          f"{s['water']:4.0f} m | {s['over30']:3.0f}% over 30 deg | {s['within50']:3.0f}% within 50 m of the OSM route, "
          f"{s['within50_any']:3.0f}% of any path | {dts:.2f} s")

conn = {}
for fc in (False, True):
    mcp = MCP_Geometric(cost, fully_connected=fc, sampling=(RES, RES))
    mcp.find_costs([(r0, c0)])
    conn[fc] = describe(np.array(mcp.traceback((r1, c1))))[0].length
print(f"4-connected moves: {conn[False]:.0f} m | 8-connected moves: {conn[True]:.0f} m "
      f"(diagonals cut the zigzag by {(1 - conn[True] / conn[False]) * 100:.1f}%)")

# ------------------------------------------------------------------ 6. direction-aware (anisotropic) cost
section("anisotropic cost with scipy.sparse.csgraph")
t = time.perf_counter()
idx = np.arange(H * W).reshape(H, W)
src_l, dst_l, w_l = [], [], []
for dr, dc in [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]:
    rs = slice(max(0, -dr), H - max(0, dr)); cs_ = slice(max(0, -dc), W - max(0, dc))   # source cells
    rd = slice(rs.start + dr, rs.stop + dr); cd = slice(cs_.start + dc, cs_.stop + dc)     # their neighbours
    step_len = RES * np.hypot(dr, dc)
    grade = (dem[rd, cd] - dem[rs, cs_]) / step_len               # signed: positive is uphill
    secs = step_len * 3.6 / tobler_kmh(grade) * (lc_mult[rs, cs_] + lc_mult[rd, cd]) / 2
    ok = np.isfinite(secs)
    src_l.append(idx[rs, cs_][ok]); dst_l.append(idx[rd, cd][ok]); w_l.append(secs[ok])
graph = sparse.csr_matrix((np.concatenate(w_l), (np.concatenate(src_l), np.concatenate(dst_l))), shape=(H * W, H * W))
t_build = time.perf_counter() - t
t = time.perf_counter()
acc_aniso, pred = csgraph.dijkstra(graph, directed=True, indices=idx[r0, c0], return_predecessors=True)
t_dij = time.perf_counter() - t
node, seq = idx[r1, c1], []
while node != idx[r0, c0]:
    seq.append(node); node = pred[node]
seq.append(node)
rc_an = np.array(np.unravel_index(np.array(seq[::-1]), (H, W))).T
aniso, st_an = describe(rc_an)
st_an.update(walk(aniso))
print(f"graph: {graph.shape[0]:,} nodes, {graph.nnz:,} directed edges, built in {t_build:.2f} s | Dijkstra {t_dij:.2f} s")
print(f"accumulated cost at the summit {acc_aniso[idx[r1, c1]] / 60:.0f} min (isotropic: {acc[r1, c1] / 60:.0f} min)")
print(f"anisotropic least-cost path: {st_an['length']:.0f} m | Tobler {st_an['minutes']:.0f} min | ascent "
      f"{st_an['ascent']:.0f} m | {st_an['over30']:.0f}% over 30 deg | {st_an['within50']:.0f}% within 50 m of the OSM "
      f"route, {st_an['within50_any']:.0f}% of any path | mean offset {st_an['mean_d']:.0f} m")
results["Anisotropic Tobler + land cover"] = (aniso, st_an)

# ------------------------------------------------------------------ 6. figures
section("figures")
hs = LightSource(azdeg=315, altdeg=45).hillshade(dem, vert_exag=1.5, dx=RES, dy=RES)
ext = (X0, X1, Y0, Y1)
zoom = (259_500, 266_000, 352_500, 357_500)
water = np.ma.masked_where(lc != 80, lc)
BLUE = ListedColormap(["#3182bd"])


def base(ax, z=True):
    ax.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
    ax.imshow(water, cmap=BLUE, extent=ext, alpha=0.9)
    if z:
        ax.set_xlim(zoom[0], zoom[1]); ax.set_ylim(zoom[2], zoom[3])
    ax.set_xticks([]); ax.set_yticks([])


def pts(ax):
    ax.plot(sx, sy, "o", ms=8, mfc="white", mec="black", zorder=6)
    ax.plot(ex, ey, "^", ms=9, mfc="white", mec="black", zorder=6)
    ax.annotate("Pen-y-Pass", (sx, sy), xytext=(6, -14), textcoords="offset points", fontsize=8,
                bbox=dict(fc="white", ec="none", alpha=0.85, pad=1))
    ax.annotate("Yr Wyddfa", (ex, ey), xytext=(-24, 10), textcoords="offset points", fontsize=8,
                bbox=dict(fc="white", ec="none", alpha=0.85, pad=1))


# Figure 1: study area
fig, ax = plt.subplots(figsize=(8.5, 7))
ax.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
im = ax.imshow(dem, cmap=SPECTRAL, extent=ext, alpha=0.55)
ax.imshow(water, cmap=BLUE, extent=ext)
for g in all_paths:
    ax.plot(*g.xy, color="black", lw=0.6, alpha=0.8)
ax.plot(*osm_line.xy, color="white", lw=4.5); ax.plot(*osm_line.xy, color="#7b3294", lw=2.5)
ax.plot(*straight.xy, color="black", lw=1.5, ls="--")
ax.plot(*shortest_line.xy, color="#d53e4f", lw=2, ls=":")
pts(ax)
ax.set_xlim(X0, X1); ax.set_ylim(Y0, Y1); ax.set_xticks([]); ax.set_yticks([])
fig.colorbar(im, ax=ax, shrink=0.7, label="elevation (m, Copernicus GLO-30)")
ax.legend(handles=[Line2D([], [], color="black", lw=0.6, label="OSM footpaths and tracks"),
                   Line2D([], [], color="#7b3294", lw=2.5, label=f"fastest OSM route without scrambles ({osm_line.length / 1000:.1f} km)"),
                   Line2D([], [], color="#d53e4f", lw=2, ls=":", label=f"shortest OSM route, over Crib Goch ({shortest_line.length / 1000:.1f} km)"),
                   Line2D([], [], color="black", lw=1.5, ls="--", label=f"straight line ({straight.length / 1000:.1f} km)"),
                   Line2D([], [], color="#3182bd", lw=6, label="permanent water (WorldCover)")],
          loc="lower left", fontsize=7.5, framealpha=0.9)
ax.set_title("Yr Wyddfa (Snowdon): 9 km x 8 km on British National Grid, 10 m cells")
fig.savefig(OUT / "01-study-area.png"); plt.close(fig)

# Figure 2: cost-surface building steps
fig, axs = plt.subplots(2, 3, figsize=(13, 8.4))
a = axs[0, 0]
a.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
im = a.imshow(slope_deg, cmap=SPECTRAL, extent=ext, vmin=0, vmax=50, alpha=0.75)
a.set_xticks([]); a.set_yticks([]); fig.colorbar(im, ax=a, shrink=0.75, label="slope (degrees)")
a.set_title("1. Slope from the DEM")
a = axs[0, 1]
s = np.linspace(0, 1.2, 200)
a.plot(np.degrees(np.arctan(s)), tobler_kmh(s), color="#d7301f", lw=2)
a.set_xlabel("slope (degrees)"); a.set_ylabel("walking speed (km/h)"); a.grid(alpha=0.3)
a.set_title("2. Tobler's hiking function (isotropic)")
for deg in (0, 10, 20, 30, 40):
    v = tobler_kmh(np.tan(np.radians(deg)))
    a.annotate(f"{v:.1f} km/h", (deg, v), xytext=(4, 4), textcoords="offset points", fontsize=7.5)
    a.plot(deg, v, "o", color="#d7301f", ms=4)
a = axs[0, 2]
im = a.imshow(slope_cost, cmap=SPECTRAL, extent=ext, norm=LogNorm(0.7, 15))
a.set_xticks([]); a.set_yticks([]); fig.colorbar(im, ax=a, shrink=0.75, label="seconds per metre")
a.set_title("3. Slope cost = 3.6 / speed")
codes = sorted(WC)
a = axs[1, 0]
a.imshow(np.ma.masked_equal(lc, 0), cmap=ListedColormap([WC[c][1] for c in codes]), extent=ext,
         norm=BoundaryNorm([c - 5 for c in codes] + [codes[-1] + 5], len(codes)), interpolation="nearest")
a.set_xticks([]); a.set_yticks([]); a.set_title("4. ESA WorldCover 2021")
present = [c for c in codes if (lc == c).any()]
a.legend(handles=[Line2D([], [], marker="s", ls="", mfc=WC[c][1], mec="grey", ms=7, label=WC[c][0]) for c in present],
         loc="lower left", fontsize=6.5, framealpha=0.9)
a = axs[1, 1]
im = a.imshow(np.where(np.isfinite(lc_mult), lc_mult, np.nan), cmap=SPECTRAL, extent=ext, vmin=1, vmax=2)
a.imshow(water, cmap=ListedColormap(["black"]), extent=ext)
a.set_xticks([]); a.set_yticks([]); fig.colorbar(im, ax=a, shrink=0.75, label="multiplier (water impassable, black)")
a.set_title("5. Land-cover multiplier")
a = axs[1, 2]
im = a.imshow(np.where(np.isfinite(cost), cost, np.nan), cmap=SPECTRAL, extent=ext, norm=LogNorm(0.7, 15))
a.imshow(water, cmap=ListedColormap(["black"]), extent=ext)
a.set_xticks([]); a.set_yticks([]); fig.colorbar(im, ax=a, shrink=0.75, label="seconds per metre")
a.set_title("6. Cost = slope cost x multiplier")
fig.suptitle("Building the cost surface, step by step", fontsize=11)
fig.tight_layout(); fig.savefig(OUT / "02-cost-surface-steps.png"); plt.close(fig)

# Figure 3: accumulated cost surfaces, isotropic and anisotropic
fig, axs = plt.subplots(1, 2, figsize=(13, 5.6))
for a, surf, line, title in ((axs[0], acc, lcp, "Isotropic: MCP_Geometric on slope + land cover"),
                             (axs[1], acc_aniso.reshape(H, W), aniso, "Anisotropic: Dijkstra on uphill/downhill edges")):
    mins = np.where(np.isfinite(surf), surf / 60, np.nan)
    a.imshow(hs, cmap="gray", extent=ext, vmin=0, vmax=1)
    im = a.imshow(mins, cmap=SPECTRAL, extent=ext, alpha=0.65, vmin=0, vmax=240)
    cs = a.contour(xs, ys, mins, levels=np.arange(30, 241, 30), colors="black", linewidths=0.6)
    a.clabel(cs, fmt="%d min", fontsize=7)
    a.imshow(water, cmap=BLUE, extent=ext)
    a.plot(*line.xy, color="white", lw=4); a.plot(*line.xy, color="#d7301f", lw=2)
    pts(a); a.set_xticks([]); a.set_yticks([])
    a.set_title(f"{title}\n{surf.flat[idx[r1, c1]] / 60:.0f} min to the summit", fontsize=9)
fig.colorbar(im, ax=axs, shrink=0.8, label="minutes from Pen-y-Pass (accumulated cost)")
fig.savefig(OUT / "03-accumulated-cost.png"); plt.close(fig)


def profile(line):
    d = np.arange(0, line.length, 10.0)
    p = np.array([line.interpolate(v).coords[0] for v in d])
    return d / 1000, z_at(p[:, 0], p[:, 1])


# Figure 4: least-cost paths vs straight line vs OSM route, with profiles
ROUTES = ((straight, "black", "--", "straight line"), (osm_line, "#7b3294", "-", "OSM walkers' route (Pyg Track)"),
          (lcp, "#f46d43", "-", "isotropic least-cost path"), (aniso, "#9e0142", "-", "anisotropic least-cost path"))
fig = plt.figure(figsize=(12, 6.2))
ax = fig.add_axes([0.0, 0.0, 0.5, 1.0]); base(ax)
for g in all_paths:
    ax.plot(*g.xy, color="black", lw=0.5, alpha=0.6)
for line, col, ls_, lab in ROUTES:
    if ls_ == "-":
        ax.plot(*line.xy, color="white", lw=4.2)
    ax.plot(*line.xy, color=col, lw=2.2 if col == "#7b3294" else 1.7, ls=ls_)
pts(ax)
ax.legend(handles=[Line2D([], [], color=c, lw=2, ls=s, label=l) for _, c, s, l in ROUTES],
          loc="lower left", fontsize=7.5, framealpha=0.9)
ax2 = fig.add_axes([0.57, 0.12, 0.42, 0.8])
for line, col, ls_, lab in ROUTES:
    d, z = profile(line)
    ax2.plot(d, z, color=col, ls=ls_, lw=1.8, label=f"{lab}: {line.length / 1000:.2f} km")
ax2.set_xlabel("distance from Pen-y-Pass (km)"); ax2.set_ylabel("elevation (m)"); ax2.grid(alpha=0.3)
ax2.legend(fontsize=8); ax2.set_title("Elevation profiles")
fig.savefig(OUT / "04-route-comparison.png"); plt.close(fig)

# Figure 5: how the weights change the route
cols = ["#5e4fa2", "#3288bd", "#d7301f", "#f46d43", "#66c2a5", "#9e0142"]
fig = plt.figure(figsize=(12.5, 6.4))
ax = fig.add_axes([0.0, 0.0, 0.5, 1.0]); base(ax)
ax.plot(*osm_line.xy, color="#fdbf2d", lw=7, alpha=0.85, solid_capstyle="round")
for (name, (line, s)), col in zip(results.items(), cols):
    ax.plot(*line.xy, color="white", lw=3.2); ax.plot(*line.xy, color=col, lw=1.7)
pts(ax)
ax.legend(handles=[Line2D([], [], color="#fdbf2d", lw=6, alpha=0.85, label="OSM walkers' route (Pyg Track)")] +
          [Line2D([], [], color=c, lw=1.7, label=n) for n, c in zip(results, cols)],
          loc="lower left", fontsize=7, framealpha=0.9)
names = list(results)
y = np.arange(len(names))
short = [n.replace("Slope + land cover + 4x above 30 deg", "Slope + LC + 4x above 30 deg")
         .replace("Anisotropic Tobler + land cover", "Anisotropic Tobler + LC") for n in names]
for i, (key, lab, fmt) in enumerate((("minutes", "walking time uphill, Tobler along the route (min)", "{:.0f}"),
                                     ("within50", "% of route within 50 m of the OSM walkers' route", "{:.0f}%"))):
    b = fig.add_axes([0.66, 0.56 - i * 0.48, 0.32, 0.38])
    vals = [results[n][1][key] for n in names]
    b.barh(y, vals, color=cols)
    for yy, v in zip(y, vals):
        b.text(v, yy, " " + fmt.format(v), va="center", fontsize=7.5)
    b.set_yticks(y, short, fontsize=7.5)
    b.invert_yaxis(); b.set_xlabel(lab); b.set_xlim(0, max(vals) * 1.25); b.grid(axis="x", alpha=0.3)
fig.savefig(OUT / "05-cost-weights.png"); plt.close(fig)
print("saved 5 figures to", OUT)