"""Boreholes on a map: AGS4 LOCA to points in British National Grid, checked against LIDAR.

Part 2 of the TechnicalGIS series "AGS data in GIS". Turns the LOCA group of the A417 Missing
Link ground investigation into a GeoDataFrame in British National Grid (EPSG:27700), shows the
two coordinate mistakes that send boreholes to the wrong place, summarises each hole's geology
from GEOL (depth to rock, the unit at rockhead, the deepest unit reached), checks every ground
level against Environment Agency LIDAR, draws two maps and writes a GeoPackage with a
ready-made HTML popup for QGIS (styled by qgis_borehole_map.py).

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/ags4-boreholes-map-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 time
from html import escape
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 pyproj import Transformer
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LightSource
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)             # 4 x 4 km around the holes, British National Grid
RES = 2.0                                               # metres per cell requested from the WCS
TILE = 1000                                             # metres per WCS request (500 x 500 cells)
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")

# One colour per kind of hole, shared with the QGIS style in qgis_borehole_map.py
TYPE_COLOURS = {"RC": "#1f78b4", "TP": "#e6550d", "CPT": "#31a354", "RO": "#756bb1", "CP": "#d62728"}


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__} | pyproj {version('pyproj')} | python-ags4 {version('python-ags4')}")

# ------------------------------------------------------------------ Step 1: LOCA to points
section("Step 1: LOCA to points in British National Grid")
ags = A.load_clean(DATA)
loca = ags["LOCA"]
holes = gpd.GeoDataFrame(loca, geometry=gpd.points_from_xy(loca.LOCA_NATE, loca.LOCA_NATN), crs="EPSG:27700")
print(f"{len(holes)} holes | CRS {holes.crs.name} ({holes.crs.to_string()})")
print(f"extent E {holes.total_bounds[0]:,.0f} to {holes.total_bounds[2]:,.0f}, "
      f"N {holes.total_bounds[1]:,.0f} to {holes.total_bounds[3]:,.0f}")
centre = holes.dissolve().centroid.to_crs("EPSG:4326").iloc[0]
print(f"centre in WGS84: {centre.y:.5f} N, {abs(centre.x):.5f} W (Crickley Hill and Birdlip)")

# ------------------------------------------------------------------ Step 2: two classic mistakes
section("Step 2: what goes wrong with the wrong CRS")
first = loca.iloc[0]
to_wgs = Transformer.from_crs("EPSG:27700", "EPSG:4326", always_xy=True)
lon, lat = to_wgs.transform(first.LOCA_NATN, first.LOCA_NATE)       # easting and northing swapped
print(f"{first.LOCA_ID} with easting and northing swapped lands at {lat:.3f} N, {lon:.3f} W: "
      f"in the sea off north Wales, {np.hypot(first.LOCA_NATE - first.LOCA_NATN, first.LOCA_NATN - first.LOCA_NATE) / 1000:,.0f} km away")
wrong = gpd.GeoDataFrame(geometry=gpd.points_from_xy(loca.LOCA_NATE, loca.LOCA_NATN), crs="EPSG:4326")
print(f"tagged as EPSG:4326 instead, x runs to {wrong.total_bounds[2]:,.0f} 'degrees'; "
      f"valid longitudes stop at 180, so nothing can be drawn")

# ------------------------------------------------------------------ Step 3: geology per hole
section("Step 3: summarise GEOL for each hole")
geol = ags["GEOL"].copy()
geol["rank"] = geol.UNIT.map(A.UNITS["rank"])
rock = geol[geol["rank"] > 0].groupby("LOCA_ID")
summary = pd.DataFrame({
    "rockhead_depth": rock.GEOL_TOP.min(),
    "rockhead_unit": rock.UNIT.first(),
    "deepest_unit": rock.UNIT.last(),
})
abbr = ags["ABBR"].query("ABBR_HDNG == 'LOCA_TYPE'").set_index("ABBR_CODE").ABBR_DESC
water = ags["WSTG"].groupby("LOCA_ID").WSTG_DPTH.min().rename("first_water_strike")
counts = pd.DataFrame({"n_samples": ags["SAMP"].groupby("LOCA_ID").size(),
                       "n_spt": ags["ISPT"].groupby("LOCA_ID").size()})


def log_html(g):
    """The hole's log as a small HTML table, one row per unit, for the QGIS popup."""
    runs = g.groupby((g.UNIT != g.UNIT.shift()).cumsum()).agg(
        unit=("UNIT", "first"), top=("GEOL_TOP", "min"), base=("GEOL_BASE", "max"))
    rows = "".join(f'<tr><td style="background:{A.UNITS.colour[r.unit]};width:10px"></td>'
                   f"<td>{r.top:.2f}&ndash;{r.base:.2f} m</td><td>{escape(A.UNITS.name[r.unit])}</td></tr>"
                   for r in runs.itertuples())
    return f'<table style="font-size:11px;border-collapse:collapse">{rows}</table>'


