#!/usr/bin/env python3
"""
import_province_boundaries.py — Pre-compute province boundaries from watershed analysis.

For each inhabited body with a heightmap row in atlas_body_heightmaps, runs D8
drainage analysis to derive drainage basin boundaries and stores them as pixel-space
polylines in atlas_province_boundaries (D-205, D-208, #907).

Algorithm:
  1. Load float32 elevation BLOB from atlas_body_heightmaps.
  2. Depression-fill: raise sinks to the lowest-outlet neighbor (iterative).
  3. D8 flow direction: assign each cell to its steepest-descent neighbor.
  4. Flow accumulation: upstream cell count per cell (topological sort).
  5. Basin labeling: seed a basin per pour-point (flow-accumulation > threshold);
     flood-fill remaining cells following flow direction.
  6. Merge small basins (< 2% area) into the largest adjacent basin.
  7. Clamp basin count to [4, 12] by iterative merging of smallest basins.
  8. Trace boundary polylines between adjacent basins.
  9. Upsert rows into atlas_province_boundaries.

Province count target: 4–12 per body (D-205). Bodies with low relief get fewer,
larger provinces; high-relief worlds get more.

Performance: ~3–5s per body on a single CPU core at canonical 512×256 resolution.
The bottleneck is the pure-Python depression-fill + flow-direction scan (O(H×W) each,
~131k cells). For a full run of ~270 inhabited bodies expect ~15–20 minutes.
Hot loops (_depression_fill, _flow_direction) are candidates for NumPy vectorization
if build time becomes a bottleneck; the current scalar implementation is correct
and deterministic, which takes priority at this stage.

Incremental: bodies that already have rows in atlas_province_boundaries are skipped
unless --force is passed.

Usage:
    tooling/planet-gen/import_province_boundaries.py
    tooling/planet-gen/import_province_boundaries.py --body GJ380c
    tooling/planet-gen/import_province_boundaries.py --force
    tooling/planet-gen/import_province_boundaries.py --dry-run

Exit codes:
    0  completed
    1  fatal error (missing DB, schema error)
"""

import argparse
import json
import sys
import time
from pathlib import Path

TOOLING_DIR = Path(__file__).resolve().parent
REPO_ROOT = (TOOLING_DIR / ".." / "..").resolve()

_venv_python = REPO_ROOT / ".venv" / "bin" / "python"
if _venv_python.exists() and Path(sys.executable).resolve() != _venv_python.resolve():
    import os
    os.execv(str(_venv_python), [str(_venv_python)] + sys.argv)

import numpy as np
import sqlite3

from atlas_common import (
    DB_PATH,
    ensure_atlas_schema,
    query_inhabited_bodies,
)

# D8 neighbor offsets: (dr, dc)
_D8 = [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]

# River threshold from D-208: cells with flow_accumulation > 200 are river cells.
# Province seeds are local flow-accumulation maxima (watershed pour points).
_FLOW_THRESHOLD = 200

# Minimum basin area as fraction of total cells before merging into neighbor.
_MIN_BASIN_FRAC = 0.02

_PROVINCE_MIN = 4
_PROVINCE_MAX = 12


def _load_elevation(body_id: str, conn: sqlite3.Connection) -> np.ndarray | None:
    """Load float32 LE elevation BLOB from atlas_body_heightmaps."""
    row = conn.execute(
        "SELECT data, width, height FROM atlas_body_heightmaps WHERE body_id = ?",
        (body_id,),
    ).fetchone()
    if not row:
        return None
    data, width, height = row
    arr = np.frombuffer(data, dtype="<f4").reshape((height, width))
    return arr.astype(np.float32)


def _depression_fill(elev: np.ndarray) -> np.ndarray:
    """Simple iterative depression fill: raise sinks to their lowest outlet.

    Uses a shallow iterative pass — good enough for province-scale basins on
    512×256 grids. Not full priority-flood (which is O(N log N)); this O(N·k)
    approach converges in ≤10 passes on real heightmaps.
    """
    H, W = elev.shape
    filled = elev.copy()
    for _ in range(10):
        changed = False
        for r in range(1, H - 1):
            for c in range(W):
                nbr_min = float("inf")
                for dr, dc in _D8:
                    nr = r + dr
                    nc = (c + dc) % W
                    if 0 <= nr < H:
                        nbr_min = min(nbr_min, filled[nr, nc])
                if filled[r, c] < nbr_min:
                    filled[r, c] = nbr_min + 1e-6
                    changed = True
        if not changed:
            break
    return filled


