#!/usr/bin/env python3
"""Compare two grids of the same area in an AO project: make a difference grid, add it to the
project as a grid, and optionally mark each place that changed as an object.

  python3 ao_change.py <project> <before.tif> <after.tif> [--mark METERS]

<project> is the project's name, or its id if several projects share the name. <before.tif>
and <after.tif> are the names of two GeoTIFF grids in that project, such as a baseline survey
and a resurvey. The difference is after minus before, on the before grid's cells, so give the
coarser grid first. With --mark, each area of at least 10 cells (--min-cells) that changed by
more than METERS gets a pin, one call per pin; running again adds them again.

Set AO_URL (such as https://ao.example.com) and AO_EMAIL; the script asks for your password.
On MCP-enabled AO, you can set AO_TOKEN to an API token instead.

This is a sample to adapt, not part of AO. Both grids need the same vertical datum and the same
corrections, such as sound velocity and tide, and heights in meters: the script can't check those.
Needs: pip install requests numpy rasterio scipy
"""

import argparse
import getpass
import hashlib
import json
import os
import re
import sys
from pathlib import Path

import numpy as np
import rasterio
import requests
from rasterio.warp import Resampling, reproject, transform
from scipy import ndimage

API = os.environ["AO_URL"].rstrip("/") + "/api/v1"
NODATA = -9999.0
session = requests.Session()
session.headers["User-Agent"] = "ao-change-script"
refresh_token = None


def use_tokens(response):
    """Keeps the tokens from a sign-in or renewal's user-context header."""
    global refresh_token
    response.raise_for_status()
    tokens = json.loads(response.headers["user-context"])
    session.headers["Authorization"] = "Bearer " + tokens["auth_token"]
    refresh_token = tokens["refresh_token"]


def sign_in():
    if os.environ.get("AO_TOKEN"):
        session.headers["Authorization"] = "Bearer " + os.environ["AO_TOKEN"]
        return
    email = os.environ["AO_EMAIL"]
    password = getpass.getpass(f"AO password for {email}: ")
    use_tokens(
        session.post(
            f"{API}/users/login",
            headers={"X-Email": email, "X-Password": password},
            timeout=60,
        )
    )


def sign_out():
    """Ends a password sign-in's session, so it doesn't count toward an account's 10."""
    if refresh_token:
        try:
            call("POST", "/users/logout", timeout=60)
        except requests.RequestException:
            pass


def call(method, path, upload=None, **options):
    """Calls AO. After a 401, renews the sign-in once and tries again."""
    options.setdefault("timeout", (30, 3600))
    for attempt in (1, 2):
        if upload:
            with open(upload, "rb") as f:
                response = session.request(
                    method, API + path, files={"file": (upload.name, f)}, **options
                )
        else:
            response = session.request(method, API + path, **options)
        if response.status_code == 401 and refresh_token and attempt == 1:
            use_tokens(
                session.post(
                    f"{API}/users/refresh-token",
                    headers={"X-Refresh-Token": refresh_token},
                    timeout=60,
                )
            )
            continue
        response.raise_for_status()
        return response


def find_project(name_or_id):
    """The project's record, which also holds its organization's id (org_id)."""
    me = call("GET", "/users/me").json()["items"]
    projects = call("GET", f"/surveys/users/{me['userid']}/preview").json()["items"][0][
        "surveys"
    ]
    if re.fullmatch(
        r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", name_or_id
    ):
        found = [p for p in projects if p["uuid"] == name_or_id]
    else:
        found = [p for p in projects if p["name"].lower() == name_or_id.lower()]
    if len(found) != 1:
        listed = "".join(
            f"\n  {p['uuid']}  {p['name']} ({p['org_name']})" for p in found
        )
        sys.exit(
            f"{len(found)} projects match {name_or_id!r}. "
            f"Give the project's id instead.{listed}"
        )
    return found[0]


def project_files(project_id):
    rows = call("GET", f"/surveys/{project_id}/files").json()["items"][0][
        "fileStatusList"
    ]
    return [f for f in rows if not f["isVirtualDir"]]


