"""Regional refinement (spec docs/superpowers/specs/2026-09-25-regional-refinement-design.md): inside drawn regions,
re-run erosion, rivers, lakes, ground and biomes on H3 cells two resolutions finer than the world, with the world
build as the boundary. Generated terrain (`idea`). Outside mapgen/ so the world build cache is unaffected."""
from __future__ import annotations

import copy
import hashlib
import json
import math
import os
import shutil
import tempfile
import threading
import time
from pathlib import Path
from types import SimpleNamespace

import h3.api.basic_int as h3
import numpy as np

import h3par

from mapgen import fields as FL, hydrology as HY, ice as IC
from mapgen.config import params
from mapgen.environment import DEFAULTS as ENV_DEFAULTS, ground, holdridge
from mapgen.graph import (OCEAN_MIN_KM2, accumulate, components, drop_smooth_cache, priority_flood, receiver_levels,
                          smooth_km, steepest_receivers)
from scipy.spatial import cKDTree

from mapgen.grid import Grid
from mapgen.sphere import latlon_to_xyz

import seafloor as SF
import servecache

MODEL = 9                      # bump when the refinement model changes (part of the cache key); 2: lake depth, shoulders;
                               # 3: halo outlets and walls; 4: 1 km hillslope (canyons stay narrow); 5: smooth bedrock
                               # means, connected sea, lake-aware incision, border report; 6: sea-floor pass (canyons,
                               # fans, ponds, volcanic fields, vents), zone fields from the world cells; 7: world lakes
                               # (deep, carved beds) fill to their surface and are ways out; no lake re-carving;
                               # 8: sea only where it reaches the open ocean (walled-off world ocean: lagoon lakes);
                               # 9: world lakes keep their world way out (no fine sill dams them above their level)
FINER = 2                      # fine resolution = world resolution + FINER (49 children per world cell)
LAPSE_C_PER_KM = 6.5
RELIEF_MIN_KM = 5.0            # procedural relief kept at wavelengths ≥ ≈ the fine cell spacing
COPY = ("continental", "age_class", "plate", "lithology", "landform", "seasonality", "P_ann", "P_jun", "P_dec", "PET",
        "dist_ocean_km", "m_o2_zones", "m_gravity_zones", "deposits", "deposit_main")
LAPSE = ("T_mean", "T_min", "T_jun", "T_dec", "biotemp")



def load_regions(path: Path) -> list[dict]:
    path = Path(path)
    return json.loads(path.read_text()).get("regions", []) if path.exists() else []


def area_cells(regions: list[dict], fine_res: int) -> list[np.ndarray]:
    """Cells of the union of all outlines, split into connected areas: outlines that overlap or touch merge (each
    outline's own cells count as one piece). Arrays, not sets: a continent has millions of cells."""
    pieces, caps = [], []
    for r in regions:
        c = np.unique(np.array(h3.polygon_to_cells(h3.LatLngPoly([tuple(p) for p in r["outline"]]), fine_res),
                               dtype=np.uint64))
        if len(c):
            pieces.append(c)
            o = np.asarray(r["outline"], dtype=np.float64)
            v = latlon_to_xyz(o[:, 0], o[:, 1]).reshape(-1, 3)
            m = v.mean(axis=0)
            m = m / max(np.linalg.norm(m), 1e-12)
            caps.append((m, float(np.max(np.arccos(np.clip(v @ m, -1, 1))))))   # a cap holding the outline
    parent = list(range(len(pieces)))

    def find(k):
        while parent[k] != k:
            parent[k] = parent[parent[k]]
            k = parent[k]
        return k
    order = sorted(range(len(pieces)), key=lambda k: len(pieces[k]))
    for a_ in range(len(order)):
        for b_ in range(a_ + 1, len(order)):
            i_, j_ = order[a_], order[b_]
            (ci, ri), (cj, rj) = caps[i_], caps[j_]
            apart = math.acos(float(np.clip(ci @ cj, -1, 1))) > ri + rj + 0.01   # caps apart (+ ≈ 2 cells): no touch
            if not apart and find(i_) != find(j_) and _touch(pieces[i_], pieces[j_]):
                parent[find(i_)] = find(j_)
    groups = {}
    for k in range(len(pieces)):
        groups.setdefault(find(k), []).append(pieces[k])
    out = [np.unique(np.concatenate(g)) if len(g) > 1 else g[0] for g in groups.values()]
    return sorted(out, key=lambda a: int(a[0]))


def _touch(small: np.ndarray, big: np.ndarray) -> bool:
    """Whether two cell sets overlap or are neighbours (the smaller one's cells and their rings against the other)."""
    if np.isin(small, big, assume_unique=True).any():
        return True
    for c in small.tolist():                                      # only the edge of the smaller one can touch
        ring = np.array(h3.grid_ring(c, 1), dtype=np.uint64)
        k = np.searchsorted(big, ring)
        k[k >= len(big)] = 0
        if (big[k] == ring).any():
            return True
    return False


def landmass(a, lat: float, lon: float) -> np.ndarray:
    """Indices of the world's land cells connected to the land cell nearest (lat, lon) (a continent or island)."""
    ids = np.asarray(a["g_ids"])
    land = ~np.asarray(a["ocean"]).astype(bool)
    L = np.where(land)[0]
    start = L[cKDTree(np.asarray(a["g_xyz"])[L]).query(latlon_to_xyz(lat, lon).reshape(3))[1]]
    seen, stack = {int(start)}, [int(start)]
    while stack:
        for n in h3.grid_ring(int(ids[stack.pop()]), 1):
            j = int(np.searchsorted(ids, np.uint64(n)))
            if j < len(ids) and ids[j] == n and land[j] and j not in seen:
                seen.add(j)
                stack.append(j)
    return np.array(sorted(seen), dtype=np.int64)


def land_outline(a, lat: float, lon: float, step_deg: float = 0.5, margin: int = 1, max_points: int = 500) -> list:
    """A simple outline polygon [[lat, lon], …] around the landmass at (lat, lon): its cells on a step_deg grid,
    grown by `margin` grid cells (the coast and a strip of sea), holes filled, traced along the grid lines."""
    from scipy.ndimage import binary_dilation, binary_fill_holes, label
    comp = landmass(a, lat, lon)
    la = np.asarray(a["g_lat"])[comp]
    lo = np.asarray(a["g_lon"])[comp]
    c0 = float(np.degrees(np.arctan2(np.sin(np.radians(lo)).mean(), np.cos(np.radians(lo)).mean())))
    rel = (lo - c0 + 180.0) % 360.0 - 180.0                       # longitudes around the landmass's middle
    while True:
        r0, q0 = np.floor(la.min() / step_deg) - margin - 1, np.floor(rel.min() / step_deg) - margin - 1
        R = int(np.ceil(la.max() / step_deg) - r0 + margin + 2)
        Q = int(np.ceil(rel.max() / step_deg) - q0 + margin + 2)
        m = np.zeros((R, Q), bool)
        m[(np.floor(la / step_deg) - r0).astype(int), (np.floor(rel / step_deg) - q0).astype(int)] = True
        m = binary_fill_holes(binary_dilation(m, np.ones((3, 3), bool), iterations=margin) if margin else m)
        lab, _ = label(m)                                         # the piece holding the landmass
        m = lab == lab[int(np.floor(la[0] / step_deg) - r0), int(np.floor(rel[0] / step_deg) - q0)]
        while True:                                               # no cells touching at a corner only
            d1 = m[:-1, :-1] & m[1:, 1:] & ~m[:-1, 1:] & ~m[1:, :-1]
            d2 = m[:-1, 1:] & m[1:, :-1] & ~m[:-1, :-1] & ~m[1:, 1:]
            if not (d1.any() or d2.any()):
                break
            m[:-1, 1:] |= d1
            m[:-1, :-1] |= d2
            m = binary_fill_holes(m)
        nxt = {}                                                  # boundary edges, inside on the left
        for r, q in zip(*np.nonzero(m)):
            if r == 0 or not m[r - 1, q]:
                nxt[(r, q + 1)] = (r, q)
            if r == R - 1 or not m[r + 1, q]:
                nxt[(r + 1, q)] = (r + 1, q + 1)
            if q == 0 or not m[r, q - 1]:
                nxt[(r, q)] = (r + 1, q)
            if q == Q - 1 or not m[r, q + 1]:
                nxt[(r + 1, q + 1)] = (r, q + 1)
        start = min(nxt)
        loop, v = [start], nxt[start]
        while v != start:
            loop.append(v)
            v = nxt[v]
        pts = [p for k, p in enumerate(loop)                      # corners only
               if (loop[k - 1][0] - p[0], loop[k - 1][1] - p[1]) != (p[0] - loop[(k + 1) % len(loop)][0],
                                                                    p[1] - loop[(k + 1) % len(loop)][1])]
        if len(pts) <= max_points:
            return [[round(float((r + r0) * step_deg), 6), round(float((q + q0) * step_deg + c0 + 180.0) % 360.0 - 180.0, 6)]
                    for r, q in pts]
        step_deg *= 1.25


