"""Geological cross-section from AGS4 borehole logs, with the ground profile from LIDAR.

Part 4 of the TechnicalGIS series "AGS data in GIS". Draws a section 3.2 km long up the Cotswold
escarpment at Crickley Hill, from the Lias mudstones of the vale to the Great Oolite on the plateau.
Every A417 borehole within 60 m of the line is projected onto it and drawn as a strip log at its true
elevation, under a ground profile sampled from Environment Agency LIDAR. Contacts between formations
are picked from the logs and correlated from hole to hole, and the script measures how much of each
correlation line rests on a nearby log and how much is inferred. Finally it writes a GeoPackage for
the QGIS elevation profile (qgis_section.py).

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/geological-cross-section-ags4-python/

Data: A417 Missing Link ground investigation (BGS deposited data items 173984 and 173975, Phase 2A;
client Highways England, contractor Geotechnical Engineering Ltd), 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 time
from importlib.metadata import version
from pathlib import Path

import numpy as np
import pandas as pd
import geopandas as gpd
import requests
import rasterio
from rasterio.transform import from_origin
from scipy.ndimage import map_coordinates
from shapely.geometry import LineString, Point
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, Rectangle

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

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

PHASE_2A = 173975                                       # BGS item: 20 more holes on the plateau east of Crickley Hill
AOI = (391_600, 212_800, 395_600, 216_800)             # same 4 x 4 km window and 2 m tiles as Part 2
RES, TILE = 2.0, 1000
# The section line, west to east, as British National Grid vertices. It follows the line of holes: up the
# landslipped slope, over the top of Crickley Hill, across the plateau and south-east to the Great Oolite.
VERTICES = [(392_030, 215_720), (392_650, 215_830), (393_060, 215_850), (393_330, 216_020), (393_800, 215_990),
            (393_990, 215_770), (393_920, 215_520), (393_820, 215_305), (394_310, 214_960)]
HALF = 60.0                                             # corridor half-width (m): holes this close are projected
VE = 5.0                                                # vertical exaggeration of the section
STEP = 2.0                                              # ground profile sample spacing (m)
SOLID_GAP = 150.0                                       # correlate holes up to this far apart with a solid line
MAX_DIP = 5.0                                           # steepest apparent dip (deg) accepted between 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)."
# Contacts to correlate: the surface is the top of the first unit with at least this rank (A.UNITS).
SURFACES = pd.DataFrame([
    ("FE",   "top of Fuller's Earth",       3, "#2f6b2f"),
    ("IOG",  "top of Inferior Oolite",      4, "#a0522d"),
    ("BDS",  "base of Inferior Oolite",     7, "#111111"),
    ("WHM",  "base of Bridport Sand",       8, "#1f4e9c"),
    ("MRB",  "top of Marlstone Rock",       9, "#7a3b8f"),
], columns=["key", "name", "rank", "colour"]).set_index("key")


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


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


print(f"geopandas {gpd.__version__} | shapely {version('shapely')} | matplotlib {matplotlib.__version__} | "
      f"python-ags4 {version('python-ags4')}")

# ------------------------------------------------------------------ Step 1: the logs, from two BGS items
section("Step 1: load the logs")
main = A.load_clean(DATA)
extra = A.numeric(A.repair(*A.merge(A.download(DATA, item=PHASE_2A)))[0])
extra["GEOL"] = extra["GEOL"].sort_values(["LOCA_ID", "GEOL_TOP"], ignore_index=True)
extra["GEOL"]["UNIT"] = A.unit_code(extra["GEOL"])
assert not set(main["LOCA"].LOCA_ID) & set(extra["LOCA"].LOCA_ID), "the two items share a hole"
loca = pd.concat([main["LOCA"], extra["LOCA"]], ignore_index=True).set_index("LOCA_ID")
geol = pd.concat([main["GEOL"], extra["GEOL"]], ignore_index=True)
geol = geol[geol.GEOL_BASE > geol.GEOL_TOP].reset_index(drop=True)
geol["rank"] = geol.UNIT.map(A.UNITS["rank"])
geol["last"] = geol.UNIT.map(A.UNITS["last"])
geol = geol.join(loca[["LOCA_NATE", "LOCA_NATN", "LOCA_GL"]], on="LOCA_ID")
geol["top_mOD"] = geol.LOCA_GL - geol.GEOL_TOP
geol["base_mOD"] = geol.LOCA_GL - geol.GEOL_BASE
print(f"item 173984: {len(main['LOCA'])} holes | Phase 2A (item {PHASE_2A}): {len(extra['LOCA'])} holes | "
      f"together {len(loca)} holes and {len(geol):,} GEOL intervals")

# ------------------------------------------------------------------ Step 2: the line and its corridor
section("Step 2: project the holes onto the section line")
line = LineString(VERTICES)
bends = [line.project(Point(v)) for v in VERTICES]
LETTERS = "ABCDEFGHI"[:len(VERTICES)]
print(f"section {LETTERS[0]}-{LETTERS[-1]}: {len(VERTICES) - 1} legs, {line.length:,.0f} m long | corridor ±{HALF:.0f} m")


def side(p, ch):
    """+1 if the point is left of the line looking along it (north-ish on the first legs), -1 if right."""
    a, b = line.interpolate(max(ch - 1, 0)), line.interpolate(min(ch + 1, line.length))
    return np.sign((b.x - a.x) * (p.y - a.y) - (b.y - a.y) * (p.x - a.x)) or 1.0


holes = loca[["LOCA_TYPE", "LOCA_NATE", "LOCA_NATN", "LOCA_GL", "LOCA_FDEP"]].copy()
pts = [Point(x, y) for x, y in zip(holes.LOCA_NATE, holes.LOCA_NATN)]
holes["chainage"] = [line.project(p) for p in pts]
holes["offset"] = [line.distance(p) * side(p, c) for p, c in zip(pts, holes.chainage)]
on = holes[(holes.offset.abs() <= HALF) & (holes.LOCA_TYPE != "CPT")].sort_values("chainage").copy()
on["item"] = np.where(on.index.isin(extra["LOCA"].LOCA_ID), PHASE_2A, 173984)
print(f"{len(on)} holes within {HALF:.0f} m (cone penetration tests left out: they have no GEOL descriptions) | "
      f"{(on.LOCA_FDEP >= 10).sum()} boreholes 10 m or deeper, {(on.LOCA_FDEP < 10).sum()} trial pits and shallow holes | "
      f"{(on.item == PHASE_2A).sum()} from Phase 2A")
print(f"distance off the line: median {on.offset.abs().median():.0f} m, {(on.offset.abs() <= 20).sum()} within 20 m")
print(on.loc[on.LOCA_FDEP >= 10, ["LOCA_TYPE", "LOCA_GL", "LOCA_FDEP", "chainage", "offset"]].round(1).to_string())

# ------------------------------------------------------------------ Step 3: ground profile from LIDAR
section("Step 3: ground profile from Environment Agency LIDAR")
X0, Y0, X1, Y1 = AOI
dtm = np.full((int((Y1 - Y0) / RES), int((X1 - X0) / RES)), np.nan, dtype="float32")
for x in range(X0, X1, TILE):
    for y in range(Y0, Y1, TILE):
        f = get(f"{WCS}&subset=E({x},{x + TILE})&subset=N({y},{y + TILE})&SCALEFACTOR={1 / RES}",
                DATA / f"ea_dtm_{x}_{y}_{TILE}m_tile_{RES:.0f}m.tif")
        with rasterio.open(f) as src:
            a = src.read(1).astype("float32")
            a[a < -1000] = np.nan
            r, c = int(round((Y1 - src.bounds.top) / RES)), int(round((src.bounds.left - X0) / RES))
            dtm[r:r + a.shape[0], c:c + a.shape[1]] = a


def lidar(x, y):
    """Bilinear sample of the 2 m DTM at British National Grid coordinates."""
    rows = (Y1 - np.asarray(y)) / RES - 0.5
    cols = (np.asarray(x) - X0) / RES - 0.5
    return map_coordinates(dtm, [rows, cols], order=1, mode="nearest")


ch = np.arange(0, line.length + STEP / 2, STEP)
xy = np.array([line.interpolate(c).coords[0] for c in ch])
ground = lidar(xy[:, 0], xy[:, 1])
# Unit normal at every sample (left of the line), for the strip of ground across the corridor.
d = np.gradient(xy, axis=0)
normal = np.c_[-d[:, 1], d[:, 0]] / np.hypot(d[:, 0], d[:, 1])[:, None]
across = np.stack([lidar(xy[:, 0] + o * normal[:, 0], xy[:, 1] + o * normal[:, 1])
                   for o in np.arange(-HALF, HALF + 1, 5)])
lo, hi = across.min(axis=0), across.max(axis=0)
print(f"{len(ch):,} samples every {STEP:.0f} m: ground {ground.min():.1f} to {ground.max():.1f} m OD | "
      f"steepest 20 m stretch {np.degrees(np.arctan(np.abs(np.convolve(np.gradient(ground, STEP), np.ones(10) / 10, 'same')).max())):.0f} deg")
print(f"across the ±{HALF:.0f} m corridor the ground varies by up to {(hi - lo).max():.1f} m (median {np.median(hi - lo):.1f} m)")

on["lidar_at_hole"] = lidar(on.LOCA_NATE, on.LOCA_NATN).round(2)
on["lidar_on_line"] = np.interp(on.chainage, ch, ground).round(2)
on["gl_minus_lidar"] = (on.LOCA_GL - on.lidar_at_hole).round(2)
on["hole_minus_line"] = (on.LOCA_GL - on.lidar_on_line).round(2)
print(f"LOCA_GL minus LIDAR at the hole: median {on.gl_minus_lidar.median():+.2f} m, "
      f"{(on.gl_minus_lidar.abs() <= 0.5).sum()} of {len(on)} within 0.5 m")
print(f"LOCA_GL minus the profile on the line: up to {on.hole_minus_line.abs().max():.1f} m "
      f"({on.hole_minus_line.abs().idxmax()}, {abs(on.loc[on.hole_minus_line.abs().idxmax(), 'offset']):.0f} m off the line)")

# ------------------------------------------------------------------ Step 4: pick the contacts
section("Step 4: pick the contacts in each log")


def in_place(g):
    """Bedrock intervals below any repeat in the sequence (CP217's slipped block), and the repeat's depth."""
    rock = g[g["rank"] > 0]
    reached, split = 0, None
    for r in rock.itertuples():
        if r.last < reached:
            split, reached = r.GEOL_TOP, 0
        reached = max(reached, r.rank)
    return (rock if split is None else rock[rock.GEOL_TOP >= split]), split


rows, below_hole, repeats = [], [], {}
for lid, g in geol[geol.LOCA_ID.isin(on.index)].groupby("LOCA_ID", sort=False):
    rock, split = in_place(g)
    if split is not None:
        repeats[lid] = split
    for key, s in SURFACES.iterrows():
        above, below = rock[rock["last"] < s["rank"]], rock[rock["rank"] >= s["rank"]]
        gl = g.LOCA_GL.iloc[0]
        if len(above) and len(below) and above.GEOL_BASE.max() <= below.GEOL_TOP.min():
            top_of_gap, depth = above.GEOL_BASE.max(), below.GEOL_TOP.min()
            rows.append(dict(LOCA_ID=lid, surface=key, depth=depth, z=gl - depth, gap=depth - top_of_gap))
        elif len(above) and not len(below) and above["rank"].max() >= s["rank"] - 3:
            below_hole.append(dict(LOCA_ID=lid, surface=key, z=gl - g.GEOL_BASE.max()))
contacts = pd.DataFrame(rows).join(on[["chainage", "offset", "LOCA_NATE", "LOCA_NATN"]], on="LOCA_ID")
below_hole = pd.DataFrame(below_hole).join(on[["chainage"]], on="LOCA_ID")
print(f"{len(contacts)} contacts picked in {contacts.LOCA_ID.nunique()} holes "
      f"({(contacts.gap > 0).sum()} fall in a gap with no recovery, up to {contacts.gap.max():.2f} m thick)")
print(contacts.groupby("surface", sort=False).agg(holes=("z", "size"), lowest=("z", "min"), highest=("z", "max"))
      .rename(index=SURFACES.name).round(1).to_string())
for lid, depth in repeats.items():
    print(f"{lid}: the sequence repeats from {depth:.2f} m; contacts are taken from the lower, in-place sequence")

# Part 3's test: is a hole's base of the Inferior Oolite far from the median of its neighbours within 100 m?
# Use every hole in both items, on the line or not, so the test has as many neighbours as possible.
allc = []
for lid, g in geol.groupby("LOCA_ID", sort=False):
    rock, _ = in_place(g)
    above, below = rock[rock["last"] < 7], rock[rock["rank"] >= 7]
    if len(above) and len(below):
        allc.append(dict(LOCA_ID=lid, E=g.LOCA_NATE.iloc[0], N=g.LOCA_NATN.iloc[0],
                         z=g.LOCA_GL.iloc[0] - below.GEOL_TOP.min()))
allc = pd.DataFrame(allc).set_index("LOCA_ID")
nb_dz = {}
for lid, c in allc.iterrows():
    nb = allc[(np.hypot(allc.E - c.E, allc.N - c.N) <= 100) & (allc.index != lid)]
    if len(nb) >= 2:
        nb_dz[lid] = c.z - nb.z.median()
nb_dz = pd.Series(nb_dz)
dropped = nb_dz[nb_dz.abs() > 8].index.tolist()
print(f"base of the Inferior Oolite in {len(allc)} holes; more than 8 m from the median of their neighbours: "
      + ", ".join(f"{h} ({nb_dz[h]:+.1f} m)" for h in dropped))

# A plane through the other contacts gives the regional dip, used to judge what projecting onto the line costs.
ok = allc.drop(index=dropped)
M = np.c_[np.ones(len(ok)), ok.E - X0, ok.N - Y0]
coef, *_ = np.linalg.lstsq(M, ok.z, rcond=None)
resid = ok.z - M @ coef
grad = coef[1:]
dip = np.degrees(np.arctan(np.hypot(*grad)))
dip_dir = (np.degrees(np.arctan2(-grad[0], -grad[1])) + 360) % 360
print(f"plane through {len(ok)} contacts: dip {dip:.1f} deg towards {dip_dir:.0f} deg (bearing) | "
      f"residuals RMS {np.sqrt((resid ** 2).mean()):.1f} m, largest {resid.abs().max():.1f} m ({resid.abs().idxmax()})")
foot = np.array([line.interpolate(c).coords[0] for c in contacts.chainage])
contacts["projection_dz"] = (grad[0] * (foot[:, 0] - contacts.LOCA_NATE) + grad[1] * (foot[:, 1] - contacts.LOCA_NATN)).round(2)
print(f"moving each contact from its hole to the line along that plane changes it by "
      f"{contacts.projection_dz.abs().median():.2f} m (median), {contacts.projection_dz.abs().max():.2f} m at most")
contacts["dropped"] = contacts.LOCA_ID.isin(dropped) & (contacts.surface == "BDS")

# ------------------------------------------------------------------ Step 5: correlate and measure the inference
section("Step 5: correlate the contacts and measure how much is inferred")
tan_max = np.tan(np.radians(MAX_DIP))


def apparent_gradient(c):
    """Rise of the fitted plane per metre along the line at chainage c (positive = rising along the line)."""
    a, b = line.interpolate(max(c - 5, 0)), line.interpolate(min(c + 5, line.length))
    u = np.array([b.x - a.x, b.y - a.y]) / np.hypot(b.x - a.x, b.y - a.y)
    return float(grad @ u)


segments, extensions, pairs = [], [], []
for key in SURFACES.index:
    c = contacts[(contacts.surface == key) & ~contacts.dropped].sort_values("chainage")
    for a, b in zip(c.itertuples(), c.iloc[1:].itertuples()):
        gap, dz = b.chainage - a.chainage, b.z - a.z
        agree = abs(dz) <= max(3.0, gap * tan_max)
        pairs.append(dict(surface=key, a=a.LOCA_ID, b=b.LOCA_ID, gap=round(gap), dz=round(dz, 1), agree=agree))
        if agree:
            segments.append(dict(surface=key, c0=a.chainage, z0=a.z, c1=b.chainage, z1=b.z,
                                 style="solid" if gap <= SOLID_GAP else "dashed"))
    # Carry the surface on from its first and last contacts at the plane's apparent dip until it meets the ground.
    if len(c):
        for end, sign in ((c.iloc[0], -1), (c.iloc[-1], 1)):
            cc, zz = end.chainage, end.z
            path = [(cc, zz)]
            while 0 < cc < line.length and len(path) < 400 / STEP:
                cc += sign * STEP
                zz += sign * STEP * apparent_gradient(cc)
                path.append((cc, zz))
                if zz >= np.interp(cc, ch, ground):
                    extensions.append(dict(surface=key, path=np.array(path), outcrop=cc))
                    break
segments = pd.DataFrame(segments)
pairs = pd.DataFrame(pairs)
outcrops = pd.DataFrame(extensions)[["surface", "outcrop"]] if extensions else pd.DataFrame(columns=["surface", "outcrop"])
print(pairs.to_string(index=False))


def summary(key):
    """How much of each surface's line rests on logs: the share drawn within 50 m (along the line) of a contact."""
    s = segments[segments.surface == key]
    e = [x["path"] for x in extensions if x["surface"] == key]
    cs = contacts[(contacts.surface == key) & ~contacts.dropped].chainage.to_numpy()
    xs = np.concatenate([np.arange(r.c0, r.c1, 1.0) for r in s.itertuples()] + [p[:, 0] for p in e]) if len(s) or e \
        else np.array([])
    near = (np.abs(xs[:, None] - cs[None, :]).min(axis=1) <= 50).mean() if len(xs) else np.nan
    p = pairs[pairs.surface == key]
    return dict(surface=SURFACES.name[key], holes=len(cs), drawn_m=len(xs),
                between_close_holes_m=round((s.c1 - s.c0)[s["style"] == "solid"].sum()),
                across_wide_gaps_m=round((s.c1 - s.c0)[s["style"] == "dashed"].sum()),
                carried_to_ground_m=sum(len(x) - 1 for x in e) * STEP,
                within_50m_of_a_log=f"{near:.0%}" if len(xs) else "-",
                neighbours_agree=f"{p.agree.sum()} of {len(p)}")


inference = pd.DataFrame([summary(k) for k in SURFACES.index])
print("\nhow much of each correlation line rests on a log:")
print(inference.to_string(index=False))
print("where each surface is carried up to the ground (chainage, m):",
      ", ".join(f"{SURFACES.name[r.surface]} {r.outcrop:.0f}" for r in outcrops.itertuples()))

# What does the log say in the dropped holes and in the repeated sequence?
steep = geol.GEOL_DESC.str.extract(r"bedded \((\d+)\s*-\s*(\d+)\s*deg\)", expand=True).astype(float)
geol["bedding_dip"] = steep[1]
geol["slickensides"] = geol.GEOL_DESC.str.contains("slickensid", case=False)
moved = geol[geol.LOCA_ID.isin(on.index) & (geol.bedding_dip.notna() | geol.slickensides)]
print(f"\nlogs on the section that record steep bedding or slickensides: {', '.join(moved.LOCA_ID.unique())}")
cp = contacts[contacts.surface == "MRB"].set_index("LOCA_ID")
if "CP217" in repeats:
    g = geol[geol.LOCA_ID == "CP217"]
    upper = g[(g.UNIT == "MRB") & (g.GEOL_TOP < repeats["CP217"])].iloc[0]
    print(f"CP217: Marlstone Rock at {upper.top_mOD:.1f} m OD in the upper sequence and {cp.z['CP217']:.1f} m OD in "
          f"the lower | neighbours on the line: " + ", ".join(f"{h} {cp.z[h]:.1f}" for h in cp.index if h != "CP217"))

# ------------------------------------------------------------------ Figures
section("Figures")
shade = LightSource(azdeg=315, altdeg=45).hillshade(np.nan_to_num(dtm, nan=np.nanmean(dtm)), vert_exag=2, dx=RES, dy=RES)
LOGS = on[on.index.isin(geol.LOCA_ID)]
UNITS_USED = [u for u in A.UNITS.index if u in set(geol[geol.LOCA_ID.isin(LOGS.index)].UNIT)]


def plan(ax, pad=180, labels=True, ticks=True):
    """GIS map of the line: LIDAR hillshade, corridor, holes in and out of it, chainage every 500 m."""
    xs_, ys_ = zip(*VERTICES)
    x0, x1, y0, y1 = min(xs_) - pad, max(xs_) + pad, min(ys_) - pad, max(ys_) + pad
    r0, r1, c0, c1 = int((Y1 - y1) / RES), int((Y1 - y0) / RES), int((x0 - X0) / RES), int((x1 - X0) / RES)
    ax.imshow(shade[r0:r1, c0:c1], cmap="gray", extent=(x0, x1, y0, y1), vmin=0, vmax=1.15)
    ax.imshow(dtm[r0:r1, c0:c1], cmap="terrain", extent=(x0, x1, y0, y1), alpha=0.25, vmin=40, vmax=340)
    corridor = line.buffer(HALF, cap_style="flat", join_style="mitre")
    ax.fill(*corridor.exterior.xy, color="#ffd400", alpha=0.25, lw=0)
    ax.plot(*line.xy, color="#c62828", lw=1.6)
    near = holes[(holes.LOCA_NATE.between(x0, x1)) & (holes.LOCA_NATN.between(y0, y1))]
    ax.scatter(near.LOCA_NATE, near.LOCA_NATN, s=6, color="#777777", lw=0, zorder=3)
    ax.scatter(on.LOCA_NATE, on.LOCA_NATN, s=14 if labels else 7, color="#111111", edgecolor="white", lw=0.4, zorder=4)
    for L_, (x, y) in zip(LETTERS, VERTICES):
        ax.annotate(L_, (x, y), xytext=(0, 6 if labels else 3), textcoords="offset points", ha="center",
                    fontsize=8 if labels else 6, weight="bold", color="#c62828",
                    bbox=dict(facecolor="white", alpha=0.8, lw=0, pad=0.6))
    if ticks:
        for c in range(500, int(line.length), 500):
            p = line.interpolate(c)
            ax.plot(p.x, p.y, marker="|", ms=7, color="#c62828", mew=1.2)
            ax.annotate(f"{c / 1000:.1f} km", (p.x, p.y), xytext=(4, -10), textcoords="offset points", fontsize=6,
                        color="#c62828")
    ax.set_xlim(x0, x1); ax.set_ylim(y0, y1); ax.set_aspect("equal")
    return x0, x1, y0, y1


def strip(ax, lid, c, width):
    """One strip log at chainage c, true elevation, coloured by unit; no recovery is hatched."""
    g = geol[geol.LOCA_ID == lid]
    for r in g.itertuples():
        ax.add_patch(Rectangle((c - width / 2, r.base_mOD), width, r.top_mOD - r.base_mOD,
                               facecolor=A.UNITS.colour[r.UNIT], edgecolor="none", zorder=3,
                               hatch="////" if r.UNIT == "NR" else None))
    gl, base = g.LOCA_GL.iloc[0], g.base_mOD.min()
    ax.add_patch(Rectangle((c - width / 2, base), width, gl - base, fill=False, edgecolor="#333333", lw=0.4, zorder=4))


def label_logs(ax, logs, min_sep, fontsize, y_fixed=None, label_len=0):
    """Hole names above the logs; where holes crowd together, names step along the line on short leaders.
    With y_fixed, names hang from that elevation instead (for the close-ups)."""
    last = -np.inf
    for lid, h in logs.sort_values("chainage").iterrows():
        x = max(h.chainage, last + min_sep)
        last = x
        if y_fixed is not None:
            ax.plot([h.chainage, h.chainage, x], [h.LOCA_GL + 0.5, y_fixed - label_len - 2, y_fixed - label_len - 1],
                    color="#777777", lw=0.4, zorder=6)
            ax.text(x, y_fixed, lid, rotation=90, ha="center", va="top", fontsize=fontsize, zorder=6)
            continue
        y = max(h.LOCA_GL, np.interp(x, ch, hi)) + 3
        if x - h.chainage > 1:
            y += 4
            ax.plot([h.chainage, x], [h.LOCA_GL + 1, y - 1], color="#777777", lw=0.4, zorder=6)
        ax.text(x, y, lid, rotation=90, ha="center", va="bottom", fontsize=fontsize, zorder=6)


def draw_section(ax, c0, c1, width, labels_min_depth=10, fontsize=6, min_sep=30, y_fixed=None, label_len=0):
    m = (ch >= c0 - 10) & (ch <= c1 + 10)
    ax.fill_between(ch[m], lo[m], hi[m], color="#c9b99a", alpha=0.45, lw=0, zorder=1)
    ax.plot(ch[m], ground[m], color="#3b3024", lw=1.0, zorder=5)
    for b, L_ in zip(bends, LETTERS):
        if c0 <= b <= c1:
            ax.axvline(b, color="#c62828", lw=0.6, ls=(0, (1, 2)), zorder=0)
            ax.text(b, 1.0, L_, transform=ax.get_xaxis_transform(), ha="center", va="bottom", fontsize=8,
                    weight="bold", color="#c62828")
    here = LOGS[(LOGS.chainage >= c0) & (LOGS.chainage <= c1)]
    for lid, h in here.iterrows():
        strip(ax, lid, h.chainage, width)
    label_logs(ax, here[here.LOCA_FDEP >= labels_min_depth], min_sep, fontsize, y_fixed, label_len)
    for key, s in SURFACES.iterrows():
        for r in segments[segments.surface == key].itertuples():
            ax.plot([r.c0, r.c1], [r.z0, r.z1], color=s.colour, lw=1.4, zorder=7,
                    ls="-" if r.style == "solid" else (0, (5, 3)))
        for e in extensions:
            if e["surface"] == key:
                ax.plot(e["path"][:, 0], e["path"][:, 1], color=s.colour, lw=1.2, ls=(0, (1, 2)), zorder=7)
        cs = contacts[(contacts.surface == key)]
        ax.scatter(cs.chainage[~cs.dropped], cs.z[~cs.dropped], s=12, color=s.colour, edgecolor="white", lw=0.5,
                   zorder=8)
        bh = below_hole[below_hole.surface == key] if len(below_hole) else below_hole
        for r in bh.itertuples():
            if c0 <= r.chainage <= c1 and loca.LOCA_FDEP[r.LOCA_ID] >= 15:
                ax.plot(r.chainage, r.z - 2, marker="v", ms=4, color=s.colour, zorder=8)
    d_ = contacts[contacts.dropped]
    ax.scatter(d_.chainage, d_.z, s=26, marker="D", color="#d62728", edgecolor="white", lw=0.6, zorder=9)
    ax.set_xlim(c0, c1)
    ax.set_aspect(VE)
    ax.set_xlabel("distance along the section (m)")
    ax.set_ylabel("elevation (m OD)")
    ax.spines[["top", "right"]].set_visible(False)


legend_units = [Patch(facecolor=A.UNITS.colour[u], edgecolor="#555555", lw=0.3, label=A.UNITS.name[u],
                      hatch="////" if u == "NR" else None) for u in UNITS_USED]
legend_lines = [Line2D([], [], color=s.colour, lw=1.4, label=s["name"]) for _, s in SURFACES.iterrows()]
legend_style = [Line2D([], [], color="#555555", lw=1.4, label=f"holes up to {SOLID_GAP:.0f} m apart"),
                Line2D([], [], color="#555555", lw=1.4, ls=(0, (5, 3)), label="across a wider gap"),
                Line2D([], [], color="#555555", lw=1.2, ls=(0, (1, 2)), label="carried to the ground at the regional dip"),
                Line2D([], [], ls="", marker="v", ms=4, color="#555555", label="surface is below the end of the hole"),
                Line2D([], [], ls="", marker="D", ms=5, color="#d62728", label="base of Inferior Oolite 11-15 m low"),
                Patch(facecolor="#c9b99a", alpha=0.45, label=f"ground across the ±{HALF:.0f} m corridor")]

# 01: the whole section, with the map as an inset
fig = plt.figure(figsize=(15, 7.4))
ax = fig.add_axes([0.05, 0.12, 0.74, 0.8])
draw_section(ax, -20, line.length + 20, width=14)
ax.set_ylim(95, 305)
ins = ax.inset_axes([0.0, 0.58, 0.34, 0.42])
plan(ins, pad=120, labels=False, ticks=False)
ins.set_xticks([]); ins.set_yticks([])
for sp in ins.spines.values():
    sp.set_edgecolor("#555555")
ins.text(0.02, 0.03, "LIDAR © EA (OGL v3)", transform=ins.transAxes, fontsize=5.5,
         bbox=dict(facecolor="white", alpha=0.75, lw=0, pad=1))
ax.set_title(f"Section {LETTERS[0]}–{LETTERS[-1]} up the Cotswold escarpment at Crickley Hill: {len(LOGS)} logs within "
             f"{HALF:.0f} m of the line, ground from LIDAR (vertical exaggeration ×{VE:.0f})", fontsize=10, loc="left", pad=16)
fig.legend(handles=legend_units, loc="upper left", bbox_to_anchor=(0.8, 0.97), fontsize=7, frameon=False,
           title="GEOL unit", title_fontsize=8, alignment="left")
fig.legend(handles=legend_lines + legend_style, loc="lower left", bbox_to_anchor=(0.8, 0.02), fontsize=7,
           frameon=False, title="correlation", title_fontsize=8, alignment="left")
fig.text(0.05, 0.02, CREDIT + f" Logs drawn at their own ground level; letters mark bends in the line.",
         fontsize=7, color="#444444")
fig.savefig(OUT / "01-cross-section.png", dpi=150)
plt.close(fig)

# 02: the section line as a GIS map
fig, ax = plt.subplots(figsize=(9, 5.6))
x0, x1, y0, y1 = plan(ax)
MAP_LABELS = {"CP217": (-8, 8), "DSRC301": (-44, -12), "DSRC109": (-40, 10), "DSRC406": (6, 12), "DSRC110": (10, 8),
              "DSRCOH410": (-10, -14), "DSRCOH306": (6, 6), "DSRCOH308": (8, -2), "OH411": (6, 2), "OH413": (-36, 6)}
for lid, (dx, dy) in MAP_LABELS.items():
    h = on.loc[lid]
    ax.annotate(lid, (h.LOCA_NATE, h.LOCA_NATN), xytext=(dx, dy), textcoords="offset points", fontsize=6,
                bbox=dict(facecolor="white", alpha=0.75, lw=0, pad=0.6), zorder=5,
                arrowprops=dict(arrowstyle="-", color="#333333", lw=0.5, shrinkA=0, shrinkB=2))
ax.set_xticks(np.arange(np.ceil(x0 / 500) * 500, x1, 500))
ax.set_yticks(np.arange(np.ceil(y0 / 500) * 500, y1, 500))
ax.tick_params(labelsize=7)
ax.set_xlabel("easting (m, EPSG:27700)", fontsize=8); ax.set_ylabel("northing (m)", fontsize=8)
sb = x0 + 80
ax.plot([sb, sb + 500], [y0 + 70] * 2, color="black", lw=2.5)
ax.text(sb + 250, y0 + 95, "500 m", ha="center", fontsize=7)
ax.annotate("N", (x1 - 90, y1 - 70), xytext=(x1 - 90, y1 - 230), ha="center", fontsize=9, weight="bold",
            arrowprops=dict(arrowstyle="-|>", color="black", lw=1.2))
ax.legend(handles=[Line2D([], [], color="#c62828", lw=1.6, label="section line"),
                   Patch(facecolor="#ffd400", alpha=0.35, label=f"corridor ±{HALF:.0f} m"),
                   Line2D([], [], ls="", marker="o", ms=4, color="#111111", label=f"projected ({len(on)})"),
                   Line2D([], [], ls="", marker="o", ms=3, color="#777777", label="other holes")],
          loc="lower left", bbox_to_anchor=(0.0, 0.1), fontsize=7, framealpha=0.9)
ax.text(0.01, 0.99, "Contains Environment Agency information © Environment Agency and database right; BGS © UKRI "
        "(OGL v3)", transform=ax.transAxes, fontsize=5.5, va="top", bbox=dict(facecolor="white", alpha=0.75, lw=0, pad=1.5))
ax.set_title(f"Section line {LETTERS[0]}–{LETTERS[-1]}: {line.length / 1000:.1f} km, {len(on)} holes within "
             f"{HALF:.0f} m (labels on holes 60 m or deeper, and CP217)", fontsize=9)
fig.savefig(OUT / "02-section-line-map.png")
plt.close(fig)

# 03: the ground profile and how the collars sit on it
fig, (a1, a2) = plt.subplots(2, 1, figsize=(11, 5.6), sharex=True, gridspec_kw={"height_ratios": [2.3, 1]})
a1.fill_between(ch, lo, hi, color="#c9b99a", alpha=0.55, lw=0, label=f"ground across the ±{HALF:.0f} m corridor")
a1.plot(ch, ground, color="#3b3024", lw=1, label="LIDAR on the line")
a1.scatter(on.chainage, on.LOCA_GL, s=10, color="#1f78b4", zorder=3, label="LOCA_GL of each projected hole")
for b, L_ in zip(bends, LETTERS):
    a1.axvline(b, color="#c62828", lw=0.6, ls=(0, (1, 2)))
    a1.text(b, 1.0, L_, transform=a1.get_xaxis_transform(), ha="center", va="bottom", fontsize=8, weight="bold",
            color="#c62828")
a1.set_ylabel("m OD")
a1.legend(fontsize=7, frameon=False, loc="lower right")
a2.axhline(0, color="#555555", lw=0.7)
a2.scatter(on.chainage, on.hole_minus_line, s=10, color="#bbbbbb", label="LOCA_GL minus the profile on the line")
a2.scatter(on.chainage, on.gl_minus_lidar, s=10, color="#1f78b4", label="LOCA_GL minus LIDAR at the hole itself")
a2.set_ylabel("m"); a2.set_xlabel("distance along the section (m)")
a2.legend(fontsize=7, frameon=False, loc="lower left", ncol=2)
for a in (a1, a2):
    a.spines[["top", "right"]].set_visible(False)
a1.set_title("Ground profile from Environment Agency LIDAR (2 m DTM), with the collars of the projected holes",
             fontsize=9, loc="left")
fig.text(0.0, -0.01, CREDIT, fontsize=6.5, color="#444444")
fig.savefig(OUT / "03-ground-profile.png")
plt.close(fig)

# 04: how much is inferred
fig, axs = plt.subplots(len(SURFACES) + 1, 1, figsize=(11, 5.8), sharex=True,
                        gridspec_kw={"height_ratios": [2.2] + [0.55] * len(SURFACES)})
a0 = axs[0]
a0.axhspan(-HALF, HALF, color="#ffd400", alpha=0.15, lw=0)
a0.axhline(0, color="#c62828", lw=1)
deep = on.LOCA_FDEP >= 10
a0.vlines(on.chainage[deep], 0, on.offset[deep], color="#555555", lw=0.8)
a0.scatter(on.chainage[deep], on.offset[deep], s=14, color="#111111", zorder=3, label="borehole 10 m or deeper")
a0.scatter(on.chainage[~deep], on.offset[~deep], s=8, color="#999999", zorder=3, label="trial pit or shallow hole")
a0.set_ylim(-HALF - 8, HALF + 8)
a0.set_ylabel("off the line (m)\nleft + / right −", fontsize=8)
pt_handles = a0.get_legend_handles_labels()[0]
a0.set_title("Where the correlation lines rest on logs. Top: how far each hole is from the line.\nBelow, for each "
             "surface: dots are logged contacts; the bar is shaded by distance to the nearest one", fontsize=9, loc="left")
xs_all = np.arange(0, line.length, 1.0)
for a, (key, s) in zip(axs[1:], SURFACES.iterrows()):
    cs = contacts[(contacts.surface == key) & ~contacts.dropped].chainage.to_numpy()
    drawn = np.zeros_like(xs_all, dtype=bool)
    for r in segments[segments.surface == key].itertuples():
        drawn |= (xs_all >= r.c0) & (xs_all <= r.c1)
    for e in extensions:
        if e["surface"] == key:
            drawn |= (xs_all >= e["path"][:, 0].min()) & (xs_all <= e["path"][:, 0].max())
    dist = np.abs(xs_all[:, None] - cs[None, :]).min(axis=1) if len(cs) else np.full_like(xs_all, np.inf)
    for lo_, hi_, alpha in [(0, 50, 0.9), (50, 150, 0.5), (150, 1e9, 0.18)]:
        m = drawn & (dist > lo_) & (dist <= hi_) if lo_ else drawn & (dist <= hi_)
        a.fill_between(xs_all, 0, 1, where=m, color=s.colour, alpha=alpha, lw=0, step="mid")
    a.scatter(cs, np.full(len(cs), 0.5), s=10, color="white", edgecolor=s.colour, zorder=3)
    dd = contacts[(contacts.surface == key) & contacts.dropped]
    a.scatter(dd.chainage, np.full(len(dd), 0.5), s=16, marker="D", color="#d62728", zorder=3)
    a.set_yticks([]); a.set_ylim(0, 1)
    a.text(-30, 0.5, s["name"], ha="right", va="center", fontsize=7)
    a.spines[["top", "right", "left"]].set_visible(False)
axs[-1].set_xlabel("distance along the section (m)")
axs[-1].set_xlim(0, line.length)
fig.legend(handles=pt_handles + [Patch(color="#555555", alpha=a, label=l) for a, l in
                    [(0.9, "within 50 m of a logged contact"), (0.5, "50 to 150 m"), (0.18, "more than 150 m")]],
           loc="lower center", ncol=5, fontsize=7, frameon=False, bbox_to_anchor=(0.5, -0.04))
fig.text(0.0, -0.07, CREDIT, fontsize=6.5, color="#444444")
fig.savefig(OUT / "04-how-much-is-inferred.png")
plt.close(fig)

# 05: close-ups of the two places the logs disagree
fig, (a1, a2) = plt.subplots(1, 2, figsize=(14, 4.6), gridspec_kw={"width_ratios": [1, 1.25]})
for a, (c0, c1, z0, z1, ll, title) in zip((a1, a2), [
        (560, 830, 135, 222, 10, "Below the escarpment: CP217 and its neighbours"),
        (1170, 1490, 180, 262, 12, "Top of Crickley Hill: the base of the Inferior Oolite")]):
    draw_section(a, c0, c1, width=7, labels_min_depth=0, fontsize=7, min_sep=11, y_fixed=z1 - 4, label_len=ll)
    a.set_ylim(z0, z1)
    a.set_aspect(2.0)
    a.set_title(title, fontsize=9, loc="left")
for r in moved.itertuples():
    c = on.chainage[r.LOCA_ID]
    if 1170 <= c <= 1490:
        mid = (r.top_mOD + r.base_mOD) / 2
        a2.plot(c + 6, mid, marker="$\u2215$" if pd.notna(r.bedding_dip) else "x", ms=9 if pd.notna(r.bedding_dip) else 6,
                mew=1.6, color="#d62728", zorder=9)
prev = -np.inf
for r in contacts[contacts.surface == "BDS"].sort_values("chainage").itertuples():
    if 1170 <= r.chainage <= 1490:
        left = r.chainage - prev < 12
        colour = "#d62728" if r.dropped else "black"
        nearby = on.chainage[(on.chainage > r.chainage + 1) & (on.chainage <= r.chainage + 15)]
        if not left and len(nearby):
            # another log sits just to the right: put the label clear of it, on a short leader
            x = nearby.max() + 8
            a2.plot([r.chainage + 1, x - 1], [r.z, r.z + 2.5], color="#777777", lw=0.4, zorder=10)
            a2.text(x, r.z + 2.5, f"{r.z:.1f}", fontsize=6.5, va="center", ha="left", color=colour, zorder=10)
        else:
            a2.text(r.chainage + (-5 if left else 5), r.z - 1.2, f"{r.z:.1f}", fontsize=6.5, va="top",
                    ha="right" if left else "left", color=colour, zorder=10)
        prev = r.chainage
if "CP217" in repeats:
    c, g = on.chainage["CP217"], geol[geol.LOCA_ID == "CP217"]
    z_split = g.LOCA_GL.iloc[0] - repeats["CP217"]
    a1.plot([c - 9, c + 9], [z_split] * 2, color="#d62728", lw=1.6, zorder=9)
    a1.annotate(f"sequence repeats\nbelow {z_split:.1f} m OD", (c + 4, z_split), xytext=(c + 35, z_split - 30),
                fontsize=7, color="#d62728", arrowprops=dict(arrowstyle="-", color="#d62728", lw=0.8), zorder=10)
    for r in cp.itertuples():
        if 560 <= r.chainage <= 830:
            a1.text(r.chainage + 4, r.z - 0.8, f"{r.z:.1f}", fontsize=6.5, va="top", zorder=10)
    a1.text(on.chainage["CP217"] + 4, upper.top_mOD - 0.8, f"{upper.top_mOD:.1f}", fontsize=6.5, va="top",
            color="#d62728", zorder=10)
handles = [Line2D([], [], ls="", marker="$\u2215$", ms=8, color="#d62728", label="log records bedding dipping 30-50°"),
           Line2D([], [], ls="", marker="x", ms=6, mew=1.6, color="#d62728", label="log records slickensides"),
           Line2D([], [], ls="", marker="D", ms=5, color="#d62728", label="base of Inferior Oolite 11-15 m low"),
           Line2D([], [], color=SURFACES.colour["BDS"], lw=1.4, label="base of Inferior Oolite"),
           Line2D([], [], color=SURFACES.colour["WHM"], lw=1.4, label="base of Bridport Sand"),
           Line2D([], [], color=SURFACES.colour["MRB"], lw=1.4, label="top of Marlstone Rock (in place)")]
fig.legend(handles=handles + legend_units, loc="lower center", ncol=5, fontsize=7, frameon=False,
           bbox_to_anchor=(0.5, -0.12))
fig.text(0.0, -0.15, CREDIT + " Vertical exaggeration ×2; labels give contact elevations in m OD.", fontsize=6.5,
         color="#444444")
fig.savefig(OUT / "05-close-ups.png")
plt.close(fig)
print("figures written:", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))