def find_file(files, name):
    """A file by its name, or by folder/name when several folders have one."""
    found = [
        f
        for f in files
        if f["name"].lower() == name.lower()
        or ((f["filePath"] or "") + f["name"]).lower() == name.lower()
    ]
    if len(found) != 1:
        listed = "".join(f"\n  {(f['filePath'] or '') + f['name']}" for f in found)
        sys.exit(f"{len(found)} files in the project match {name!r}.{listed}")
    return found[0]


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def fetch(f):
    """Downloads a project file into this folder, keeping its project folder, unless the
    same file is already here."""
    target = Path(f["filePath"] or "", f["name"])
    if not target.resolve().is_relative_to(Path.cwd().resolve()):
        sys.exit(f"{target} is outside this folder.")
    if target.exists() and sha256(target) == f["hash"]:
        return target
    target.parent.mkdir(parents=True, exist_ok=True)
    partial = target.with_name(target.name + ".part")
    with (
        call("GET", f"/file-management/geo/{f['fileUuid']}", stream=True) as response,
        open(partial, "wb") as out,
    ):
        for chunk in response.iter_content(1024 * 1024):
            out.write(chunk)
    if sha256(partial) != f["hash"]:
        partial.unlink()
        sys.exit(f"The download of {f['name']} doesn't match AO's copy: run again.")
    partial.replace(target)
    print("downloaded", target)
    return target


def difference(before_path, after_path, out_path):
    """Writes after minus before on the before grid's cells, as one band of Float32 values.
    The after grid is resampled onto those cells first, so the CRS and cell size match: averaged
    when its cells are smaller, interpolated otherwise."""
    with rasterio.open(before_path) as before, rasterio.open(after_path) as after:
        for grid in (before, after):
            if grid.crs is None:
                sys.exit(f"{grid.name} has no CRS: AO and this script need one.")
        before_values = before.read(1, masked=True).astype("float32").filled(np.nan)
        after_values = np.full(before_values.shape, np.nan, dtype="float32")
        finer = after.crs == before.crs and abs(after.res[0]) < abs(before.res[0])
        reproject(
            source=rasterio.band(after, 1),
            destination=after_values,
            src_nodata=after.nodata,
            dst_nodata=np.nan,
            dst_transform=before.transform,
            dst_crs=before.crs,
            resampling=Resampling.average if finer else Resampling.bilinear,
        )
        epsg = before.crs.to_epsg()
        if epsg is None:
            sys.exit(
                f"{before_path}'s CRS has no EPSG code, which AO needs to show it."
            )
        change = after_values - before_values  # NaN where either grid has no data
        compared = np.count_nonzero(~np.isnan(change))
        if compared == 0:
            sys.exit("The two grids don't overlap, or one of them has no data.")
        profile = {
            "driver": "GTiff",
            "height": before.height,
            "width": before.width,
            "count": 1,
            "dtype": "float32",
            "crs": before.crs,
            "transform": before.transform,
            "nodata": NODATA,
            "compress": "deflate",
        }
        with rasterio.open(out_path, "w", **profile) as out:
            out.write(np.where(np.isnan(change), NODATA, change).astype("float32"), 1)
        print(
            f"compared {compared} cells: change from {np.nanmin(change):+.2f} "
            f"to {np.nanmax(change):+.2f}"
        )
        return change, before.transform, before.crs, epsg


def upload(project_id, path):
    """Adds a file to the project. Returns its id, and whether it's new to the project."""
    record = {
        "name": path.name,
        "filePath": "",
        "hash": sha256(path),
        "on_conflict": "SKIP",
    }
    answer = call(
        "POST", f"/surveys/{project_id}/files", json={"files": {"files": [record]}}
    ).json()
    status = answer["items"][0]["fileCreationStatus"]
    if not status:
        sys.exit(
            f"A different file named {path.name} is already in the project: "
            "rename or delete it, then run again."
        )
    if status[0]["status"] == "dup":
        print(path.name, "is already in the project")
        return status[0]["fid"], False
    call("POST", f"/file-management/geo/{status[0]['fid']}", upload=path)
    print("uploaded", path.name)
    return status[0]["fid"], True


