settled-reach/tooling/domains/atlas/planet/render_heightmap.py

"""
render_heightmap.py
-------------------
Renders an annotated equirectangular heightmap PNG from a terrain dict.

This is the PRIMARY output of the planet generator pipeline.
The globe render is a separate downstream step that reads the same terrain dict.

Equirectangular projection:
    X axis: longitude  0°→360°  (left to right)
    Y axis: latitude  +90°→-90° (top to bottom, north pole at row 0)

Each terrain grid cell maps to a block of output pixels via bicubic upscale.
All rendering is in float32; final conversion to uint8 at save time.

Output layers (composited in order):
    1. Biome colour     — smooth-blended from Whittaker grid, not hard-snapped
    2. Elevation shading — subtle darkening in valleys, lightening on peaks
    3. Hillshade        — surface normal lighting pass (makes terrain 3D-readable)
    4. Coastline        — 1px dark border at sea level threshold
    5. Rivers           — anti-aliased polylines from river list
    6. Lat/lon grid     — every 30°, semi-transparent
    7. Title panel      — body metadata strip at top
    8. Legend           — biome colour swatches at bottom

Geographic only. No settlements, roads, or cultural data.
Those live in a separate JSON sidecar and are overlaid by the atlas app.

Usage:
    from tooling.domains.atlas.planet.render_heightmap import render_heightmap
    from tooling.domains.atlas.planet.planet_simulation import simulate
    from tooling.domains.atlas.planet.body_definition_parser import parse_system

    defs    = parse_system("index.md")
    terrain = simulate(defs[0])
    img     = render_heightmap(defs[0], terrain)
    img.save("GJ144d_heightmap.png")
"""

import numpy as np
from PIL import Image, ImageDraw, ImageFont
from scipy.ndimage import binary_dilation

from tooling.domains.atlas.planet.biome_config import (
    BIOME_PALETTE as _BIOME_PALETTE_CFG,
    RIVER_RGB as _RIVER_RGB_CFG,
    COAST_RGB as _COAST_RGB_CFG,
    build_biome_rgb,
)
from tooling.core import console

# ---------------------------------------------------------------------------
# Output resolution
# ---------------------------------------------------------------------------

OUT_W = 1024
OUT_H = 512
UI_SCALE = OUT_W / 1024   # 1.0 at default 1024 — all pixel sizes scale with this

# ---------------------------------------------------------------------------
# Biome colour palette
# Indices match planet_simulation.WHITTAKER_TABLE class IDs.
# Extended exotic classes appended at end.
# ---------------------------------------------------------------------------

# Biome palette loaded from biomes.toml via biome_config.
# Per-planet overrides can patch _BIOME_PALETTE_CFG before rendering.
BIOME_PALETTE = _BIOME_PALETTE_CFG

def _build_biome_rgb(mode: str = "cartographic") -> dict:
    return build_biome_rgb(mode)

RENDER_MODE = "cartographic"
BIOME_RGB   = _build_biome_rgb(RENDER_MODE)

def _ocean_arrays(mode: str = "cartographic"):
    return (
        np.array(BIOME_PALETTE[0][mode], dtype=np.float32),
        np.array(BIOME_PALETTE[1][mode], dtype=np.float32),
        np.array(BIOME_PALETTE[2][mode], dtype=np.float32),
    )

OCEAN_DEEP, OCEAN_MID, OCEAN_SHALLOW = _ocean_arrays(RENDER_MODE)

RIVER_RGB = _RIVER_RGB_CFG
COAST_RGB = _COAST_RGB_CFG

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _upscale(grid: np.ndarray, order: int = 1) -> np.ndarray:
    """
    Upscale a (GRID_H, GRID_W) float32 grid to (OUT_H, OUT_W).
    order=1 → bilinear (smooth, good for continuous fields)
    order=0 → nearest  (sharp, good for integer class grids)
    """
    from scipy.ndimage import zoom
    zy = OUT_H / grid.shape[0]
    zx = OUT_W / grid.shape[1]
    return zoom(grid.astype(np.float32), (zy, zx), order=order).astype(np.float32)


def _upscale_int(grid: np.ndarray) -> np.ndarray:
    """Nearest-neighbour upscale for integer class grids (biome, etc)."""
    from scipy.ndimage import zoom
    zy = OUT_H / grid.shape[0]
    zx = OUT_W / grid.shape[1]
    return zoom(grid.astype(np.int32), (zy, zx), order=0).astype(np.int8)