holes = holes.join(summary, on="LOCA_ID").join(water, on="LOCA_ID").join(counts, on="LOCA_ID")
holes[["n_samples", "n_spt"]] = holes[["n_samples", "n_spt"]].fillna(0).astype(int)
holes["type_name"] = holes.LOCA_TYPE.map(abbr)
holes["rockhead_mOD"] = (holes.LOCA_GL - holes.rockhead_depth).round(2)
holes["rockhead_name"] = holes.rockhead_unit.map(A.UNITS.name)
holes["log_html"] = holes.LOCA_ID.map(geol.groupby("LOCA_ID").apply(log_html, include_groups=False))
print(holes.type_name.value_counts().to_string())
print(f"\n{holes.rockhead_depth.notna().sum()} holes reach bedrock; depth to rock "
      f"{holes.rockhead_depth.min():.2f} to {holes.rockhead_depth.max():.2f} m (median {holes.rockhead_depth.median():.2f} m)")
print("unit at rockhead:", holes.rockhead_unit.value_counts().to_dict())
print(holes.loc[holes.LOCA_ID.isin(["DSRC415", "RC516", "CP217"]),
                ["LOCA_ID", "type_name", "LOCA_GL", "LOCA_FDEP", "rockhead_depth", "rockhead_mOD", "rockhead_name",
                 "first_water_strike", "n_samples"]].to_string(index=False))

# ------------------------------------------------------------------ Step 4: ground level against LIDAR
section("Step 4: check LOCA_GL against 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
transform = from_origin(X0, Y1, RES, RES)
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 | "
      f"no-data cells: {np.isnan(dtm).sum():,}")
profile = dict(driver="GTiff", width=dtm.shape[1], height=dtm.shape[0], count=1, dtype="float32",
               crs="EPSG:27700", transform=transform, nodata=np.nan, compress="deflate")
with rasterio.open(OUT / "a417_dtm_2m.tif", "w", **profile) as dst:
    dst.write(dtm, 1)
ls = LightSource(azdeg=315, altdeg=45)
shade = ls.hillshade(np.nan_to_num(dtm, nan=np.nanmean(dtm)), vert_exag=2, dx=RES, dy=RES)
with rasterio.open(OUT / "a417_hillshade.tif", "w", **{**profile, "dtype": "uint8", "nodata": None}) as dst:
    dst.write((shade * 255).astype("uint8"), 1)

col = ((holes.LOCA_NATE - X0) / RES).astype(int)
row = ((Y1 - holes.LOCA_NATN) / RES).astype(int)
holes["lidar_gl"] = dtm[row, col].round(2)
holes["gl_minus_lidar"] = (holes.LOCA_GL - holes.lidar_gl).round(2)
d = holes.gl_minus_lidar
print(f"LOCA_GL minus LIDAR: median {d.median():+.2f} m | middle half {d.quantile(.25):+.2f} to {d.quantile(.75):+.2f} m | "
      f"within 0.5 m: {(d.abs() <= 0.5).sum()} of {d.notna().sum()} | within 1 m: {(d.abs() <= 1).sum()}")
odd = holes.loc[d.abs() > 1, ["LOCA_ID", "type_name", "LOCA_CLST", "LOCA_GL", "lidar_gl", "gl_minus_lidar"]]
print("more than 1 m apart:")
print(odd.sort_values("gl_minus_lidar").to_string(index=False))

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.8), gridspec_kw={"width_ratios": [1.1, 1]})
for t, c in TYPE_COLOURS.items():
    s = holes[holes.LOCA_TYPE == t]
    ax1.scatter(s.lidar_gl, s.LOCA_GL, s=12, color=c, label=f"{t} {abbr.get(t, '')}", edgecolor="white", lw=0.3)
lim = [holes.lidar_gl.min() - 5, holes.lidar_gl.max() + 5]
ax1.plot(lim, lim, color="#555555", lw=0.8)
ax1.set(xlim=lim, ylim=lim, xlabel="LIDAR DTM at the hole (m OD)", ylabel="LOCA_GL in the AGS file (m OD)",
        title="Ground level: AGS against LIDAR")
ax1.legend(fontsize=7, frameon=False)
ax2.hist(d.clip(-3, 3), bins=np.arange(-3, 3.01, 0.2), color="#3f7f9f")
ax2.axvline(0, color="#555555", lw=0.8)
ax2.set(xlabel="LOCA_GL minus LIDAR (m, clipped at ±3)", ylabel="holes",
        title=f"{(d.abs() <= 0.5).sum()} of {d.notna().sum()} holes agree within 0.5 m")
for ax in (ax1, ax2):
    ax.spines[["top", "right"]].set_visible(False)
fig.savefig(OUT / "01-ground-level-check.png")
plt.close(fig)

# ------------------------------------------------------------------ Step 5: maps
section("Step 5: maps")
extent = (X0, X1, Y0, Y1)