# ------------------------------------------------------------------ Step 6: GeoPackage for the QGIS elevation profile
section("Step 6: write a GeoPackage for QGIS")
gpkg = OUT / "a417_section.gpkg"
gpkg.unlink(missing_ok=True)
gpd.GeoDataFrame({"name": [f"{LETTERS[0]}-{LETTERS[-1]}"], "length_m": [round(line.length, 1)]},
                 geometry=[line], crs="EPSG:27700").to_file(gpkg, layer="section_line", driver="GPKG")
gpd.GeoDataFrame({"half_width_m": [HALF]}, geometry=[line.buffer(HALF, cap_style="flat", join_style="mitre")],
                 crs="EPSG:27700").to_file(gpkg, layer="corridor", driver="GPKG")
gpd.GeoDataFrame({"letter": list(LETTERS), "chainage": np.round(bends, 1)},
                 geometry=[Point(v) for v in VERTICES], crs="EPSG:27700").to_file(gpkg, layer="bends", driver="GPKG")
cols = ["LOCA_TYPE", "LOCA_GL", "LOCA_FDEP", "chainage", "offset", "item", "lidar_at_hole", "gl_minus_lidar"]
gpd.GeoDataFrame(on[cols].round(2).reset_index(), geometry=[Point(x, y, z) for x, y, z in
                 zip(on.LOCA_NATE, on.LOCA_NATN, on.LOCA_GL)], crs="EPSG:27700").to_file(gpkg, layer="holes", driver="GPKG")