def plateau_outline(p: dict, seed: int, radius_km: float, margin_km: float = 300.0, n: int = 96) -> list:
    """An outline [[lat, lon], …] around a sunken plateau (mapgen.plateaus' noise-warped ellipse) grown by margin_km:
    per bearing, the farthest point still inside the plateau plus the margin (star-shaped: a simple polygon)."""
    from mapgen import plateaus as PL
    from mapgen.sphere import east_north, great_circle_point, xyz_to_latlon
    c = latlon_to_xyz(*p["center"])
    e, nn = east_north(c[None])
    a, _ = PL.semi_axes(p)
    steps = np.linspace(0.0, a * (1.0 + PL.EDGE_WARP) / (1.0 - PL.EDGE_WARP), 400)
    out = []
    for az in np.linspace(0.0, 2.0 * np.pi, n, endpoint=False):
        t = np.cos(az) * nn[0] + np.sin(az) * e[0]
        pts = np.cos(steps / radius_km)[:, None] * c + np.sin(steps / radius_km)[:, None] * t
        inside = np.flatnonzero(PL.rho(pts, p, seed, radius_km) < 1.0)
        r = float(steps[inside[-1]]) if len(inside) else 0.0
        la, lo = xyz_to_latlon(great_circle_point(c, t, r + margin_km, radius_km))
        out.append([round(float(la), 4), round(float(lo), 4)])
    return out


def region_key(region: dict) -> str:
    """A region's key in the build record: its id, or its outline for hand-written regions without one."""
    return str(region.get("id") or "outline-" + outline_hash(region))


def outline_hash(region: dict) -> str:
    return hashlib.sha1(json.dumps(region["outline"]).encode()).hexdigest()[:12]


def region_areas(regions: list[dict], areas: list[np.ndarray], fine_res: int) -> dict:
    """{region id: {"outline": outline hash, "area": hash of the area holding it (None: too small for a cell)}}."""
    out = {}
    for r in regions:
        cells = h3.polygon_to_cells(h3.LatLngPoly([tuple(p) for p in r["outline"]]), fine_res)
        c = np.uint64(cells[0]) if cells else None
        hit = None
        for a in areas:
            k = int(np.searchsorted(a, c)) if c is not None else len(a)
            if k < len(a) and a[k] == c:
                hit = area_hash(a)
                break
        out[region_key(r)] = {"outline": outline_hash(r), "area": hit}
    return out


def read_record(regions_root: Path, key: str):
    """What the last successful build made of which outline (build_areas writes it), or None."""
    try:
        rec = json.loads((results_dir(regions_root, key) / "outlines.json").read_text())
    except (OSError, ValueError):
        return None
    ok = isinstance(rec, dict) and isinstance(rec.get("regions"), dict) and all(
        isinstance(v, dict) and "outline" in v and "area" in v for v in rec["regions"].values())
    return rec if ok else None                                  # a damaged record counts as none (a rebuild)


def area_hash(cells: np.ndarray) -> str:
    return hashlib.sha1(np.asarray(cells, dtype=np.uint64).tobytes()).hexdigest()[:12]


AREAS = "areas"                # <regions_root>/areas/<key>/: results by content, shared between worlds
KEY_SKIP = ("era_mask",)       # era bookkeeping (and zone_* weights), not terrain: equal ground, equal key
KEY_RINGS = 3                  # world cells around an area whose values its refinement reads (relief raster,
                               # mean correction, river crossings)


def areas_dir(regions_root: Path) -> Path:
    return Path(regions_root) / AREAS


def world_dirs(root: Path, res: int, log=print) -> list:
    """(era, world folder) of every world the viewer can show: the base build first (named after the first era
    without events, else "base"), then each era with events that has been built (out/r<res>/eras/<era>/)."""
    from mapgen import config as C, eras as ER
    root = Path(root)
    _, tect = C.load(root)
    order = (tect.get("eras") or {}).get("order", [])
    plain = [n for n in order if not C.era_events(tect, n)]
    out = [(plain[0] if plain else "base", root / "out" / f"r{res}")]
    for n in order:
        if not C.era_events(tect, n):
            continue
        d = ER.era_dir(root, res, n)
        if (d / "cells.npz").exists():
            out.append((n, d))
        else:
            log(f"refine: era {n} is not built (run mapgen.py build): its areas wait")
    return out


def cell_parents(cells, res: int) -> np.ndarray:
    """h3.cell_to_parent for an array of cells (all of resolution ≥ res), by the index bits: the resolution field
    set to res, the digits below it set to 7 (unused)."""
    cells = np.asarray(cells, dtype=np.uint64)
    if len(cells) and int(((cells >> np.uint64(52)) & np.uint64(15)).min()) < res:
        raise ValueError(f"cell_parents: a cell is coarser than resolution {res}")
    unused = 0
    for r in range(res + 1, 16):
        unused |= 7 << ((15 - r) * 3)
    out = (cells & np.uint64(~(15 << 52) & (2**64 - 1))) | np.uint64(res << 52)
    return out | np.uint64(unused)


def world_near(world: "WorldCells", cells, rings: int = KEY_RINGS) -> np.ndarray:
    """Indices of the world cells under the area and `rings` rings around them."""
    par = np.unique(cell_parents(cells, world.res))
    near = np.unique(np.fromiter((n for c in par.tolist() for n in h3.grid_disk(c, rings)), dtype=np.uint64))
    k = np.minimum(np.searchsorted(world.ids, near), len(world.ids) - 1)
    return k[world.ids[k] == near].astype(np.int64)


def area_key(world: "WorldCells", cells, cfg: dict, plateaus: list) -> str:
    """Content key of an area's refinement: the model, its cells, the config and every world value it reads (the
    world cells under and around it, their river levels, the plateaus there). An era that left those values as they
    were gives the same key: the area is built once for both."""
    import tiles
    idx = world_near(world, cells)
    h = hashlib.sha256(f"m{MODEL}:t{tiles.VERSION}:f{FINER}".encode())
    h.update(np.asarray(cells, dtype=np.uint64).tobytes())
    h.update(json.dumps(cfg, sort_keys=True, default=str).encode())
    for k in sorted(world.a.keys()):
        if k in KEY_SKIP or k.startswith("zone_"):
            continue
        v = world.a.peek(k)
        if v.shape[:1] == world.ids.shape:
            h.update(k.encode())
            h.update(np.ascontiguousarray(v[idx]).tobytes())
    h.update(np.ascontiguousarray(world.river_level[idx]).tobytes())
    pid = np.asarray(world.a.peek("plateau_id"))[idx] if "plateau_id" in world.a else np.zeros(0, np.int64)
    near = [int(i) for i in np.unique(pid) if 0 <= i < len(plateaus)]
    h.update(json.dumps([plateaus[i] for i in near], sort_keys=True, default=str).encode())
    return h.hexdigest()[:16]


def _write_json(path: Path, obj) -> None:
    fd, tmp = tempfile.mkstemp(dir=path.parent, suffix=".json")
    with os.fdopen(fd, "w") as fh:
        json.dump(obj, fh)
    os.replace(tmp, path)


def _remove(p: Path) -> None:
    if p.is_symlink():
        p.unlink(missing_ok=True)
    elif p.is_dir():
        shutil.rmtree(p, ignore_errors=True)


def _link(set_dir: Path, name: str, target: Path) -> Path:
    """set_dir/name → target as a relative symlink, replaced atomically (an old real folder goes first)."""
    link = set_dir / name
    rel = os.path.relpath(target, set_dir)
    if link.is_symlink() and os.readlink(link) == rel:
        return link
    tmp = set_dir / f".{name}.link-{os.getpid()}"
    tmp.unlink(missing_ok=True)
    os.symlink(rel, tmp)
    if link.exists() and not link.is_symlink():
        shutil.rmtree(link)
    os.replace(tmp, link)
    return link


def _rings(ids) -> tuple[np.ndarray, np.ndarray]:
    """Neighbours of each cell (h3.grid_ring order), flat, with per-cell counts (pentagons have five)."""
    return h3par.run(h3par.rings, ids)