def visualize(project, file_id, epsg, name):
    """Visualizes the difference as a Grid with a custom vertical offset of 0, so the legend,
    the cursor's readout and profiles show the change itself."""
    params = {
        "name": name,
        "org_uuid": project["org_id"],
        "epsg": str(epsg),
        "data_product_type": "model",
        "tileset_type": "3DTILES",
        "data_product_subtype": "jet",
        "vertical_epsg": None,
        "custom_vertical_offset": 0,
        "elev_units": "m",
        "z_positive_down": False,
    }
    call(
        "POST",
        f"/surveys/{project['uuid']}/files/{file_id}/render",
        json={"params": params},
    )
    print(f"visualizing {name!r}: see the project's Visualizations tab")


def mark(project, change, grid_transform, crs, meters, label, limit, min_cells):
    """Adds a pin where each area that changed by more than `meters` changed the most.
    Areas of fewer than `min_cells` cells are skipped: they're usually noise."""
    areas, count = ndimage.label(
        np.nan_to_num(np.abs(change)) > meters, structure=np.ones((3, 3))
    )
    if count == 0:
        print(f"no area changed by more than {meters} m")
        return
    peaks, small = [], 0
    for index, cells in enumerate(ndimage.find_objects(areas), start=1):
        inside = areas[cells] == index
        if inside.sum() < min_cells:
            small += 1
            continue
        values = np.where(inside, change[cells], 0)
        row, col = np.unravel_index(np.argmax(np.abs(values)), values.shape)
        peaks.append(
            (
                values[row, col],
                int(inside.sum()),
                cells[0].start + row,
                cells[1].start + col,
            )
        )
    if small:
        print(f"skipped {small} areas of fewer than {min_cells} cells")
    peaks.sort(key=lambda peak: -abs(peak[0]))
    for value, cells, row, col in peaks[:limit]:
        x, y = rasterio.transform.xy(grid_transform, row, col)
        lon, lat = (v[0] for v in transform(crs, "EPSG:4326", [x], [y]))
        call(
            "POST",
            "/marker",
            json={
                "name": f"Change {value:+.2f} m",
                "type": "Pin",
                "org_id": project["org_id"],
                "survey_uuid": project["uuid"],
                "marker_coordinates": f"POINT({lon:.8f} {lat:.8f})",
                "label": "Change",
                "comments": f"{value:+.2f} m at the most, over {cells} cells: {label}",
                "source": "ao_change.py",
            },
        )
    print(f"added {min(limit, len(peaks))} objects", end="")
    print(f" (the {limit} largest of {len(peaks)} areas)" if len(peaks) > limit else "")


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("project")
    parser.add_argument("before")
    parser.add_argument("after")
    parser.add_argument(
        "--mark",
        type=float,
        metavar="METERS",
        help="add a pin for each area that changed by more than this",
    )
    parser.add_argument(
        "--limit", type=int, default=50, help="add at most this many pins (default 50)"
    )
    parser.add_argument(
        "--min-cells",
        type=int,
        default=10,
        help="skip changed areas smaller than this many cells (default 10)",
    )
    options = parser.parse_args()
    project = find_project(options.project)
    files = project_files(project["uuid"])
    before = fetch(find_file(files, options.before))
    after = fetch(find_file(files, options.after))
    out_path = Path(f"{after.stem}_minus_{before.stem}.tif")
    change, grid_transform, crs, epsg = difference(before, after, out_path)
    file_id, new = upload(project["uuid"], out_path)
    if new:
        visualize(project, file_id, epsg, out_path.stem)
    if options.mark is not None:
        mark(
            project,
            change,
            grid_transform,
            crs,
            options.mark,
            out_path.stem,
            options.limit,
            options.min_cells,
        )


if __name__ == "__main__":
    sign_in()
    try:
        main()
    finally:
        sign_out()