def _flow_direction(filled: np.ndarray) -> np.ndarray:
    """D8 flow direction: index into _D8 (0–7), or -1 for no outflow (edge/flat)."""
    H, W = filled.shape
    fdir = np.full((H, W), -1, dtype=np.int8)
    for r in range(H):
        for c in range(W):
            best_drop = 0.0
            best_k = -1
            for k, (dr, dc) in enumerate(_D8):
                nr = r + dr
                nc = (c + dc) % W
                if nr < 0 or nr >= H:
                    continue
                drop = filled[r, c] - filled[nr, nc]
                if drop > best_drop:
                    best_drop = drop
                    best_k = k
            fdir[r, c] = best_k
    return fdir


def _flow_accumulation(fdir: np.ndarray) -> np.ndarray:
    """Flow accumulation via topological sort of the D8 DAG."""
    H, W = fdir.shape
    in_degree = np.zeros((H, W), dtype=np.int32)

    for r in range(H):
        for c in range(W):
            k = int(fdir[r, c])
            if k < 0:
                continue
            dr, dc = _D8[k]
            nr = r + dr
            nc = (c + dc) % W
            if 0 <= nr < H:
                in_degree[nr, nc] += 1

    from collections import deque
    queue = deque()
    for r in range(H):
        for c in range(W):
            if in_degree[r, c] == 0:
                queue.append((r, c))

    accum = np.ones((H, W), dtype=np.int32)
    while queue:
        r, c = queue.popleft()
        k = int(fdir[r, c])
        if k < 0:
            continue
        dr, dc = _D8[k]
        nr = r + dr
        nc = (c + dc) % W
        if 0 <= nr < H:
            accum[nr, nc] += accum[r, c]
            in_degree[nr, nc] -= 1
            if in_degree[nr, nc] == 0:
                queue.append((nr, nc))

    return accum


def _label_basins(fdir: np.ndarray, accum: np.ndarray) -> np.ndarray:
    """Label each cell with a basin ID via pour-point flood fill.

    Pour points are local flow-accumulation maxima above the river threshold.
    Each pour point seeds a basin; remaining cells are labeled by tracing
    flow direction back to their pour-point seed.
    """
    H, W = fdir.shape
    labels = np.full((H, W), -1, dtype=np.int32)

    # Seed one label per local accum maximum above threshold.
    # Use a simple scan: a cell is a local maximum if no neighbor has higher accum.
    pour_pts: list[tuple[int, int]] = []
    for r in range(H):
        for c in range(W):
            if accum[r, c] <= _FLOW_THRESHOLD:
                continue
            is_max = True
            for dr, dc in _D8:
                nr = r + dr
                nc = (c + dc) % W
                if 0 <= nr < H and accum[nr, nc] > accum[r, c]:
                    is_max = False
                    break
            if is_max:
                pour_pts.append((r, c))

    # If no pour points (e.g. flat/ocean world), create a single basin.
    if not pour_pts:
        labels[:] = 0
        return labels

    for basin_id, (r, c) in enumerate(pour_pts):
        labels[r, c] = basin_id

    # BFS flood: for each unlabeled cell, follow flow direction until a labeled
    # cell is reached; assign that label back along the path.

    def _trace(r0: int, c0: int) -> int:
        path: list[tuple[int, int]] = []
        r, c = r0, c0
        for _ in range(H * W):
            if labels[r, c] >= 0:
                lbl = labels[r, c]
                for pr, pc in path:
                    labels[pr, pc] = lbl
                return lbl
            path.append((r, c))
            k = int(fdir[r, c])
            if k < 0:
                # No outflow — assign basin 0
                lbl = 0
                for pr, pc in path:
                    labels[pr, pc] = lbl
                return lbl
            dr, dc = _D8[k]
            nr = r + dr
            nc = (c + dc) % W
            if nr < 0 or nr >= H:
                lbl = 0
                for pr, pc in path:
                    labels[pr, pc] = lbl
                return lbl
            r, c = nr, nc
        # Cycle guard
        lbl = 0
        for pr, pc in path:
            labels[pr, pc] = lbl
        return lbl

    for r in range(H):
        for c in range(W):
            if labels[r, c] < 0:
                _trace(r, c)

    return labels