def subgrid(cells: np.ndarray, radius_km: float):
    """The area cells plus a one-cell halo ring, as a Grid whose neighbours stay inside that set (numpy: a continent's
    millions of cells must not become Python sets and lists)."""
    cells = np.asarray(cells, dtype=np.uint64)
    if len(cells) > 1 and not (cells[1:] > cells[:-1]).all():
        cells = np.unique(cells)
    nb_a, cnt_a = _rings(cells)
    u = np.unique(nb_a)
    p = np.minimum(np.searchsorted(cells, u), len(cells) - 1)
    ring = u[cells[p] != u]                                         # the halo: neighbours outside the area
    del u, p
    ids = np.sort(np.concatenate([cells, ring]), kind="stable")    # sorted: the area and its ring
    nb_r, cnt_r = _rings(ring)                                      # the area's own rings are known: only the ring's
    pc, pr = np.searchsorted(ids, cells), np.searchsorted(ids, ring)
    order = np.empty(len(ids), np.int64)                            # each id's ring, in ids order
    order[pc] = np.arange(len(cells))
    order[pr] = len(cells) + np.arange(len(ring))
    cnt = np.concatenate([cnt_a, cnt_r])
    start = np.concatenate([[0], np.cumsum(cnt)[:-1]])
    counts = cnt[order]
    pool = np.concatenate([nb_a, nb_r])
    del nb_a, nb_r
    first = np.repeat(start[order] - np.concatenate([[0], np.cumsum(counts)[:-1]]), counts)
    flat = pool[first + np.arange(int(counts.sum()))]
    del pool, first
    k = np.searchsorted(ids, flat)
    k[k >= len(ids)] = 0
    keep = ids[k] == flat                                          # neighbours inside the set (in ring order)
    per = np.add.reduceat(keep.astype(np.int64), np.concatenate([[0], np.cumsum(counts)[:-1]])) if len(ids) else counts
    ptr = np.concatenate([[0], np.cumsum(per)]).astype(np.int64)
    idx = k[keep].astype(np.int64)
    del flat, k, keep
    res = h3.get_resolution(int(ids[0]))
    ll, area = h3par.run(h3par.centres_areas, ids)
    g = Grid(res, radius_km, ids, ll[:, 0], ll[:, 1], latlon_to_xyz(ll[:, 0], ll[:, 1]), area * radius_km ** 2, ptr, idx)
    halo = np.zeros(len(ids), bool)
    halo[pr] = True
    return g, halo


class LazyArrays(dict):
    """An npz file's arrays, each read on first use (a refinement needs a few of the world's fields, not all)."""
    def __init__(self, path: Path):
        super().__init__()
        self._npz = np.load(path)
        self._keys = set(self._npz.files)

    def __missing__(self, k):
        if k not in self._keys:
            raise KeyError(k)
        v = self[k] = self._npz[k]
        return v

    def __contains__(self, k):
        return k in self._keys

    def keys(self):
        return self._keys

    def __iter__(self):
        return iter(self._keys)

    def __len__(self):
        return len(self._keys)

    def peek(self, k):
        """The array without keeping it (a key hashes every field once)."""
        return dict.__getitem__(self, k) if dict.__contains__(self, k) else self._npz[k]


class WorldCells:
    def __init__(self, out_dir: Path):
        self.a = LazyArrays(Path(out_dir) / "cells.npz")
        self.meta = json.loads((Path(out_dir) / "cells_meta.json").read_text())
        self.ids = self.a["g_ids"]
        self.res = h3.get_resolution(int(self.ids[0]))
        self._levels = None
        self._tree = None

    @property
    def tree(self) -> cKDTree:
        if self._tree is None:
            self._tree = cKDTree(np.asarray(self.a["g_xyz"], dtype=np.float64))
        return self._tree

    @property
    def river_level(self) -> np.ndarray:
        """The world rivers' graded water levels (rivers.py, spec §10): how deep a river may cut where it leaves."""
        if self._levels is None:
            import rivers as RV
            self._levels = RV.RiverNet(self.a, self.meta["legends"], float(self.meta["radius_km"])).level
        return self._levels

    def index(self, fine_ids) -> np.ndarray:
        return np.searchsorted(self.ids, cell_parents(fine_ids, self.res))


def initial_fields(g, halo, world: WorldCells, src) -> dict:
    p = world.index(g.ids)
    a = world.a
    z = src.relief_at(g.lat, g.lon, RELIEF_MIN_KM)
    dz_km = (z - a["z_surface_m"][p]) / 1000.0
    out = {"parent": p, "elevation_m": z, "ocean_world": a["ocean"][p].astype(bool)}
    for k in COPY:
        if k in a:
            out[k] = a[k][p]
    for k in LAPSE:
        out[k] = a[k][p] - LAPSE_C_PER_KM * dz_km
    out["T_range"] = a["T_range"][p]
    out["biotemp"] = np.clip(out["biotemp"], 0.0, 30.0)
    return out


EROSION = {"steps": 8, "dt_myr": 1.0, "k": 0.004, "m": 0.8, "hillslope_km": 1.0, "sediment_fill": 0.15,
           "area_per_km3": 2000.0}   # water flux → equivalent drainage area (0.5 m/yr runoff)
K_MULT = {0: 0.0, 1: 1.5, 2: 1.2, 3: 1.0, 4: 1.0, 5: 0.6, 6: 0.9, 7: 0.8}   # age class → erodibility (as the world)


def k_mult(age_class) -> np.ndarray:
    """K_MULT per cell (1.0 for classes it lacks)."""
    c = np.asarray(age_class)
    c = np.trunc(c).astype(np.int64) if c.dtype.kind == "f" else c.astype(np.int64)
    table = np.array([K_MULT.get(i, 1.0) for i in range(max(K_MULT) + 1)], dtype=np.float64)
    ok = (c >= 0) & (c < len(table))
    return np.where(ok, table[np.where(ok, c, 0)], 1.0)


def runoff_km3(g, f) -> np.ndarray:
    p, pet = f["P_ann"], np.maximum(f["PET"], 1e-6)
    aet = p / np.sqrt(1.0 + (p / pet) ** 2)                                 # Pike (1964), as the world
    return np.maximum(p - aet, 0.0) * g.area_km2 * 1e-6


def crossings(g, halo, world):
    """World river segments crossing the area outline: inflow (km³/yr) at the first inside cell, and the halo cells
    where they leave, with the world river's level there."""
    a = world.a
    riv = np.where(a["river"].astype(bool))[0]
    near = np.zeros(len(world.ids), bool)                # only segments near the area can cross its outline
    near[world_near(world, g.ids)] = True
    riv = riv[near[riv] | near[np.asarray(a["recv"])[riv]]]
    res = h3.get_resolution(int(g.ids[0]))
    pos = {int(c): k for k, c in enumerate(g.ids)}
    inflow, exits, levels = np.zeros(g.n), [], []
    for i in riv:
        j = int(a["recv"][i])
        if j == i:
            continue
        pa, pb = a["g_xyz"][i], a["g_xyz"][j]
        n = max(2, int(np.linalg.norm(pb - pa) * g.radius_km / 1.0))           # 1 km steps
        t = np.linspace(0, 1, n)[:, None]
        pts = pa * (1 - t) + pb * t
        pts /= np.linalg.norm(pts, axis=1, keepdims=True)
        lat, lon = np.degrees(np.arcsin(pts[:, 2])), np.degrees(np.arctan2(pts[:, 1], pts[:, 0]))
        k = [pos.get(h3.latlng_to_cell(float(la), float(lo), res), -1) for la, lo in zip(lat, lon)]
        inside = [kk >= 0 and not halo[kk] for kk in k]
        for s in range(1, len(k)):
            if not inside[s - 1] and inside[s]:
                inflow[k[s]] += float(a["discharge_km3_yr"][i])
            if inside[s - 1] and not inside[s] and k[s] >= 0:
                exits.append(k[s])
                lv = world.river_level
                li = lv[i] if np.isfinite(lv[i]) else a["z_surface_m"][i]
                lj = lv[j] if np.isfinite(lv[j]) else (0.0 if a["ocean"][j] else a["z_surface_m"][j])
                levels.append(float(min(li, lj)))
    return inflow, np.array(exits, dtype=np.int64), np.array(levels)


WALL_M = 1.0e6                 # halo cells that are no way out: walls for routing