# ---------------------------------------------------------------------------
# Layer 1 + 2 + 3: Biome colour + elevation shading + hillshade
# ---------------------------------------------------------------------------

def _render_surface(terrain: dict) -> np.ndarray:
    """
    Returns (OUT_H, OUT_W, 3) float32 RGB in [0, 1].

    Compositing order:
        biome_colour × elevation_shade × hillshade_factor
    """
    elevation   = _upscale(terrain["elevation"], order=1)
    hillshade   = _upscale(terrain["hillshade"],  order=1)
    biome_up    = _upscale_int(terrain["biome"])
    surf_water  = _upscale(terrain["surface_water"].astype(np.float32),
                           order=0) > 0.5
    sea_level   = terrain["sea_level"]

    H, W = elevation.shape

    # ── Biome base colour ─────────────────────────────────────────────────
    # Clamp biome index, look up palette
    # Build lookup array from active BIOME_RGB dict for vectorised indexing
    max_id   = max(BIOME_RGB.keys())
    pal_arr  = np.zeros((max_id + 1, 3), dtype=np.float32)
    for k, v in BIOME_RGB.items():
        pal_arr[k] = v
    biome_clamped = np.clip(biome_up, 0, max_id)
    rgb = pal_arr[biome_clamped].astype(np.float32) / 255.0

    # ── Ocean depth blending ───────────────────────────────────────────────
    # Override flat ocean biome with smooth depth gradient
    if surf_water.any():
        depth = np.clip((sea_level - elevation) / (sea_level + 1e-9), 0, 1)
        deep_col    = OCEAN_DEEP    / 255.0
        mid_col     = OCEAN_MID     / 255.0
        shallow_col = OCEAN_SHALLOW / 255.0

        # Three-stop blend: 0=shallow, 0.5=mid, 1=deep
        t1 = np.clip(depth * 2.0, 0, 1)           # 0→0.5 depth: shallow→mid
        t2 = np.clip((depth - 0.5) * 2.0, 0, 1)   # 0.5→1 depth: mid→deep
        ocean_rgb = (shallow_col * (1 - t1)[..., None]
                     + mid_col   * (t1 * (1 - t2))[..., None]
                     + deep_col  * t2[..., None])
        rgb = np.where(surf_water[..., None], ocean_rgb, rgb)

    # ── Elevation shading on land ──────────────────────────────────────────
    # Slight darkening in lowlands, brightening on ridges
    elev_norm = np.where(
        ~surf_water,
        np.clip((elevation - sea_level) / (1.0 - sea_level + 1e-9), 0, 1),
        0.0)
    elev_shade = 0.88 + 0.18 * elev_norm    # [0.88, 1.06] — clamp below
    rgb = np.where(~surf_water[..., None],
                   np.clip(rgb * elev_shade[..., None], 0, 1),
                   rgb)

    # ── Hillshade ──────────────────────────────────────────────────────────
    # Apply only on land — ocean gets its own depth shading
    # Blend factor: 0.55 hillshade + 0.45 flat (keeps colours readable)
    hs_blend = 0.55 * hillshade + 0.45
    rgb = np.where(~surf_water[..., None],
                   np.clip(rgb * hs_blend[..., None], 0, 1),
                   rgb)

    return rgb.astype(np.float32)


# ---------------------------------------------------------------------------
# Layer 4: Coastline
# ---------------------------------------------------------------------------

def _render_coastline(terrain: dict,
                      rgb: np.ndarray) -> np.ndarray:
    """Draw a 1–2px dark border at the sea level threshold."""
    surf_water = _upscale(terrain["surface_water"].astype(np.float32),
                          order=0) > 0.5

    # Dilate water mask by 1px, XOR with original → coastline ring
    dilated   = binary_dilation(surf_water, iterations=2)
    coastline = dilated & ~surf_water

    coast_col = np.array(COAST_RGB, dtype=np.float32) / 255.0
    out = rgb.copy()
    out[coastline] = coast_col
    return out


# ---------------------------------------------------------------------------
# Layer 5: Lat/lon grid
# ---------------------------------------------------------------------------

