#!/usr/bin/env python3
"""Download a USACE channel's public condition surveys from eHydro and grid them onto shared
cells, ready for ao_shoaling.py or your own analysis.

  python3 ehydro_fetch.py <channel area> [--since DATE] [--types CS,PR,AD,BD] [--list]

<channel area> is the channel's eHydro id, such as CESWG_NR_02_BTB (Galveston District, Neches
River). Find it on the eHydro map (https://www.arcgis.com/apps/mapviewer/index.html?layers=f9d6ff0437a34a6a9d1b0d81cc2ed4cf)
or with --list and a district code, such as --list CESWG.

Writes a folder named for the channel area:
  surveys/<survey>/survey.zip   each survey as published by USACE
  grids/<survey>.tif            depth below the survey's vertical datum, in feet, one band of
                                Float32, NoData -9999, on the same cells for every survey
  surveys.csv                   one row per survey: name, type, date, datum, soundings, grid
  channel.geojson               the channel's outline (WGS 84), for AO
  channel.tif                   the channel's cells (1 inside), on the same cells as the grids
  channel.json                  the channel's maintained and authorized depths, cells and CRS

Survey types: CS condition, PR condition of part of the channel, BD before dredging, AD after
dredging. XC surveys combine several PR surveys under one date, so they're skipped unless you
ask for them with --types.

This is a sample to adapt, not part of AO. eHydro surveys use the vertical datum named on each
plot sheet, usually MLLW: check it before comparing surveys, and mix datums only after
converting them. Galveston District surveys before 2016 use MLT.
Needs: pip install requests numpy rasterio scipy pyproj
"""

import argparse
import csv
import json
import re
import sys
import zipfile
from datetime import date, datetime, timezone
from pathlib import Path

import numpy as np
import rasterio
import requests
from pyproj import CRS
from rasterio.features import rasterize
from rasterio.transform import from_origin
from scipy.interpolate import LinearNDInterpolator
from scipy.spatial import cKDTree

ARCGIS = "https://services7.arcgis.com/n1YM8pTrFmm7L4hs/arcgis/rest/services"
SURVEYS = f"{ARCGIS}/eHydro_Survey_Data/FeatureServer/0/query"
CHANNEL_AREAS = f"{ARCGIS}/National_Channel_Framework/FeatureServer/1/query"
CHANNEL_REACHES = f"{ARCGIS}/National_Channel_Framework/FeatureServer/2/query"
NODATA = -9999.0
session = requests.Session()
session.headers["User-Agent"] = "ehydro-fetch-script"


def query(url, **params):
    """All the records an ArcGIS query returns, page by page."""
    params = {"f": "json", "returnGeometry": "false", "outFields": "*", **params}
    records, offset = [], 0
    while True:
        response = session.get(
            url, params={**params, "resultOffset": offset}, timeout=120
        )
        response.raise_for_status()
        answer = response.json()
        if "error" in answer:
            sys.exit(f"eHydro said: {answer['error'].get('message')}")
        page = answer.get("features", [])
        records += page
        if not answer.get("exceededTransferLimit") or not page:
            return records
        offset += len(page)


def day(milliseconds):
    return datetime.fromtimestamp(milliseconds / 1000, timezone.utc).date()


def list_channels(district):
    """Prints a district's channel areas with their number of surveys."""
    rows = query(
        SURVEYS,
        where=f"usacedistrictcode='{district}'",
        outFields="channelareaidfk,sdsfeaturename",
    )
    counts = {}
    for row in rows:
        key = (
            row["attributes"]["channelareaidfk"],
            row["attributes"]["sdsfeaturename"],
        )
        counts[key] = counts.get(key, 0) + 1
    for (area, name), n in sorted(counts.items()):
        print(f"{area:22} {n:5} surveys  {name}")


def find_surveys(area, since, until, types):
    rows = query(
        SURVEYS, where=f"channelareaidfk='{area}'", orderByFields="surveydatestart"
    )
    surveys = []
    for row in rows:
        a = row["attributes"]
        when = day(a["surveydatestart"])
        if since <= when <= until and a["surveytype"] in types:
            surveys.append(
                {
                    "survey": a["sourcedatalocation"]
                    .rsplit("/", 1)[-1]
                    .removesuffix(".ZIP")
                    .removesuffix(".zip"),
                    "type": a["surveytype"],
                    "date": when.isoformat(),
                    "url": a["sourcedatalocation"],
                    "projection": a["sourceprojection"],
                }
            )
    return surveys


def epsg_for(projection):
    """The EPSG code for eHydro's projection name, such as NAD_1983_StatePlane_Texas_South_Central_FIPS_4204_Feet."""
    try:
        code = CRS.from_user_input(projection).to_epsg(min_confidence=70)
    except Exception:
        code = None
    if not code:
        sys.exit(f"No EPSG code found for {projection}: give one with --epsg.")
    return code