def basemap(ax):
    ax.imshow(shade, cmap="gray", extent=extent, vmin=0, vmax=1.15)
    ax.imshow(dtm, cmap="terrain", extent=extent, alpha=0.25, vmin=np.nanmin(dtm) - 40, vmax=np.nanmax(dtm) + 40)
    ax.set_xlim(X0, X1); ax.set_ylim(Y0, Y1)
    ax.set_xticks(range(X0, X1 + 1, 1000), [f"{v / 1000:.0f}" for v in range(X0, X1 + 1, 1000)], fontsize=7)
    ax.set_yticks(range(Y0, Y1 + 1, 1000), [f"{v / 1000:.0f}" for v in range(Y0, Y1 + 1, 1000)], fontsize=7)
    ax.set_xlabel("easting (km, EPSG:27700)", fontsize=8); ax.set_ylabel("northing (km)", fontsize=8)
    ax.text(0.01, 0.01, "Contains Environment Agency information © Environment Agency and database right; "
            "BGS © UKRI (OGL v3)", transform=ax.transAxes, fontsize=5.5, color="#333333",
            bbox=dict(facecolor="white", alpha=0.7, lw=0, pad=1.5))


fig, ax = plt.subplots(figsize=(7.5, 7.5))
basemap(ax)
for t, c in TYPE_COLOURS.items():
    s = holes[holes.LOCA_TYPE == t]
    ax.scatter(s.geometry.x, s.geometry.y, s=22, color=c, edgecolor="white", lw=0.6,
               label=f"{abbr.get(t, t)} ({len(s)})", zorder=3)
placed = []  # label the deepest holes first, skipping any within 180 m of a label already placed
for _, h in holes[holes.LOCA_FDEP >= 40].sort_values("LOCA_FDEP", ascending=False).iterrows():
    x, y = h.geometry.x, h.geometry.y
    if any(np.hypot(x - px, y - py) < 180 for px, py in placed):
        continue
    placed.append((x, y))
    ax.annotate(h.LOCA_ID, (x, y), xytext=(4, 3), textcoords="offset points", fontsize=6,
                bbox=dict(facecolor="white", alpha=0.75, lw=0, pad=0.8), zorder=4)
ax.legend(loc="upper right", fontsize=7, framealpha=0.9)
ax.set_title("A417 Missing Link: 112 exploratory holes (labels on the deepest, 40 m or more)")
fig.savefig(OUT / "02-borehole-map.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(7.5, 7.5))
basemap(ax)
none = holes[holes.rockhead_unit.isna()]
ax.scatter(none.geometry.x, none.geometry.y, s=14, facecolor="white", edgecolor="#555555", lw=0.6, zorder=3)
order = [u for u in A.UNITS.index if u in set(holes.rockhead_unit.dropna())]
for u in order:
    s = holes[holes.rockhead_unit == u]
    ax.scatter(s.geometry.x, s.geometry.y, s=30, color=A.UNITS.colour[u], edgecolor="#222222", lw=0.6, zorder=3)
handles = [Line2D([], [], marker="o", ls="", markersize=6, markerfacecolor=A.UNITS.colour[u], markeredgecolor="#222222",
                  label=f"{A.UNITS.name[u]} ({(holes.rockhead_unit == u).sum()})") for u in order]
handles.append(Line2D([], [], marker="o", ls="", markersize=5, markerfacecolor="white", markeredgecolor="#555555",
                      label=f"bedrock not reached ({len(none)})"))
ax.legend(handles=handles, loc="upper right", fontsize=6.5, framealpha=0.9, title="first bedrock unit", title_fontsize=7)
ax.set_title("The formation at rockhead in each hole")
fig.savefig(OUT / "03-rockhead-unit-map.png")
plt.close(fig)

# ------------------------------------------------------------------ Step 6: GeoPackage for QGIS
section("Step 6: write a GeoPackage for QGIS")
keep = ["LOCA_ID", "LOCA_TYPE", "type_name", "LOCA_CLST", "LOCA_GL", "LOCA_FDEP", "lidar_gl", "gl_minus_lidar",
        "rockhead_depth", "rockhead_mOD", "rockhead_unit", "rockhead_name", "deepest_unit", "first_water_strike",
        "n_samples", "n_spt", "LOCA_BGSI", "LOCA_REM", "log_html", "geometry"]
gpkg = OUT / "a417_boreholes.gpkg"
gpkg.unlink(missing_ok=True)
holes[keep].to_file(gpkg, layer="boreholes", driver="GPKG")
check = gpd.read_file(gpkg, layer="boreholes")
print(f"wrote {gpkg.as_posix()}: layer 'boreholes', {len(check)} points, {len(check.columns) - 1} fields, "
      f"CRS {check.crs.to_epsg()}")
print("rasters for the QGIS basemap: a417_dtm_2m.tif, a417_hillshade.tif")
print("figures written:", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))