"""Turn a lake or seabed survey into a 16-bit PNG heightmap for the relief map generator.

    pip install rasterio numpy shapely
    python bathymetry.py SOURCE OUT.png [--width 900] [--surface METRES] [--bbox WEST SOUTH EAST NORTH]

SOURCE is anything GDAL reads: a GeoTIFF, an ESRI ASCII grid (.asc), an ArcInfo grid (the folder), a NetCDF file.
The values can be heights (a lake floor at 1,300 m, a seabed at -4,000 m) or depths; the PNG just runs from
black at the lowest value to white at the highest, so with heights the deepest point is black: tick "Dark is
high" in the generator to make it the tallest layer.

- Cells with no data (land, islands, gaps in the survey) come out transparent: off the map. With the generator's
  shape set to "outline" the map takes the shape of the water.
- --surface: for a grid that has land as well as water (GEBCO, a merged model). Cells higher than this are left
  out, and so is any low ground that isn't joined to the deepest point: --surface 0 keeps the sea and drops the
  land; for a lake, give its water level.
- --bbox: read only this rectangle (degrees for a latitude/longitude grid, else the grid's own units). Worth it
  for a big file, or one read over the web like GEBCO.
- Latitude/longitude grids are reprojected to UTM, so the pixels are square on the ground.
- The picture is averaged down to --width pixels on its long side (the generator keeps at most 300 cells, so
  300, 600 or 900 are good choices) and cropped to where there is data.
"""

from __future__ import annotations

import argparse
import struct
import zlib

import numpy as np
import rasterio
from rasterio.crs import CRS
from rasterio.features import geometry_mask, shapes
from rasterio.transform import Affine, array_bounds
from rasterio.warp import Resampling, calculate_default_transform, reproject
from rasterio.windows import from_bounds
from shapely.geometry import Point, shape


def utm(lon: float, lat: float) -> CRS:
    """The UTM zone a point lies in: a flat map in metres."""
    return CRS.from_epsg((32600 if lat >= 0 else 32700) + int((lon + 180) // 6) + 1)


def water(z: np.ndarray, surface: float) -> np.ndarray:
    """Cells below `surface` that are joined to the deepest one (the lake or sea, not a low valley next to it)."""
    below = ~np.isnan(z) & (z <= surface)
    if not below.any():
        raise SystemExit(f"nothing in the grid is lower than {surface:g}")
    r, c = np.unravel_index(np.nanargmin(np.where(below, z, np.nan)), z.shape)
    deepest = Point(c + 0.5, r + 0.5)
    body = next(shape(g) for g, _ in shapes(below.astype("uint8"), mask=below) if shape(g).contains(deepest))
    return ~geometry_mask([body], z.shape, Affine.identity())


def read(source: str, surface: float | None = None, bbox: tuple | None = None) -> tuple[np.ndarray, CRS, Affine]:
    """The grid as float32 heights with NaN where there is no data, in a CRS with square metre pixels."""
    with rasterio.open(source) as src:
        win = from_bounds(*bbox, src.transform).round_offsets().round_lengths() if bbox else None
        z = src.read(1, masked=True, window=win).astype("float32").filled(np.nan)
        crs, transform = src.crs, src.window_transform(win) if win else src.transform
        if crs is None:
            raise SystemExit(f"{source} has no coordinate system: set one in QGIS first")
        if surface is not None:
            z[~water(z, surface)] = np.nan
        if crs.is_geographic:
            w, s, e, n = array_bounds(*z.shape, transform)
            dst = utm((w + e) / 2, (s + n) / 2)
            t, cols, rows = calculate_default_transform(crs, dst, z.shape[1], z.shape[0], w, s, e, n)
            out = np.full((rows, cols), np.nan, "float32")
            reproject(z, out, src_transform=transform, src_crs=crs, src_nodata=np.nan, dst_transform=t,
                      dst_crs=dst, dst_nodata=np.nan, resampling=Resampling.bilinear)
            z, crs, transform = out, dst, t
    return z, crs, transform


def fit(z: np.ndarray, crs: CRS, transform: Affine, width: int) -> tuple[np.ndarray, float]:
    """Crop to the data and average down to `width` pixels on the long side; returns the grid and its pixel size."""
    rows, cols = np.nonzero(~np.isnan(z))
    if not len(rows):
        raise SystemExit("the grid has no data (check --surface)")
    r0, r1, c0, c1 = rows.min(), rows.max() + 1, cols.min(), cols.max() + 1
    z = z[r0:r1, c0:c1]
    transform = transform * Affine.translation(c0, r0)
    pixel = max(abs(transform.a) * z.shape[1], abs(transform.e) * z.shape[0]) / width
    w = max(1, round(abs(transform.a) * z.shape[1] / pixel))
    h = max(1, round(abs(transform.e) * z.shape[0] / pixel))
    out = np.full((h, w), np.nan, "float32")
    dst = Affine(pixel, 0, transform.c, 0, -pixel, transform.f)
    reproject(z, out, src_transform=transform, src_crs=crs, src_nodata=np.nan, dst_transform=dst, dst_crs=crs,
              dst_nodata=np.nan, resampling=Resampling.average)
    return out, pixel


def png16(z: np.ndarray, path: str) -> tuple[float, float]:
    """Save heights as a 16-bit grey + alpha PNG, lowest = black, highest = white, no data transparent."""
    ok = ~np.isnan(z)
    lo, hi = float(z[ok].min()), float(z[ok].max())
    px = np.empty((*z.shape, 2), ">u2")                            # PNG wants big-endian samples
    px[..., 0] = np.where(ok, np.round((np.nan_to_num(z, nan=lo) - lo) / (hi - lo) * 65535), 0)
    px[..., 1] = np.where(ok, 65535, 0)
    rows = px.reshape(z.shape[0], -1)
    raw = b"".join(b"\0" + row.tobytes() for row in rows)          # filter type 0 on every row

    def chunk(kind: bytes, body: bytes) -> bytes:
        return struct.pack(">I", len(body)) + kind + body + struct.pack(">I", zlib.crc32(kind + body))

    head = struct.pack(">IIBBBBB", z.shape[1], z.shape[0], 16, 4, 0, 0, 0)   # 16 bits, grey + alpha
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", head) + chunk(b"IDAT", zlib.compress(raw, 9))
                + chunk(b"IEND", b""))
    return lo, hi


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("source")
    ap.add_argument("out")
    ap.add_argument("--width", type=int, default=900, help="pixels on the long side (default 900)")
    ap.add_argument("--surface", type=float, help="keep only the water below this height (0 for a sea)")
    ap.add_argument("--bbox", type=float, nargs=4, metavar=("WEST", "SOUTH", "EAST", "NORTH"),
                    help="read only this rectangle")
    a = ap.parse_args()
    z, pixel = fit(*read(a.source, a.surface, a.bbox), a.width)
    lo, hi = png16(z, a.out)
    print(f"{z.shape[1]} x {z.shape[0]} px of {pixel:.1f} m ({z.shape[1] * pixel / 1000:.2f} x "
          f"{z.shape[0] * pixel / 1000:.2f} km), black = {lo:.1f}, white = {hi:.1f}, "
          f"{np.isnan(z).mean():.0%} transparent")