"""Interpolated stratum surfaces from AGS4: a 3D ground model of the A417 escarpment, with its evidence.

Part 5 of the TechnicalGIS series "AGS data in GIS". Picks four geological surfaces from the A417
Missing Link borehole logs (rockhead, the top and base of the Inferior Oolite and the base of the
Bridport Sand), kriges them onto a 10 m grid under Environment Agency LIDAR ground, and then tests the
result: leave-one-out cross-validation at the boreholes, kriging standard deviation maps, every borehole
that stops short of a surface used as a one-sided check, and a check that no surface crosses another.
The two places Part 3 found where the logs disagree with their neighbours (the slipped block in CP217 and
the low block at RC507-RC509) are handled explicitly rather than smoothed away.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/geological-ground-model-surfaces-python/

Data: A417 Missing Link ground investigation (BGS deposited data items 173984 and 173975; client
Highways England), from the BGS AGS download service under the Open Government Licence v3.0. Contains
British Geological Survey materials © UKRI 2026. Environment Agency LIDAR Composite DTM 1 m via the
Defra WCS (Open Government Licence v3).
"""
import json
import time
from importlib.metadata import version
from pathlib import Path

import numpy as np
import pandas as pd
import requests
import rasterio
from rasterio.transform import from_origin
from scipy.optimize import curve_fit
from scipy.spatial import cKDTree
from pykrige.ok import OrdinaryKriging
import pyvista as pv
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LightSource
from matplotlib.lines import Line2D
from matplotlib.patches import Patch

import a417_ags as A

plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9,
                     "figure.facecolor": "white", "axes.facecolor": "white"})
pd.set_option("display.width", 130)

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

ITEMS = [173984, 173975]                                # the main GI, plus 20 extra holes on the escarpment crest
AOI = (391_600, 212_800, 395_600, 216_800)             # same 4 x 4 km window as Parts 2 and 3
RES = 5.0                                               # LIDAR download resolution (m)
CELL = 10.0                                             # model grid (m)
REACH = 300.0                                           # model only within this distance of a borehole
VE = 3.0
SLIPPED = {"CP217": 19.10}                              # Part 3: a slipped block repeats the Lias above this depth
CAMBER = ["RC507", "RC508", "RC509"]                    # Part 3: base of Inferior Oolite 11-15 m below neighbours
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")
CREDIT = "BGS AGS data © UKRI; LIDAR © Environment Agency (both OGL v3)."

# Model units, youngest first, and the GEOL units that belong to each.
GROUPS = {"SUP": ["TOPS", "MGR", "HEAD", "SLIP"], "GOG": ["WHL", "HMB", "GOG", "FE"],
          "IOG": ["SALS", "ASLS", "BLPL", "IOG"], "BDS": ["BDS"], "LIAS": ["WHM", "MRB", "DYS", "CHAM"]}
ORDER = list(GROUPS)
GROUP = {u: g for g, us in GROUPS.items() for u in us}
UNIT = pd.DataFrame({
    "name": ["Superficial deposits", "Great Oolite Group", "Inferior Oolite Group", "Bridport Sand Formation",
             "Lias mudstones (Whitby Mudstone and older)"],
    "colour": ["#c98c5a", "#f7f1bd", "#d98f3e", "#bfae6e", "#7d93ab"]}, index=ORDER)
# Each modelled surface is the base of one group and the top of the next.
SURFACES = {"rockhead": "base of the superficial deposits", "top_io": "top of the Inferior Oolite",
            "base_io": "base of the Inferior Oolite", "base_bds": "base of the Bridport Sand"}
BELOW = {"top_io": ("GOG", "IOG"), "base_io": ("IOG", "BDS"), "base_bds": ("BDS", "LIAS")}


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


def get(url, path, tries=6):
    """Download once and cache, backing off when the Defra WCS is busy."""
    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)
    return path


print(f"pykrige {version('pykrige')} | pyvista {pv.__version__} | numpy {np.__version__} | "
      f"python-ags4 {version('python-ags4')}")

# ------------------------------------------------------------------ Step 1: logs to model units
section("Step 1: logs to model units")
paths = [p for item in ITEMS for p in A.download(DATA, item)]
tables, headings = A.repair(*A.merge(paths))
ags = A.numeric(tables)
loca = ags["LOCA"].set_index("LOCA_ID")
geol = ags["GEOL"].sort_values(["LOCA_ID", "GEOL_TOP"], ignore_index=True)
geol["UNIT"] = A.unit_code(geol)
geol = geol[geol.GEOL_BASE > geol.GEOL_TOP].join(loca[["LOCA_NATE", "LOCA_NATN", "LOCA_GL"]], on="LOCA_ID")
geol["group"] = geol.UNIT.map(GROUP)                     # NR (no recovery) stays blank and is skipped
for lid, depth in SLIPPED.items():
    moved = (geol.LOCA_ID == lid) & (geol.GEOL_TOP < depth) & geol.group.notna()
    print(f"{lid}: {moved.sum()} intervals above {depth} m (a repeated Lias sequence) counted with the landslide mass")
    geol.loc[moved, "group"] = "SUP"
print(f"{loca.shape[0]} holes from items {', '.join(map(str, ITEMS))}; "
      f"{geol.group.notna().groupby(geol.LOCA_ID).any().sum()} have at least one classified interval")

# Collapse each log to a run of model units, top to bottom.
runs = []
for lid, g in geol[geol.group.notna()].groupby("LOCA_ID", sort=False):
    g = g.assign(run=(g.group != g.group.shift()).cumsum())
    r = g.groupby("run").agg(group=("group", "first"), top=("GEOL_TOP", "min"), base=("GEOL_BASE", "max"))
    runs.append(r.assign(LOCA_ID=lid))
runs = pd.concat(runs, ignore_index=True).join(loca[["LOCA_NATE", "LOCA_NATN", "LOCA_GL", "LOCA_FDEP"]],
                                               on="LOCA_ID")
runs["rank"] = runs.group.map(ORDER.index)
backwards = runs.groupby("LOCA_ID")["rank"].apply(lambda s: (s.diff() < 0).any())
print("holes whose model units step back to a younger unit:", ", ".join(backwards[backwards].index) or "none")

# ------------------------------------------------------------------ Step 2: pick the contacts
section("Step 2: pick the contacts, and what the other holes say")
picks, bounds = [], []
for lid, r in runs.groupby("LOCA_ID", sort=False):
    h = r.iloc[0]
    xy = dict(LOCA_ID=lid, E=h.LOCA_NATE, N=h.LOCA_NATN, GL=h.LOCA_GL)
    ranks, bottom = r["rank"].to_numpy(), h.LOCA_GL - r.base.max()
    # Rockhead: the first bedrock run. A hole that never leaves the superficial deposits only tells us
    # the superficial deposits are at least that thick.
    rock = r[r.group != "SUP"]
    if rock.empty:
        bounds.append(dict(xy, surface="rockhead", kind="deeper than", z=bottom))
    else:
        sup = r[r.group == "SUP"]
        depth = (sup.base.max() + rock.top.iloc[0]) / 2 if len(sup) else rock.top.iloc[0]
        picks.append(dict(xy, surface="rockhead", depth=depth, z=h.LOCA_GL - depth))
    first_rock = rock["rank"].iloc[0] if len(rock) else None
    for s, (up, down) in BELOW.items():
        a, b = r[r.group == up], r[r.group == down]
        if len(a) and len(b):                           # seen in the hole: the midpoint of any unlogged gap
            depth = (a.base.max() + b.top.min()) / 2
            picks.append(dict(xy, surface=s, depth=depth, z=h.LOCA_GL - depth))
        elif first_rock is not None and first_rock >= ORDER.index(down):
            # the hole starts in older rock, so this surface has been eroded and lay above rockhead
            bounds.append(dict(xy, surface=s, kind="above", z=h.LOCA_GL - rock.top.iloc[0]))
        elif len(rock) and ranks.max() <= ORDER.index(up):
            bounds.append(dict(xy, surface=s, kind="below", z=bottom))   # stops in older bedrock above it