# For the elevation profile, each interval is a 3D point at its base with an extrusion height equal to its thickness.
iv = geol[geol.LOCA_ID.isin(on.index)]
gpd.GeoDataFrame({"LOCA_ID": iv.LOCA_ID, "unit": iv.UNIT, "unit_name": iv.UNIT.map(A.UNITS.name),
                  "colour": iv.UNIT.map(A.UNITS.colour), "top_mOD": iv.top_mOD.round(2), "base_mOD": iv.base_mOD.round(2),
                  "thickness": (iv.top_mOD - iv.base_mOD).round(2), "GEOL_DESC": iv.GEOL_DESC},
                 geometry=[Point(r.LOCA_NATE, r.LOCA_NATN, r.base_mOD) for r in iv.itertuples()],
                 crs="EPSG:27700").to_file(gpkg, layer="intervals", driver="GPKG")
gpd.GeoDataFrame(contacts.assign(surface_name=contacts.surface.map(SURFACES.name)).drop(columns=["LOCA_NATE", "LOCA_NATN"]),
                 geometry=[Point(r.LOCA_NATE, r.LOCA_NATN, r.z) for r in contacts.itertuples()],
                 crs="EPSG:27700").to_file(gpkg, layer="contacts", driver="GPKG")
xs_, ys_ = zip(*VERTICES)
cx0, cy1 = int((min(xs_) - 300 - X0) / RES), int((Y1 - max(ys_) - 300) / RES)
cx1, cy0 = int((max(xs_) + 300 - X0) / RES), int((Y1 - min(ys_) + 300) / RES)
clip = dtm[cy1:cy0, cx0:cx1]
with rasterio.open(OUT / "a417_section_dtm_2m.tif", "w", driver="GTiff", width=clip.shape[1], height=clip.shape[0],
                   count=1, dtype="float32", crs="EPSG:27700", nodata=np.nan, compress="deflate",
                   transform=from_origin(X0 + cx0 * RES, Y1 - cy1 * RES, RES, RES)) as dst:
    dst.write(clip, 1)
pd.DataFrame({"chainage": ch, "easting": xy[:, 0].round(2), "northing": xy[:, 1].round(2), "ground_mOD": ground.round(2),
              "corridor_min": lo.round(2), "corridor_max": hi.round(2)}).to_csv(OUT / "ground_profile.csv", index=False)
for layer in ["section_line", "corridor", "bends", "holes", "intervals", "contacts"]:
    g = gpd.read_file(gpkg, layer=layer)
    print(f"  {layer:13s} {len(g):5d} {g.geom_type.iloc[0]}{' Z' if g.has_z.all() else ''}")
print(f"wrote {gpkg.as_posix()}, a417_section_dtm_2m.tif ({clip.shape[1]} x {clip.shape[0]} cells) and ground_profile.csv")