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

#!/usr/bin/env python3
"""
Planet Generator — CLI entry point.

Two input modes:
  1. Body definition JSON:  generate body_def.json --output-dir ./output
  2. System index.md:       generate --system wiki/star-systems/GJ-144/index.md

Outputs per body into {output-dir}/{body_id}/:
    heightmap.png      — clean equirectangular cartographic map (no chrome)
    globe.png          — 512×512 sphere render
    body.json          — body descriptor + rendering metadata
    terrain.npz        — compressed terrain grids for downstream generators
    rivers.json        — river polylines in grid coords

Optional (spike/review only):
    heightmap_chrome.png — heightmap with title bar + legend overlay
"""

import argparse
import json
import os
import time

# Venv bootstrap — re-exec into .venv/bin/python if not already there.
from pathlib import Path

from tooling.core import config, console
from tooling.core.errors import ReachError
PLANET_DIR = Path(__file__).resolve().parent
WORKTREE_ROOT = config.repo_root()
# The venv re-exec that used to sit here is gone (T-1288). It relaunched the
# script under .venv/bin/python so numpy would resolve when run by path.
# reach declares numpy and Pillow itself, so its own environment already has
# them — and an os.execv into a different interpreter, carrying reach's
# argv, would have relaunched something that is not this command at all.

import numpy as np

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



def _build_markers(body_def: dict, terrain: dict) -> dict:
    """Extract geographic markers from terrain data."""
    from scipy.ndimage import label, center_of_mass

    grid_h = terrain["_grid_h"]
    grid_w = terrain["_grid_w"]
    sea_level = terrain["sea_level"]
    elevation = terrain["elevation"]
    surface_water = terrain["surface_water"]
    biome = terrain["biome"]

    markers = {
        "grid": {"w": grid_w, "h": grid_h},
        "rivers": [],
        "oceans": [],
        "mountain_ranges": [],
        "roads": [],
        "cities": [],
        "railroads": [],
        "pois": [],
    }

    # ── Rivers ───────────────────────────────────────────────────────────
    for i, path in enumerate(terrain.get("rivers", [])):
        markers["rivers"].append({
            "id": f"river_{i}",
            "name": None,  # named by copy team or procedural namer
            "path": path,
        })

    # ── Oceans / seas ────────────────────────────────────────────────────
    # Label connected water bodies and record their center + area
    if surface_water.any():
        water_labels, n_bodies = label(surface_water)
        total_cells = grid_h * grid_w
        for lbl in range(1, n_bodies + 1):
            mask = water_labels == lbl
            area_cells = int(mask.sum())
            area_frac = area_cells / total_cells
            if area_frac < 0.005:
                continue  # skip tiny puddles
            cy, cx = center_of_mass(mask)
            kind = "ocean" if area_frac > 0.10 else "sea" if area_frac > 0.02 else "lake"
            markers["oceans"].append({
                "id": f"water_{lbl}",
                "name": None,
                "kind": kind,
                "center": [int(cy), int(cx)],
                "area_fraction": round(area_frac, 4),
            })

    # ── Mountain ranges ──────────────────────────────────────────────────
    # High-elevation connected regions on land
    land = ~surface_water
    elev_norm = np.where(land,
                         (elevation - sea_level) / (1.0 - sea_level + 1e-9),
                         0.0)
    mountains = land & (elev_norm > 0.55)
    if mountains.any():
        mtn_labels, n_ranges = label(mountains)
        for lbl in range(1, n_ranges + 1):
            mask = mtn_labels == lbl
            area = int(mask.sum())
            if area < 20:
                continue  # skip tiny peaks
            cy, cx = center_of_mass(mask)
            # Find the ridge line: cells with highest elevation in the range
            ys, xs = np.where(mask)
            peak_idx = np.argmax(elevation[ys, xs])
            markers["mountain_ranges"].append({
                "id": f"range_{lbl}",
                "name": None,
                "center": [int(cy), int(cx)],
                "peak": [int(ys[peak_idx]), int(xs[peak_idx])],
                "area_cells": area,
            })

    return markers


