settled-reach/tooling/domains/atlas/planet/sol_data/earth.py

"""
Earth (GJ0d) terrain builder.

Data sources:
  - Elevation: ETOPO 2022 60 arc-second (NOAA) — GeoTIFF
  - Temperature: WorldClim v2.1 annual mean (10 arc-min) — GeoTIFF
  - Precipitation: WorldClim v2.1 annual total (10 arc-min) — GeoTIFF
  - Rivers: Natural Earth 10m rivers — GeoJSON

All sources are equirectangular with col 0 = 180°W. ETOPO and WorldClim
use col 0 = 180°W natively. Natural Earth uses -180 to 180 longitude.
"""

import json
import zipfile
import numpy as np
from pathlib import Path

from tooling.domains.atlas.planet.sol_data.download import ensure_cached
from tooling.domains.atlas.planet.sol_data.shared import (
    GRID_W, GRID_H,
    load_tiff_as_array,
    resample_to_grid, normalize_01, compute_sea_level,
    compute_hillshade, assemble_terrain,
)
from tooling.core import console

# ─── Data source URLs ───────────────────────────────────────────────────────

# ETOPO 2022 60 arc-second — surface elevation (ice surface, not bedrock)
# ~130 MB GeoTIFF, 21600 x 10800, int16 metres
ETOPO_URL = "https://www.ngdc.noaa.gov/mgg/global/relief/ETOPO2022/data/60s/60s_surface_elev_gtif/ETOPO_2022_v1_60s_N90W180_surface.tif"
ETOPO_FILE = "ETOPO_2022_v1_60s_N90W180_surface.tif"

# WorldClim v2.1 — 10 arc-minute resolution (migrated to geodata.ucdavis.edu)
# Temperature: mean annual, °C × 10 (int16), in a zip
WCLIM_TEMP_URL = "https://geodata.ucdavis.edu/climate/worldclim/2_1/base/wc2.1_10m_tavg.zip"
WCLIM_TEMP_FILE = "wc2.1_10m_tavg.zip"

# Precipitation: annual total mm (int16), in a zip
WCLIM_PREC_URL = "https://geodata.ucdavis.edu/climate/worldclim/2_1/base/wc2.1_10m_prec.zip"
WCLIM_PREC_FILE = "wc2.1_10m_prec.zip"

# Natural Earth 10m rivers — GeoJSON from GitHub
RIVERS_URL = "https://raw.githubusercontent.com/nvkelso/natural-earth-vector/master/geojson/ne_10m_rivers_lake_centerlines.geojson"
RIVERS_FILE = "ne_10m_rivers_lake_centerlines.geojson"

# Earth physical constants
EARTH_OCEAN_FRACTION = 0.71
EARTH_MIN_ELEV_M = -10994.0  # Mariana Trench
EARTH_MAX_ELEV_M = 8849.0    # Everest


# ─── River filtering ────────────────────────────────────────────────────────

# Rivers to include (smart scatter: 1-2 per continent + Rhine)
INCLUDED_RIVERS = {
    # Europe
    "Danube", "Volga", "Rhine",
    # North America
    "Mississippi", "St. Lawrence",
    # South America
    "Amazon", "Paraná",
    # Africa
    "Nile", "Congo",
    # West Asia
    "Tigris",
    # East/South Asia
    "Yangtze", "Ganges", "Mekong",
    # Australia
    "Murray",
}

# Fuzzy matching — some NE names differ slightly
RIVER_NAME_ALIASES = {
    "Parana": "Paraná",
    "Chang Jiang": "Yangtze",
    "Huang He": "Yellow",
    "Ganga": "Ganges",
    "Nil": "Nile",
    "Danau": "Danube",
    "Donau": "Danube",
    "Rhin": "Rhine",
    "Rhein": "Rhine",
    "Saint Lawrence": "St. Lawrence",
    "St Lawrence": "St. Lawrence",
    "Río Paraná": "Paraná",
    "Rio Parana": "Paraná",
}


def _match_river_name(feature_name: str) -> str:
    """Check if a Natural Earth river name matches our included set."""
    if not feature_name:
        return None
    name = feature_name.strip()
    # Direct match
    if name in INCLUDED_RIVERS:
        return name
    # Alias match
    if name in RIVER_NAME_ALIASES:
        alias = RIVER_NAME_ALIASES[name]
        if alias in INCLUDED_RIVERS:
            return alias
    # Substring match (e.g. "Mississippi River" contains "Mississippi")
    for included in INCLUDED_RIVERS:
        if included.lower() in name.lower() or name.lower() in included.lower():
            return included
    return None


# ─── Data loaders ───────────────────────────────────────────────────────────

def _load_etopo() -> np.ndarray:
    """Load ETOPO 2022 elevation data, return raw metres array."""
    path = ensure_cached(ETOPO_URL, ETOPO_FILE)
    console.event(f"    loading ETOPO: {path}")
    try:
        arr = load_tiff_as_array(str(path))
    except Exception as e:
        raise RuntimeError(
            f"Failed to load ETOPO GeoTIFF: {e}\n"
            f"If PIL can't read this TIFF, install Pillow with TIFF support "
            f"or convert to raw binary."
        ) from e
    console.event(f"    ETOPO shape: {arr.shape}, range: [{arr.min():.0f}, {arr.max():.0f}] m")
    return arr


