"""Read and validate AGS4 ground investigation data with Python: the A417 Missing Link.

Part 1 of the TechnicalGIS series "AGS data in GIS". Finds the 112 exploratory holes of the
A417 Missing Link ground investigation in the British Geological Survey (BGS) AGS index,
downloads them as AGS4 files, explains how an AGS4 file is put together, reads it into pandas
with python-ags4 (maintained by the Association of Geotechnical and Geoenvironmental
Specialists), validates it, fixes every error the checker finds and writes a clean file that
passes. It finishes by checking the geology itself: is every borehole's sequence in the right
stratigraphic order?

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/read-validate-ags4-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.
"""
from importlib.metadata import version
from pathlib import Path

import numpy as np
import pandas as pd
from python_ags4 import AGS4
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
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)
pd.set_option("display.max_colwidth", 60)

DATA = Path("data"); DATA.mkdir(exist_ok=True)
OUT = Path("output"); OUT.mkdir(exist_ok=True)
CLEAN = OUT / "clean" / "A417_missing_link_GI.ags"


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


def error_counts(path):
    """Run the AGS checker and count the problems under each rule (ignoring its summaries)."""
    errors = AGS4.check_file(str(path))
    return {rule: len(v) for rule, v in errors.items()
            if rule not in ("Metadata", "Summary of data", "General")}


print(f"python-ags4 {version('python-ags4')} | pandas {pd.__version__}")

# ------------------------------------------------------------------ Step 1: find the boreholes
section("Step 1: the A417 Missing Link in the BGS AGS borehole index")
index = A.borehole_index()
print(index.proj_name.iloc[0])
print(f"{len(index)} exploratory holes | deepest {index.loca_fdep.max():.1f} m | "
      f"E {index.x.min():,.0f} to {index.x.max():,.0f} | N {index.y.min():,.0f} to {index.y.max():,.0f}")
print(index[["bgs_loca_id", "loca_id", "x", "y", "loca_fdep"]].head(3).to_string(index=False))

# ------------------------------------------------------------------ Step 2: download
section("Step 2: download as AGS4, 40 holes per request")
paths = A.download(DATA)
for p in paths:
    print(f"{p.name}: {p.stat().st_size / 1024:,.0f} kB, {sum(1 for _ in open(p, encoding='utf-8')):,} lines")

# ------------------------------------------------------------------ Step 3: anatomy of an AGS4 file
section("Step 3: what an AGS4 file looks like")
lines = paths[0].read_text(encoding="utf-8").splitlines()
start = lines.index('"GROUP","LOCA"')
for line in lines[start:start + 6]:
    print(line[:96] + (" ..." if len(line) > 96 else ""))

# ------------------------------------------------------------------ Step 4: read into pandas
section("Step 4: read every group into pandas")
tables, headings = AGS4.AGS4_to_dataframe(str(paths[0]))
print(f"{len(tables)} groups in {paths[0].name}: {' '.join(tables)}")
print(tables["LOCA"][["HEADING", "LOCA_ID", "LOCA_TYPE", "LOCA_NATE", "LOCA_NATN", "LOCA_GL"]].head(4)
      .to_string(index=False))

merged, merged_headings = A.merge(paths)
std, _ = AGS4.AGS4_to_dataframe(str(Path(AGS4.__file__).parent / "Standard_dictionary_v4_0_3.ags"))
group_desc = std["DICT"].query("DICT_TYPE == 'GROUP'").set_index("DICT_GRP").DICT_DESC
counts = pd.Series({g: int((df.HEADING == "DATA").sum()) for g, df in merged.items()
                    if g not in ("PROJ", "TRAN", "UNIT", "TYPE", "DICT", "ABBR", "FILE")})
print("\nAll three files stacked, DATA rows per group:")
for g, n in counts.sort_values(ascending=False).items():
    print(f"  {g}  {n:>5,}  {group_desc.get(g, '')}")