def outlets(g, halo, world, parent, z=None) -> np.ndarray:
    """Halo cells where water may leave the area: where the world's own flow leaves it (a world cell with children
    inside drains into this halo cell's world cell), or the world sea. Elsewhere (e.g. upstream of an inflowing
    river) the halo is a wall, so inflows cross the area instead of turning back out."""
    a = world.a
    inside = np.unique(parent[~halo])
    targets = np.setdiff1d(np.asarray(a["recv"])[inside], inside)
    out = halo & (np.isin(parent, targets) | np.asarray(a["ocean"])[parent].astype(bool))
    if z is not None:   # and wherever the ground outside lies below the area cells next to it (water runs out there)
        inner_nb = ~halo[g.dst]
        low_in = np.full(g.n, np.inf)
        np.minimum.at(low_in, g.src[inner_nb], np.asarray(z)[g.dst[inner_nb]])
        out |= halo & (np.asarray(z) < low_in)
    if not out.any():                                   # a closed basin in the world: its lowest border cell drains
        k = np.where(halo)[0]
        out[k[np.argmin(np.asarray(a["z_surface_m"])[parent[k]])]] = True
    return out


def erode_area(g, z, halo, sea, water, kmult, P=EROSION, outlet=None, progress=None):
    """Stream-power incision (implicit, as the world build) driven by water flux, with the halo and the sea fixed;
    water leaves through the sea and the outlet halo cells only (default: the whole halo)."""
    z = np.asarray(z, dtype=np.float64).copy()
    fixed = halo | sea
    sinks = sea | (halo if outlet is None else outlet)
    wall = halo & ~sinks
    ar = np.arange(g.n)
    for step in range(int(P["steps"])):
        if progress:
            progress("erosion", 0.1 + 0.5 * step / int(P["steps"]))
        base = np.where(sea, 0.0, z)                          # rivers cut to sea level, never toward the seabed
        zf = priority_flood(g, np.where(wall, WALL_M, base), sinks)
        zb = np.where(fixed, base, z + P["sediment_fill"] * (zf - z))
        recv, _, dist = steepest_receivers(g, zf)
        recv = np.where(fixed, ar, recv)
        levels = receiver_levels(recv)
        q = accumulate(recv, levels, water)
        F = P["k"] * kmult * P["dt_myr"] * (q * P["area_per_km3"]) ** P["m"] / np.where(np.isfinite(dist), dist, 1.0)
        F[fixed] = 0.0
        floor = zb.copy()                                     # the level of the way out each cell drains to
        for lv in levels[1:]:
            floor[lv] = floor[recv[lv]]
        zn = zb.copy()
        for lv in levels[1:]:
            r = recv[lv]
            cut = np.minimum(zb[lv], (zb[lv] + F[lv] * zn[r]) / (1.0 + F[lv]))
            zn[lv] = np.maximum(cut, np.minimum(floor[lv], zb[lv]))   # never below it (hollows are no base level)
        zn = smooth_km(g, zn, P["hillslope_km"], keep=True)
        z = np.where(fixed, z, np.minimum(z, zn))
    drop_smooth_cache(g)
    return z


def mean_correction(xyz, z, parent, halo, target, world, iters=12, hold=None, memo=None) -> np.ndarray:
    """A smooth height correction (the compact 4-of-5 kernel over nearby world cell centres, evaluated at each fine
    cell) that makes each world cell's area children average the target (the world's bedrock): no block steps
    between world cells. Tapers to 0 over the two cell rows next to the halo (held at world values), so exits and
    the edge meet the world without a step. Cells in `hold` (e.g. the fixed sea) count in the means but are not moved.
    memo: a dict shared by calls on the same cells (the kernel and the taper are computed once)."""
    z = np.asarray(z, dtype=np.float64)
    inner = ~halo
    wx = np.asarray(world.a["g_xyz"], dtype=np.float64)
    if memo:
        near, idx, w, has, taper = memo["geometry"]
    else:
        pids = np.unique(parent)
        ring = world.tree.query(wx[pids], k=7)[1].ravel()               # the parents and their world neighbours
        near = np.unique(np.concatenate([pids, ring]))
        nw = max(1, h3par.workers())                                    # KD-tree queries on threads: same answers
        d, idx = cKDTree(wx[near]).query(np.asarray(xyz, dtype=np.float64), k=5, workers=nw)
        w = np.clip(1.0 - d[:, :4] / np.maximum(d[:, 4:5], 1e-12), 0.0, None) ** 2
        w = w / np.maximum(w.sum(1, keepdims=True), 1e-300)
        idx = idx[:, :4]
        has = np.bincount(parent[inner], minlength=len(wx)) > 0
        xyz = np.asarray(xyz, dtype=np.float64)
        step = np.median(cKDTree(xyz).query(xyz, k=2, workers=nw)[0][:, 1])   # the fine cell spacing
        taper = (np.clip((cKDTree(xyz[halo]).query(xyz, workers=nw)[0] / step - 1.0) / 2.0, 0.0, 1.0) if halo.any()
                 else 1.0)
        if memo is not None:
            memo["geometry"] = (near, idx.astype(np.int32), w, has, taper)
    corr = np.zeros(len(z))
    for _ in range(iters):
        zc = z + corr
        sums = np.bincount(parent[inner], weights=zc[inner], minlength=len(wx))
        cnt = np.bincount(parent[inner], minlength=len(wx))
        resid = np.where(has, np.asarray(target, dtype=np.float64) - sums / np.maximum(cnt, 1), 0.0)
        corr += taper * (resid[near][idx] * w).sum(1)
        corr[halo] = 0.0
        if hold is not None:
            corr[hold] = 0.0
    return corr


RIVER_MIN_KM3_YR = 0.2
LAKE_MIN_DEPTH_M = 15.0        # only hollows deeper than this hold lakes (5 km cells: shallow dips are noise)


def _p_for_runoff(target_mm, pet):
    """Precipitation (mm) whose Pike runoff is target_mm (bisection): lets the world hydrology code carry inflows."""
    lo, hi = np.zeros_like(target_mm), target_mm + pet * 4 + 10
    for _ in range(60):
        mid = (lo + hi) / 2
        run = mid - mid / np.sqrt(1 + (mid / pet) ** 2)
        lo, hi = np.where(run < target_mm, mid, lo), np.where(run < target_mm, hi, mid)
    return (lo + hi) / 2


def _sea(g, z, ocean_world, halo, min_km2: float = OCEAN_MIN_KM2) -> np.ndarray:
    """Sea: cells at or below 0 m under the world's ocean that reach it beyond the area (through the halo), or a
    separate basin as big as the world counts as sea (its own rule). World ocean that the fine heights wall off is
    no sea: a lagoon, which the lakes take (author 2026-10-03). Inland dry basins stay land."""
    wet = np.asarray(z) <= 0
    lab = components(g, wet)
    ow = wet & np.asarray(ocean_world)
    seas = np.unique(lab[ow & np.asarray(halo)])
    area = np.bincount(lab[ow], weights=np.asarray(g.area_km2)[ow], minlength=int(lab.max()) + 1 if lab.size else 0)
    seas = np.union1d(seas[seas >= 0], np.flatnonzero(area >= min_km2))
    return wet & np.isin(lab, seas)


OUTFLOW_TOL_M = 1.0            # a world lake may stand this much above its world level before its way out is cut
OUTFLOW_DROP_M = 0.5           # the cut way out starts this far below the world lake's level


