"""SPT N-values from AGS4: where the A417 landslide deposits stop and the escarpment's bedrock starts.

Part 6 of the TechnicalGIS series "AGS data in GIS". Reads the ISPT group of the A417 Missing Link
ground investigation, separates full tests from refusals, finds the GEOL unit each test was made in,
and presents the N-values four ways: against elevation by unit, against height above or below the
logged base of the landslide deposits, as coloured discs on 3D borehole logs over Environment Agency
LIDAR terrain, and as a map of the depth to the first refusal. Finally it writes the tests to a
GeoPackage for QGIS.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/spt-n-values-ags-python/

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 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
import pyvista as pv
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LightSource, to_rgb
from matplotlib.lines import Line2D

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)             # the same 4 x 4 km LIDAR window as Parts 2 and 3
RES = 5.0
VE = 3.0                                                # vertical exaggeration for the 3D view
SCARP = (391_950, 393_700, 215_250, 216_350)            # Crickley Hill: the part of the route with most SPTs
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")
BGS = "BGS AGS data © UKRI (OGL v3)."
BGS_EA = "BGS AGS data © UKRI; LIDAR © Environment Agency (both OGL v3)."

# SPT N bands. For sands and gravels these are the BS EN ISO 14688-2 / BS 5930 density terms.
BANDS = [0, 4, 10, 30, 50, np.inf]
BAND_LABELS = ["N < 4", "4 to 10", "10 to 30", "30 to 50", "refusal (N > 50)"]
BAND_COLOURS = ["#fde725", "#7ad151", "#22a884", "#2a788e", "#440154"]


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


def unit_legend(units):
    return [Line2D([], [], ls="", marker="o", mfc=A.UNITS.colour[u], mec="#333333", label=A.UNITS.name[u])
            for u in units]


print(f"pyvista {pv.__version__} | vtk {version('vtk')} | geopandas {gpd.__version__} | "
      f"python-ags4 {version('python-ags4')}")

# ------------------------------------------------------------------ Step 1: read the SPTs
section("Step 1: SPT N-values and refusals")
ags = A.load_clean(DATA)
loca = ags["LOCA"].set_index("LOCA_ID")
geol = ags["GEOL"].copy()
geol["UNIT"] = geol.UNIT.fillna("NR")
spt = ags["ISPT"][["LOCA_ID", "ISPT_TOP", "ISPT_REP"]].copy()
spt["ISPT_TOP"] = pd.to_numeric(spt.ISPT_TOP)
spt["N"] = pd.to_numeric(spt.ISPT_REP)
spt["refusal"] = spt.N > 50
spt["pen_mm"] = np.where(spt.refusal, 50 * 300 / spt.N, 300.0)    # penetration implied by 50 blows
spt = spt.join(loca[["LOCA_NATE", "LOCA_NATN", "LOCA_GL"]], on="LOCA_ID")
spt["mOD"] = spt.LOCA_GL - spt.ISPT_TOP
print(f"{len(spt)} tests in {spt.LOCA_ID.nunique()} holes, {spt.ISPT_TOP.min():.1f} to {spt.ISPT_TOP.max():.1f} m deep")
print(f"N from {spt.N.min():.0f} to {spt.N.max():.0f}; {spt.refusal.sum()} tests ({spt.refusal.mean():.0%}) above 50")
pens = np.sort(spt.loc[spt.refusal, "pen_mm"].unique())
print(f"penetration implied by those values for 50 blows: {pens.min():.1f} to {pens.max():.1f} mm, e.g. "
      + ", ".join(f"{p:.0f}" for p in pens[np.linspace(0, len(pens) - 1, 6).astype(int)]))


# ------------------------------------------------------------------ Step 2: which unit was each test in?
section("Step 2: the GEOL unit at each test")
g = geol[["LOCA_ID", "GEOL_TOP", "GEOL_BASE", "UNIT"]]
m = spt.reset_index().merge(g, on="LOCA_ID", how="left")
m = m[(m.GEOL_TOP <= m.ISPT_TOP) & (m.ISPT_TOP < m.GEOL_BASE)].drop_duplicates("index").set_index("index")
spt["UNIT"] = m.UNIT.reindex(spt.index).fillna("NR")
order = [u for u in A.UNITS.index if u in set(spt.UNIT)]
table = (spt.groupby("UNIT").agg(tests=("N", "size"), holes=("LOCA_ID", "nunique"), median_N=("N", "median"),
                                 refusals=("refusal", "mean")).reindex(order))
table["refusals"] = (table.refusals * 100).round().astype(int).astype(str) + "%"
print(table.to_string())

# ------------------------------------------------------------------ Step 3: N against elevation, by unit
section("Step 3: N against elevation, by unit")
facets = ["MGR", "HEAD", "SLIP", "FE", "BLPL", "BDS", "WHM", "DYS"]
fig, axes = plt.subplots(1, len(facets) + 1, figsize=(15.5, 5.8), sharey=True)
for ax, u in zip(axes, facets + ["other"]):
    d = spt[spt.UNIT == u] if u != "other" else spt[~spt.UNIT.isin(facets)]
    ax.axvspan(50, 2000, color="#eeeeee", zorder=0)
    for refused, marker in [(False, "o"), (True, ">")]:
        e = d[d.refusal == refused]
        ax.scatter(e.N, e.mOD, c=[A.UNITS.colour[x] for x in e.UNIT], marker=marker, s=22 if refused else 16,
                   edgecolor="#333333", lw=0.4, zorder=3)
    ax.set_xscale("log")
    ax.set_xlim(2, 1500)
    ax.set_xticks([3, 10, 50, 500], ["3", "10", "50", "500"])
    name = "Other units" if u == "other" else A.UNITS.name[u].replace(" Formation", "")
    ax.set_title(f"{name}\n{len(d)} tests, median {d.N.median():.0f}", fontsize=7.5)
    ax.spines[["top", "right"]].set_visible(False)
    ax.grid(axis="y", color="#dddddd", lw=0.5)
    print(f"{name:<22} {len(d):>4} tests  median N {d.N.median():>4.0f}  "
          f"{d.mOD.min():.0f} to {d.mOD.max():.0f} m OD")
axes[0].set_ylabel("elevation of test (m OD)")
fig.supxlabel("SPT N (log scale). Grey band: refusal, 50 blows for less than 300 mm, reported as an extrapolated N",
              fontsize=8.5)
fig.suptitle("A417 Missing Link: SPT N against elevation, by geological unit", weight="bold", x=0.01, ha="left")
fig.legend(handles=[Line2D([], [], ls="", marker="o", mfc="white", mec="#333", label="full 300 mm test drive"),
                    Line2D([], [], ls="", marker=">", mfc="white", mec="#333", label="refusal")],
           loc="upper right", ncol=2, frameon=False, fontsize=8)
fig.text(0.01, -0.02, BGS, fontsize=7, color="#444444")
fig.savefig(OUT / "02-n-vs-elevation-by-unit.png")
plt.close(fig)

# ------------------------------------------------------------------ Step 4: across the base of the landslide deposits
section("Step 4: across the base of the landslide deposits")
bedrock = A.UNITS.index[A.UNITS["rank"] > 0]
rows = []
for lid, h in geol.groupby("LOCA_ID"):
    s = h[h.UNIT == "SLIP"]
    if s.empty:
        continue
    under = h[(h.GEOL_TOP >= s.GEOL_BASE.max() - 0.01) & h.UNIT.isin(bedrock)]
    if len(under):
        rows.append((lid, s.GEOL_BASE.max(), under.UNIT.iloc[0]))
base = pd.DataFrame(rows, columns=["LOCA_ID", "slip_base", "below"]).set_index("LOCA_ID")
b = spt.join(base, on="LOCA_ID", how="inner")
b["rel"] = b.slip_base - b.ISPT_TOP                      # metres above (+) or below (-) the logged base
above, below = b[b.rel > 0], b[b.rel <= 0]
print(f"{len(base)} holes log landslide deposits directly on bedrock; {b.LOCA_ID.nunique()} of them have SPTs "
      f"({len(b)} tests)")
print(f"base of the deposits: {base.slip_base.min():.1f} to {base.slip_base.max():.1f} m deep, "
      f"median {base.slip_base.median():.1f} m")
print(f"above the base: {len(above)} tests, median N {above.N.median():.0f}, refusals {above.refusal.mean():.0%}")
print(f"below the base: {len(below)} tests, median N {below.N.median():.0f}, refusals {below.refusal.mean():.0%}")
first_below = below.sort_values("ISPT_TOP").groupby("LOCA_ID").first()
print(f"first test below the base refused in {first_below.refusal.sum()} of {len(first_below)} holes; "
      f"bedrock beneath: " + ", ".join(f"{A.UNITS.name[u]} {n}" for u, n in base.below.value_counts().items()))

fig, ax = plt.subplots(figsize=(7.5, 6))
ax.axvspan(50, 2000, color="#eeeeee", zorder=0)
ax.axhline(0, color="black", lw=1.2)
for refused, marker in [(False, "o"), (True, ">")]:
    e = b[b.refusal == refused]
    ax.scatter(e.N, e.rel, c=[A.UNITS.colour[x] for x in e.UNIT], marker=marker, s=24, edgecolor="#333333",
               lw=0.4, zorder=3)
lo, hi = np.floor(b.rel.min() / 2) * 2, np.ceil(b.rel.max() / 2) * 2
b["band"] = pd.cut(b.rel, np.arange(lo, hi + 2, 2))
med = b.groupby("band", observed=True).N.median()
ax.plot(med.values, [i.mid for i in med.index], color="#b2182b", lw=2, marker="s", ms=4, zorder=4)
ax.set_xscale("log")
ax.set_xlim(2, 1500)
ax.set_xticks([3, 10, 50, 300, 1000], ["3", "10", "50", "300", "1000"])
ax.set_ylim(lo - 0.5, hi + 0.5)
ax.set_ylabel("metres above (+) or below (−) the logged base of the landslide deposits")
ax.set_xlabel("SPT N (log scale); grey = refusal, triangles = refused tests")
ax.text(2.3, 0.4, "base of landslide deposits", fontsize=8, va="bottom")
ax.set_title(f"{b.LOCA_ID.nunique()} holes through the landslide deposits into bedrock: {len(b)} SPTs",
             fontsize=9, loc="left")
units = [u for u in A.UNITS.index if u in set(b.UNIT)]
ax.legend(handles=unit_legend(units) + [Line2D([], [], color="#b2182b", lw=2, marker="s", ms=4,
                                               label="median N in 2 m bands")],
          fontsize=7, frameon=False, loc="center left", bbox_to_anchor=(1, 0.5))
ax.spines[["top", "right"]].set_visible(False)
fig.text(0.01, -0.02, BGS, fontsize=7, color="#444444")
fig.savefig(OUT / "03-n-across-slip-base.png")
plt.close(fig)

# ------------------------------------------------------------------ Step 5: N discs on 3D logs over LIDAR
section("Step 5: N on the 3D logs")
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 = src.bounds.left, src.bounds.top
xs = left - X0 + RES / 2 + np.arange(dtm.shape[1]) * RES
ys = top - Y0 - RES / 2 - np.arange(dtm.shape[0]) * RES
z = np.nan_to_num(dtm, nan=np.nanmin(dtm))
win = (SCARP[0] - X0, SCARP[1] - X0, SCARP[2] - Y0, SCARP[3] - Y0)
ci = (xs >= win[0]) & (xs <= win[1])
ri = (ys >= win[2]) & (ys <= win[3])
gx, gy = np.meshgrid(xs[ci], ys[ri])
terrain = pv.StructuredGrid(gx, gy, z[np.ix_(ri, ci)] * VE)

spt["band"] = pd.cut(spt.N, BANDS, labels=False, right=False)
holes = loca.loc[spt.LOCA_ID.unique()]
holes = holes[(holes.LOCA_NATE - X0).between(win[0], win[1]) & (holes.LOCA_NATN - Y0).between(win[2], win[3])]


def coloured(mesh, colour):
    mesh.cell_data["rgb"] = np.tile((np.array(to_rgb(colour)) * 255).astype(np.uint8), (mesh.n_cells, 1))
    return mesh


logs = []
for r in geol[geol.LOCA_ID.isin(holes.index)].join(loca[["LOCA_NATE", "LOCA_NATN", "LOCA_GL"]],
                                                    on="LOCA_ID").itertuples():
    t, bse = r.LOCA_GL - r.GEOL_TOP, r.LOCA_GL - r.GEOL_BASE
    if t > bse:
        logs.append(coloured(pv.Cylinder(center=(r.LOCA_NATE - X0, r.LOCA_NATN - Y0, (t + bse) / 2 * VE),
                                         direction=(0, 0, 1), radius=3.5, height=(t - bse) * VE, resolution=16),
                             A.UNITS.colour[r.UNIT]))
discs = [coloured(pv.Cylinder(center=(r.LOCA_NATE - X0, r.LOCA_NATN - Y0, r.mOD * VE), direction=(0, 0, 1),
                              radius=11, height=1.4 * VE, resolution=32), BAND_COLOURS[int(r.band)])
         for r in spt[spt.LOCA_ID.isin(holes.index)].itertuples()]
rings = [pv.Disc(center=(loca.at[lid, "LOCA_NATE"] - X0, loca.at[lid, "LOCA_NATN"] - Y0,
                         (loca.at[lid, "LOCA_GL"] - d) * VE + 0.5), inner=13, outer=22, normal=(0, 0, 1), c_res=48)
         for lid, d in base.slip_base.items() if lid in holes.index]
print(f"{len(holes)} holes with SPTs on Crickley Hill: {len(logs)} log intervals, {len(discs)} SPT discs, "
      f"{len(rings)} rings at the base of the landslide deposits")
logs, discs, rings = pv.merge(logs), pv.merge(discs), pv.merge(rings)

pl = pv.Plotter(off_screen=True, window_size=(1900, 1150))
pl.set_background("#dfe9f3", top="#9fbbd8")
pl.add_mesh(terrain, scalars=terrain.points[:, 2] / VE, cmap="gist_earth", clim=[60, 330], opacity=0.28,
            show_scalar_bar=False, ambient=0.5)
pl.add_mesh(terrain.contour(np.arange(110, 300, 10) * VE, scalars=terrain.points[:, 2]), color="#5e584c",
            line_width=1, opacity=0.6)
pl.add_mesh(logs, scalars="rgb", rgb=True, ambient=0.5)
pl.add_mesh(discs, scalars="rgb", rgb=True, ambient=0.5, smooth_shading=True)
pl.add_mesh(rings, color="black", ambient=0.6)
cx, cy = (win[0] + win[1]) / 2, (win[2] + win[3]) / 2
pl.camera_position = [(cx - 1100, cy - 1700, 420 * VE), (cx + 60, cy + 60, 175 * VE), (0, 0, 1)]
pl.camera.zoom(1.45)
img = pl.screenshot(return_img=True)
pl.close()

h, w = img.shape[:2]
fig = plt.figure(figsize=(12, 12 * h / w * 0.9))
ax = fig.add_axes([0, 0.04, 0.82, 0.9])
ax.imshow(img)
ax.set_axis_off()
fig.text(0.01, 0.97, "SPT N on the A417 borehole logs, looking north-east up Crickley Hill", weight="bold",
         fontsize=11, va="top")
fig.text(0.01, 0.01, f"{BGS_EA} Vertical exaggeration x{VE:.0f}; 10 m ground contours. "
         "Thin columns are the logs, coloured by unit as in Part 3.", fontsize=7, color="#444444")
fig.legend(handles=[Line2D([], [], ls="", marker="o", ms=9, mfc=c, mec="#333", label=l)
                    for c, l in zip(BAND_COLOURS, BAND_LABELS)]
           + [Line2D([], [], ls="", marker="o", ms=11, mfc="none", mec="black", mew=2.5,
                     label="base of landslide\ndeposits")],
           title="SPT N", loc="center left", bbox_to_anchor=(0.81, 0.5), frameon=False, fontsize=8)
fig.savefig(OUT / "01-spt-n-3d.png", dpi=150)
plt.close(fig)

# ------------------------------------------------------------------ Step 6: depth to the first refusal
section("Step 6: depth to the first refusal")
first = spt[spt.refusal].groupby("LOCA_ID").ISPT_TOP.min()
none = spt.groupby("LOCA_ID").ISPT_TOP.max()[~spt.groupby("LOCA_ID").refusal.any()]
print(f"{len(first)} holes reached refusal, at {first.min():.1f} to {first.max():.1f} m (median {first.median():.1f} m); "
      f"{len(none)} holes never did (deepest test {none.min():.1f} to {none.max():.1f} m)")
geo = first.to_frame("depth").join(loca[["LOCA_NATE", "LOCA_NATN"]])
for name, cond in [("plateau, east of 393,100 E", geo.LOCA_NATE >= 393_100),
                   ("escarpment, west of 393,100 E", geo.LOCA_NATE < 393_100)]:
    d = geo[cond & geo.LOCA_NATE.between(SCARP[0], SCARP[1]) & geo.LOCA_NATN.between(SCARP[2], SCARP[3])]
    print(f"  {name}: {len(d)} holes, median depth to refusal {d.depth.median():.1f} m")

hs = LightSource(azdeg=315, altdeg=45).hillshade(z, dx=RES, dy=RES, vert_exag=1)
fig, ax = plt.subplots(figsize=(9, 6))
ax.imshow(hs, cmap="gray", extent=(left, left + dtm.shape[1] * RES, top - dtm.shape[0] * RES, top), alpha=0.8)
ax.contour(xs + X0, ys + Y0, z, levels=np.arange(110, 300, 10), colors="#6b6253", linewidths=0.4)
p = loca.loc[first.index]
sc = ax.scatter(p.LOCA_NATE, p.LOCA_NATN, c=first.values, cmap="plasma_r", vmin=0, vmax=25, s=70,
                edgecolor="black", lw=0.6, zorder=3)
q = loca.loc[none.index]
ax.scatter(q.LOCA_NATE, q.LOCA_NATN, marker="o", s=70, facecolor="white", edgecolor="black", lw=0.8, zorder=3,
           label="no refusal (label: depth of deepest test, m)")
for lid, d in none.items():
    ax.annotate(f">{d:.0f}", (loca.at[lid, "LOCA_NATE"], loca.at[lid, "LOCA_NATN"]), xytext=(5, 3),
                textcoords="offset points", fontsize=7)
ax.set_xlim(SCARP[0], SCARP[1])
ax.set_ylim(SCARP[2], SCARP[3])
ax.set_aspect("equal")
ax.ticklabel_format(useOffset=False, style="plain")
ax.set_xlabel("Easting (m, British National Grid)")
ax.set_ylabel("Northing (m, British National Grid)")
fig.colorbar(sc, ax=ax, shrink=0.7, label="depth of first SPT refusal (m below ground)")
ax.legend(loc="lower right", fontsize=7.5)
ax.set_title("Depth to the first SPT refusal on Crickley Hill", fontsize=10, loc="left")
ax.text(0, -0.17, BGS_EA + " Hillshade and 10 m contours from the 5 m LIDAR DTM.", fontsize=7,
        color="#444444", transform=ax.transAxes)
fig.savefig(OUT / "04-depth-to-refusal-map.png")
plt.close(fig)

# ------------------------------------------------------------------ Step 7: SPTs for QGIS
section("Step 7: write the tests to a GeoPackage")
pts = gpd.GeoDataFrame(
    spt[["LOCA_ID", "ISPT_TOP", "mOD", "N", "refusal", "pen_mm", "UNIT"]]
    .assign(unit_name=spt.UNIT.map(A.UNITS.name), band=[BAND_LABELS[int(i)] for i in spt.band]),
    geometry=gpd.points_from_xy(spt.LOCA_NATE, spt.LOCA_NATN, spt.mOD), crs="EPSG:27700")
gpkg = OUT / "a417_spt.gpkg"
pts.to_file(gpkg, layer="spt", driver="GPKG")
check = gpd.read_file(gpkg, layer="spt")
print(f"{gpkg.name}: {len(check)} PointZ features in {check.crs.to_string()}, fields "
      + ", ".join(c for c in check.columns if c != "geometry"))
print("figures:", ", ".join(sorted(p.name for p in OUT.glob("0*.png"))))