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

#!/usr/bin/env python3
"""
sol_import.py — Import real-world data for the Sol system (GJ-0).

Produces the same output format as generate.py (heightmap.png, globe.png,
markers.json, terrain.npz) by constructing terrain dicts from real
planetary science data instead of procedural simulation.

Usage:
    reach atlas planet sol-import                        # All Sol bodies
    reach atlas planet sol-import --body GJ0d            # Earth only
    reach atlas planet sol-import --body GJ0d --body GJ0e  # Earth + Mars
    reach atlas planet sol-import --download-only        # Fetch data, skip rendering
    reach atlas planet sol-import --heightmap-size 2048x1024 --globe-size 1024

Data is cached in tooling/domains/atlas/planet/sol_data/.cache/ after first download.
"""

import argparse
import json
import time

from pathlib import Path

import numpy as np

from tooling.core import config, console
from tooling.core.errors import ReachError
from tooling.domains.atlas.planet.generate import _build_markers
from tooling.domains.atlas.planet.planet_simulation import simulate
from tooling.domains.atlas.planet.render_heightmap import render_heightmap

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.


# Per-body importers (lazy-loaded)
SOL_INDEX = WORKTREE_ROOT / "wiki" / "star-systems" / "GJ-0" / "index.md"
SOL_OVERRIDES = PLANET_DIR / "sol_overrides.json"
SOL_BODIES_DIR = WORKTREE_ROOT / "wiki" / "star-systems" / "GJ-0" / "bodies"
SOL_MARKERS_DIR = PLANET_DIR / "sol_markers"

# Bodies that use real-world data (keyed by body_id → importer module)
REAL_DATA_BODIES = {
    "GJ0b":   "mercury",
    "GJ0c":   "venus",
    "GJ0d":   "earth",
    "GJ0d-1": "luna",
    "GJ0e":   "mars",
    "GJ0f-1": "io_moon",
    "GJ0f-2": "ice_moons",
    "GJ0f-3": "ice_moons",
    "GJ0f-4": "ice_moons",
    "GJ0g-1": "titan",
    "GJ0g-2": "ice_moons",
}

# Bodies that fall through to procedural simulation
PROCEDURAL_BODIES = {"GJ0e-1", "GJ0e-2"}

# Non-renderable body types
SKIP_TYPES = {"asteroid_belt", "oort_cloud"}


def _load_importer(module_name: str):
    """Lazy-import a sol_data.* module."""
    import importlib
    return importlib.import_module(f"tooling.domains.atlas.planet.sol_data.{module_name}")