fig, ax = plt.subplots(figsize=(7.5, 3.4))
s = counts.sort_values()
ax.barh([f"{g}  {group_desc.get(g, '')}" for g in s.index], s.values, color="#3f7f9f")
for y, v in enumerate(s.values):
    ax.text(v + 40, y, f"{v:,}", va="center", fontsize=8)
ax.set_xlabel("DATA rows")
ax.set_title("A417 Missing Link: what is in the AGS data (112 exploratory holes)")
ax.spines[["top", "right"]].set_visible(False)
ax.set_xlim(0, s.max() * 1.12)
fig.savefig(OUT / "01-group-rows.png")
plt.close(fig)

# ------------------------------------------------------------------ Step 5: LOCA, GEOL and SAMP
section("Step 5: LOCA, GEOL and SAMP")
num = A.numeric(merged)
loca, geol, samp = num["LOCA"], num["GEOL"], num["SAMP"]
print("LOCA_TYPE:", loca.LOCA_TYPE.value_counts().to_dict())
print("LOCA_CLST:", loca.LOCA_CLST.value_counts().to_dict())
print(f"ground level {loca.LOCA_GL.min():.1f} to {loca.LOCA_GL.max():.1f} m OD | "
      f"final depth {loca.LOCA_FDEP.min():.1f} to {loca.LOCA_FDEP.max():.1f} m")
print(f"\nGEOL: {len(geol):,} strata in {geol.LOCA_ID.nunique()} holes. One borehole, top to bottom:")
one = geol[geol.LOCA_ID == "RC516"].sort_values("GEOL_TOP")
print(one[["GEOL_TOP", "GEOL_BASE", "GEOL_GEOL", "GEOL_BGS", "GEOL_DESC"]].to_string(index=False, max_colwidth=50))
print(f"\nSAMP: {len(samp):,} samples. SAMP_TYPE:", samp.SAMP_TYPE.value_counts().head(8).to_dict())
abbr = merged["ABBR"].query("HEADING == 'DATA' and ABBR_HDNG == 'SAMP_TYPE'").set_index("ABBR_CODE").ABBR_DESC
print("  " + " | ".join(f"{c} = {abbr.get(c, '?')}" for c in samp.SAMP_TYPE.value_counts().head(8).index))

# ------------------------------------------------------------------ Step 6: validate
section("Step 6: validate the downloads with the AGS checker")
before = pd.DataFrame({p.name: error_counts(p) for p in paths}).fillna(0).astype(int)
before["total"] = before.sum(axis=1)
print(before.to_string())
errors = AGS4.check_file(str(paths[0]))
for rule in ["AGS Format Rule 1", "AGS Format Rule 9", "AGS Format Rule 10a", "AGS Format Rule 10c"]:
    print(f"{rule}: {errors[rule][0]['desc']} (line {errors[rule][0]['line']})")

# ------------------------------------------------------------------ Step 7: fix
section("Step 7: merge, repair and write a clean file")
raw_text = "\n".join(p.read_text(encoding="utf-8") for p in paths)
print(f"double-encoded degree signs to repair: {raw_text.count('Β°')}")
merged_only = OUT / "merged_only" / "A417_merged.ags"
A.write(*A.merge(paths), merged_only, DATA / "bgs_173984" / "batch01_BGS_download_metadata.txt")
tables, headings = A.repair(*A.merge(paths))
A.write(tables, headings, CLEAN, DATA / "bgs_173984" / "batch01_BGS_download_metadata.txt")
print(f"wrote {CLEAN.as_posix()} ({CLEAN.stat().st_size / 1024:,.0f} kB) with FILE/BGSFileSet01 beside it")
print("CRLF line endings:", CLEAN.read_bytes().count(b"\r\n") == CLEAN.read_bytes().count(b"\n"))

