"""Download, merge and repair the A417 Missing Link ground investigation AGS4 data.

Shared by every part of the TechnicalGIS series "AGS data in GIS"; each part keeps an identical copy.

The British Geological Survey (BGS) holds AGS files donated by the ground investigation industry.
Its OGC API lists every borehole it can export (collection "agsboreholeindex"), and its AGS
File Utilities API exports those boreholes as AGS4 files, one ZIP per request. A request URL
can only carry so many borehole IDs, so we ask for them in batches of 40.

repair() applies the fixes worked through step by step in Part 1, so Parts 2 and 3 can start
from a clean, merged set of tables with load_clean().

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.
"""
import io
import time
import zipfile
from pathlib import Path

import pandas as pd
import requests
from python_ags4 import AGS4

ITEM = 173984                       # BGS deposited data item: HE551505 A417 Missing Link GI
INDEX = "https://ogcapi.bgs.ac.uk/collections/agsboreholeindex/items"
EXPORT = "https://agsapi.bgs.ac.uk/v7/ags_export/"
BATCH = 40
UA = {"User-Agent": "TechnicalGIS-tutorial (https://www.technicalgis.com)"}

# Geological units met in the A417 logs, youngest at the top. Bedrock codes are BGS Lexicon codes
# (GEOL_BGS); rank gives the expected order down a borehole, and an undivided group spans the ranks of its
# formations (rank to last). Superficial deposits have rank 0.
UNITS = pd.DataFrame([
    # code,  name,                                   group,                   rank, colour
    ("TOPS", "Topsoil",                              "Superficial",            0, "#5b4636"),
    ("MGR",  "Made ground",                          "Superficial",            0, "#9e9e9e"),
    ("HEAD", "Head",                                 "Superficial",            0, "#d9b77e"),
    ("SLIP", "Landslide deposits",                   "Superficial",            0, "#c98c5a"),
    ("WHL",  "White Limestone Formation",            "Great Oolite Group",     1, "#f7f1bd"),
    ("HMB",  "Hampen Formation",                     "Great Oolite Group",     2, "#e4d79c"),
    ("GOG",  "Great Oolite Group (undivided)",       "Great Oolite Group",     1, "#eee3a6"),
    ("FE",   "Fuller's Earth Formation",             "Great Oolite Group",     3, "#9fb08f"),
    ("SALS", "Salperton Limestone Formation",        "Inferior Oolite Group",  4, "#f4cd72"),
    ("ASLS", "Aston Limestone Formation",            "Inferior Oolite Group",  5, "#e9b157"),
    ("BLPL", "Birdlip Limestone Formation",          "Inferior Oolite Group",  6, "#d98f3e"),
    ("IOG",  "Inferior Oolite Group (undivided)",    "Inferior Oolite Group",  4, "#e6a855"),
    ("BDS",  "Bridport Sand Formation",              "Lias Group",             7, "#d2c48c"),
    ("WHM",  "Whitby Mudstone Formation",            "Lias Group",             8, "#7d93ab"),
    ("MRB",  "Marlstone Rock Formation",             "Lias Group",             9, "#b0703f"),
    ("DYS",  "Dyrham Formation",                     "Lias Group",            10, "#a3b8cc"),
    ("CHAM", "Charmouth Mudstone Formation",         "Lias Group",            11, "#55698a"),
    ("NR",   "No recovery or not classified",        "Other",                 0, "#f2f2f2"),
], columns=["code", "name", "group", "rank", "colour"]).set_index("code")
UNITS["last"] = UNITS["rank"]
UNITS.loc[["GOG", "IOG"], "last"] = [3, 6]

# Where GEOL_BGS is blank, the project's own code (GEOL_GEOL) still tells us the unit.
FROM_GEOL = {"TOP": "TOPS", "MG": "MGR", "HDD": "HEAD", "GOG": "GOG", "IOG": "IOG", "SLIP": "SLIP"}


def borehole_index(item=ITEM):
    """Every borehole in one BGS deposited data item, from the BGS OGC API (GeoJSON)."""
    r = requests.get(INDEX, params={"f": "json", "limit": 1000, "item_index_id": item}, headers=UA, timeout=120)
    r.raise_for_status()
    rows = [f["properties"] for f in r.json()["features"]]
    return pd.DataFrame(rows).sort_values("bgs_loca_id").reset_index(drop=True)