def _apply_named_features(markers: dict, body_id: str) -> dict:
    """Overlay named features from sol_markers/ onto auto-detected markers."""
    features_map = {
        "GJ0d":   "earth_features.json",
        "GJ0e":   "mars_features.json",
        "GJ0d-1": "luna_features.json",
    }
    outer_bodies = {"GJ0f-1", "GJ0f-2", "GJ0f-3", "GJ0f-4",
                    "GJ0g-1", "GJ0g-2"}

    filename = features_map.get(body_id)
    if not filename and body_id in outer_bodies:
        filename = "outer_features.json"

    if not filename:
        return markers

    features_path = SOL_MARKERS_DIR / filename
    if not features_path.exists():
        return markers

    with open(features_path) as f:
        features = json.load(f)

    body_features = features.get(body_id, features)

    # Name auto-detected oceans by matching center coordinates
    if "oceans" in body_features:
        for named_ocean in body_features["oceans"]:
            best_match = None
            best_dist = float("inf")
            nc = named_ocean["center"]
            for detected in markers["oceans"]:
                dc = detected["center"]
                dist = (dc[0] - nc[0])**2 + (dc[1] - nc[1])**2
                if dist < best_dist:
                    best_dist = dist
                    best_match = detected
            if best_match and best_dist < 2500:  # within ~50 cells
                best_match["name"] = named_ocean["name"]

    # Name auto-detected mountain ranges by matching peak coordinates
    if "mountain_ranges" in body_features:
        for named_range in body_features["mountain_ranges"]:
            best_match = None
            best_dist = float("inf")
            nc = named_range.get("peak", named_range.get("center", [0, 0]))
            for detected in markers["mountain_ranges"]:
                dp = detected.get("peak", detected.get("center", [0, 0]))
                dist = (dp[0] - nc[0])**2 + (dp[1] - nc[1])**2
                if dist < best_dist:
                    best_dist = dist
                    best_match = detected
            if best_match and best_dist < 1600:  # within ~40 cells
                best_match["name"] = named_range["name"]

    # Name rivers by matching start/end coordinates
    if "rivers" in body_features:
        for named_river in body_features["rivers"]:
            best_match = None
            best_dist = float("inf")
            nc = named_river.get("mouth", named_river.get("center", [0, 0]))
            for detected in markers["rivers"]:
                if not detected["path"]:
                    continue
                # Check last point (mouth) of river path
                dp = detected["path"][-1]
                dist = (dp[0] - nc[0])**2 + (dp[1] - nc[1])**2
                if dist < best_dist:
                    best_dist = dist
                    best_match = detected
            if best_match and best_dist < 900:
                best_match["name"] = named_river["name"]

    # Add cities as POIs
    if "cities" in body_features:
        for city in body_features["cities"]:
            markers["cities"].append({
                "id": f"city_{city['name'].lower().replace(' ', '_')}",
                "name": city["name"],
                "center": city["center"],
                "population": city.get("population"),
            })

    # Add POIs
    if "pois" in body_features:
        for poi in body_features["pois"]:
            markers["pois"].append(poi)

    return markers