def _render_grid(rgb: np.ndarray) -> np.ndarray:
    """Draw lat/lon lines every 30° as semi-transparent overlays."""
    out  = rgb.copy()
    col  = np.array([255, 255, 255], dtype=np.float32) / 255.0
    alpha = 0.12   # very subtle

    # Latitude lines (horizontal) every 30°: at 1/6, 2/6, 3/6, 4/6, 5/6 of height
    for frac in [1/6, 2/6, 3/6, 4/6, 5/6]:
        y = int(frac * OUT_H)
        y0 = max(0, y - 1);  y1 = min(OUT_H - 1, y + 1)
        out[y0:y1, :] = out[y0:y1, :] * (1 - alpha) + col * alpha

    # Longitude lines (vertical) every 30°
    for frac in [1/6, 2/6, 3/6, 4/6, 5/6]:
        x = int(frac * OUT_W)
        x0 = max(0, x - 1);  x1 = min(OUT_W - 1, x + 1)
        out[:, x0:x1] = out[:, x0:x1] * (1 - alpha) + col * alpha

    return out


# ---------------------------------------------------------------------------
# Layer 7: Title panel
# ---------------------------------------------------------------------------

def _load_font(size: int):
    try:
        return ImageFont.load_default(size=size)
    except TypeError:
        return ImageFont.load_default()


def _render_title(img: Image.Image, body_def: dict) -> Image.Image:
    """Draw metadata strip at top of image."""
    panel_h = int(52 * UI_SCALE)
    panel   = Image.new("RGBA", (OUT_W, panel_h), (12, 15, 22, 210))

    img_rgba = img.convert("RGBA")
    img_rgba.paste(panel, (0, 0), panel)
    img_out  = img_rgba.convert("RGB")
    draw     = ImageDraw.Draw(img_out)

    name   = body_def.get("name") or body_def.get("id", "Unknown")
    bid    = body_def.get("id", "")
    pclass = body_def.get("planet_class", "").replace("_", " ")
    star   = body_def.get("star", {})
    orbit  = body_def.get("orbit", {})
    phys   = body_def.get("physical", {})
    env    = body_def.get("environment", {})

    star_str  = f"{star.get('type','?')}-type"
    dist_str  = f"{orbit.get('distance_au', 0):.2f} AU"
    grav_str  = f"{phys.get('gravity_g', '?')}g"
    atmo_str  = phys.get("atmosphere", "?")
    hydro_str = env.get("hydrosphere", "?")

    px = int(14 * UI_SCALE)
    py = int(7  * UI_SCALE)
    lh = int(17 * UI_SCALE)

    title_col = (200, 210, 228)
    sub_col   = (130, 145, 168)
    dim_col   = (75,  88, 110)

    line1 = f"{name.upper()}  ·  {bid}  ·  {pclass}"
    line2 = f"{star_str}  ·  {dist_str}  ·  {grav_str}  ·  atmo: {atmo_str}  ·  hydro: {hydro_str}"
    line3 = "HEIGHTMAP  ·  Settled Reach"

    draw.text((px, py),        line1, fill=title_col, font=_load_font(int(14 * UI_SCALE)))
    draw.text((px, py + lh),   line2, fill=sub_col,   font=_load_font(int(12 * UI_SCALE)))
    draw.text((px, py + lh*2), line3, fill=dim_col,   font=_load_font(int(11 * UI_SCALE)))

    return img_out


# ---------------------------------------------------------------------------
# Layer 8: Legend
# ---------------------------------------------------------------------------

def _biome_legend_items(terrain: dict) -> list:
    """
    Return list of (label, RGB) for biome classes actually present
    in this terrain — no phantom legend entries.
    """
    biome = terrain["biome"]
    present = set(np.unique(biome).tolist())

    LABELS = {
        0:  "ocean deep",     1:  "ocean",         2:  "coastal water",
        3:  "coast",          5:  "trop. rainforest", 6:  "trop. forest",
        7:  "savanna",        8:  "grassland",      9:  "forest",
        10: "temp. rainforest", 11: "boreal",       12: "shrubland",
        13: "temperate desert", 14: "desert",       15: "hot desert",
        16: "tundra",        17: "ice / snow",      18: "mountain rock",
        19: "lava field",    20: "chemosyn. mat",   21: "thermophilic",
        22: "sulfuric scrub", 23: "crypto. crust",  25: "ash field",
    }

    items = []
    # Fixed display order — most common first, exotic last
    order = [0, 1, 2, 3, 7, 8, 5, 6, 9, 10, 11, 12, 13, 14, 15, 16, 17,
             18, 19, 20, 21, 22, 23, 25]
    for cls_id in order:
        if cls_id in present and cls_id in LABELS:
            rgb = BIOME_RGB.get(cls_id, (128, 128, 128))
            items.append((LABELS[cls_id], rgb))

    # Always include river swatch if rivers exist
    if terrain.get("rivers"):
        items.append(("river", RIVER_RGB))

    return items