def _generate_body(body_def: dict, hmap_w: int, hmap_h: int,
                   globe_size: int, render_mode: str, output_dir: str,
                   chrome: bool = False):
    """Generate all outputs for a single body definition."""
    body_id = body_def["id"]
    if body_def.pop("_flat_output", False):
        body_dir = output_dir  # output directly, no body_id subdir
    else:
        body_dir = os.path.join(output_dir, body_id)
    os.makedirs(body_dir, exist_ok=True)

    planet_class = body_def.get("planet_class", "unknown")
    name = body_def.get("name") or body_id

    console.event(f"  {body_id} ({name}) — {planet_class}")

    t0 = time.time()

    # ── 1. Simulate ──────────────────────────────────────────────────────
    terrain = simulate(body_def)
    is_gas = not terrain
    t_sim = time.time()

    if is_gas:
        console.event(f"    simulate: gas giant ({t_sim - t0:.1f}s)")
    else:
        console.event(f"    simulate: {t_sim - t0:.1f}s  "
              f"sea={terrain['sea_level']:.3f}  "
              f"land={int((~terrain['surface_water']).sum())}")

    # ── 2. Render heightmap ──────────────────────────────────────────────
    t_hmap = t_sim
    if not is_gas:
        # Clean heightmap (no title/legend)
        hmap_img = render_heightmap(body_def, terrain,
                                    out_w=hmap_w, out_h=hmap_h,
                                    render_mode=render_mode, chrome=False)
        hmap_img.save(os.path.join(body_dir, "heightmap.png"))

        # Chrome version for review (optional — saved to /tmp, not shipped)
        if chrome:
            hmap_chrome = render_heightmap(body_def, terrain,
                                           out_w=hmap_w, out_h=hmap_h,
                                           render_mode=render_mode, chrome=True)
            chrome_path = f"/tmp/{body_id}_heightmap_chrome.png"
            hmap_chrome.save(chrome_path)
            console.event(f"    chrome: {chrome_path}")

        t_hmap = time.time()
        console.event(f"    heightmap: {t_hmap - t_sim:.1f}s  {hmap_w}×{hmap_h}")

    # ── 3. Render globe ──────────────────────────────────────────────────
    try:
        from tooling.domains.atlas.planet.planet_renderer import render_globe
        globe_img = render_globe(body_def, terrain, size=globe_size)
        globe_img.save(os.path.join(body_dir, "globe.png"))
        t_globe = time.time()
        console.event(f"    globe: {t_globe - t_hmap:.1f}s  {globe_size}×{globe_size}")
    except Exception as e:
        console.event(f"    globe: FAILED — {e}")
        t_globe = time.time()

    # ── 4. Write data files ──────────────────────────────────────────────
    # Body definition lives in the index.md frontmatter — no body.json needed.

    if not is_gas:
        # terrain.npz — grids for downstream generators
        save_dict = {}
        for key in ("elevation", "temperature", "moisture", "hillshade",
                    "biome", "surface_water", "river_grid"):
            if key in terrain:
                save_dict[key] = terrain[key]
        save_dict["sea_level"] = np.array([terrain["sea_level"]])
        np.savez_compressed(os.path.join(body_dir, "terrain.npz"), **save_dict)

        # markers.json — named geographic and cultural features.
        markers = _build_markers(body_def, terrain)
        with open(os.path.join(body_dir, "markers.json"), "w") as f:
            json.dump(markers, f, indent=2)

    elapsed = time.time() - t0
    console.event(f"    total: {elapsed:.1f}s → {body_dir}/")


def main(argv: list[str] | None = None):
    parser = argparse.ArgumentParser(
        description="Planet generator — heightmap + globe from body definitions")

    # Input modes
    parser.add_argument("body_def", nargs="?",
                        help="Path to body definition JSON file")
    parser.add_argument("--system",
                        help="Path to system index.md — generates all bodies")
    parser.add_argument("--overrides",
                        help="Path to per-body overrides JSON (used with --system)")

    # Output
    parser.add_argument("--output-dir", default=".",
                        help="Root output directory (bodies get subdirs)")

    # Rendering
    parser.add_argument("--heightmap-size", default="1024x512",
                        help="Heightmap output resolution (WxH)")
    parser.add_argument("--globe-size", type=int, default=512,
                        help="Globe output resolution (square, locked at 512)")
    parser.add_argument("--render-mode", choices=["cartographic", "photographic"],
                        default="cartographic")
    parser.add_argument("--chrome", action="store_true",
                        help="Also render heightmap with title/legend (review only, not shipped)")

    args = parser.parse_args(argv)

    # Parse heightmap size
    try:
        hw, hh = args.heightmap_size.lower().split("x")
        hmap_w, hmap_h = int(hw), int(hh)
    except ValueError:
        raise ReachError(f"invalid heightmap size '{args.heightmap_size}'", fix="pass --heightmap-size as WxH, e.g. 1024x512")

    # ── Collect body definitions ─────────────────────────────────────────
    body_defs = []

    if args.system:
        # Read from system index.md → parse bodies table
        from tooling.domains.atlas.planet.body_definition_parser import parse_system
        overrides = {}
        if args.overrides:
            with open(args.overrides) as f:
                overrides = json.load(f)
        body_defs = parse_system(args.system, overrides=overrides)
        console.event(f"System: {args.system} — {len(body_defs)} bodies")
    elif args.body_def:
        input_path = args.body_def
        if input_path.endswith(".json"):
            with open(input_path) as f:
                body_defs = [json.load(f)]
        elif input_path.endswith(".md"):
            # Read body definition from frontmatter
            import yaml
            with open(input_path) as f:
                content = f.read()
            if content.startswith("---"):
                fm_end = content.index("---", 3)
                fm = yaml.safe_load(content[3:fm_end])
                if "id" in fm and "planet_class" in fm:
                    body_defs = [fm]
                    fm["_flat_output"] = True  # no body_id subdir
                    # Output alongside the body index.md if no --output-dir
                    if args.output_dir == ".":
                        args.output_dir = str(Path(input_path).parent)
                else:
                    raise ReachError(f"{input_path} frontmatter missing 'id' or 'planet_class'", fix="re-scaffold the system: reach atlas planet scaffold <system index>")
            else:
                raise ReachError(f"{input_path} has no YAML frontmatter", fix="re-scaffold the system: reach atlas planet scaffold <system index>")
        else:
            raise ReachError(f"unrecognized input format: {input_path}", fix="pass a body definition .json or a body index.md")
    else:
        parser.error("Provide a body_def (.json or .md) or --system index.md")

    # ── Generate ─────────────────────────────────────────────────────────
    t_total = time.time()
    for bd in body_defs:
        _generate_body(bd, hmap_w, hmap_h, args.globe_size,
                       args.render_mode, args.output_dir, chrome=args.chrome)

    elapsed = time.time() - t_total
    console.event(f"  All done: {len(body_defs)} bodies in {elapsed:.1f}s")


if __name__ == "__main__":
    main()