picks, bounds = pd.DataFrame(picks), pd.DataFrame(bounds)
picks["camber"] = picks.LOCA_ID.isin(CAMBER)
print("contacts seen in the logs (holes):", picks.groupby("surface").size().reindex(SURFACES).to_dict())
print("one-sided checks from holes that miss a surface:")
print(bounds.groupby(["surface", "kind"]).size().unstack(fill_value=0).reindex(SURFACES).to_string())
print(picks[picks.surface != "rockhead"].round(2).sort_values(["surface", "E"]).to_string(index=False))

# ------------------------------------------------------------------ Step 3: ground and the model grid
section("Step 3: LIDAR ground and the model grid")
X0, Y0, X1, Y1 = AOI
f = get(f"{WCS}&subset=E({X0},{X1})&subset=N({Y0},{Y1})&SCALEFACTOR={1 / RES}",
        DATA / f"ea_dtm_{X0}_{Y0}_{X1 - X0}m_{RES:.0f}m.tif")
with rasterio.open(f) as src:
    dtm5 = src.read(1).astype("float64")
    dtm5[dtm5 < -1000] = np.nan
k = int(CELL / RES)
ground = dtm5.reshape(dtm5.shape[0] // k, k, dtm5.shape[1] // k, k).mean(axis=(1, 3))   # 5 m to 10 m
ny, nx = ground.shape
transform = from_origin(X0, Y1, CELL, CELL)
gx, gy = np.meshgrid(CELL / 2 + np.arange(nx) * CELL, (Y1 - Y0) - CELL / 2 - np.arange(ny) * CELL)  # local metres
holes = runs.drop_duplicates("LOCA_ID")
hx, hy = holes.LOCA_NATE.to_numpy() - X0, holes.LOCA_NATN.to_numpy() - Y0
dist, _ = cKDTree(np.c_[hx, hy]).query(np.c_[gx.ravel(), gy.ravel()])
inside = (dist.reshape(ny, nx) <= REACH) & np.isfinite(ground)
print(f"grid {nx} x {ny} cells of {CELL:.0f} m; {inside.sum():,} cells ({inside.sum() * CELL**2 / 1e6:.2f} km²) "
      f"lie within {REACH:.0f} m of one of the {len(holes)} logged holes")
picks["x"], picks["y"] = picks.E - X0, picks.N - Y0
bounds["x"], bounds["y"] = bounds.E - X0, bounds.N - Y0
row = np.clip(((Y1 - picks.N) // CELL).astype(int), 0, ny - 1)
col = np.clip(((picks.E - X0) // CELL).astype(int), 0, nx - 1)
gl_minus_dtm = (picks.GL - ground[row, col])[picks.surface == "rockhead"]
print(f"logged ground level minus LIDAR at the holes: median {gl_minus_dtm.median():+.2f} m, "
      f"90% within {gl_minus_dtm.abs().quantile(0.9):.2f} m")


def dedupe(p, tol=2.0):
    """Average holes drilled within tol metres of each other (a repeat hole would make kriging singular)."""
    p = p.reset_index(drop=True)
    label = np.arange(len(p))
    for i, j in sorted(cKDTree(p[["x", "y"]]).query_pairs(tol)):
        label[label == label[j]] = label[i]
    num = p.select_dtypes("number").columns
    out = p.groupby(label).agg({**{c: "mean" for c in num}, "LOCA_ID": "+".join, "surface": "first"})
    return out.reset_index(drop=True)


def experimental(x, y, v, edges):
    """Semivariance in distance bins: half the mean squared difference of every pair of holes."""
    i, j = np.triu_indices(len(v), 1)
    d = np.hypot(x[i] - x[j], y[i] - y[j])
    g = 0.5 * (v[i] - v[j]) ** 2
    b = np.digitize(d, edges) - 1
    ok = (b >= 0) & (b < len(edges) - 1)
    n = np.bincount(b[ok], minlength=len(edges) - 1)
    lag = np.bincount(b[ok], d[ok], len(edges) - 1) / np.maximum(n, 1)
    gam = np.bincount(b[ok], g[ok], len(edges) - 1) / np.maximum(n, 1)
    return lag[n > 0], gam[n > 0], n[n > 0]


def spherical(h, psill, rng, nugget):
    h = np.minimum(np.asarray(h, float) / rng, 1.0)
    return nugget + psill * (1.5 * h - 0.5 * h**3)


def fit_spherical(x, y, v, edges):
    """Fit a spherical variogram to the binned semivariances, weighting each bin by its number of pairs."""
    lag, gam, n = experimental(x, y, v, edges)
    p, _ = curve_fit(spherical, lag, gam, p0=[gam.max(), edges[-1] / 3, 0], sigma=1 / np.sqrt(n),
                     bounds=([0, edges[1], 0], [10 * gam.max(), edges[-1], gam.max()]))
    return dict(psill=p[0], range=p[1], nugget=p[2]), (lag, gam, n)


def krige(x, y, v, params, trend=False):
    """Ordinary kriging with a fixed spherical variogram, optionally of the residuals from a fitted plane."""
    X = np.c_[np.ones_like(x), x, y]
    coef = np.linalg.lstsq(X, v, rcond=None)[0] if trend else np.zeros(3)
    ok = OrdinaryKriging(x, y, v - X @ coef, variogram_model="spherical", variogram_parameters=params)

    def predict(px, py):
        r, ss = ok.execute("points", np.asarray(px, float), np.asarray(py, float))
        return np.asarray(r) + np.c_[np.ones_like(px), px, py] @ coef, np.sqrt(np.clip(np.asarray(ss), 0, None))
    return predict, coef


def residuals(x, y, v):
    X = np.c_[np.ones_like(x), x, y]
    coef = np.linalg.lstsq(X, v, rcond=None)[0]
    return v - X @ coef, coef


def loocv(x, y, v, params, trend=False, extra=((), (), ())):
    """Leave each hole out in turn, rebuild from the rest (same variogram), and predict it."""
    pred, sd = np.empty(len(v)), np.empty(len(v))
    for i in range(len(v)):
        keep = np.arange(len(v)) != i
        f, _ = krige(np.r_[x[keep], extra[0]], np.r_[y[keep], extra[1]], np.r_[v[keep], extra[2]], params, trend)
        p, s = f(x[i:i + 1], y[i:i + 1])
        pred[i], sd[i] = p[0], s[0]
    return pred, sd


TOL = 1.0                                               # metres of slack for logged levels and 10 m cells


def honour(name, x, y, v, params, b, trend=False):
    """Krige, then turn every one-sided check the surface breaks into a data point at its limit, and repeat.

    b has columns x, y, kind ('above' or 'below': the surface must lie above or below b.limit) and limit.
    """
    added = np.zeros(len(b), bool)
    first = None
    for _ in range(20):
        ex = b[added]
        f, coef = krige(np.r_[x, ex.x], np.r_[y, ex.y], np.r_[v, ex.limit], params, trend)
        m, _ = f(b.x.values, b.y.values)
        broken = np.where(b.kind == "above", m < b.limit - TOL, m > b.limit + TOL) & ~added
        first = broken.sum() if first is None else first
        if not broken.any():
            break
        added |= broken
    ex = b[added]
    passes[name] = (len(b), first, added.sum())
    print(f"  one-sided checks: the first pass broke {first} of {len(b)}; {added.sum()} holes added at their "
          f"limits; now broken: {broken.sum()}")
    return f, coef, (ex.x.values, ex.y.values, ex.limit.values), b.LOCA_ID[added].tolist()


def scores(obs, pred, sd):
    e = pred - obs
    return dict(n=len(e), RMSE=np.sqrt(np.mean(e**2)), mean_error=e.mean(), MAE=np.abs(e).mean(),
                mean_SE=np.mean(e / sd), var_SE=np.var(e / sd))


def on_grid(f):
    p, s = f(gx[inside], gy[inside])
    a, b = np.full(ground.shape, np.nan), np.full(ground.shape, np.nan)
    a[inside], b[inside] = p, s
    return a, b


cv, vgm, passes = {}, {}, {}
LAGS = np.arange(0, 2050, 100.0)                         # 100 m bins out to 2 km

# ------------------------------------------------------------------ Step 4: rockhead
section("Step 4: rockhead = LIDAR ground minus kriged superficial thickness")


def at(grid, E, N):
    """Value of a model grid at British National Grid points (nearest cell)."""
    r = np.clip(((Y1 - np.asarray(N)) // CELL).astype(int), 0, ny - 1)
    c = np.clip(((np.asarray(E) - X0) // CELL).astype(int), 0, nx - 1)
    return grid[r, c]


def show(name):
    p = vgm[name][0]
    print(f"  variogram: spherical, partial sill {p['psill']:.1f} m², range {p['range']:.0f} m, "
          f"nugget {p['nugget']:.1f} m²")


def checks(surface, limit, flip=False):
    """The one-sided checks for a surface, as limits on the quantity being kriged."""
    b = bounds[bounds.surface == surface].copy()
    b["kind"] = b.kind.replace({"deeper than": "above"})
    if flip:
        b["kind"] = b.kind.map({"above": "below", "below": "above"})
    b["limit"] = limit(b)
    return b.reset_index(drop=True)


sup = dedupe(picks[picks.surface == "rockhead"])
xs_, ys_, v_ = sup.x.values, sup.y.values, sup.depth.values
vgm["Superficial thickness"] = fit_spherical(xs_, ys_, v_, LAGS)
show("Superficial thickness")
print(f"  {len(sup)} thicknesses (after averaging repeat holes): median {sup.depth.median():.2f} m, "
      f"max {sup.depth.max():.2f} m ({sup.LOCA_ID[sup.depth.idxmax()]})")
f_sup, _, ex_sup, _ = honour("Rockhead", xs_, ys_, v_, vgm["Superficial thickness"][0], checks("rockhead", lambda b: b.GL - b.z))
p, s = loocv(xs_, ys_, v_, vgm["Superficial thickness"][0], extra=ex_sup)
sup["pred"], sup["sd"] = np.clip(p, 0, None), s
cv["Superficial thickness"] = scores(sup.depth, sup.pred, sup.sd)
thick_sup, sd_sup = on_grid(f_sup)
thick_sup = np.clip(thick_sup, 0, None)
rockhead = ground - thick_sup

# ------------------------------------------------------------------ Step 5: base of the Inferior Oolite
section("Step 5: base of the Inferior Oolite: regional dip plus kriged residuals")
bio_all = dedupe(picks[picks.surface == "base_io"])
bio = bio_all[~bio_all.LOCA_ID.isin(CAMBER)].reset_index(drop=True)
res, coef = residuals(bio.x.values, bio.y.values, bio.z.values)
grad = np.hypot(coef[1], coef[2])
azimuth = (np.degrees(np.arctan2(-coef[1], -coef[2])) + 360) % 360      # direction the surface falls towards
print(f"{len(bio)} holes (RC507-RC509 held back). Fitted plane dips {np.degrees(np.arctan(grad)):.2f}° "
      f"({grad * 1000:.1f} m per km) towards {azimuth:.0f}°; residuals {res.min():+.1f} to {res.max():+.1f} m")
vgm["Base of Inferior Oolite"] = fit_spherical(bio.x.values, bio.y.values, res, LAGS)
show("Base of Inferior Oolite")
b_bio = checks("base_io", lambda b: b.z)
f_bio, coef, ex_bio, added_bio = honour("Base of Inferior Oolite", bio.x.values, bio.y.values, bio.z.values, vgm["Base of Inferior Oolite"][0],
                                        b_bio, trend=True)
print("  added:", ", ".join(added_bio) or "none")
p, s = loocv(bio.x.values, bio.y.values, bio.z.values, vgm["Base of Inferior Oolite"][0], True, ex_bio)
bio["pred"], bio["sd"] = p, s
cv["Base of Inferior Oolite"] = scores(bio.z, bio.pred, bio.sd)
base_io, sd_bio = on_grid(f_bio)

# The three low holes, against a surface built without them.
cam = bio_all[bio_all.LOCA_ID.isin(CAMBER)].reset_index(drop=True)
cam["regional"], cam["sd"] = f_bio(cam.x.values, cam.y.values)
cam["offset"] = cam.z - cam.regional
print("RC507-RC509 against the surface built from the other holes:")
print(cam[["LOCA_ID", "z", "regional", "sd", "offset"]].round(2).to_string(index=False))
# ...and what happens to everyone else if they are simply kriged in with the rest.
res_w, _ = residuals(bio_all.x.values, bio_all.y.values, bio_all.z.values)
vgm_with = fit_spherical(bio_all.x.values, bio_all.y.values, res_w, LAGS)[0]
f_with, _ = krige(np.r_[bio_all.x, ex_bio[0]], np.r_[bio_all.y, ex_bio[1]], np.r_[bio_all.z, ex_bio[2]], vgm_with,
                  trend=True)
p, s = loocv(bio_all.x.values, bio_all.y.values, bio_all.z.values, vgm_with, True, ex_bio)
others = ~bio_all.LOCA_ID.isin(CAMBER)
cv_with = scores(bio_all.z[others], p[others], s[others])
print(f"kriging them in with the rest instead: LOOCV RMSE at the other {others.sum()} holes "
      f"{cv_with['RMSE']:.2f} m (without them {cv['Base of Inferior Oolite']['RMSE']:.2f} m); the variogram "
      f"nugget rises from {vgm['Base of Inferior Oolite'][0]['nugget']:.1f} to {vgm_with['nugget']:.1f} m²")
base_io_with, _ = on_grid(f_with)
base_io_regional = base_io.copy()

# Deliberate handling: a block bounded by the nearest-hole rule. Each cell nearer to one of the three
# than to any other hole that reached the contact, and within 60 m of it, drops by that hole's offset.
d_good, _ = cKDTree(bio[["x", "y"]]).query(np.c_[gx.ravel(), gy.ravel()])
d_cam, i_cam = cKDTree(cam[["x", "y"]]).query(np.c_[gx.ravel(), gy.ravel()])
block = ((d_cam < d_good) & (d_cam <= 60)).reshape(ny, nx) & inside
base_io[block] += cam.offset.to_numpy()[i_cam.reshape(ny, nx)[block]]
print(f"low block: {block.sum()} cells ({block.sum() * CELL**2 / 1e4:.1f} ha) lowered by "
      f"{-cam.offset.max():.1f} to {-cam.offset.min():.1f} m")

# ------------------------------------------------------------------ Step 6: hang the other surfaces from it
section("Step 6: top of the Inferior Oolite and base of the Bridport Sand from thicknesses")
tio = dedupe(picks[picks.surface == "top_io"])
both = tio.LOCA_ID.isin(bio_all.LOCA_ID)
tio["base_here"] = np.where(both, tio.LOCA_ID.map(bio_all.set_index("LOCA_ID").z), at(base_io, tio.E, tio.N))
tio["thick"] = tio.z - tio.base_here
print("Inferior Oolite thickness = top seen in the log minus the base (seen if the hole goes through, else modelled)")
print(tio[["LOCA_ID", "z", "base_here", "thick"]].assign(through=both).round(2).to_string(index=False))
vgm["Inferior Oolite thickness"] = fit_spherical(tio.x.values, tio.y.values, tio.thick.values,
                                                 np.arange(0, 2050, 250.0))
show("Inferior Oolite thickness")
f_tio, _, ex_tio, added_tio = honour("Top of Inferior Oolite", tio.x.values, tio.y.values, tio.thick.values, vgm["Inferior Oolite thickness"][0],
                                     checks("top_io", lambda b: b.z - at(base_io, b.E, b.N)))
p, s = loocv(tio.x.values, tio.y.values, tio.thick.values, vgm["Inferior Oolite thickness"][0], extra=ex_tio)
tio["pred"], tio["sd"] = p + tio.base_here, s
cv["Top of Inferior Oolite"] = scores(tio.z, tio.pred, tio.sd)
thick_io, sd_tio = on_grid(f_tio)
top_io = base_io + thick_io
sd_tio = np.hypot(sd_tio, sd_bio)

bds = dedupe(picks[picks.surface == "base_bds"])
bds["z_top"] = bds.LOCA_ID.map(bio_all.set_index("LOCA_ID").z)
bds["thick"] = bds.z_top - bds.z
print(f"\nBridport Sand, {len(bds)} complete penetrations: " +
      ", ".join(f"{a} {b:.1f} m" for a, b in zip(bds.LOCA_ID, bds.thick)))
# Five thicknesses cannot give a variogram of their own: borrow the range of the base of the Inferior Oolite
# and use their variance as the sill.
vgm["Bridport Sand thickness"] = (dict(psill=bds.thick.var(), range=vgm["Base of Inferior Oolite"][0]["range"],
                                       nugget=0.0), experimental(bds.x.values, bds.y.values, bds.thick.values, LAGS))
show("Bridport Sand thickness")
f_bds, _, ex_bds, added_bds = honour("Base of Bridport Sand", bds.x.values, bds.y.values, bds.thick.values, vgm["Bridport Sand thickness"][0],
                                     checks("base_bds", lambda b: at(base_io, b.E, b.N) - b.z, flip=True))
print("  added (holes that end in Bridport Sand below the first model):", ", ".join(added_bds) or "none")
p, s = loocv(bds.x.values, bds.y.values, bds.thick.values, vgm["Bridport Sand thickness"][0], extra=ex_bds)
bds["pred"], bds["sd"] = bds.z_top - p, s
cv["Base of Bridport Sand"] = scores(bds.z, bds.pred, bds.sd)
thick_bds, sd_t = on_grid(f_bds)
base_bds = base_io - np.clip(thick_bds, 0, None)
sd_bds = np.hypot(sd_bio, sd_t)
print(f"modelled Bridport Sand thickness {np.nanmin(thick_bds[inside]):.1f} to {np.nanmax(thick_bds[inside]):.1f} m")
# ------------------------------------------------------------------ Step 7: the checks
section("Step 7: does the model hold together?")
print("leave-one-out cross-validation (prediction minus log, m; standardised errors should average 0, variance 1):")
print(pd.DataFrame(cv).T.round(2).to_string())

sep = (top_io - base_io)[inside]
sep2 = (base_io - base_bds)[inside]
print(f"\nno-crossing: top minus base of Inferior Oolite {np.nanmin(sep):.1f} to {np.nanmax(sep):.1f} m "
      f"({(sep < 0).sum()} cells cross); base of Inferior Oolite minus base of Bridport Sand {np.nanmin(sep2):.1f} "
      f"to {np.nanmax(sep2):.1f} m; rockhead minus ground {np.nanmax((rockhead - ground)[inside]):.1f} m at most")
raw = {"rockhead": rockhead, "top_io": top_io, "base_io": base_io, "base_bds": base_bds}
r = ((Y1 - bounds.N) // CELL).astype(int); c = ((bounds.E - X0) // CELL).astype(int)
bounds["model"] = [raw[s][i, j] for s, i, j in zip(bounds.surface, r, c)]
bounds["margin"] = np.where(bounds.kind == "below", bounds.z - bounds.model, bounds.model - bounds.z)
bounds.loc[bounds.kind == "deeper than", "margin"] = (bounds.z - bounds.model)[bounds.kind == "deeper than"]
bounds["ok"] = bounds.margin >= -0.5                    # half a metre for the 10 m cells and logged levels
tab = bounds.groupby(["surface", "kind"]).agg(holes=("ok", "size"), agree=("ok", "sum"),
                                              worst_m=("margin", "min")).reindex(SURFACES, level=0)
print("\nholes that miss a surface: does the model agree with them?")
print(tab.round(1).to_string())
print("disagreements:")
print(bounds[~bounds.ok][["LOCA_ID", "surface", "kind", "z", "model", "margin"]].round(1).to_string(index=False))

# Stack the surfaces: where a younger surface is lower, the older one has been eroded and takes its level.
stack = {"ground": ground.copy(), "rockhead": rockhead}
stack["top_io"] = np.fmin(top_io, stack["rockhead"])
stack["base_io"] = np.fmin(base_io, stack["top_io"])
stack["base_bds"] = np.fmin(base_bds, stack["base_io"])
for g in stack.values():
    g[~inside] = np.nan
MODEL_BASE = float(np.floor((np.nanmin(stack["base_bds"]) - 10) / 10) * 10)
present = {u: np.nanmean((stack[a] - stack[b])[inside] > 0.01) for u, a, b in
           [("SUP", "ground", "rockhead"), ("GOG", "rockhead", "top_io"), ("IOG", "top_io", "base_io"),
            ("BDS", "base_io", "base_bds")]}
present["LIAS"] = 1.0
print(f"\nmodel base {MODEL_BASE:.0f} m OD. Share of the model footprint where each unit is present: " +
      ", ".join(f"{UNIT.name[u].split(' (')[0]} {v:.0%}" for u, v in present.items()))

# ------------------------------------------------------------------ Step 8: write the model
section("Step 8: write the surfaces")
FILES = {"ground": "ground_mOD", "rockhead": "rockhead_mOD", "top_io": "top_inferior_oolite_mOD",
         "base_io": "base_inferior_oolite_mOD", "base_bds": "base_bridport_sand_mOD"}
sds = {"rockhead": sd_sup, "top_io": sd_tio, "base_io": sd_bio, "base_bds": sd_bds}


def write_tif(name, a):
    with rasterio.open(MODEL / f"{name}.tif", "w", driver="GTiff", width=nx, height=ny, count=1, dtype="float32",
                       crs="EPSG:27700", transform=transform, nodata=np.nan, compress="deflate") as dst:
        dst.write(np.where(inside, a, np.nan).astype("float32"), 1)


for s, name in FILES.items():
    write_tif(name, stack[s])
for s, a in sds.items():
    write_tif(f"sd_{FILES[s]}", a)
write_tif("low_block", block.astype("float32"))
tops = list(FILES.values())
meta = {"model_base_mOD": MODEL_BASE, "cell_m": CELL, "crs": "EPSG:27700",
        "units": [{"code": u, "name": UNIT.name[u], "colour": UNIT.colour[u], "top": tops[i],
                   "base": tops[i + 1] if i + 1 < len(tops) else None} for i, u in enumerate(ORDER)],
        "credit": CREDIT}
(MODEL / "model.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")
pd.concat([sup.assign(surface="superficial thickness", obs=sup.depth), bio.assign(obs=bio.z),
           tio.assign(obs=tio.z), bds.assign(obs=bds.z)])[["LOCA_ID", "surface", "E", "N", "obs", "pred", "sd"]] \
    .round(2).to_csv(OUT / "loocv.csv", index=False)
print("wrote", ", ".join(sorted(p.name for p in MODEL.glob("*"))), "and loocv.csv")

# ------------------------------------------------------------------ Step 9: figures
section("Step 9: figures")
EXT = [X0, X1, Y0, Y1]
hill = LightSource(azdeg=315, altdeg=40).hillshade(np.nan_to_num(dtm5, nan=np.nanmean(dtm5)), vert_exag=2,
                                                    dx=RES, dy=RES)
# The route corridor, vale to plateau, for the cross-section.
LINE = np.array([(392_050, 215_640), (392_650, 215_880), (393_080, 215_870), (393_330, 216_030),
                 (393_700, 216_000), (394_000, 215_770), (394_240, 215_150), (394_680, 213_850)], float)
seg = np.diff(LINE, axis=0)
seglen = np.hypot(*seg.T)
chain0 = np.r_[0, np.cumsum(seglen)]


def project(E, N):
    """Chainage along LINE and distance off it for each point."""
    P = np.c_[E, N][:, None, :]
    t = np.clip(((P - LINE[:-1]) * seg).sum(-1) / seglen**2, 0, 1)
    foot = LINE[:-1] + t[..., None] * seg
    d = np.hypot(*(P - foot).transpose(2, 0, 1))
    k = d.argmin(axis=1)
    return chain0[k] + t[np.arange(len(k)), k] * seglen[k], d[np.arange(len(k)), k]


def frame(ax, ext=None, grid=True):
    ext = ext or EXT
    ax.set_xlim(ext[0], ext[1]); ax.set_ylim(ext[2], ext[3])
    ax.set_aspect("equal")
    ax.ticklabel_format(useOffset=False, style="plain")
    ax.tick_params(labelsize=6.5)
    for lab in ax.get_xticklabels() + ax.get_yticklabels():
        lab.set_rotation(0)


def footprint(ax, colour="#333333"):
    ax.contour(gx + X0, gy + Y0, inside.astype(float), levels=[0.5], colors=colour, linewidths=0.8)


role = pd.Series("logged, rockhead only", index=holes.LOCA_ID)
role[bounds.loc[bounds.kind == "deeper than", "LOCA_ID"]] = "ends in superficial deposits"
role[picks.loc[picks.surface == "top_io", "LOCA_ID"]] = "top of Inferior Oolite seen"
role[picks.loc[picks.surface == "base_io", "LOCA_ID"]] = "base of Inferior Oolite seen"
role[CAMBER] = "RC507-RC509 (low block)"
STYLE = {"logged, rockhead only": dict(marker="o", s=10, c="#ffffff", edgecolors="#333333", lw=0.6),
         "ends in superficial deposits": dict(marker="o", s=10, c="#c98c5a", edgecolors="#333333", lw=0.6),
         "top of Inferior Oolite seen": dict(marker="^", s=30, c="#f4cd72", edgecolors="black", lw=0.7),
         "base of Inferior Oolite seen": dict(marker="s", s=24, c="#1f3b73", edgecolors="white", lw=0.5),
         "RC507-RC509 (low block)": dict(marker="s", s=30, c="#b2182b", edgecolors="white", lw=0.5)}

# Figure 2: the data.
fig, ax = plt.subplots(figsize=(8.6, 8.2))
ax.imshow(hill, cmap="gray", extent=EXT, vmin=0, vmax=1.15)
footprint(ax)
ax.plot(*LINE.T, color="#b2182b", lw=1.4, ls="--")
ax.annotate("section, figure 7", LINE[0], xytext=(-8, -16), textcoords="offset points", fontsize=7, color="#b2182b")
for k, st in STYLE.items():
    sel = holes[holes.LOCA_ID.map(role) == k]
    ax.scatter(sel.LOCA_NATE, sel.LOCA_NATN, label=f"{k} ({len(sel)})", zorder=3, **st)
for lid in ["CP217", "OH413", "DSRC404", "OH411", "DSRCOH400", "DSRCOH306"]:
    h = holes.set_index("LOCA_ID").loc[lid]
    ax.annotate(lid, (h.LOCA_NATE, h.LOCA_NATN), xytext=(5, 3), textcoords="offset points", fontsize=7,
                bbox=dict(boxstyle="round,pad=0.15", fc="white", ec="none", alpha=0.8))
ax.plot([], [], color="#333333", lw=0.8, label=f"model footprint (within {REACH:.0f} m of a hole)")
ax.legend(loc="lower left", fontsize=7, framealpha=0.9)
frame(ax)
ax.set_title(f"{len(holes)} logged holes and what each one tells the model", fontsize=10)
fig.text(0.01, 0.01, f"{CREDIT} Hillshade from LIDAR DTM. British National Grid (m).", fontsize=6.5,
         color="#444444")
fig.savefig(OUT / "02-borehole-data-map.png")
plt.close(fig)

# Figure 3: variograms and cross-validation.
names = list(vgm)
cvd = {"Superficial thickness": (sup.depth, sup.pred, "superficial thickness (m)"),
       "Base of Inferior Oolite": (bio.z, bio.pred, "base of Inferior Oolite (m OD)"),
       "Inferior Oolite thickness": (tio.z, tio.pred, "top of Inferior Oolite (m OD)"),
       "Bridport Sand thickness": (bds.z, bds.pred, "base of Bridport Sand (m OD)")}
cvname = dict(zip(names, cv))
fig, axes = plt.subplots(2, 4, figsize=(13, 6.4))
for i, n in enumerate(names):
    p, (lag, gam, cnt) = vgm[n]
    ax = axes[0, i]
    ax.scatter(lag, gam, s=8 + 40 * cnt / cnt.max(), color="#4a6fa5", label="pairs of holes, binned")
    h = np.linspace(0, 2000, 200)
    ax.plot(h, spherical(h, p["psill"], p["range"], p["nugget"]), color="black", lw=1.2, label="spherical model")
    ax.set_xlim(0, 2000); ax.set_ylim(0, None)
    ax.set_title(("Residual of the base of the Inferior Oolite" if n.startswith("Base of") else n), fontsize=8.5)
    ax.set_xlabel("separation (m)", fontsize=7.5); ax.set_ylabel("semivariance (m²)", fontsize=7.5)
    ax.text(0.97, 0.05, f"range {p['range']:.0f} m\nnugget {p['nugget']:.1f} m²", transform=ax.transAxes,
            ha="right", fontsize=7)
    obs, pred, lab = cvd[n]
    sc = cv[cvname[n]]
    ax = axes[1, i]
    lo, hi = min(obs.min(), pred.min()), max(obs.max(), pred.max())
    pad = (hi - lo) * 0.06
    ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], color="#999999", lw=0.8)
    ax.scatter(obs, pred, s=14, color="#d98f3e", edgecolors="#5b3a12", lw=0.4)
    ax.set_xlim(lo - pad, hi + pad); ax.set_ylim(lo - pad, hi + pad); ax.set_aspect("equal")
    ax.set_xlabel(f"logged {lab}", fontsize=7.5); ax.set_ylabel("predicted with that hole left out", fontsize=7.5)
    ax.text(0.04, 0.96, f"n = {sc['n']}\nRMSE {sc['RMSE']:.1f} m\nmean error {round(sc['mean_error'], 1) + 0.0:+.1f} m\n"
            f"variance of standardised\nerrors {sc['var_SE']:.2f}", transform=ax.transAxes, va="top", fontsize=7)
    worst = np.argsort(-np.abs(np.asarray(pred) - np.asarray(obs)))[:2]
    lids = {"Superficial thickness": sup, "Base of Inferior Oolite": bio, "Inferior Oolite thickness": tio,
            "Bridport Sand thickness": bds}[n].LOCA_ID.values
    for w in worst:
        ax.annotate(lids[w], (np.asarray(obs)[w], np.asarray(pred)[w]), xytext=(4, -9), textcoords="offset points",
                    fontsize=6.5)
for ax in axes.ravel():
    ax.tick_params(labelsize=6.5)
    ax.spines[["top", "right"]].set_visible(False)
axes[0, 0].legend(fontsize=6.5, loc="upper left", frameon=False)
fig.suptitle("Top: how alike two holes are against how far apart they are. Bottom: leave-one-out cross-validation",
             fontsize=10, x=0.01, ha="left")
fig.text(0.01, -0.01, CREDIT, fontsize=6.5, color="#444444")
fig.tight_layout()
fig.savefig(OUT / "03-variograms-cross-validation.png")
plt.close(fig)

# Figure 4: surfaces and their kriging standard deviation.
rows, cols = np.where(inside)
ZEXT = [X0 + cols.min() * CELL - 100, X0 + (cols.max() + 1) * CELL + 100,
        Y1 - (rows.max() + 1) * CELL - 100, Y1 - rows.min() * CELL + 100]
crop = [Y1 - ZEXT[3], Y1 - ZEXT[2], ZEXT[0] - X0, ZEXT[1] - X0]
show_ = {"superficial thickness (m)": (thick_sup, sd_sup, "YlOrBr", None),
         "base of Inferior Oolite (m OD)": (base_io, sd_bio, "viridis", stack["base_io"] < base_io - 0.01),
         "top of Inferior Oolite (m OD)": (top_io, sd_tio, "viridis", stack["top_io"] < top_io - 0.01)}
fig, axes = plt.subplots(2, 3, figsize=(13, 9.6))
for i, (lab, (val, sd, cmap, eroded)) in enumerate(show_.items()):
    for j, (arr, cm, title) in enumerate([(val, cmap, lab), (sd, "magma_r", "kriging standard deviation (m)")]):
        ax = axes[j, i]
        ax.imshow(hill, cmap="gray", extent=EXT, vmin=0, vmax=1.3)
        a = np.where(inside, arr, np.nan)
        im = ax.imshow(a, cmap=cm, extent=EXT, alpha=0.85,
                       vmin=0 if j else np.nanpercentile(a, 1), vmax=np.nanpercentile(a, 99))
        if j == 0 and i > 0:
            cs = ax.contour(gx + X0, gy + Y0, a, levels=np.arange(150, 320, 10), colors="black", linewidths=0.5)
            ax.clabel(cs, fontsize=5.5, fmt="%d")
            ax.contourf(gx + X0, gy + Y0, np.where(inside, eroded, np.nan).astype(float), levels=[0.5, 1.5],
                        hatches=["////"], colors="none")
        if j == 0 and i == 1:
            ax.contour(gx + X0, gy + Y0, block.astype(float), levels=[0.5], colors="#b2182b", linewidths=1.2)
        footprint(ax)
        pts = {0: sup, 1: bio, 2: tio}[i]
        ax.scatter(pts.E, pts.N, s=5, c="black", zorder=3)
        frame(ax, ZEXT)
        cb = fig.colorbar(im, ax=ax, shrink=0.75, pad=0.02)
        cb.ax.tick_params(labelsize=6.5)
        ax.set_title(title if j else title + (" (hatched: eroded, above rockhead)" if i else ""), fontsize=8)
fig.suptitle("The kriged surfaces (top) and how uncertain each one is (bottom). Dots: holes that log the contact; "
             "red outline: the low block", fontsize=10, x=0.01, ha="left")
fig.text(0.01, 0.0, f"{CREDIT} Hillshade from LIDAR DTM. British National Grid (m).", fontsize=6.5, color="#444444")
fig.tight_layout()
fig.savefig(OUT / "04-surfaces-and-uncertainty.png")
plt.close(fig)

# Figure 5: the low block at RC507-RC509.
c0 = cam[["E", "N"]].mean().to_numpy()
ZB = [c0[0] - 280, c0[0] + 280, c0[1] - 200, c0[1] + 200]
fig = plt.figure(figsize=(13, 5.2))
ax = fig.add_axes([0.03, 0.1, 0.36, 0.8])
ax.imshow(hill, cmap="gray", extent=EXT, vmin=0, vmax=1.3)
im = ax.imshow(np.where(inside, base_io, np.nan), cmap="viridis", extent=EXT, alpha=0.8, vmin=185, vmax=235)
cs = ax.contour(gx + X0, gy + Y0, np.where(inside, base_io, np.nan), levels=np.arange(180, 240, 2.5),
                colors="black", linewidths=0.4)
ax.contour(gx + X0, gy + Y0, block.astype(float), levels=[0.5], colors="#b2182b", linewidths=1.4)
near = bio_all[(bio_all.E.between(*ZB[:2])) & (bio_all.N.between(*ZB[2:]))]
for r in near.itertuples():
    lowb = r.LOCA_ID in CAMBER
    ax.scatter(r.E, r.N, s=26, marker="s", c="#b2182b" if lowb else "#1f3b73", edgecolors="white", lw=0.5, zorder=3)
    ax.annotate(f"{r.LOCA_ID} {r.z:.0f}" if lowb else f"{r.z:.0f}", (r.E, r.N),
                xytext={"RC507": (-14, -14), "RC508": (6, -12), "RC509": (6, 5)}.get(
                    r.LOCA_ID, (-18, 6) if np.hypot(cam.E - r.E, cam.N - r.N).min() < 40 else (4, 4)),
                textcoords="offset points", fontsize=6.5, color="#b2182b" if lowb else "black",
                bbox=dict(fc="white", ec="none", alpha=0.75, pad=0.5))
frame(ax, ZB)
fig.colorbar(im, ax=ax, shrink=0.8, pad=0.02, label="base of Inferior Oolite (m OD)").ax.tick_params(labelsize=6.5)
ax.set_title("Final surface; red outline: the low block. Labels: logged level (m OD)", fontsize=8)
# profile through the three holes
u = np.r_[np.cos(np.radians(90 - azimuth)), np.sin(np.radians(90 - azimuth))]
ends = np.array([c0 - 280 * u, c0 + 280 * u])
t = np.linspace(0, 560, 113)
pe, pn = ends[0][0] + t * u[0], ends[0][1] + t * u[1]
ax.plot(*ends.T, color="white", lw=2.2); ax.plot(*ends.T, color="black", lw=1, ls="--")
ax2 = fig.add_axes([0.47, 0.13, 0.3, 0.74])
for g, lab, st in [(base_io_regional, "built without RC507-RC509", dict(color="#1f3b73", lw=1.4)),
                   (base_io_with, "all holes kriged together", dict(color="#7a7a7a", lw=1.2, ls=":")),
                   (base_io, "final: regional surface plus the low block", dict(color="#b2182b", lw=1.6))]:
    ax2.plot(t, at(g, pe, pn), label=lab, **st)
sd_line = at(sd_bio, pe, pn)
ax2.fill_between(t, at(base_io_regional, pe, pn) - sd_line, at(base_io_regional, pe, pn) + sd_line,
                 color="#1f3b73", alpha=0.12, lw=0, label="± one kriging standard deviation")
off = (np.c_[bio_all.E, bio_all.N] - ends[0]) @ np.c_[-u[1], u[0]].T
along = (np.c_[bio_all.E, bio_all.N] - ends[0]) @ u
sel = (np.abs(off.ravel()) < 120) & (along >= 0) & (along <= 560)
for r, a_ in zip(bio_all[sel].itertuples(), along[sel]):
    lowb = r.LOCA_ID in CAMBER
    ax2.scatter(a_, r.z, marker="s", s=28, c="#b2182b" if lowb else "#1f3b73", edgecolors="white", zorder=3)
    if lowb:
        ax2.annotate(r.LOCA_ID, (a_, r.z), xytext=(5, -3), textcoords="offset points", fontsize=6.5, color="#b2182b")
ax2.set_xlabel(f"distance along the profile, towards {azimuth:.0f}° (m); holes within 120 m", fontsize=7.5)
ax2.set_ylabel("base of Inferior Oolite (m OD)", fontsize=7.5)
ax2.legend(fontsize=6.5, loc="lower left", frameon=False)
ax2.tick_params(labelsize=6.5); ax2.spines[["top", "right"]].set_visible(False)
ax2.set_title("Profile down the regional dip, through the three holes", fontsize=8)
ax3 = fig.add_axes([0.83, 0.13, 0.15, 0.74])
vals = [cv["Base of Inferior Oolite"]["RMSE"], cv_with["RMSE"]]
ax3.bar([0, 1], vals, color=["#1f3b73", "#9a9a9a"], width=0.6)
for i, v in enumerate(vals):
    ax3.text(i, v + 0.1, f"{v:.2f} m", ha="center", fontsize=7.5)
ax3.set_xticks([0, 1], ["held out,\nthen a block", "kriged in\nwith the rest"], fontsize=7)
ax3.set_ylabel(f"LOOCV RMSE at the other {others.sum()} holes (m)", fontsize=7.5)
ax3.tick_params(labelsize=6.5); ax3.spines[["top", "right"]].set_visible(False)
ax3.set_title("Error elsewhere", fontsize=8)
fig.text(0.01, 0.0, f"{CREDIT} Hillshade from LIDAR DTM. British National Grid (m).", fontsize=6.5, color="#444444")
fig.savefig(OUT / "05-low-block.png")
plt.close(fig)

# Figure 6: thicknesses, crossing and the one-sided checks.
fig = plt.figure(figsize=(13, 5.4))
ax = fig.add_axes([0.02, 0.1, 0.34, 0.82])
ax.imshow(hill, cmap="gray", extent=EXT, vmin=0, vmax=1.3)
im = ax.imshow(np.where(inside, thick_io, np.nan), cmap="cividis", extent=EXT, alpha=0.85)
footprint(ax)
ax.scatter(tio.E, tio.N, s=18, marker="^", c="#f4cd72", edgecolors="black", lw=0.6, zorder=3)
for r in tio[tio.LOCA_ID.isin(bio_all.LOCA_ID)].itertuples():
    ax.annotate(f"{r.thick:.0f}", (r.E, r.N), xytext=(4, 2), textcoords="offset points", fontsize=6.5,
                bbox=dict(fc="white", ec="none", alpha=0.75, pad=0.4))
frame(ax, ZEXT)
fig.colorbar(im, ax=ax, shrink=0.8, pad=0.02, label="Inferior Oolite thickness (m)").ax.tick_params(labelsize=6.5)
ax.set_title("Kriged Inferior Oolite thickness; labels: holes that go right through", fontsize=8)
ax2 = fig.add_axes([0.45, 0.13, 0.22, 0.75])
ax2.hist(sep, bins=40, color="#d98f3e", alpha=0.85, label="top minus base of Inferior Oolite")
ax2.hist(sep2, bins=40, color="#b8a96a", alpha=0.85, label="base of Inferior Oolite minus\nbase of Bridport Sand")
ax2.axvline(0, color="#b2182b", lw=1)
ax2.text(1, ax2.get_ylim()[1] * 0.97, "crossing\nwould show\nleft of here", fontsize=6.5, color="#b2182b", va="top")
ax2.set_xlim(-8, None)
ax2.set_xlabel("vertical gap between surfaces, every model cell (m)", fontsize=7.5)
ax2.set_ylabel("cells", fontsize=7.5)
ax2.legend(fontsize=6.5, frameon=False, loc="upper right")
ax2.tick_params(labelsize=6.5); ax2.spines[["top", "right"]].set_visible(False)
ax2.set_title(f"No crossing: {(sep < 0).sum() + (sep2 < 0).sum()} of {2 * len(sep):,} cell gaps below zero",
              fontsize=8)
ax3 = fig.add_axes([0.75, 0.13, 0.23, 0.75])
sumr = bounds.groupby("surface").agg(n=("ok", "size"), ok=("ok", "sum")).reindex(SURFACES)
first = {"rockhead": passes["Rockhead"], "top_io": passes["Top of Inferior Oolite"],
         "base_io": passes["Base of Inferior Oolite"], "base_bds": passes["Base of Bridport Sand"]}
yy = np.arange(len(sumr))[::-1]
ax3.barh(yy + 0.2, [first[s][0] - first[s][1] for s in sumr.index], 0.38, color="#bdbdbd",
         label="agree with the first model")
ax3.barh(yy - 0.2, sumr.ok, 0.38, color="#1f3b73", label="agree with the final model (0.5 m slack)")
for y_, s in zip(yy, sumr.index):
    ax3.text(first[s][0] - first[s][1] + 1, y_ + 0.2, f"{first[s][0] - first[s][1]} of {first[s][0]}",
             fontsize=6.5, va="center")
    ax3.text(sumr.ok[s] + 1, y_ - 0.2, f"{sumr.ok[s]} of {sumr.n[s]}", fontsize=6.5, va="center")
ax3.set_yticks(yy, ["Rockhead", "Top of Inferior\nOolite", "Base of Inferior\nOolite", "Base of Bridport\nSand"],
               fontsize=7)
ax3.set_xlabel("holes that stop short of, or start below, the surface", fontsize=7.5)
ax3.legend(fontsize=6.5, frameon=False, loc="upper right")
ax3.tick_params(labelsize=6.5); ax3.spines[["top", "right"]].set_visible(False)
ax3.set_xlim(0, sumr.n.max() * 1.25)
ax3.set_title("One-sided checks", fontsize=8)
fig.text(0.01, 0.0, f"{CREDIT} Hillshade from LIDAR DTM. British National Grid (m).", fontsize=6.5, color="#444444")
fig.savefig(OUT / "06-thickness-and-checks.png")
plt.close(fig)

# Figure 7: the section.
STEP = 10.0
ch = np.arange(0, chain0[-1], STEP)
k_ = np.clip(np.searchsorted(chain0, ch, side="right") - 1, 0, len(seg) - 1)
fr = (ch - chain0[k_]) / seglen[k_]
se, sn = LINE[k_, 0] + fr * seg[k_, 0], LINE[k_, 1] + fr * seg[k_, 1]
prof = {s: at(stack[s], se, sn) for s in stack}
prof["model_base"] = np.where(np.isfinite(prof["ground"]), MODEL_BASE, np.nan)
SEQ = [("SUP", "ground", "rockhead"), ("GOG", "rockhead", "top_io"), ("IOG", "top_io", "base_io"),
       ("BDS", "base_io", "base_bds"), ("LIAS", "base_bds", "model_base")]
VS = 5.0
fig, ax = plt.subplots(figsize=(14, 4.4))
for u_, a_, b_ in SEQ:
    ax.fill_between(ch, prof[b_], prof[a_], color=UNIT.colour[u_], lw=0, label=UNIT.name[u_])
ax.plot(ch, at(ground, se, sn), color="black", lw=0.9, label="LIDAR ground")
sdb = at(sd_bio, se, sn)
ax.plot(ch, prof["base_io"] - sdb, color="#1f3b73", lw=0.7, ls="--")
ax.plot(ch, prof["base_io"] + sdb, color="#1f3b73", lw=0.7, ls="--",
        label="base of Inferior Oolite ± one kriging SD")
cE, dE = project(holes.LOCA_NATE.to_numpy(), holes.LOCA_NATN.to_numpy())
near = holes.assign(ch=cE, off=dE)[dE <= 75]
w = 14
for h in near.itertuples():
    for r in runs[runs.LOCA_ID == h.LOCA_ID].itertuples():
        ax.add_patch(plt.Rectangle((h.ch - w / 2, r.LOCA_GL - r.base), w, r.base - r.top, fc=UNIT.colour[r.group],
                                   ec="black", lw=0.5, zorder=4))
for lid in ["CP217", "RC508", "DSRCOH306", "OH411", "OH413", "DSRCOH400"]:
    if lid in set(near.LOCA_ID):
        h = near.set_index("LOCA_ID").loc[lid]
        ax.annotate(lid, (h.ch, h.LOCA_GL + 4), ha="center", fontsize=6.5, rotation=90, va="bottom")
ax.set_xlim(0, ch[-1]); ax.set_ylim(MODEL_BASE, np.nanmax(prof["ground"]) + 75)   # headroom so the legend clears the hole labels
ax.set_aspect(VS)
ax.set_xlabel("distance along the section (m)", fontsize=8); ax.set_ylabel("m OD", fontsize=8)
ax.tick_params(labelsize=7)
for cx_ in chain0[1:-1]:
    ax.axvline(cx_, color="#555555", lw=0.4, ls=":")
ax.legend(fontsize=7, ncol=4, loc="upper right", frameon=True, framealpha=0.9)
ax.set_title(f"Section along the route corridor, vale (west) to plateau (south-east), {VS:.0f}x vertical "
             f"exaggeration. Logs within 75 m, collapsed to model units", fontsize=9, loc="left")
fig.text(0.01, 0.0, f"{CREDIT} Dotted lines: bends in the section.", fontsize=6.5, color="#444444")
fig.savefig(OUT / "07-section.png")
plt.close(fig)

# Figure 1: the model in 3D, with the superficial deposits lifted off.
X3, Y3 = np.stack([gx, gx], -1), np.stack([gy, gy], -1)
lower = {"SUP": "rockhead", "GOG": "top_io", "IOG": "base_io", "BDS": "base_bds", "LIAS": None}
upper = {"SUP": "ground", "GOG": "rockhead", "IOG": "top_io", "BDS": "base_io", "LIAS": "base_bds"}
pl = pv.Plotter(off_screen=True, window_size=(1800, 1150))
pl.set_background("white")
for u_ in ORDER:
    top = np.nan_to_num(stack[upper[u_]], nan=MODEL_BASE)
    base = np.full_like(top, MODEL_BASE) if lower[u_] is None else np.nan_to_num(stack[lower[u_]], nan=MODEL_BASE)
    g = pv.StructuredGrid(X3, Y3, np.stack([base, top], -1) * VE)
    g.point_data["keep"] = np.stack([inside & (top - base > 0.05)] * 2, -1).ravel(order="F").astype(float)
    solid = g.threshold(0.5, scalars="keep", all_scalars=True)
    if solid.n_cells == 0 or u_ == "SUP":
        continue
    pl.add_mesh(solid.extract_surface(algorithm=None), color=UNIT.colour[u_], smooth_shading=True, split_sharp_edges=True,
                ambient=0.35, diffuse=0.7, specular=0.0)
gnd = pv.StructuredGrid(gx, gy, np.nan_to_num(stack["ground"], nan=MODEL_BASE) * VE + 2)
gnd.point_data["keep"] = inside.ravel(order="F").astype(float)
pl.add_mesh(gnd.threshold(0.5, scalars="keep", all_scalars=True).extract_surface(algorithm=None), color="#c98c5a",
            opacity=0.18)
pins = pv.merge([pv.Line((e - X0, n - Y0, z * VE), (e - X0, n - Y0, (z + 18) * VE))
                 for e, n, z in zip(holes.LOCA_NATE, holes.LOCA_NATN, holes.LOCA_GL)])
pl.add_mesh(pins, color="black", line_width=2)
pl.camera_position = [(-1500, -2600, 2400 * VE / 3 + 900), (2100, 2300, 160 * VE), (0, 0, 1)]
pl.camera.zoom(1.35)
img = pl.screenshot(return_img=True)
pl.close()


def compose(img, units, path, title, note=""):
    """Put a PyVista screenshot beside a proper legend (PyVista's own legend tints the text)."""
    ink = np.where((img[..., :3] < 245).any(axis=2))                 # trim the empty white margins
    img = img[max(ink[0].min() - 10, 0):ink[0].max() + 10, max(ink[1].min() - 10, 0):ink[1].max() + 10]
    h, w = img.shape[:2]
    fig = plt.figure(figsize=(11, 11 * h / w * 0.86))
    ax = fig.add_axes([0, 0.035, 0.8, 0.905])
    ax.imshow(img)
    ax.set_axis_off()
    fig.text(0.01, 0.975, title, fontsize=11, weight="bold", va="top")
    if note:
        fig.text(0.01, 0.005, note, fontsize=7, color="#444444")
    handles = [Patch(facecolor=UNIT.colour[u], edgecolor="#555555", lw=0.4, label=UNIT.name[u]) for u in units]
    handles.append(Patch(facecolor="#c98c5a", alpha=0.25, edgecolor="#555555", lw=0.4,
                         label="LIDAR ground (superficial\ndeposits lifted off)"))
    handles.append(Line2D([], [], color="black", lw=1.2, label="logged hole"))
    fig.legend(handles=handles, loc="center left", bbox_to_anchor=(0.79, 0.5), fontsize=7.5, frameon=False,
               title="Model unit", title_fontsize=8, alignment="left")
    fig.savefig(path, dpi=150)
    plt.close(fig)


compose(img, ["GOG", "IOG", "BDS", "LIAS"], OUT / "01-ground-model-3d.png",
        "A417 ground model: kriged stratum surfaces under LIDAR ground, looking north-east up Crickley Hill",
        f"{CREDIT} {VE:.0f}x vertical exaggeration; model base {MODEL_BASE:.0f} m OD.")
print("wrote", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))