def _get(url, params, tries=5):
    for attempt in range(tries):
        try:
            r = requests.get(url, params=params, headers=UA, timeout=300)
            r.raise_for_status()
            return r
        except requests.RequestException:
            if attempt == tries - 1:
                raise
            time.sleep(10 * (attempt + 1))


def download(data_dir, item=ITEM):
    """Download the item's boreholes as AGS4 files, once. Returns the .ags paths in order."""
    raw = Path(data_dir) / f"bgs_{item}"
    raw.mkdir(parents=True, exist_ok=True)
    if not list(raw.glob("*.ags")):
        ids = borehole_index(item).bgs_loca_id.tolist()
        for i in range(0, len(ids), BATCH):
            r = _get(EXPORT, {"bgs_loca_id": ";".join(ids[i:i + BATCH])})
            with zipfile.ZipFile(io.BytesIO(r.content)) as z:
                for name in z.namelist():
                    (raw / f"batch{i // BATCH + 1:02d}_{Path(name).name}").write_bytes(z.read(name))
            time.sleep(1)
    return sorted(raw.glob("*.ags"))


def merge(paths):
    """Read several AGS4 files from one project and stack them, group by group.

    Each table keeps its UNIT and TYPE rows (from the first file) followed by every DATA row;
    rows repeated in every file (PROJ, TRAN, ABBR, DICT and so on) are kept once.
    """
    tables, headings = {}, {}
    for p in paths:
        t, h = AGS4.AGS4_to_dataframe(str(p))
        for group, df in t.items():
            if group not in tables:
                tables[group], headings[group] = df, list(h[group])
                continue
            headings[group] += [c for c in h[group] if c not in headings[group]]
            tables[group] = pd.concat([tables[group], df[df.HEADING == "DATA"]], ignore_index=True)
    for group, df in tables.items():
        df = df.reindex(columns=headings[group]).fillna("")
        tables[group] = df.drop_duplicates(ignore_index=True)
    return tables, headings