def polygons(geometry):
    """Esri rings as GeoJSON polygons: clockwise rings are outlines, the others are holes."""
    shapes = []
    for ring in geometry["rings"]:
        x, y = np.array(ring).T
        if np.sum((x[1:] - x[:-1]) * (y[1:] + y[:-1])) > 0:
            shapes.append({"type": "Polygon", "coordinates": [ring]})
        elif shapes:
            shapes[-1]["coordinates"].append(ring)
    return shapes


def channel(area, epsg, folder):
    """The channel's outline in the surveys' CRS, its depths, and its outline in WGS 84 for AO."""
    where = f"channelareaidpk='{area}'"
    found = query(CHANNEL_AREAS, where=where, returnGeometry="true", outSR=epsg)
    if not found:
        sys.exit(f"The National Channel Framework has no channel area {area}.")
    outline = [shape for record in found for shape in polygons(record["geometry"])]
    response = session.get(
        CHANNEL_AREAS,
        params={
            "where": where,
            "outFields": "channelareaidpk",
            "outSR": 4326,
            "f": "geojson",
        },
        timeout=120,
    )
    response.raise_for_status()
    (folder / "channel.geojson").write_text(response.text)
    reaches = query(CHANNEL_REACHES, where=f"channelareaidfk='{area}'")
    depths = {
        "maintained_ft": max((r["attributes"]["depthmaintained"] or 0) for r in reaches)
        if reaches
        else None,
        "authorized_ft": max((r["attributes"]["depthauthorized"] or 0) for r in reaches)
        if reaches
        else None,
    }
    return outline, depths


def fetch(survey, folder):
    """Downloads a survey's zip into its own folder, unless it's already there."""
    target = folder / "surveys" / survey["survey"] / "survey.zip"
    if target.exists() and zipfile.is_zipfile(target):
        return target
    target.parent.mkdir(parents=True, exist_ok=True)
    partial = target.with_name("survey.zip.part")
    with session.get(survey["url"], stream=True, timeout=(30, 3600)) as response:
        response.raise_for_status()
        with open(partial, "wb") as out:
            for chunk in response.iter_content(1024 * 1024):
                out.write(chunk)
    if not zipfile.is_zipfile(partial):
        partial.unlink()
        sys.exit(f"The download of {survey['survey']} isn't a zip file: run again.")
    partial.replace(target)
    print(f"downloaded {survey['survey']} ({target.stat().st_size / 1e6:.0f} MB)")
    return target


def soundings(path):
    """The survey's thinned soundings (the .XYZ file) as x, y, depth below the datum in feet."""
    with zipfile.ZipFile(path) as z:
        names = [
            n
            for n in z.namelist()
            if n.upper().endswith(".XYZ")
            and not re.search(r"_(A|FULL)\.XYZ$", n.upper())
        ]
        if not names:
            return None, ""
        text = z.read(names[0]).decode("latin-1")
    rows = [line.replace(",", " ").split()[:3] for line in text.splitlines()]
    rows = [r for r in rows if len(r) == 3 and re.fullmatch(r"-?[\d.]+", r[2])]
    points = np.array(rows, dtype=float)
    # some files list each sounding twice: keep one per position
    _, first = np.unique(points[:, :2].round(2), axis=0, return_index=True)
    points = points[np.sort(first)]
    # eHydro XYZ files hold elevations, negative below the datum; a file of mostly positive
    # values holds depths already
    if np.median(points[:, 2]) < 0:
        points[:, 2] = -points[:, 2]
    datum = re.search(r"_(MLLW|MLT|MLW|NAVD88|LWRP)", names[0].upper())
    return points, datum.group(1) if datum else ""