def _load_worldclim_temperature() -> np.ndarray:
    """
    Load WorldClim v2.1 annual mean temperature.
    Returns temperature in Kelvin at native resolution.
    """
    zip_path = ensure_cached(WCLIM_TEMP_URL, WCLIM_TEMP_FILE)
    console.event(f"    loading WorldClim temperature: {zip_path}")

    # The zip contains monthly TIFFs (tavg_01.tif to tavg_12.tif).
    # Compute annual mean from all 12 months.
    cache_dir = zip_path.parent
    monthly_sum = None
    count = 0

    with zipfile.ZipFile(zip_path) as zf:
        tif_names = sorted([n for n in zf.namelist() if n.endswith(".tif")])
        for tif_name in tif_names:
            extracted = cache_dir / Path(tif_name).name
            if not extracted.exists():
                zf.extract(tif_name, cache_dir)
                # Handle nested paths in zip
                nested = cache_dir / tif_name
                if nested != extracted and nested.exists():
                    nested.rename(extracted)
            try:
                arr = load_tiff_as_array(str(extracted))
            except Exception:
                # Try the nested path
                nested = cache_dir / tif_name
                if nested.exists():
                    arr = load_tiff_as_array(str(nested))
                else:
                    continue
            # Replace nodata with NaN
            arr[arr < -999] = np.nan
            if monthly_sum is None:
                monthly_sum = arr.copy()
            else:
                monthly_sum += arr
            count += 1

    if count == 0:
        raise RuntimeError("No temperature TIFFs found in WorldClim archive")

    # Annual mean (WorldClim tavg is °C × 10)
    temp_C = (monthly_sum / count) / 10.0
    # Convert to Kelvin
    temp_K = temp_C + 273.15
    # Replace NaN (ocean/nodata) with a reasonable ocean temperature
    temp_K = np.nan_to_num(temp_K, nan=288.0)

    console.event(f"    WorldClim temp shape: {temp_K.shape}, "
          f"range: [{np.nanmin(temp_K):.0f}, {np.nanmax(temp_K):.0f}] K")
    return temp_K


def _load_worldclim_precipitation() -> np.ndarray:
    """
    Load WorldClim v2.1 annual precipitation (sum of 12 months).
    Returns precipitation in mm/year at native resolution.
    """
    zip_path = ensure_cached(WCLIM_PREC_URL, WCLIM_PREC_FILE)
    console.event(f"    loading WorldClim precipitation: {zip_path}")

    cache_dir = zip_path.parent
    annual_sum = None

    with zipfile.ZipFile(zip_path) as zf:
        tif_names = sorted([n for n in zf.namelist() if n.endswith(".tif")])
        for tif_name in tif_names:
            extracted = cache_dir / Path(tif_name).name
            if not extracted.exists():
                zf.extract(tif_name, cache_dir)
                nested = cache_dir / tif_name
                if nested != extracted and nested.exists():
                    nested.rename(extracted)
            try:
                arr = load_tiff_as_array(str(extracted))
            except Exception:
                nested = cache_dir / tif_name
                if nested.exists():
                    arr = load_tiff_as_array(str(nested))
                else:
                    continue
            arr[arr < -999] = 0.0
            if annual_sum is None:
                annual_sum = arr.copy()
            else:
                annual_sum += arr

    if annual_sum is None:
        raise RuntimeError("No precipitation TIFFs found in WorldClim archive")

    console.event(f"    WorldClim precip shape: {annual_sum.shape}, "
          f"range: [{annual_sum.min():.0f}, {annual_sum.max():.0f}] mm/yr")
    return annual_sum


def _load_rivers_geojson() -> list:
    """
    Load Natural Earth rivers GeoJSON and extract polylines for included rivers.
    Returns list of (name, [(row, col), ...]) in grid coordinates.
    """
    path = ensure_cached(RIVERS_URL, RIVERS_FILE)
    console.event(f"    loading rivers: {path}")

    with open(path) as f:
        geojson = json.load(f)

    rivers = []
    for feature in geojson.get("features", []):
        props = feature.get("properties", {})
        fname = props.get("name") or props.get("name_en") or ""
        matched = _match_river_name(fname)
        if not matched:
            continue

        geom = feature.get("geometry", {})
        geom_type = geom.get("type", "")
        coords_list = []

        if geom_type == "LineString":
            coords_list = [geom["coordinates"]]
        elif geom_type == "MultiLineString":
            coords_list = geom["coordinates"]
        else:
            continue

        for coords in coords_list:
            path_grid = []
            for lon, lat in coords:
                # Convert lon/lat to grid coordinates
                # Grid: row 0 = 90°N, row 255 = 90°S
                #        col 0 = 180°W, col 511 = 180°E
                row = int((90.0 - lat) / 180.0 * GRID_H)
                col = int((lon + 180.0) / 360.0 * GRID_W)
                row = max(0, min(GRID_H - 1, row))
                col = max(0, min(GRID_W - 1, col))
                # Deduplicate: skip if same grid cell as previous point.
                # Natural Earth has hundreds of lon/lat points per river,
                # many of which land on the same 512x256 cell. Without
                # dedup, the renderer sees len(path)=300 and draws width 6.
                if path_grid and path_grid[-1] == (row, col):
                    continue
                path_grid.append((row, col))
            if len(path_grid) >= 2:
                rivers.append((matched, path_grid))

    # Deduplicate: keep longest segment per river name
    by_name = {}
    for name, path in rivers:
        if name not in by_name or len(path) > len(by_name[name]):
            by_name[name] = path

    console.event(f"    matched {len(by_name)} rivers: {', '.join(sorted(by_name.keys()))}")
    return [(name, path) for name, path in by_name.items()]


