"""3D borehole logs from AGS4: strata cylinders coloured by GEOL, in PyVista and QGIS.

Part 3 of the TechnicalGIS series "AGS data in GIS". Turns every GEOL interval of the A417 Missing
Link ground investigation into a short cylinder at the hole's British National Grid position and its
true elevation (ground level from LOCA_GL minus depth), coloured by formation, and draws them under
Environment Agency LIDAR terrain. Then it checks whether neighbouring logs tell the same story at the
base of the Inferior Oolite, finds a block that does not, and shows what its logs record. Finally it
writes LineStringZ intervals to a GeoPackage for the QGIS 3D map view (qgis_3d_logs.py).

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/3d-borehole-logs-pyvista-qgis/

Data: A417 Missing Link ground investigation (BGS deposited data item 173984; 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 itertools
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 shapely.geometry import LineString, Point
import pyvista as pv
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import to_rgb
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", 120)

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

AOI = (391_600, 212_800, 395_600, 216_800)             # same 4 x 4 km window as Part 2
RES = 5.0                                               # 5 m terrain is plenty under 112 holes
VE = 3.0                                                # vertical exaggeration for the 3D views
CLUSTER = (393_320, 216_060, 250)                       # centre and radius (m) of the close-up at the top of Crickley Hill
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). "
          f"Vertical exaggeration x{VE:.0f}.")


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"pyvista {pv.__version__} | vtk {version('vtk')} | geopandas {gpd.__version__} | "
      f"python-ags4 {version('python-ags4')}")

# ------------------------------------------------------------------ Step 1: intervals with elevations
section("Step 1: GEOL intervals with elevations")
ags = A.load_clean(DATA)
loca = ags["LOCA"].set_index("LOCA_ID")
geol = ags["GEOL"].copy()
geol["UNIT"] = geol.UNIT.fillna("NR")
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
geol["thick"] = geol.GEOL_BASE - geol.GEOL_TOP
geol = geol[geol.thick > 0].reset_index(drop=True)
print(f"{len(geol)} intervals in {geol.LOCA_ID.nunique()} holes | {geol.thick.sum():,.0f} m of logged ground | "
      f"elevations {geol.base_mOD.min():.1f} to {geol.top_mOD.max():.1f} m OD")
used = [u for u in A.UNITS.index if u in set(geol.UNIT)]
print(geol.groupby("UNIT").agg(intervals=("LOCA_ID", "size"), holes=("LOCA_ID", "nunique"), metres=("thick", "sum"))
      .reindex(used).round(1).to_string())

# ------------------------------------------------------------------ Step 2: terrain
section("Step 2: LIDAR terrain")
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:
    dtm = src.read(1).astype("float32")
    dtm[dtm < -1000] = np.nan
    left, top, dtm_transform = src.bounds.left, src.bounds.top, src.transform
print(f"DTM {dtm.shape[1]} x {dtm.shape[0]} cells of {RES:.0f} m, {np.nanmin(dtm):.1f} to {np.nanmax(dtm):.1f} m OD")

# Local coordinates keep the numbers small for the graphics card: metres east and north of the AOI corner.
xs = left - X0 + RES / 2 + np.arange(dtm.shape[1]) * RES
ys = top - Y0 - RES / 2 - np.arange(dtm.shape[0]) * RES
gx, gy = np.meshgrid(xs, ys)
z = np.nan_to_num(dtm, nan=np.nanmin(dtm))
terrain = pv.StructuredGrid(gx, gy, z * VE)
terrain["elevation"] = z.ravel(order="F")


# ------------------------------------------------------------------ Step 3: one cylinder per interval
section("Step 3: build the strata cylinders")


def cylinders(intervals, radius):
    """One cylinder per GEOL interval, with an RGB colour per face so the whole set is a single mesh."""
    pieces = []
    for r in intervals.itertuples():
        c = pv.Cylinder(center=(r.LOCA_NATE - X0, r.LOCA_NATN - Y0, (r.top_mOD + r.base_mOD) / 2 * VE),
                        direction=(0, 0, 1), radius=radius, height=r.thick * VE, resolution=24)
        rgb = (np.array(to_rgb(A.UNITS.colour[r.UNIT])) * 255).astype(np.uint8)
        c.cell_data["rgb"] = np.tile(rgb, (c.n_cells, 1))
        pieces.append(c)
    return pv.merge(pieces)


logs = cylinders(geol, radius=22)
print(f"{len(geol)} cylinders merged into one mesh: {logs.n_points:,} points, {logs.n_cells:,} faces")


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=A.UNITS.colour[u], edgecolor="#555555", lw=0.4, label=A.UNITS.name[u]) for u in units]
    fig.legend(handles=handles, loc="center left", bbox_to_anchor=(0.79, 0.5), fontsize=7.5, frameon=False,
               title="GEOL unit", title_fontsize=8, alignment="left")
    fig.savefig(path, dpi=150)
    plt.close(fig)


# ------------------------------------------------------------------ Step 4: the overview
section("Step 4: render the overview")
pl = pv.Plotter(off_screen=True, window_size=(1800, 1150))
pl.set_background("white")
pl.add_mesh(terrain, scalars="elevation", cmap="gist_earth", clim=[60, 330], opacity=0.4, smooth_shading=True,
            specular=0.0, ambient=0.3, show_scalar_bar=False)
pl.add_mesh(terrain.contour(np.arange(90, 310, 10) * VE, scalars=terrain.points[:, 2]), color="#6f6a5f",
            line_width=1, opacity=0.5)
pl.add_mesh(logs, scalars="rgb", rgb=True, smooth_shading=True, ambient=0.35)
for lid in ["CP217", "RC516", "DSRC301"]:
    p = loca.loc[lid]
    pl.add_point_labels([(p.LOCA_NATE - X0, p.LOCA_NATN - Y0, (p.LOCA_GL + 30) * VE)], [lid], font_size=20,
                        point_size=1, shape_opacity=0.8, always_visible=True, text_color="black",
                        shape_color="white")
pl.camera_position = [(-900, -1200, 1950 * VE / 3 + 1500), (2200, 2700, 200 * VE), (0, 0, 1)]
pl.camera.zoom(1.5)
img = pl.screenshot(return_img=True)
pl.close()
compose(img, used, OUT / "01-3d-borehole-logs.png",
        "A417 Missing Link: 112 borehole logs under LIDAR terrain, looking north-east up Crickley Hill",
        CREDIT + " 10 m ground contours.")
print("rendered 01-3d-borehole-logs.png")

# ------------------------------------------------------------------ Step 5: do neighbouring logs agree?
section("Step 5: do neighbouring logs agree? The base of the Inferior Oolite")
rank = geol.UNIT.map(A.UNITS["rank"])
contacts = []
for lid, g in geol[rank > 0].groupby("LOCA_ID", sort=False):
    bds = g[g.UNIT == "BDS"]
    if bds.empty:
        continue
    b = bds.iloc[0]
    above = g[(g.UNIT.isin(["SALS", "ASLS", "BLPL", "IOG"])) & (g.GEOL_TOP < b.GEOL_TOP)]
    if above.empty:
        continue
    contacts.append(dict(LOCA_ID=lid, E=b.LOCA_NATE, N=b.LOCA_NATN, GL=b.LOCA_GL, depth=b.GEOL_TOP, z=b.top_mOD,
                         above=above.iloc[-1].GEOL_DESC, below=b.GEOL_DESC))
contacts = pd.DataFrame(contacts).set_index("LOCA_ID")
print(f"{len(contacts)} holes log Inferior Oolite limestone resting on Bridport Sand; the contact lies at "
      f"{contacts.z.min():.1f} to {contacts.z.max():.1f} m OD")
lith = contacts.below.str.extract(r"\b(SILT|CLAY|MUDSTONE|SILTSTONE|SAND|SANDSTONE|LIMESTONE)\b", expand=False)
print("first Bridport Sand material below the contact:", lith.value_counts().to_dict())

pairs = []
for a, b in itertools.combinations(contacts.itertuples(), 2):
    s = np.hypot(a.E - b.E, a.N - b.N)
    if s <= 100:
        pairs.append(dict(a=a.Index, b=b.Index, metres_apart=round(s), dz=round(a.z - b.z, 1)))
pairs = pd.DataFrame(pairs).sort_values("metres_apart")
pairs["agree_3m"] = pairs.dz.abs() <= 3
print(f"\n{len(pairs)} pairs of holes within 100 m of each other; {pairs.agree_3m.sum()} agree within 3 m")
print(pairs.to_string(index=False))

# A hole is "offset" if its contact sits more than 8 m from the median of its neighbours within 100 m.
off = {}
for lid, c in contacts.iterrows():
    nb = contacts[(np.hypot(contacts.E - c.E, contacts.N - c.N) <= 100) & (contacts.index != lid)]
    if len(nb) >= 2:
        off[lid] = c.z - nb.z.median()
off = pd.Series(off)
offset = off[off.abs() > 8].index.tolist()
print("\ncontact minus the median of its neighbours within 100 m (holes with 2+ neighbours):")
print(off.round(1).sort_values().to_string())
rest = pairs[~pairs.a.isin(offset) & ~pairs.b.isin(offset)]
print(f"offset by more than 8 m: {', '.join(offset)}. Without them, {len(rest)} pairs agree to within "
      f"{rest.dz.abs().max():.1f} m (median {rest.dz.abs().median():.1f} m)")

# What do the offset holes record that their neighbours do not?
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)
cx, cy, cr = CLUSTER
in_cluster = contacts[np.hypot(contacts.E - cx, contacts.N - cy) <= cr].index.tolist()
ev = geol[geol.LOCA_ID.isin(in_cluster) & (geol.bedding_dip.notna() | geol.slickensides)]
print(f"\n{len(in_cluster)} holes reach the contact within {cr} m of the close-up centre; their logs record:")
for r in ev.itertuples():
    what = f"bedding dipping up to {r.bedding_dip:.0f} deg" if pd.notna(r.bedding_dip) else "slickensides"
    print(f"  {r.LOCA_ID:8s} {r.GEOL_TOP:6.2f}-{r.GEOL_BASE:6.2f} m  {A.UNITS.name[r.UNIT]:28s} {what}")

# Strip logs of the cluster at true elevation, west to east.
cl = contacts.loc[in_cluster].sort_values("E")
fig, ax = plt.subplots(figsize=(10, 5.2))
for i, (lid, c) in enumerate(cl.iterrows()):
    for r in geol[geol.LOCA_ID == lid].itertuples():
        ax.bar(i, r.thick, bottom=r.base_mOD, width=0.62, color=A.UNITS.colour[r.UNIT], edgecolor="none")
    ax.plot([i - 0.42, i + 0.42], [c.z, c.z], color="black", lw=1.6)
    ax.text(i + 0.36, c.z + 0.8, f"{c.z:.1f}", fontsize=7, ha="left", va="bottom")
    e = ev[ev.LOCA_ID == lid]
    for r in e[e.bedding_dip.notna()].itertuples():
        ax.plot(i, (r.top_mOD + r.base_mOD) / 2, marker="$\u2215$", ms=11, color="#b2182b")
    for r in e[e.slickensides].itertuples():
        ax.plot(i, (r.top_mOD + r.base_mOD) / 2, marker="x", ms=7, mew=2, color="#b2182b")
names = [f"{lid}\n{'offset' if lid in offset else ''}" for lid in cl.index]
ax.set_xticks(range(len(cl)), names, fontsize=7, rotation=0)
for t, lid in zip(ax.get_xticklabels(), cl.index):
    t.set_color("#b2182b" if lid in offset else "black")
ax.set_ylim(170, 245)
ax.set_xlim(-0.8, len(cl) - 0.2)
ax.set_ylabel("elevation (m OD)")
ax.set_title(f"Top of Crickley Hill: {len(cl)} holes within {cr} m, west to east. Black bars mark the base of the "
             "Inferior Oolite (m OD)", fontsize=9)
ax.spines[["top", "right"]].set_visible(False)
handles = [Patch(facecolor=A.UNITS.colour[u], label=A.UNITS.name[u])
           for u in used if u in set(geol.loc[geol.LOCA_ID.isin(cl.index), "UNIT"])]
handles += [Line2D([], [], color="black", lw=1.6, label="base of Inferior Oolite"),
            Line2D([], [], ls="", marker="$\u2215$", ms=9, color="#b2182b", label="log records bedding dipping 30-50°"),
            Line2D([], [], ls="", marker="x", ms=6, mew=1.6, color="#b2182b", label="log records slickensides")]
ax.legend(handles=handles, fontsize=6.5, loc="center left", bbox_to_anchor=(1.0, 0.5), frameon=False)
fig.savefig(OUT / "02-contact-check-strip-logs.png")
plt.close(fig)

# The same holes in 3D, with the contact drawn as a black disc and the ground as 2 m contours.
box = (cx - X0 - 330, cx - X0 + 330, cy - Y0 - 330, cy - Y0 + 330, -1e6, 1e6)
near = loca[np.hypot(loca.LOCA_NATE - cx, loca.LOCA_NATN - cy) <= 330].index
close_logs = cylinders(geol[geol.LOCA_ID.isin(near)], radius=5)
local = terrain.clip_box(box, invert=False)
pl = pv.Plotter(off_screen=True, window_size=(1800, 1150))
pl.set_background("white")
pl.add_mesh(local.contour(np.arange(200, 300, 2) * VE, scalars=local.points[:, 2]), color="#8f8a7e", line_width=1.2)
pl.add_mesh(close_logs, scalars="rgb", rgb=True, smooth_shading=True, ambient=0.25)
for lid, c in contacts.loc[contacts.index.isin(near)].iterrows():
    pl.add_mesh(pv.Disc(center=(c.E - X0, c.N - Y0, c.z * VE), inner=0, outer=11, normal=(0, 0, 1), c_res=32),
                color="#b2182b" if lid in offset else "black")
labels = contacts.loc[contacts.index.isin(near)]
pl.add_point_labels(np.c_[labels.E - X0, labels.N - Y0, (labels.GL + 6) * VE],
                    [f"{lid} {c.z:.0f}" for lid, c in labels.iterrows()], font_size=15, point_size=1,
                    shape_opacity=0.75, always_visible=True, text_color="black", shape_color="white")
pl.camera_position = [(cx - X0 - 150, cy - Y0 - 900, 360 * VE), (cx - X0, cy - Y0, 205 * VE), (0, 0, 1)]
img = pl.screenshot(return_img=True)
pl.close()
compose(img, [u for u in used if u in set(geol.loc[geol.LOCA_ID.isin(near), "UNIT"])],
        OUT / "03-contact-check-3d.png",
        "Top of Crickley Hill, looking north: discs mark the base of the Inferior Oolite (red = offset holes)",
        CREDIT + " 2 m ground contours. Labels give the contact in m OD.")
print("rendered 03-contact-check-3d.png")

# ------------------------------------------------------------------ Step 6: LineStringZ for QGIS 3D
section("Step 6: write LineStringZ intervals for QGIS")
lines = gpd.GeoDataFrame(
    {"LOCA_ID": geol.LOCA_ID, "unit": geol.UNIT, "unit_name": geol.UNIT.map(A.UNITS.name),
     "colour": geol.UNIT.map(A.UNITS.colour), "top_depth": geol.GEOL_TOP, "base_depth": geol.GEOL_BASE,
     "top_mOD": geol.top_mOD.round(2), "base_mOD": geol.base_mOD.round(2), "GEOL_DESC": geol.GEOL_DESC},
    geometry=[LineString([(r.LOCA_NATE, r.LOCA_NATN, r.top_mOD), (r.LOCA_NATE, r.LOCA_NATN, r.base_mOD)])
              for r in geol.itertuples()], crs="EPSG:27700")
collars = gpd.GeoDataFrame(loca[["LOCA_TYPE", "LOCA_GL", "LOCA_FDEP"]].reset_index(),
                           geometry=[Point(r.LOCA_NATE, r.LOCA_NATN, r.LOCA_GL) for r in loca.itertuples()],
                           crs="EPSG:27700")
contact_pts = gpd.GeoDataFrame(contacts[["depth", "z"]].assign(offset=contacts.index.isin(offset)).reset_index(),
                               geometry=[Point(c.E, c.N, c.z) for c in contacts.itertuples()], crs="EPSG:27700")
gpkg = OUT / "a417_logs_3d.gpkg"
gpkg.unlink(missing_ok=True)
lines.to_file(gpkg, layer="intervals", driver="GPKG")
collars.to_file(gpkg, layer="collars", driver="GPKG")
contact_pts.to_file(gpkg, layer="base_inferior_oolite", driver="GPKG")
check = gpd.read_file(gpkg, layer="intervals")
print(f"wrote {gpkg.as_posix()}: 'intervals' {len(check)} {check.geom_type.iloc[0]}s with z: {check.has_z.all()}, "
      f"'collars' {len(collars)} points, 'base_inferior_oolite' {len(contact_pts)} points")
with rasterio.open(OUT / "a417_dtm_5m.tif", "w", driver="GTiff", width=dtm.shape[1], height=dtm.shape[0], count=1,
                   dtype="float32", crs="EPSG:27700", transform=dtm_transform, nodata=np.nan,
                   compress="deflate") as dst:
    dst.write(dtm, 1)
print("terrain for QGIS 3D: a417_dtm_5m.tif")
print("figures written:", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))