def world_outflows(g, z, halo, sinks, parent, a) -> np.ndarray:
    """z with each world lake's way out kept: a lake that the fine relief dams above its world level (a sill the world
    cells' means hide) gets a channel cut along its world drainage path, down to where that path already drains lower
    (the sea, a lower lake, a way out of the area), never below sea level. The channel follows the least digging (Dijkstra on the height above
    the lake's level) and falls evenly from just below the lake's level. Lakes that drain low enough stay as they are."""
    import heapq
    z = np.asarray(z, dtype=np.float64).copy()
    sinks = np.asarray(sinks, bool)
    lake_w, ocean_w = np.asarray(a["lake"]).astype(bool), np.asarray(a["ocean"]).astype(bool)
    lid, lev_w, recv = np.asarray(a["lake_id"]), np.asarray(a["lake_level_m"], np.float64), np.asarray(a["recv"])
    inner_lake = ~halo & lake_w[parent]
    ids = np.unique(lid[parent[inner_lake]])
    ids = ids[ids >= 0]
    if not ids.size:
        return z
    inside_w = np.zeros(len(lake_w), bool)
    inside_w[parent[~halo]] = True
    levels = {int(L): float(np.nanmax(lev_w[lid == L])) for L in ids}
    zf = None
    for L in sorted(levels, key=levels.get):
        lvl = levels[L]
        if not np.isfinite(lvl):
            continue
        own = lid == L
        water = inner_lake & own[parent] & (z < lvl)
        if not water.any():
            continue
        if zf is None:
            zf = priority_flood(g, z, sinks)
        if zf[water].min() <= lvl + OUTFLOW_TOL_M:
            continue
        exits = np.flatnonzero(own & ~own[recv])                          # the world path: lake exit → sea or lower
        path = set(exits.tolist())
        for c in exits:
            c = int(recv[c])
            while c not in path and inside_w[c]:
                path.add(c)
                if ocean_w[c] or (lake_w[c] and lid[c] != L) or recv[c] == c:
                    break
                c = int(recv[c])
        corridor = np.isin(parent, list(path)) & (~halo | sinks)
        seeds = np.flatnonzero(corridor & (zf <= lvl - OUTFLOW_DROP_M) & ~(own[parent] & (z < lvl)))
        goal = corridor & own[parent] & (z < lvl)
        if not seeds.size or not goal.any():
            continue
        cost = np.full(g.n, np.inf)
        prev = np.full(g.n, -1, np.int64)
        cost[seeds] = 0.0
        heap = [(0.0, int(s)) for s in seeds]
        heapq.heapify(heap)
        hit = -1
        while heap:
            c0, i = heapq.heappop(heap)
            if c0 > cost[i]:
                continue
            if goal[i]:
                hit = i
                break
            for j in g.nbr_idx[g.nbr_ptr[i]:g.nbr_ptr[i + 1]]:
                if not corridor[j]:
                    continue
                cj = c0 + max(z[j] - lvl, 0.0) + 1e-3
                if cj < cost[j]:
                    cost[j], prev[j] = cj, i
                    heapq.heappush(heap, (cj, int(j)))
        if hit < 0:
            continue
        chain = [hit]                                                     # lake → … → a cell that drains lower
        while prev[chain[-1]] >= 0:
            chain.append(int(prev[chain[-1]]))
        chain = np.array(chain[1:], dtype=np.int64)
        if not chain.size:
            continue
        top = lvl - OUTFLOW_DROP_M
        end = min(max(float(zf[chain[-1]]), 0.0), top)                    # the sea's surface, not its bed
        z[chain] = np.minimum(z[chain], top + (end - top) * np.arange(1, len(chain) + 1) / len(chain))
        zf = None
    return z


def _lake_water(g, z, parent, world, halo) -> tuple[np.ndarray, np.ndarray]:
    """(water, surface): cells below a world lake's surface connected to its cells (as _sea for the world ocean), and
    that lake's surface (m; NaN elsewhere). The world's sea is never lake water."""
    a = world.a
    sea = _sea(g, z, np.asarray(a["ocean"]).astype(bool)[parent], halo)
    lake = np.asarray(a["lake"]).astype(bool)[parent]
    lev = np.asarray(a["lake_level_m"] if "lake_level_m" in a else a["z_filled_m"], dtype=np.float64)[parent]
    water, surf = np.zeros(g.n, bool), np.full(g.n, np.nan)
    for L in np.unique(lev[lake]):
        wet = (np.asarray(z) < L) & ~sea
        lab = components(g, wet)
        own = np.unique(lab[wet & lake & (lev == L)])
        m = wet & np.isin(lab, own[own >= 0]) & ~water
        water |= m
        surf[m] = L
    return water, surf


ZONE_KEYS = ("gravity_g", "o2_fraction", "fire_reactivity", "plant_height_x")


def zone_fields(world, parent, z_surface, sea_p, scale_height_m) -> dict:
    """Air, gravity and fire of the fine cells: their world cell's values (config zones and era zone events alike);
    the pressure falls from the world cell's sea-level pressure with the fine cell's own height."""
    out = {k: np.asarray(world.a[k])[parent] for k in ZONE_KEYS}
    out["pressure_bar"] = sea_p * FL.pressure(1.0, z_surface, scale_height_m)
    out["po2_bar"] = np.asarray(out["o2_fraction"], dtype=np.float64) * out["pressure_bar"]
    return out


def refine_area(g, halo, world, src, cfg, progress=None, plateaus=None) -> dict:
    tell = progress or (lambda stage, f: None)
    tell("fields", 0.05)
    f = initial_fields(g, halo, world, src)
    p = f["parent"]
    bedrock = np.asarray(world.a["elevation_eroded_m"], dtype=np.float64)   # ice comes back on top later
    mc = None if servecache.low_memory() else {}                        # the correction's kernel, built once
    z0 = f["elevation_m"] + mean_correction(g.xyz, f["elevation_m"], p, halo, bedrock, world, memo=mc)
    inflow, exits, exit_lev = crossings(g, halo, world)
    lw, surf = _lake_water(g, z0, p, world, halo)                       # halo under a world lake stands at its surface
    on_lake = halo & lw                                                 # (water fills it, not drains via its bed) and
    z0 = np.where(on_lake, surf, z0)                                    # is a way out (into the world lake)
    np.minimum.at(z0, exits, exit_lev)                                  # world exits are the lowest ways out
    sea = _sea(g, z0, f["ocean_world"], halo)
    water = runoff_km3(g, f) + inflow
    kmult = k_mult(f["age_class"])
    out_cells = outlets(g, halo, world, p, z0)
    out_cells[exits] = True
    out_cells |= on_lake
    z0 = world_outflows(g, z0, halo, sea | out_cells, p, world.a)       # no fine sill dams a world lake (MODEL 9)
    z = erode_area(g, z0, halo, sea, water, kmult, outlet=out_cells, progress=tell)
    z = z + mean_correction(g.xyz, z, p, halo, bedrock, world, hold=sea, memo=mc)   # means back, smoothly
    del mc
    ocean = _sea(g, z, f["ocean_world"], halo)                         # land that sank to the sea joins it
    tell("sea floor", 0.62)                                              # canyons, fans, ponds, volcanoes, vents
    z, sea_floor, vents = SF.run(g, z, ocean, halo, p, world.a, plateaus or [], int(cfg["build"]["seed"]))
    out_cells |= outlets(g, halo, world, p, z)                          # the raised area may now spill over more halo
    z = world_outflows(g, z, halo, ocean | out_cells, p, world.a)       # again on the final heights (means put back)
    # temperatures follow the final heights (lapse rate from the world's bedrock height, as the world climate)
    dz_km = (z - bedrock[p]) / 1000.0
    for k in LAPSE:
        f[k] = np.asarray(world.a[k])[p] - LAPSE_C_PER_KM * dz_km
    f["biotemp"] = np.clip(f["biotemp"], 0.0, 30.0)
    tell("rivers", 0.65)
    # the world hydrology code, on the sub-grid: outlets = the ways out; inflow carried as extra runoff
    pet = np.maximum(f["PET"], 1e-6)
    run_mm = np.where(ocean, 0.0, np.maximum(f["P_ann"] - f["P_ann"] / np.sqrt(1 + (f["P_ann"] / pet) ** 2), 0.0))
    extra_mm = inflow / np.maximum(g.area_km2 * 1e-6, 1e-12)
    p_hy = np.where(extra_mm > 0, _p_for_runoff(run_mm + extra_mm, pet), f["P_ann"])
    p_hy = np.where(halo, 0.0, p_hy)                                     # no rain from outside the area
    hcfg = copy.deepcopy(cfg)
    hcfg.setdefault("hydrology", {})["river_min_km3_yr"] = RIVER_MIN_KM3_YR
    hcfg["hydrology"]["min_depth_m"] = LAKE_MIN_DEPTH_M
    ctx = SimpleNamespace(grid=g, cfg=hcfg, seed=int(cfg["build"]["seed"]), data={
        "elevation_eroded_m": np.where(halo & ~out_cells, WALL_M, z), "P_ann": p_hy, "PET": pet, "ocean": ocean | out_cells,
        "lake_carve": np.zeros(g.n, bool)})                               # lake beds come carved from the world
    ctx.need = lambda *keys: [ctx.data[k] for k in keys]
    hy = HY.run(ctx)
    hy["runoff_mm"] = run_mm
    river = hy["river"] & ~halo
    tell("ground", 0.85)
    # environment: life zones and ground from the refined fields (lithology, landform from the world cells)
    P = params(cfg, "environment", ENV_DEFAULTS)
    H = float(cfg["planet"]["scale_height_km"]) * 1000.0
    sea_p = (np.asarray(world.a["pressure_bar"], dtype=np.float64)[p]
             / FL.pressure(1.0, np.asarray(world.a["z_surface_m"])[p], H))   # the world cell's sea-level air
    wet = np.sqrt(sea_p / float(cfg["planet"]["sea_level_pressure_bar"]))  # dense air: rain acts wetter (eras)
    zone, _ = holdridge(f["biotemp"], f["P_ann"] * wet, f["T_min"])
    gr = ground(g, z, f["lithology"], f["T_mean"], f["T_min"], f["P_ann"], f["dist_ocean_km"], river, hy["strahler"],
                hy["discharge_km3_yr"], hy["lake"] & ~halo, hy["salt_flat"] & ~halo, P, ocean)
    ctx.data.update({"T_mean": f["T_mean"], "T_jun": f["T_jun"], "T_dec": f["T_dec"], "P_ann": f["P_ann"], "ocean": ocean,
                     "elevation_eroded_m": z})                                # real heights again (no routing walls)
    ic = IC.run(ctx)
    fl = zone_fields(world, p, ic["z_surface_m"], sea_p, H)
    exit_level = np.full(g.n, np.nan)
    np.fmin.at(exit_level, exits, exit_lev)
    out = {"g_ids": g.ids, "g_lat": g.lat, "g_lon": g.lon, "g_xyz": g.xyz, "g_area_km2": g.area_km2,
           "parent": p, "halo": halo, "elevation_eroded_m": z, "ocean": ocean, "holdridge": zone, "ground": gr,
           **{k: v for k, v in hy.items() if k not in ("river", "elevation_eroded_m", "lake_cut_m")}, "river": river, **ic, **fl,
           **{k: f[k] for k in (*COPY, *LAPSE, "T_range")}}
    out["lake"] = hy["lake"] & ~halo
    out["inflow_km3_yr"] = inflow
    out["outlet"] = out_cells
    out["exit_level_m"] = exit_level
    out["z_filled_m"] = np.where(halo & ~out_cells, z, out["z_filled_m"])   # no routing walls in the saved result
    out.update(sea_floor)
    out["vents"] = vents                                                 # per vent (build_areas: vents.npz)
    return out