# ─── Main builder ───────────────────────────────────────────────────────────

def build_terrain(body_def: dict) -> dict:
    """
    Build Earth terrain dict from real-world data.

    Returns the same dict format as planet_simulation.simulate().
    """
    from tooling.domains.atlas.planet.planet_simulation import compute_biome

    console.event("  Earth: loading real-world data...")

    # ── 1. Elevation ────────────────────────────────────────────────────
    etopo_raw = _load_etopo()

    # ETOPO 2022 N90W180 is already col 0 = 180°W — no shift needed
    # Resample to grid
    elevation_m = resample_to_grid(etopo_raw, GRID_H, GRID_W, order=1)

    # Normalise to [0, 1]
    elevation = normalize_01(elevation_m, EARTH_MIN_ELEV_M, EARTH_MAX_ELEV_M)

    # Sea level: Earth's ocean fraction is ~0.71
    sea_level = compute_sea_level(elevation, EARTH_OCEAN_FRACTION)
    surface_water = elevation < sea_level

    console.event(f"    elevation: sea_level={sea_level:.4f}, "
          f"ocean={surface_water.sum()}/{GRID_H*GRID_W} cells")

    # ── 2. Temperature ──────────────────────────────────────────────────
    temp_raw_K = _load_worldclim_temperature()

    # WorldClim uses col 0 = 180°W — no shift needed
    temperature_K = resample_to_grid(temp_raw_K, GRID_H, GRID_W, order=1)

    # Fill ocean areas with latitude-dependent ocean temperature
    v = np.linspace(0, 1, GRID_H, dtype=np.float32)
    lat_abs = np.abs(v - 0.5) * 2.0  # 0 at equator, 1 at poles
    ocean_temp = 301.0 - lat_abs[:, np.newaxis] * 30.0  # ~28°C equator, ~-2°C poles
    temperature_K = np.where(surface_water, ocean_temp, temperature_K)

    console.event(f"    temperature: [{temperature_K.min():.0f}, {temperature_K.max():.0f}] K")

    # ── 3. Moisture ─────────────────────────────────────────────────────
    precip_raw = _load_worldclim_precipitation()

    # WorldClim uses col 0 = 180°W — no shift needed
    precip = resample_to_grid(precip_raw, GRID_H, GRID_W, order=1)

    # Normalise to [0, 1] — global max is ~10000 mm/yr (tropical rainforest)
    moisture = normalize_01(precip, 0.0, 6000.0)
    # Ocean moisture = high (drives adjacent land humidity)
    moisture = np.where(surface_water, 0.9, moisture)

    console.event(f"    moisture: [{moisture.min():.2f}, {moisture.max():.2f}]")

    # ── 4. Biome classification ─────────────────────────────────────────
    # Use the existing Whittaker table with real temperature and moisture
    biome = compute_biome(body_def, elevation, sea_level, surface_water,
                          temperature_K, moisture)
    n_biomes = len(np.unique(biome))
    console.event(f"    biomes: {n_biomes} classes present")

    # ── 5. Hillshade ────────────────────────────────────────────────────
    hillshade = compute_hillshade(elevation)

    # ── 6. Rivers ───────────────────────────────────────────────────────
    named_rivers = _load_rivers_geojson()

    # Clip rivers: stop each path when it hits surface water.
    # Rivers like the Amazon/Nile/Rhine otherwise draw through seas.
    clipped = []
    for name, path in named_rivers:
        clipped_path = []
        for r, c in path:
            if surface_water[r, c]:
                break
            clipped_path.append((r, c))
        if len(clipped_path) >= 2:
            clipped.append((name, clipped_path))

    n_orig = len(named_rivers)
    n_kept = len(clipped)
    console.event(f"    rivers: {n_kept}/{n_orig} kept after water clipping")
    named_rivers = clipped
    rivers = [path for _, path in named_rivers]

    # ── 7. Assemble ─────────────────────────────────────────────────────
    terrain = assemble_terrain(
        elevation=elevation,
        temperature_K=temperature_K,
        moisture=moisture,
        biome=biome,
        surface_water=surface_water,
        hillshade=hillshade,
        rivers=rivers,
        sea_level=sea_level,
    )

    # Store river names for the marker overlay
    terrain["_river_names"] = {i: name for i, (name, _) in enumerate(named_rivers)}

    return terrain