def repair(tables, headings):
    """Apply the Part 1 fixes to merged A417 tables, in place, and return them."""
    # Rule 1: a degree sign double-encoded as "Β°" (Greek capital beta + degree); the logs write "deg".
    for df in tables.values():
        df.replace("Β°", "deg", regex=True, inplace=True)

    # Rule 10a: three DICT headings defined twice; keep the later, archive descriptions.
    d = tables["DICT"]
    tables["DICT"] = d[~d.duplicated(["HEADING", "DICT_TYPE", "DICT_GRP", "DICT_HDNG"], keep="last")
                       | (d.HEADING != "DATA")].reset_index(drop=True)

    # Rules 7, 9, 19a, 19b: rename the BGS identifiers to valid user-defined headings and define them.
    new_dict = []
    for group, old, new, desc in [("PROJ", "BGS_PROJ_ID", "PROJ_BGSI", "BGS export project identifier"),
                                  ("LOCA", "BGS_LOCA_ID", "LOCA_BGSI", "BGS borehole index identifier")]:
        tables[group] = tables[group].rename(columns={old: new})
        tables[group].loc[tables[group].HEADING == "TYPE", new] = "X"
        headings[group] = [new if c == old else c for c in headings[group]]
        new_dict.append({"HEADING": "DATA", "DICT_TYPE": "HEADING", "DICT_GRP": group, "DICT_HDNG": new,
                         "DICT_STAT": "OTHER", "DICT_DTYP": "X", "DICT_DESC": desc})

    # The *_TEXT columns hold the depths and diameters exactly as originally written ("0.50"), where the
    # plain columns were reformatted ("0.5"). Put the original text back, restore the standard data type
    # and drop the extra column.
    std, _ = AGS4.AGS4_to_dataframe(str(Path(AGS4.__file__).parent / "Standard_dictionary_v4_0_3.ags"))
    dtype = std["DICT"].query("DICT_TYPE == 'HEADING'").set_index(["DICT_GRP", "DICT_HDNG"]).DICT_DTYP
    for group in list(tables):
        df = tables[group]
        for col in [c for c in headings[group] if c.endswith("_TEXT")]:
            base = col[:-5]
            data = df.HEADING == "DATA"
            same = pd.to_numeric(df.loc[data, base]).round(6).equals(pd.to_numeric(df.loc[data, col]).round(6))
            assert same, f"{group}.{col} disagrees with {base}"
            df.loc[data, base] = df.loc[data, col]
            df.loc[df.HEADING == "TYPE", base] = dtype[(group, base)]
            tables[group] = df = df.drop(columns=col)
            headings[group].remove(col)

    # Rule 10c: in the child groups LOCA_ID was exported with TYPE "X", so it no longer matches the
    # parent key in LOCA (TYPE "ID"). Restore the standard type.
    for group, df in tables.items():
        if group != "LOCA" and "LOCA_ID" in df:
            df.loc[df.HEADING == "TYPE", "LOCA_ID"] = "ID"

    tables["DICT"] = pd.concat([tables["DICT"], pd.DataFrame(new_dict)], ignore_index=True).fillna("")

    # Rule 7 (only checkable once the headings above are defined): headings must follow the order of
    # the standard dictionary, with the file's own DICT headings after them in DICT order.
    own = tables["DICT"].query("HEADING == 'DATA' and DICT_TYPE == 'HEADING'")
    defined = pd.concat([std["DICT"].query("DICT_TYPE == 'HEADING'"), own])
    for group in tables:
        listed = defined[defined.DICT_GRP == group].DICT_HDNG.drop_duplicates().tolist()
        order = {h: i for i, h in enumerate(listed)}
        headings[group] = ["HEADING"] + sorted(headings[group][1:], key=lambda h: order.get(h, len(order)))
        tables[group] = tables[group][headings[group]]

    tables["TRAN"].loc[tables["TRAN"].HEADING == "DATA", "TRAN_DESC"] += \
        "; three BGS exports merged and repaired (TechnicalGIS AGS data in GIS, Part 1)"
    return tables, headings


def write(tables, headings, ags_path, metadata):
    """Write a clean AGS4 file with the FILE folder AGS Rule 20 expects beside it."""
    ags_path = Path(ags_path)
    ags_path.parent.mkdir(parents=True, exist_ok=True)
    AGS4.dataframe_to_AGS4(tables, headings, str(ags_path), warnings=False)
    fset = ags_path.parent / "FILE" / "BGSFileSet01"
    fset.mkdir(parents=True, exist_ok=True)
    (fset / "BGS_download_metadata.txt").write_bytes(Path(metadata).read_bytes())
    return ags_path


def numeric(tables):
    """DATA rows only, with numbers as numbers. python-ags4 reads the TYPE row to decide; it leaves
    "U" (variable decimal places) as text, so convert those columns too."""
    out = {}
    for g, df in tables.items():
        types = df[df.HEADING == "TYPE"].iloc[0]
        n = AGS4.convert_to_numeric(df)
        for col in types.index[types == "U"]:
            n[col] = pd.to_numeric(n[col], errors="coerce")
        out[g] = n
    return out


def unit_code(geol):
    """One unit code per GEOL row: the BGS Lexicon code, else the project code, else NR."""
    code = geol.GEOL_BGS.where(geol.GEOL_BGS != "", geol.GEOL_GEOL.map(FROM_GEOL))
    return code.fillna("NR")


def in_order(codes):
    """True if a top-to-bottom list of bedrock unit codes never steps back to a younger unit."""
    reached = 0
    for code in codes:
        if UNITS.at[code, "last"] < reached:
            return False
        reached = max(reached, UNITS.at[code, "rank"])
    return True


def load_clean(data_dir):
    """Download (once), merge and repair the A417 data. Returns numeric DataFrames by group."""
    tables, headings = repair(*merge(download(data_dir)))
    out = numeric(tables)
    out["GEOL"] = out["GEOL"].sort_values(["LOCA_ID", "GEOL_TOP"], ignore_index=True)
    out["GEOL"]["UNIT"] = unit_code(out["GEOL"])
    return out