def stats(a) -> dict:
    """Build report: size, rivers, largest lakes and the water balance (sources vs outflow + lake evaporation)."""
    inner = ~a["halo"]
    q, recv = a["discharge_km3_yr"], a["recv"]
    sink = a["ocean"] | a["outlet"]
    out_q = q[~sink & sink[recv]].sum()                                   # flow into the sea or out of the area
    ar = np.arange(len(q))
    terminal = q[~sink & (recv == ar)].sum()                              # endorheic basins keep (evaporate) theirs
    lost = a.get("lake_loss_km3_yr", np.zeros(len(q)))[~sink].sum()
    src = (a["runoff_mm"][inner & ~a["ocean"]] * a["g_area_km2"][inner & ~a["ocean"]] * 1e-6).sum() + a["inflow_km3_yr"].sum()
    lakes = np.bincount(np.maximum(a["lake_id"][a["lake"]], 0), weights=a["g_area_km2"][a["lake"]]) if a["lake"].any() else np.zeros(1)
    # border: where world rivers leave, the refined water level just inside vs the world river's level there
    lev = a.get("exit_level_m", np.full(len(q), np.nan))
    errs = [float(a["z_filled_m"][(recv == e) & inner].min() - lev[e]) for e in np.where(np.isfinite(lev))[0]
            if ((recv == e) & inner).any()]
    return {"cells": int(inner.sum()), "rivers": int(a["river"].sum()),
            "largest_lakes_km2": sorted(lakes.tolist(), reverse=True)[:3],
            "water_balance_error": float(abs(out_q + terminal + lost - src) / max(src, 1e-9)),
            "exits": int(np.isfinite(lev).sum()), "exits_reached": len(errs),
            "exit_level_error_m": max(errs, key=abs) if errs else None,
            "canyon_max_m": float(a["canyon_m"][inner].max()) if "canyon_m" in a and inner.any() else 0.0,
            "fan_max_m": float(a["fan_m"][inner].max()) if "fan_m" in a and inner.any() else 0.0}


def results_dir(regions_root: Path, key: str) -> Path:
    return Path(regions_root) / f"{key}-m{MODEL}"


def built_areas(regions_root: Path, key: str) -> list[Path]:
    d = results_dir(regions_root, key)
    return sorted(p for p in d.glob("*") if (p / "cells.npz").exists()) if d.exists() else []


def regions_fingerprint(regions_root: Path, key: str, dirs: list | None = None) -> str:
    """Changes whenever any built result changes (model, cells, rebuild): tile URLs must follow. dirs: a listing
    already made (built_areas)."""
    h = hashlib.sha1(f"m{MODEL}".encode())
    for p in (built_areas(regions_root, key) if dirs is None else dirs):
        st = (p / "cells.npz").stat()
        h.update(f"{p.name}:{st.st_size}:{st.st_mtime_ns}".encode())
    return h.hexdigest()[:6]


def _save(d: Path, arrays: dict, meta: dict, vents: dict | None = None) -> None:
    d.mkdir(parents=True, exist_ok=True)
    (d / "cells.npz").unlink(missing_ok=True)                 # cells.npz last: its presence means "complete"
    fd, tmp = tempfile.mkstemp(dir=d, suffix=".json")
    with os.fdopen(fd, "w") as fh:
        fh.write(json.dumps(meta, indent=1) + "\n")
    os.replace(tmp, d / "meta.json")
    if vents is not None:
        fd, tmp = tempfile.mkstemp(dir=d, suffix=".npz")
        with os.fdopen(fd, "wb") as fh:
            np.savez_compressed(fh, **vents)
        os.replace(tmp, d / "vents.npz")
    fd, tmp = tempfile.mkstemp(dir=d, suffix=".npz")
    with os.fdopen(fd, "wb") as fh:
        np.savez_compressed(fh, **arrays)
    os.replace(tmp, d / "cells.npz")


def _meta(d: Path):
    try:
        return json.loads((d / "meta.json").read_text())
    except (OSError, ValueError):
        return None


def build_areas(root: Path, res: int, regions_path: Path, log=print, regions_root: Path | None = None,
                progress=None) -> list[dict]:
    """Refine every connected area of regions_path's outlines for every world the viewer can show (the base build
    and each built era). Results are keyed by content (area_key) under <regions_root>/areas/<key>/, so an area an era
    left as it was is built once; each world's set folder <regions_root>/<world fingerprint>-m<MODEL>/ links its areas
    by cell hash (what RegionSet reads) and records its outlines (outlines.json) and era (era.json). Other worlds'
    sets and unused areas go only after every area built (a failed build keeps the map as it was).
    progress(stage, fraction) reports a running build (fraction over all areas still to build)."""
    import gc

    import serve
    import tiles
    from mapgen import config as C
    root = Path(root)
    cfg, tect = C.load(root)
    plateaus = tect.get("plateau", [])
    regions_root = Path(regions_root or (root / "out" / f"r{res}" / "regions"))
    store = areas_dir(regions_root)
    store.mkdir(parents=True, exist_ok=True)
    regions = load_regions(regions_path)                     # once: the records must describe what was built
    areas = area_cells(regions, res + FINER)
    try:
        memo = json.loads((store / "keys.json").read_text())
    except (OSError, ValueError):
        memo = {}
    if progress:
        progress("keys", 0.0)
    sets = []
    for era, out_dir in world_dirs(root, res, log):
        world = serve.World(root, res, out=out_dir)
        src = tiles.TileSource(world, int(cfg["build"]["seed"]), regions_dir=regions_root, with_regions=False)
        wc = WorldCells(out_dir)
        keys = {}
        for cells in areas:
            m = f"{src.fingerprint}:m{MODEL}:{area_hash(cells)}"      # a new model: new keys
            if m not in memo:
                memo[m] = area_key(wc, cells, cfg, plateaus)
            keys[area_hash(cells)] = memo[m]
        sets.append((era, world, src, wc, keys))
    todo, queued = [], set()
    for era, world, src, wc, keys in sets:
        for cells in areas:
            k = keys[area_hash(cells)]
            done = (store / k / "cells.npz").exists() and _meta(store / k) is not None
            if not done and k not in queued:
                queued.add(k)
                todo.append((cells, k, era, world, src, wc))
    for n, (cells, k, era, world, src, wc) in enumerate(todo):
        tell = (lambda stage, f, n=n: progress(stage, (n + f) / len(todo))) if progress else None
        R = float(world.cells_meta["radius_km"])
        km2 = len(cells) * float(np.mean([h3.cell_area(int(c), unit="rads^2") for c in cells[:200]])) * R ** 2
        if km2 > 2.0e6:
            log(f"refine: area {area_hash(cells)} is {km2:,.0f} km² (over ≈ 2,000,000): this will take a while")
        t0 = time.time()
        g, halo = subgrid(cells, R)
        a = refine_area(g, halo, wc, src, cfg, progress=tell, plateaus=plateaus)
        vents = a.pop("vents", None)
        meta = {"hash": area_hash(cells), "key": k, "model": MODEL, "era": era, "world": src.fingerprint,
                "area_km2": km2, "vents": 0 if vents is None else int(len(vents["cell"])), "seconds": round(time.time() - t0, 1), **stats(a)}
        if tell:
            tell("saving", 0.95)
        _save(store / k, a, meta, vents)
        log(f"refine: area {meta['hash']} ({era}): {meta['cells']} cells, {meta['rivers']} river cells, "
            f"{meta['seconds']} s")
        del a, g, halo                                        # the next area must not share memory with this one
        gc.collect()
        trim_memory()
    metas, built_now, keep_sets = [], {t[1] for t in todo}, set()
    for era, world, src, wc, keys in sets:                   # the records: links, outlines, era
        sd = results_dir(regions_root, src.fingerprint)
        sd.mkdir(parents=True, exist_ok=True)
        keep_sets.add(sd.name)
        for cells in areas:
            h, k = area_hash(cells), keys[area_hash(cells)]
            link = _link(sd, h, store / k)
            metas.append({**_meta(store / k), "dir": str(link), "era": era, "built": k in built_now})
        _write_json(sd / "outlines.json", {"regions": region_areas(regions, areas, res + FINER)})
        _write_json(sd / "era.json", {"era": era})
    fps = {s[2].fingerprint for s in sets}
    _write_json(store / "keys.json", {m: v for m, v in memo.items()
                                      if m.split(":", 1)[0] in fps and m.split(":")[1] == f"m{MODEL}"})
    # only after every area built: other worlds' sets, links of redrawn or deleted regions, unused areas go
    names = {area_hash(c) for c in areas}
    for d in regions_root.iterdir():
        if d.name != AREAS and d.name not in keep_sets and (d.is_dir() or d.is_symlink()):
            _remove(d)
    used = set()
    for s in keep_sets:
        for p in (regions_root / s).iterdir():
            if p.name.startswith(servecache.PREFIX) or p.name.startswith(".") or not (p.is_dir() or p.is_symlink()):
                continue                                      # serve caches: the server prunes them
            if p.name in names and p.is_symlink():
                used.add(os.path.basename(os.readlink(p)))
            else:
                _remove(p)
    for p in store.iterdir():
        if p.is_dir() and p.name not in used:
            _remove(p)
    return metas