def _generate_body(body_def: dict, hmap_w: int, hmap_h: int,
                   globe_size: int, render_mode: str, output_dir: Path,
                   download_only: bool = False):
    """Generate all outputs for a single Sol body."""
    body_id = body_def["id"]
    body_type = body_def.get("body_type", "planet")
    planet_class = body_def.get("planet_class", "unknown")
    name = body_def.get("name") or body_id

    # Skip non-renderable types
    if body_type in SKIP_TYPES:
        console.event(f"  {body_id} ({name}) — skipped ({body_type})")
        return

    body_dir = output_dir / body_id
    body_dir.mkdir(parents=True, exist_ok=True)

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

    t0 = time.time()

    # ── 1. Build terrain ────────────────────────────────────────────────
    terrain = {}
    is_gas = planet_class in ("gas_giant",) or body_type == "gas_giant"

    if is_gas:
        # Gas giants: no terrain, renderer handles bands procedurally
        terrain = {}
        console.event("    terrain: gas giant (procedural bands)")
    elif body_id in REAL_DATA_BODIES:
        # Real-world data import
        module_name = REAL_DATA_BODIES[body_id]
        console.event(f"    importing real data via sol_data.{module_name}...")
        importer = _load_importer(module_name)
        terrain = importer.build_terrain(body_def)
        if download_only:
            console.event("    download complete, skipping render")
            return
    elif body_id in PROCEDURAL_BODIES:
        # Fall through to standard procedural simulation
        console.event("    procedural simulation (irregular body)...")
        terrain = simulate(body_def)
    else:
        console.event(f"no importer for {body_id}, using procedural", level="warn")
        terrain = simulate(body_def)

    t_terrain = time.time()

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

    # ── 2. Render heightmap ─────────────────────────────────────────────
    t_hmap = t_terrain
    if terrain:
        hmap_img = render_heightmap(body_def, terrain,
                                    out_w=hmap_w, out_h=hmap_h,
                                    render_mode=render_mode, chrome=False)
        hmap_img.save(str(body_dir / "heightmap.png"))
        t_hmap = time.time()
        console.event(f"    heightmap: {t_hmap - t_terrain:.1f}s  {hmap_w}x{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(str(body_dir / "globe.png"))
        t_globe = time.time()
        console.event(f"    globe: {t_globe - t_hmap:.1f}s  {globe_size}x{globe_size}")
    except Exception as e:
        console.event(f"    globe: FAILED — {e}")
        t_globe = time.time()

    # ── 4. Write data files ─────────────────────────────────────────────
    if terrain:
        # terrain.npz
        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(str(body_dir / "terrain.npz"), **save_dict)

        # markers.json — auto-detected + named features overlay
        markers = _build_markers(body_def, terrain)
        markers = _apply_named_features(markers, body_id)
        with open(body_dir / "markers.json", "w") as f:
            json.dump(markers, f, indent=2)

    # ── 5. Write index.md frontmatter ───────────────────────────────────
    _write_index_md(body_def, body_dir)

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


def _write_index_md(body_def: dict, body_dir: Path):
    """Write body index.md with YAML frontmatter."""
    import yaml

    # Strip internal fields
    bd = {k: v for k, v in body_def.items()
          if not k.startswith("_") and k != "wiki"}

    fm = yaml.dump(bd, default_flow_style=False, sort_keys=False,
                   allow_unicode=True)

    name = body_def.get("name") or body_def["id"]
    planet_class = body_def.get("planet_class", "unknown")
    system_link = "[GJ-0](../../index.md)"

    md = f"""---
{fm.rstrip()}
---

# {name}

{planet_class.replace('_', ' ').title()} {'planet' if body_def.get('body_type') == 'planet' else body_def.get('body_type', 'body')}.

**System:** {system_link}

## Visual

![Heightmap](heightmap.png)

![Globe](globe.png)
"""
    with open(body_dir / "index.md", "w") as f:
        f.write(md)


def main(argv: list[str] | None = None):
    parser = argparse.ArgumentParser(
        description="Sol system (GJ-0) real-world terrain importer")

    parser.add_argument("--body", action="append", default=None,
                        help="Specific body ID(s) to generate (repeatable)")
    parser.add_argument("--download-only", action="store_true",
                        help="Download source data without rendering")
    parser.add_argument("--output-dir", default=None,
                        help="Override output directory")
    parser.add_argument("--heightmap-size", default="1024x512",
                        help="Heightmap resolution (WxH)")
    parser.add_argument("--globe-size", type=int, default=512,
                        help="Globe resolution (square)")
    parser.add_argument("--render-mode", choices=["cartographic", "photographic"],
                        default="cartographic")

    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")

    output_dir = Path(args.output_dir) if args.output_dir else SOL_BODIES_DIR

    # Parse body definitions from GJ-0 index.md
    from tooling.domains.atlas.planet.body_definition_parser import parse_system

    overrides = {}
    if SOL_OVERRIDES.exists():
        with open(SOL_OVERRIDES) as f:
            overrides = json.load(f)

    body_defs = parse_system(str(SOL_INDEX), overrides=overrides)
    console.event(f"Sol system: {len(body_defs)} bodies parsed")

    # Filter to requested bodies
    if args.body:
        requested = set(args.body)
        body_defs = [bd for bd in body_defs if bd["id"] in requested]
        if not body_defs:
            raise ReachError(f"no matching bodies for {args.body}", fix="pass a Sol body id from wiki/star-systems/GJ-0/bodies/ (e.g. GJ0c), or omit --body")

    # Generate
    t_total = time.time()
    failed = []
    for bd in body_defs:
        try:
            _generate_body(bd, hmap_w, hmap_h, args.globe_size,
                           args.render_mode, output_dir,
                           download_only=args.download_only)
        except Exception as e:
            console.event(f"  FAILED: {bd['id']} — {e}")
            failed.append(bd["id"])

    elapsed = time.time() - t_total
    n_ok = len(body_defs) - len(failed)
    console.event(f"  Done: {n_ok}/{len(body_defs)} bodies in {elapsed:.1f}s")
    if failed:
        console.event(f"  Failed: {', '.join(failed)}")


if __name__ == "__main__":
    main()