stages = pd.DataFrame({"downloads": before["total"],
                       "merged and rewritten": pd.Series(error_counts(merged_only)),
                       "repaired": pd.Series(error_counts(CLEAN))}).fillna(0).astype(int)
stages = stages.loc[stages.max(axis=1).sort_values(ascending=False).index]
print(stages.to_string())
errors = AGS4.check_file(str(CLEAN))
print("left in the clean file:", [k for k in errors if k.startswith("AGS Format Rule")] or "no rule errors")
for e in errors.get("FYI (Related to Rule 16)", [])[:2]:
    print("  FYI:", e["desc"])

fig, ax = plt.subplots(figsize=(7.5, 3.4))
show = stages.drop(index="AGS Format Rule 2a")
x = np.arange(len(show))
for i, (col, colour) in enumerate(zip(show.columns, ["#c0504d", "#e0a03a", "#3f7f9f"])):
    ax.bar(x + (i - 1) * 0.27, show[col], 0.27, label=col, color=colour)
ax.set_xticks(x, [r.replace("AGS Format Rule ", "Rule ").replace("FYI (Related to Rule 16)", "Rule 16\n(FYI only)")
                  for r in show.index], fontsize=8)
ax.set_ylabel("messages from the checker")
ax.set_title(f"AGS checker messages at each stage (plus Rule 2a line endings: "
             f"{stages.loc['AGS Format Rule 2a', 'downloads']:,} -> 0)", fontsize=9)
ax.legend(frameon=False)
ax.spines[["top", "right"]].set_visible(False)
fig.savefig(OUT / "02-validation-stages.png")
plt.close(fig)

clean = A.load_clean(DATA)
reread = A.numeric(AGS4.AGS4_to_dataframe(str(CLEAN))[0])
same = all(reread[g].reset_index(drop=True).equals(A.numeric(tables)[g].reset_index(drop=True)) for g in tables)
print(f"clean file reads back identically: {same} | load_clean() groups: {len(clean)}")

# ------------------------------------------------------------------ Step 8: check the geology
section("Step 8: check the logs make geological sense")
geol, loca = clean["GEOL"], clean["LOCA"].set_index("LOCA_ID")
units = A.UNITS
print("Units, youngest first:")
print(geol.UNIT.value_counts().reindex(units.index).dropna().astype(int)
      .rename(lambda c: f"{c:<5}{units.name[c]}").to_string())

problems = []
for lid, g in geol.groupby("LOCA_ID"):
    if (g.GEOL_BASE <= g.GEOL_TOP).any():
        problems.append((lid, "base not below top"))
    step = g.GEOL_TOP.values[1:] - g.GEOL_BASE.values[:-1]
    if (np.abs(step) > 0.005).any():
        problems.append((lid, f"gap or overlap of {step[np.abs(step) > 0.005][0]:+.2f} m"))
    if g.GEOL_BASE.max() > loca.LOCA_FDEP[lid] + 0.005:
        problems.append((lid, "strata below final depth"))
    rank = g.UNIT.map(units["rank"])
    rock = g.UNIT[rank > 0]
    if not A.in_order(rock):
        problems.append((lid, "bedrock out of order: " + " > ".join(rock[rock != rock.shift()])))
    if len(rock) and (g.GEOL_TOP[(rank == 0) & (g.UNIT != "NR")] > g.GEOL_TOP[rank > 0].min()).any():
        problems.append((lid, "superficial deposits below bedrock"))
n_rock = geol[geol.UNIT.map(units["rank"]) > 0].LOCA_ID.nunique()
print(f"\n{geol.LOCA_ID.nunique()} holes with GEOL, {n_rock} reach bedrock. Holes needing a second look:")
for lid, msg in problems:
    print(f"  {lid}: {msg}")

cp = geol[geol.LOCA_ID == "CP217"]
print("\nCP217, top to bottom:")
print(cp[["GEOL_TOP", "GEOL_BASE", "GEOL_GEOL", "UNIT", "GEOL_DESC"]].to_string(index=False, max_colwidth=58))