NET_KEYS = ("g_xyz", "g_ids", "river", "recv", "ocean", "lake", "z_surface_m", "z_filled_m", "discharge_km3_yr",
            "lithology", "age_class", "landform", "ground", "P_ann")   # what rivers.RiverNet reads
BLEND_KM = 10.0
LAKE_FLOOR_M = 30.0            # lake water reaches this far below a lake's deepest cell (procedural detail)


def shoulders(a, radius_km: float = 12742.0) -> np.ndarray:
    """Heights with each river cell raised to its valley shoulders (the mean of its dry neighbours), so smooth
    interpolation keeps the plateau and the valley shape can be cut in at sub-cell scale."""
    z = np.asarray(a["z_surface_m"], dtype=np.float64).copy()
    riv = np.asarray(a["river"]).astype(bool) & ~np.asarray(a["halo"]).astype(bool)
    dry = (~np.asarray(a["river"]).astype(bool) & ~np.asarray(a["lake"]).astype(bool) & ~np.asarray(a["ocean"]).astype(bool)
           & ~np.asarray(a["halo"]).astype(bool))
    if not riv.any() or not dry.any():
        return z
    tree = cKDTree(a["g_xyz"][dry])
    spacing = np.sqrt(np.mean(a["g_area_km2"])) * 1.07 / radius_km               # neighbour distance on the unit sphere
    d, k = tree.query(a["g_xyz"][riv], k=min(6, int(dry.sum())), distance_upper_bound=1.5 * spacing)
    d, k = np.atleast_2d(d), np.atleast_2d(k)
    ok = np.isfinite(d)
    zd = np.where(ok, z[dry][np.minimum(k, dry.sum() - 1)], 0.0)
    mean = np.where(ok.any(1), zd.sum(1) / np.maximum(ok.sum(1), 1), -np.inf)
    z[riv] = np.maximum(z[riv], mean)
    return z


def touches_trees(trees, spacing_km: float, radius_km: float, z, x, y, margin_km: float = 0.0) -> bool:
    """Whether tile z/x/y comes within (2 cell spacings + margin) of the cells in any of the KD-trees."""
    if not trees:
        return False
    span = 180.0 / 2 ** z
    f = np.array([0.0, 0.5, 1.0])
    LA, LO = np.meshgrid(90.0 - (y + f) * span, -180.0 + (x + f) * span, indexing="ij")
    p = latlon_to_xyz(LA, LO).reshape(-1, 3)
    r = float(np.max(np.linalg.norm(p - p[4], axis=1))) + (2 * spacing_km + margin_km) / radius_km
    return any(t.query_ball_point(p[4], r, return_length=True) > 0 for t in trees)


def trim_memory() -> None:
    """Give freed memory back to the system (glibc keeps it otherwise)."""
    try:
        import ctypes
        ctypes.CDLL("libc.so.6").malloc_trim(0)
    except (OSError, AttributeError):
        pass


def _compute_region_set(dirs, legends: dict, radius_km: float, spill=None):
    """Everything a RegionSet serves, computed from its areas' results (as the old RegionSet.__init__ did):
    (arrays, meta) for the serve cache. Keys: 'a.<field>' (the concatenated inner cells, sorted by H3 id, recv
    remapped), 'halo_xyz', 'area' (index into meta['area_names'] per fine cell), 'z_env', 'lake_level',
    'lake_floor', 'net.<state>' (rivers.STATE) when a river net exists. spill: a folder to keep the arrays in
    (memory maps; same values) instead of RAM."""
    R = float(radius_km)
    keep = servecache.spiller(spill)
    files = [np.load(p / "cells.npz") for p in dirs]               # read key by key: each array once, no copies kept
    try:
        halo = [f["halo"].astype(bool) for f in files]
        inner = [~h for h in halo]
        ids = [f["g_ids"] for f in files]
        recv_p = [f["recv"].astype(np.int64) for f in files]
        river_p = [f["river"].astype(bool) for f in files]
        xyz_p = [f["g_xyz"] for f in files]
        halo_xyz = np.concatenate([x[h] for x, h in zip(xyz_p, halo)])
        order = np.argsort(np.concatenate([i_[m] for i_, m in zip(ids, inner)]))
        area_of = np.concatenate([np.full(int(m.sum()), p, np.int32) for p, m in enumerate(inner)])[order]
        # halo cells rivers leave through (for the river net): per part, their indices
        ends = []
        for m, h, r, rv in zip(inner, halo, recv_p, river_p):
            out = np.where(m & rv & h[r])[0]
            ends.append((out, *np.unique(r[out], return_inverse=True)) if len(out) else None)
        keys = set.intersection(*(set(f.files) for f in files))
        a, end_vals = {}, {}
        for k in sorted(keys):
            vs = [f[k] for f in files]
            if any(len(v) != len(i_) for v, i_ in zip(vs, ids)):
                continue
            a[k] = keep(f"a.{k}", np.concatenate([v[m] for v, m in zip(vs, inner)])[order])
            if k in NET_KEYS:
                end_vals[k] = [v[e[1]] if e is not None else v[:0] for v, e in zip(vs, ends)]
            del vs
    finally:
        for f in files:
            f.close()
    # receivers index into each part's own arrays: remap to the sorted inner cells (halo/outside → self)
    inv = np.empty_like(order)
    inv[order] = np.arange(len(order))
    recv_all, off = [], 0
    for r, m in zip(recv_p, inner):
        new = -np.ones(len(m), np.int64)
        new[np.where(m)[0]] = np.arange(off, off + m.sum())
        rr = new[r[m]]
        recv_all.append(np.where(rr >= 0, rr, np.arange(off, off + m.sum())))
        off += m.sum()
    a["recv"] = keep("a.recv", inv[np.concatenate(recv_all)][order])
    ids_all = a["g_ids"]
    z_env = shoulders(a, R)                                                    # plateau across river cells
    lake = np.asarray(a["lake"]).astype(bool)
    lake_level = np.where(lake, a["z_filled_m"], np.nan)
    # each lake cell's bed (+ LAKE_FLOOR_M for the fine detail): ground far below it lies past a dam or a cliff,
    # not under water (no water walls hanging over a drop)
    lake_floor = np.where(lake, lake_level - np.asarray(a["depression_depth_m"], dtype=np.float64) - LAKE_FLOOR_M,
                          np.nan)
    # the river net runs on into the halo cells where rivers leave, so they meet the world's rivers at the edge
    # (built from the few arrays it needs: a copy of every field would double the memory)
    n = len(ids_all)
    net_recv = a["recv"].copy()
    extra, base = [], n
    for part, e in enumerate(ends):
        if e is None:
            continue
        out, h, k = e
        net_recv[np.searchsorted(ids_all, ids[part][out])] = base + k
        extra.append(part)
        base += len(h)
    nkeys = [k for k in NET_KEYS if k in a]
    net_a = {k: np.concatenate([a[k], *[end_vals[k][p] for p in extra]]) for k in nkeys if k != "recv"}
    net_a["recv"] = np.concatenate([net_recv, np.arange(n, base)])
    net_a["river"] = np.concatenate([a["river"].astype(bool), np.zeros(base - n, bool)])
    net_a["lake"] = np.concatenate([a["lake"].astype(bool), np.zeros(base - n, bool)])
    ground = np.concatenate([z_env, net_a["z_surface_m"][n:]])
    arrays = {f"a.{k}": v for k, v in a.items()}
    arrays.update(halo_xyz=keep("halo_xyz", halo_xyz), area=keep("area", area_of), z_env=keep("z_env", z_env),
                  lake_level=keep("lake_level", lake_level), lake_floor=keep("lake_floor", lake_floor))
    meta = {"area_names": [p.name for p in dirs], "res": h3.get_resolution(int(ids_all[0])),
            "spacing_km": float(np.sqrt(np.mean(a["g_area_km2"])) * 1.07), "net": None}
    try:
        import rivers as RV
        net = RV.RiverNet(net_a, legends, R, levels=net_a["z_surface_m"], ground=ground)
        st, meta["net"] = net.state()
        arrays.update({f"net.{k}": v for k, v in st.items()})
    except Exception:
        meta["net"] = None
    return arrays, meta


