#!/usr/bin/env python3
"""
Download ETOPO 2022 30-arc-second elevation data and convert to ruaumoko-compatible
binary format (int16 little-endian, 21601 lat x 43200 lon, south-to-north).

Output: ~1.74 GiB binary file.

Usage: python3 build_elevation.py [output_path]
  Default output: /srv/ruaumoko-dataset
"""

import sys
import os
import struct
import tempfile
import numpy as np

CELLS_PER_DEGREE = 120
NUM_LATS = 180 * CELLS_PER_DEGREE + 1  # 21601
NUM_LONS = 360 * CELLS_PER_DEGREE      # 43200
EXPECTED_SIZE = NUM_LATS * NUM_LONS * 2  # 1,866,326,400

ETOPO_URL = "https://www.ngdc.noaa.gov/thredds/fileServer/global/ETOPO2022/30s/30s_surface_elev_netcdf/ETOPO_2022_v1_30s_N90W180_surface.nc"


def download_etopo(output_path):
    """Download ETOPO 2022 NetCDF and convert to ruaumoko binary format."""
    try:
        import xarray as xr
    except ImportError:
        print("ERROR: xarray is required. Install with: pip install xarray netcdf4")
        sys.exit(1)

    # Check if we can download directly or need a local file
    nc_path = os.environ.get("ETOPO_NC_PATH")
    if nc_path and os.path.exists(nc_path):
        print(f"Using local ETOPO file: {nc_path}")
    else:
        print(f"Downloading ETOPO 2022 30-second data (~1.1 GB)...")
        print(f"  URL: {ETOPO_URL}")
        print(f"  (Set ETOPO_NC_PATH env var to use a pre-downloaded file)")

        import urllib.request
        with tempfile.NamedTemporaryFile(suffix=".nc", delete=False) as f:
            nc_path = f.name

        try:
            urllib.request.urlretrieve(ETOPO_URL, nc_path, _progress)
            print()
        except Exception as e:
            os.unlink(nc_path)
            print(f"\nDownload failed: {e}")
            print("\nAlternative: manually download ETOPO 2022 30s NetCDF from:")
            print("  https://www.ncei.noaa.gov/products/etopo-global-relief-model")
            print(f"  Then set ETOPO_NC_PATH=/path/to/file.nc and re-run")
            sys.exit(1)

    print(f"Opening NetCDF dataset...")
    ds = xr.open_dataset(nc_path)

    # ETOPO 2022 30s has:
    # - lat: -90 to +90, 21601 points (south to north)
    # - lon: -180 to +180, 43201 points
    # We need:
    # - lat: -90 to +90, 21601 points (south to north) ← same
    # - lon: 0 to 360 (exclusive), 43200 points ← need to shift and drop last

    z = ds["z"]  # elevation variable
    print(f"  Shape: {z.shape}")
    print(f"  Lat range: {float(z.lat.min())} to {float(z.lat.max())}")
    print(f"  Lon range: {float(z.lon.min())} to {float(z.lon.max())}")

    # Sort latitude south-to-north (should already be, but ensure)
    z = z.sortby("lat")

    # Shift longitude from [-180, 180] to [0, 360)
    print("Shifting longitude to [0, 360)...")
    z = z.assign_coords(lon=(z.lon % 360))
    z = z.sortby("lon")

    data = z.values
    print(f"  Raw shape after sort: {data.shape}")

    # Handle longitude dimension: drop last col if it wraps (43201 → 43200)
    if data.shape[1] == NUM_LONS + 1:
        data = data[:, :NUM_LONS]
    elif data.shape[1] != NUM_LONS:
        print(f"ERROR: unexpected lon dimension: {data.shape[1]}, expected {NUM_LONS} or {NUM_LONS+1}")
        sys.exit(1)

    # Handle latitude dimension: ETOPO 2022 is cell-centered (21600 rows),
    # ruaumoko expects grid-registered (21601 rows including both poles).
    # Pad by duplicating edge rows for the poles.
    if data.shape[0] == NUM_LATS - 1:
        print(f"  Padding latitude from {data.shape[0]} to {NUM_LATS} (adding north pole row)")
        north_pole = data[-1:, :]    # duplicate +89.99... as +90
        data = np.concatenate([data, north_pole], axis=0)
    elif data.shape[0] != NUM_LATS:
        print(f"ERROR: unexpected lat dimension: {data.shape[0]}, expected {NUM_LATS} or {NUM_LATS-1}")
        sys.exit(1)

    print(f"Final grid shape: {data.shape}")
    print(f"Elevation range: {data.min():.1f} to {data.max():.1f} metres")

    # Write as int16 little-endian
    print(f"Writing to {output_path}...")
    elev_int16 = np.clip(data, -32768, 32767).astype(np.dtype("<i2"))
    elev_int16.tofile(output_path)

    actual_size = os.path.getsize(output_path)
    print(f"Written {actual_size:,} bytes (expected {EXPECTED_SIZE:,})")
    if actual_size == EXPECTED_SIZE:
        print("SUCCESS")
    else:
        print("WARNING: size mismatch!")

    ds.close()

    # Spot check
    verify(output_path)


def verify(path):
    """Quick spot-check of the elevation dataset."""
    data = np.memmap(path, dtype="<i2", mode="r", shape=(NUM_LATS, NUM_LONS))

    tests = [
        ("Mt Everest (~28.0N, 86.9E)", 28.0, 86.9, 8000, 9000),
        ("Dead Sea (~31.5N, 35.5E)", 31.5, 35.5, -500, 0),
        ("Pacific Ocean (~0N, 180E)", 0.0, 180.0, -6000, 0),
        ("Auburn AU (~-34.0S, 138.7E)", -34.03, 138.69, 200, 400),
    ]

    print("\n  Spot-check:")
    for name, lat, lon, lo, hi in tests:
        lat_idx = int((lat + 90) * CELLS_PER_DEGREE)
        lon_idx = int(lon * CELLS_PER_DEGREE)
        val = int(data[lat_idx, lon_idx])
        ok = "OK" if lo <= val <= hi else "FAIL"
        print(f"    {name}: {val}m [{ok}] (expected {lo}-{hi})")


_last_pct = -1
def _progress(block_num, block_size, total_size):
    global _last_pct
    if total_size > 0:
        pct = int(block_num * block_size * 100 / total_size)
        if pct != _last_pct and pct % 5 == 0:
            _last_pct = pct
            print(f"  {pct}%...", end="", flush=True)


if __name__ == "__main__":
    output = sys.argv[1] if len(sys.argv) > 1 else "/srv/ruaumoko-dataset"
    download_etopo(output)