def strip(ax, lid, x, width=0.6, elevation=True):
    g = geol[geol.LOCA_ID == lid]
    gl = loca.LOCA_GL[lid] if elevation else 0          # depth mode: plot -depth
    for _, r in g.iterrows():
        ax.add_patch(plt.Rectangle((x - width / 2, gl - r.GEOL_BASE), width, r.GEOL_BASE - r.GEOL_TOP,
                                   facecolor=units.colour[r.UNIT], edgecolor="#333333", linewidth=0.3))


# Figure 3: deep cored holes from the top of the escarpment to the vale, in metres above Ordnance Datum
deep = loca[(loca.LOCA_TYPE == "RC") & (loca.LOCA_FDEP >= 25)].sort_values("LOCA_GL", ascending=False)
pick = deep.iloc[np.linspace(0, len(deep) - 1, 10).round().astype(int)]
fig, ax = plt.subplots(figsize=(8.5, 5))
for i, lid in enumerate(pick.index):
    strip(ax, lid, i)
    ax.text(i, loca.LOCA_GL[lid] + 3, lid, ha="center", fontsize=7, rotation=90, va="bottom")
ax.set_xlim(-0.7, len(pick) - 0.3)
ax.set_ylim(min(loca.LOCA_GL[pick.index] - loca.LOCA_FDEP[pick.index]) - 5, loca.LOCA_GL.max() + 22)
ax.set_xticks([])
ax.set_ylabel("elevation (m above Ordnance Datum)")
ax.set_title("Ten cored boreholes from the escarpment top to the vale, logged units at true elevation")
present = [c for c in units.index if c in set(geol[geol.LOCA_ID.isin(pick.index)].UNIT)]
ax.legend(handles=[Patch(facecolor=units.colour[c], edgecolor="#333333", label=f"{c} {units.name[c]}")
                   for c in present], loc="center left", bbox_to_anchor=(1.01, 0.5), fontsize=7, frameon=False)
ax.spines[["top", "right", "bottom"]].set_visible(False)
fig.savefig(OUT / "03-strip-logs.png")
plt.close(fig)

# Figure 4: CP217, where the Lias sequence repeats under landslide deposits
fig, ax = plt.subplots(figsize=(6.5, 5.2))
strip(ax, "CP217", 0, width=0.5, elevation=False)
runs = cp.groupby((cp.UNIT != cp.UNIT.shift()).cumsum()).agg(UNIT=("UNIT", "first"), top=("GEOL_TOP", "min"),
                                                           base=("GEOL_BASE", "max"))
for _, r in runs.iterrows():
    ax.text(0.32, -(r.top + r.base) / 2, f"{units.name[r.UNIT]}  {r.top:.2f}-{r.base:.2f} m", va="center", fontsize=7)
for d in (9.5, 19.1):
    ax.hlines(-d, -2.1, 0.25, color="#c0504d", lw=0.8, ls="--")
ax.text(-2.1, -9.3, "base of landslide deposits, 9.50 m", va="bottom", fontsize=7, color="#c0504d")
ax.text(-2.1, -18.9, "Whitby Mudstone again, 19.10 m", va="bottom", fontsize=7, color="#c0504d")
ax.set_xlim(-2.2, 2.0)
ax.set_ylim(-loca.LOCA_FDEP["CP217"] - 1, 1)
ax.set_xticks([])
ax.set_yticks(-np.arange(0, 35, 5), [f"{d}" for d in range(0, 35, 5)])
ax.set_ylabel("depth below ground (m)")
ax.set_title(f"CP217 (ground level {loca.LOCA_GL['CP217']:.2f} m OD): a repeated Lias sequence", fontsize=9)
ax.spines[["top", "right", "bottom"]].set_visible(False)
fig.savefig(OUT / "04-cp217-repeat.png")
plt.close(fig)
print("\nfigures written:", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))