def grid(points, transform, inside, cell, max_gap):
    """Depths on the shared cells: interpolated across a surface through the soundings (a TIN,
    as USACE builds), and NoData where the nearest sounding is farther than max_gap."""
    rows, cols = np.nonzero(inside)
    x = transform.c + (cols + 0.5) * cell
    y = transform.f - (rows + 0.5) * cell
    near = (
        (points[:, 0] > x.min() - max_gap)
        & (points[:, 0] < x.max() + max_gap)
        & (points[:, 1] > y.min() - max_gap)
        & (points[:, 1] < y.max() + max_gap)
    )
    points = points[near]
    depths = np.full(inside.shape, NODATA, dtype="float32")
    if len(points) < 3:
        return depths
    gap, _ = cKDTree(points[:, :2]).query(
        np.column_stack([x, y]), distance_upper_bound=max_gap
    )
    covered = np.isfinite(gap)
    values = LinearNDInterpolator(points[:, :2], points[:, 2])(x[covered], y[covered])
    keep = np.isfinite(values)
    depths[rows[covered][keep], cols[covered][keep]] = values[keep]
    return depths


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "area", help="eHydro channel area id, or a district code with --list"
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="list the surveys (or a district's channel areas) and stop",
    )
    parser.add_argument(
        "--since",
        type=date.fromisoformat,
        default=date(2016, 1, 1),
        help="first survey date (default 2016-01-01)",
    )
    parser.add_argument(
        "--until",
        type=date.fromisoformat,
        default=date.today(),
        help="last survey date (default today)",
    )
    parser.add_argument(
        "--types",
        default="CS,PR,AD,BD",
        help="survey types to keep (default CS,PR,AD,BD)",
    )
    parser.add_argument(
        "--pick",
        type=Path,
        help="a file of survey names, one per line: keep only these",
    )
    parser.add_argument(
        "--cell",
        type=float,
        default=10,
        help="grid cell size in the survey CRS's units (default 10 ft)",
    )
    parser.add_argument(
        "--max-gap",
        type=float,
        default=300,
        help="leave cells farther than this from a sounding empty (default 300 ft: condition surveys often run cross-section lines 200 ft apart, and USACE's own surfaces bridge gaps up to 1,000 ft)",
    )
    parser.add_argument(
        "--epsg",
        type=int,
        help="the surveys' EPSG code, if it can't be found from eHydro's projection name",
    )
    options = parser.parse_args()

    if options.list and "_" not in options.area:
        return list_channels(options.area)
    surveys = find_surveys(
        options.area, options.since, options.until, set(options.types.split(","))
    )
    if options.pick:
        wanted = {
            line.strip().removesuffix(".ZIP")
            for line in options.pick.read_text().splitlines()
            if line.strip()
        }
        surveys = [s for s in surveys if s["survey"] in wanted]
    if not surveys:
        sys.exit("No surveys match.")
    if options.list:
        for s in surveys:
            print(f"{s['date']}  {s['type']:3} {s['survey']}")
        return
    projections = {s["projection"] for s in surveys}
    if len(projections) > 1 and not options.epsg:
        sys.exit(
            f"The surveys use different projections ({', '.join(sorted(projections))}): reproject them first."
        )
    epsg = options.epsg or epsg_for(projections.pop())

    folder = Path(options.area)
    (folder / "grids").mkdir(parents=True, exist_ok=True)
    outline, depths = channel(options.area, epsg, folder)
    left, bottom, right, top = rasterio.features.bounds(
        {"type": "GeometryCollection", "geometries": outline}
    )
    cell = options.cell
    transform = from_origin(
        np.floor(left / cell) * cell, np.ceil(top / cell) * cell, cell, cell
    )
    shape = (
        int(np.ceil((transform.f - bottom) / cell)),
        int(np.ceil((right - transform.c) / cell)),
    )
    inside = rasterize(outline, out_shape=shape, transform=transform).astype(bool)
    with rasterio.open(
        folder / "channel.tif",
        "w",
        driver="GTiff",
        height=shape[0],
        width=shape[1],
        count=1,
        dtype="uint8",
        crs=CRS.from_epsg(epsg).to_wkt(),
        transform=transform,
        nodata=0,
        compress="deflate",
    ) as tif:
        tif.write(inside.astype("uint8"), 1)
    (folder / "channel.json").write_text(
        json.dumps(
            {"area": options.area, "epsg": epsg, "cell": cell, **depths}, indent=2
        )
        + "\n"
    )
    print(
        f"{options.area}: {len(surveys)} surveys, {inside.sum()} cells of {cell:g} ft inside the channel"
    )

    profile = {
        "driver": "GTiff",
        "height": shape[0],
        "width": shape[1],
        "count": 1,
        "dtype": "float32",
        "crs": CRS.from_epsg(epsg).to_wkt(),
        "transform": transform,
        "nodata": NODATA,
        "compress": "deflate",
    }
    with open(folder / "surveys.csv", "w", newline="") as out:
        table = csv.writer(out)
        table.writerow(
            ["survey", "type", "date", "datum", "soundings", "cells", "grid"]
        )
        for s in surveys:
            points, datum = soundings(fetch(s, folder))
            if points is None:
                print(f"{s['survey']} has no XYZ file: skipped")
                continue
            depths_grid = grid(points, transform, inside, cell, options.max_gap)
            path = folder / "grids" / f"{s['survey']}.tif"
            with rasterio.open(path, "w", **profile) as tif:
                tif.write(depths_grid, 1)
            cells = int((depths_grid != NODATA).sum())
            table.writerow(
                [
                    s["survey"],
                    s["type"],
                    s["date"],
                    datum,
                    len(points),
                    cells,
                    path.relative_to(folder),
                ]
            )
            print(f"{s['date']} {s['type']:3} {cells:7} cells  {s['survey']}")
    with open(folder / "surveys.csv") as written:
        datums = {r["datum"] for r in csv.DictReader(written)} - {""}
    if len(datums) > 1:
        print(
            f"Note: the file names give different vertical datums ({', '.join(sorted(datums))}): convert before comparing."
        )
    print(f"done: see {folder}/surveys.csv")


if __name__ == "__main__":
    main()