CODE_FILES = (Path(__file__), Path(__file__).with_name("rivers.py"))   # what computes a region set's served arrays


def code_key() -> str:
    """Hash of the code that computes a region set's arrays: new code, new serve cache (stored river nets, shoulders
    and lake floors are values of this code, not only of the areas' results)."""
    h = hashlib.sha1()
    for p in CODE_FILES:
        h.update(Path(p).read_bytes())
    return h.hexdigest()[:8]


class _AreaTrees:
    """Per-area KD-trees of a RegionSet, built on first use (one area: the set's main tree)."""
    def __init__(self, rs):
        self.rs, self._built = rs, {}

    def __contains__(self, name) -> bool:
        return name in self.rs.area_names

    def __getitem__(self, name):
        if name not in self._built:
            rs = self.rs
            if len(rs.area_names) == 1:
                self._built[name] = rs.tree
            else:
                k = rs.area_names.index(name)
                self._built[name] = cKDTree(np.asarray(rs.arrays["g_xyz"], dtype=np.float64)[np.asarray(rs.area) == k])
        return self._built[name]


class RegionSet:
    """All refined areas built for the current world: fine cells for tiles, 3D heights and the inspector. The arrays
    come from the serve cache (memory-mapped; computed and written on first load); search trees are built on first
    use (warm() builds them before render workers fork, so they share them)."""
    @classmethod
    def none(cls) -> "RegionSet":
        """No refined areas (e.g. for a build, which needs only the world's relief)."""
        rs = cls.__new__(cls)
        rs.fingerprint, rs.area_names, rs.empty, rs.R, rs.cache = "", [], True, 0.0, None
        rs.area_keys = {}
        rs._tree = rs._halo_tree = None
        rs._lock = threading.Lock()
        rs.area_trees = _AreaTrees(rs)
        return rs

    def __init__(self, regions_root: Path, key: str, legends: dict, radius_km: float, use_cache: bool = True):
        self.R = float(radius_km)
        dirs = built_areas(regions_root, key)
        self.fingerprint = regions_fingerprint(regions_root, key, dirs)
        self.area_names = [p.name for p in dirs]                  # area hashes
        self.area_keys = {p.name: (Path(os.readlink(p)).name if p.is_symlink() else p.name) for p in dirs}
        self.empty = not dirs
        self.cache = None
        self._tree = self._halo_tree = None
        self._lock = threading.Lock()
        self.area_trees = _AreaTrees(self)
        if self.empty:
            return
        spill = None
        if servecache.low_memory() and use_cache:                               # fields on disk while the cache is written
            spill = Path(tempfile.mkdtemp(prefix=".spill-", dir=results_dir(regions_root, key)))
        build = lambda: _compute_region_set(dirs, legends, self.R, spill=spill)
        try:
            if use_cache:
                self.cache = servecache.cache_dir(results_dir(regions_root, key), f"{self.fingerprint}-{code_key()}")
                arrays, meta = servecache.load_or_build(self.cache, build)
            else:
                arrays, meta = build()
            self._from_cache(arrays, meta)
        finally:
            if spill is not None:
                shutil.rmtree(spill, ignore_errors=True)
        trim_memory()                                                 # the loading's scratch memory

    def _from_cache(self, arrays: dict, meta: dict) -> None:
        import rivers as RV
        self.arrays = {k[2:]: v for k, v in arrays.items() if k.startswith("a.")}
        self.ids = self.arrays["g_ids"]
        self.res = int(meta["res"])
        self.spacing_km = float(meta["spacing_km"])
        self.area = arrays["area"]
        self._halo_xyz = arrays["halo_xyz"]
        self.z_env, self.lake_level, self.lake_floor = arrays["z_env"], arrays["lake_level"], arrays["lake_floor"]
        st = {k[4:]: v for k, v in arrays.items() if k.startswith("net.")}
        self.net = RV.RiverNet.from_state(st, meta["net"]) if meta.get("net") else None

    @property
    def tree(self):
        if self._tree is None and not self.empty:
            with self._lock:
                if self._tree is None:
                    self._tree = cKDTree(np.asarray(self.arrays["g_xyz"], dtype=np.float64))
        return self._tree

    @property
    def halo_tree(self):
        if self._halo_tree is None and not self.empty:
            with self._lock:
                if self._halo_tree is None:
                    self._halo_tree = cKDTree(np.asarray(self._halo_xyz, dtype=np.float64))
        return self._halo_tree

    def warm(self) -> None:
        """Build the search trees now (before forking render workers, so they share them)."""
        if self.empty:
            return
        self.tree
        self.halo_tree
        if self.net is not None:
            self.net.tree

    def nearest(self, xyz, k=1):
        d, i = self.tree.query(xyz, k=k)
        return i, d * self.R

    def edge_km(self, xyz):
        return self.halo_tree.query(xyz)[0] * self.R

    def weight(self, xyz):
        """0 outside an area, rising to 1 over BLEND_KM inside its edge."""
        if self.empty:
            return np.zeros(len(xyz))
        d_in = self.tree.query(xyz, distance_upper_bound=1.5 * self.spacing_km / self.R)[0] * self.R   # inf: far outside
        e = self.halo_tree.query(xyz, distance_upper_bound=(2 * BLEND_KM + 2 * self.spacing_km) / self.R)[0] * self.R
        depth = np.full(len(d_in), -np.inf)
        m = np.isfinite(d_in)
        depth[m] = (e[m] - d_in[m]) / 2                                  # ≈ distance inside the area edge; 0 on it
        t = np.clip(depth / BLEND_KM, 0.0, 1.0)                          # (bounded queries: deep inside → inf → 1)
        return t * t * (3 - 2 * t)

    def touches(self, z, x, y, areas=None, margin_km: float = 0.0) -> bool:
        """Whether tile z/x/y (with a margin for borders) can show any refined area (or one of `areas`)."""
        if self.empty or (areas is not None and not areas):
            return False
        trees = [self.tree] if areas is None else [self.area_trees[h] for h in areas if h in self.area_trees]
        return touches_trees(trees, self.spacing_km, self.R, z, x, y, margin_km)

    def index_of(self, lat, lon):
        if self.empty:
            return None
        c = np.uint64(h3.latlng_to_cell(float(lat), float(lon), self.res))
        i = int(np.searchsorted(self.ids, c))
        return i if i < len(self.ids) and self.ids[i] == c else None