def _merge_small_basins(
    labels: np.ndarray, target_min: int, target_max: int
) -> np.ndarray:
    """Merge tiny basins into their largest neighbor until count is in [target_min, target_max]."""
    H, W = labels.shape
    labels = labels.copy()

    def _basin_sizes() -> dict[int, int]:
        ids, counts = np.unique(labels, return_counts=True)
        return dict(zip(ids.tolist(), counts.tolist()))

    def _neighbors(basin_id: int) -> set[int]:
        mask = labels == basin_id
        # Dilate mask by 1 pixel in each direction, find adjacent basin IDs.
        nbrs: set[int] = set()
        rs, cs = np.where(mask)
        for r, c in zip(rs.tolist(), cs.tolist()):
            for dr, dc in _D8:
                nr = r + dr
                nc = (c + dc) % W
                if 0 <= nr < H:
                    nbr_id = int(labels[nr, nc])
                    if nbr_id != basin_id:
                        nbrs.add(nbr_id)
        return nbrs

    total = H * W
    for _ in range(200):
        sizes = _basin_sizes()
        n_basins = len(sizes)
        if n_basins <= target_max and all(
            v / total >= _MIN_BASIN_FRAC for v in sizes.values()
        ):
            break
        if n_basins <= target_min:
            break

        # Find the smallest basin
        smallest_id = min(sizes, key=lambda b: sizes[b])
        smallest_frac = sizes[smallest_id] / total

        if n_basins <= target_max and smallest_frac >= _MIN_BASIN_FRAC:
            break

        # Merge into its largest neighbor
        nbrs = _neighbors(smallest_id)
        if not nbrs:
            break
        merge_into = max(nbrs, key=lambda b: sizes.get(b, 0))
        labels[labels == smallest_id] = merge_into

    # Re-number contiguously from 0
    unique_ids = sorted(np.unique(labels).tolist())
    remap = {old: new for new, old in enumerate(unique_ids)}
    new_labels = np.zeros_like(labels)
    for old, new in remap.items():
        new_labels[labels == old] = new
    return new_labels


def _trace_boundary(labels: np.ndarray, basin_id: int) -> list[list[int]]:
    """Trace the outer boundary of a basin as a pixel-space polyline.

    Returns a list of [row, col] points forming the boundary polygon.
    Uses a simple contour walk: find all boundary cells (cells adjacent to a
    different basin), then sort them by angle from centroid to approximate a
    closed polygon.
    """
    H, W = labels.shape
    mask = labels == basin_id

    # Boundary cells: in this basin AND adjacent to a different basin
    boundary: list[tuple[int, int]] = []
    rs, cs = np.where(mask)
    for r, c in zip(rs.tolist(), cs.tolist()):
        on_boundary = False
        for dr, dc in _D8:
            nr = r + dr
            nc = (c + dc) % W
            if nr < 0 or nr >= H:
                on_boundary = True
                break
            if labels[nr, nc] != basin_id:
                on_boundary = True
                break
        if on_boundary:
            boundary.append((r, c))

    if not boundary:
        return []

    # Sort by angle from centroid — produces a rough polygon outline.
    arr = np.array(boundary, dtype=np.float32)
    centroid_r = float(np.mean(arr[:, 0]))
    centroid_c = float(np.mean(arr[:, 1]))
    angles = np.arctan2(arr[:, 0] - centroid_r, arr[:, 1] - centroid_c)
    order = np.argsort(angles)

    # Subsample if very large — keep at most 500 points for storage efficiency.
    pts = [boundary[i] for i in order.tolist()]
    if len(pts) > 500:
        step = len(pts) // 500
        pts = pts[::step]

    return [[r, c] for r, c in pts]


def compute_province_boundaries(
    body_id: str, elevation: np.ndarray
) -> list[dict]:
    """Run full watershed analysis; return list of basin dicts.

    Each dict:
        basin_id: int
        path:     JSON-serialisable [[row, col], ...]
        area_pct: float
    """
    H, W = elevation.shape
    total_cells = H * W

    filled = _depression_fill(elevation)
    fdir = _flow_direction(filled)
    accum = _flow_accumulation(fdir)
    labels = _label_basins(fdir, accum)
    labels = _merge_small_basins(labels, _PROVINCE_MIN, _PROVINCE_MAX)

    unique_ids = sorted(np.unique(labels).tolist())
    basins = []
    for basin_id in unique_ids:
        count = int(np.sum(labels == basin_id))
        area_pct = count / total_cells
        path = _trace_boundary(labels, basin_id)
        if not path:
            continue
        basins.append({
            "basin_id": basin_id,
            "path": path,
            "area_pct": area_pct,
        })

    return basins