def _render_legend(img: Image.Image, terrain: dict) -> Image.Image:
    """Draw biome legend strip at bottom of image."""
    items = _biome_legend_items(terrain)
    if not items:
        return img

    draw   = ImageDraw.Draw(img)
    sw     = int(14 * UI_SCALE)    # swatch width
    sh     = int(12 * UI_SCALE)    # swatch height
    pad_x  = int(14 * UI_SCALE)
    leg_y  = OUT_H - int(34 * UI_SCALE)
    font   = _load_font(int(10 * UI_SCALE))
    gap    = int(6  * UI_SCALE)
    # Scale step to fit all items within image width
    max_items = len(items)
    max_step = (OUT_W - 2 * pad_x) // max(max_items, 1)
    step = min(int(108 * UI_SCALE), max_step)

    # Max label width = step minus swatch minus gaps
    max_label_px = step - sw - gap - int(4 * UI_SCALE)

    lx = pad_x
    for label, rgb in items:
        draw.rectangle([(lx, leg_y), (lx + sw, leg_y + sh)], fill=rgb)
        # Truncate label to fit allocated space
        display = label
        while display and draw.textlength(display, font=font) > max_label_px:
            display = display[:-1]
        if display != label and display:
            display = display[:-1] + "…"
        draw.text((lx + sw + gap, leg_y), display,
                  fill=(185, 192, 205), font=font)
        lx += step

    return img


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

def render_heightmap(body_def: dict,
                     terrain: dict,
                     out_w: int = OUT_W,
                     out_h: int = OUT_H,
                     render_mode: str = "cartographic",
                     chrome: bool = True) -> Image.Image:
    """
    Render an annotated equirectangular heightmap PNG.

    Parameters
    ----------
    body_def : dict   — body definition from body_definition_parser
    terrain  : dict   — terrain dict from planet_simulation.simulate()
    out_w, out_h      — output resolution (default 1024×512)

    Returns
    -------
    PIL.Image.Image  RGB
    """
    global OUT_W, OUT_H, UI_SCALE
    OUT_W    = out_w
    OUT_H    = out_h
    UI_SCALE = out_w / 1024

    # Set active colour mode for this render
    global BIOME_RGB, OCEAN_DEEP, OCEAN_MID, OCEAN_SHALLOW, RENDER_MODE
    RENDER_MODE = render_mode
    BIOME_RGB   = _build_biome_rgb(render_mode)
    OCEAN_DEEP, OCEAN_MID, OCEAN_SHALLOW = _ocean_arrays(render_mode)

    # Guard: require simulation data
    required = ("elevation", "biome", "surface_water", "hillshade", "sea_level")
    missing  = [k for k in required if k not in terrain]
    if missing:
        raise ValueError(f"terrain dict missing keys: {missing}")

    # 1+2+3: surface colour with elevation shading and hillshade
    rgb = _render_surface(terrain)

    # 4: coastline
    rgb = _render_coastline(terrain, rgb)

    # 5: lat/lon grid
    rgb = _render_grid(rgb)

    # Convert to PIL for text rendering
    img = Image.fromarray(
        (rgb * 255).clip(0, 255).astype(np.uint8), mode="RGB")

    if chrome:
        # 7: title panel
        img = _render_title(img, body_def)
        # 8: legend
        img = _render_legend(img, terrain)

    return img


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys
    import json
    import time

    if len(sys.argv) < 2:
        console.event("Usage: python -m tooling.domains.atlas.planet.render_heightmap body_def.json [--large]")
        sys.exit(1)

    with open(sys.argv[1]) as f:
        bd = json.load(f)

    # --large flag renders at 4096×2048 for high-res review
    large = "--large" in sys.argv
    w, h  = (4096, 2048) if large else (OUT_W, OUT_H)

    from tooling.domains.atlas.planet.planet_simulation import simulate

    console.event(f"Simulating: {bd['id']}  ({bd['planet_class']})")
    t0      = time.time()
    terrain = simulate(bd)
    sim_t   = time.time() - t0

    if not terrain:
        console.event("Gas giant — no heightmap.")
        sys.exit(0)

    console.event(f"Rendering heightmap {w}×{h}…")
    t1  = time.time()
    img = render_heightmap(bd, terrain, out_w=w, out_h=h)
    ren_t = time.time() - t1

    out = f"/mnt/user-data/outputs/{bd['id']}_heightmap.png"
    img.save(out, format="PNG")
    console.event(f"Saved: {out}")
    console.event(f"  simulate={sim_t:.1f}s  render={ren_t:.1f}s  total={sim_t+ren_t:.1f}s")