def import_body_provinces(
    body_id: str,
    conn: sqlite3.Connection,
    force: bool,
    dry_run: bool,
    verbose: bool,
) -> dict:
    """Import province boundary rows for one body.

    Returns dict:
        status:   'imported' | 'skipped' | 'no_heightmap' | 'error'
        imported: count of basins written
        message:  detail on error/skip
    """
    if not force:
        existing = conn.execute(
            "SELECT COUNT(*) FROM atlas_province_boundaries WHERE body_id = ?",
            (body_id,),
        ).fetchone()[0]
        if existing > 0:
            return {"status": "skipped", "imported": 0,
                    "message": f"already has {existing} rows"}

    elevation = _load_elevation(body_id, conn)
    if elevation is None:
        return {"status": "no_heightmap", "imported": 0,
                "message": "no row in atlas_body_heightmaps"}

    try:
        basins = compute_province_boundaries(body_id, elevation)
    except Exception as exc:
        return {"status": "error", "imported": 0, "message": str(exc)}

    if not basins:
        return {"status": "error", "imported": 0,
                "message": "no basins produced from watershed analysis"}

    if verbose:
        areas = [f"{b['basin_id']}:{b['area_pct']:.1%}" for b in basins]
        print(f"    {body_id}: {len(basins)} basins — {', '.join(areas)}")

    if not dry_run:
        with conn:  # per-body transaction (BEGIN/COMMIT): rolls back this body on exception, keeps prior commits
            conn.execute(
                "DELETE FROM atlas_province_boundaries WHERE body_id = ?",
                (body_id,),
            )
            for b in basins:
                conn.execute(
                    """INSERT INTO atlas_province_boundaries
                           (body_id, basin_id, path, area_pct)
                       VALUES (?, ?, ?, ?)""",
                    (body_id, b["basin_id"], json.dumps(b["path"]), b["area_pct"]),
                )

    return {"status": "imported", "imported": len(basins)}


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Pre-compute province boundaries from watershed analysis (D-205, #907)"
    )
    parser.add_argument("--db", default=str(DB_PATH), help="Path to systems.db")
    parser.add_argument("--body", help="Process only this body_id")
    parser.add_argument("--force", action="store_true",
                        help="Re-import even if rows already exist")
    parser.add_argument("--dry-run", action="store_true",
                        help="Analyse without writing to DB")
    parser.add_argument("--verbose", action="store_true",
                        help="Print per-body detail")
    args = parser.parse_args()

    db_path = Path(args.db)
    if not db_path.exists():
        print(f"error: {db_path} not found", file=sys.stderr)
        sys.exit(1)

    print("\n  Province Boundary Import (#907)")
    print(f"  DB: {db_path}")
    if args.dry_run:
        print("  Mode: DRY RUN (no DB writes)")
    if args.force:
        print("  Force: enabled (will overwrite existing rows)")
    print()

    conn = sqlite3.connect(str(db_path))
    conn.execute("PRAGMA foreign_keys=ON")
    ensure_atlas_schema(conn)

    bodies = query_inhabited_bodies(conn)
    if args.body:
        bodies = [b for b in bodies if b["body_id"] == args.body]
        if not bodies:
            print(f"error: body '{args.body}' not found or has no terrain_reference",
                  file=sys.stderr)
            conn.close()
            sys.exit(1)

    print(f"  {len(bodies)} inhabited bodies with terrain_reference\n")

    t_total = time.time()
    n_imported = 0
    n_skipped = 0
    n_no_hmap = 0
    n_errors = 0

    for i, body_info in enumerate(bodies):
        body_id = body_info["body_id"]
        t0 = time.time()

        result = import_body_provinces(body_id, conn, args.force, args.dry_run, args.verbose)
        elapsed = time.time() - t0
        status = result["status"]

        if status == "imported":
            n_imported += 1
            print(f"  [{i+1}/{len(bodies)}] {body_id:20s} {result['imported']} basins ({elapsed:.1f}s)")
        elif status == "skipped":
            n_skipped += 1
            if args.verbose:
                print(f"  [{i+1}/{len(bodies)}] {body_id:20s} skipped ({result['message']})")
        elif status == "no_heightmap":
            n_no_hmap += 1
            if args.verbose:
                print(f"  [{i+1}/{len(bodies)}] {body_id:20s} no heightmap — skipping")
        elif status == "error":
            n_errors += 1
            print(f"  [{i+1}/{len(bodies)}] {body_id:20s} ERROR: {result.get('message', '')}")

    conn.close()

    elapsed_total = time.time() - t_total
    print(f"\n  Done in {elapsed_total:.1f}s")
    print(f"  imported={n_imported}  skipped={n_skipped}  "
          f"no_heightmap={n_no_hmap}  errors={n_errors}")

    if n_errors > 0:
        print(f"\n  {n_errors} error(s) — check output above", file=sys.stderr)


if __name__ == "__main__":
    main()