"""Stage `render`: equirectangular rasters, cells.npz, metadata, previews, contact sheet."""
from __future__ import annotations

import colorsys
import json

import numpy as np
from PIL import Image, ImageDraw
from scipy.spatial import cKDTree

from . import geo, projections, viewer_export
from . import plateaus as PL
from .crust import AGE_NAMES
from .environment import LEGENDS, REGIONS, ZONES
from .ice import ICE_NAMES
from .noise import fbm
from .seabed import MINERAL_NAMES, SEABED_NAMES
from .sphere import east_north, latlon_to_xyz

CONTINUOUS = {
    "elevation": ("z_surface_m", 0.5, -12000.0, "m"),
    "T_mean": ("T_mean", 0.01, -100.0, "degC"),
    "T_range": ("T_range", 0.01, 0.0, "degC"),
    "P_ann": ("P_ann", 0.5, 0.0, "mm/yr"),
    "P_jun": ("P_jun", 0.5, 0.0, "mm/yr"),
    "P_dec": ("P_dec", 0.5, 0.0, "mm/yr"),
    "po2": ("po2_bar", 2e-4, 0.0, "bar"),
    "gravity": ("gravity_g", 1e-4, 0.0, "g"),
    "pressure": ("pressure_bar", 2e-4, 0.0, "bar"),
    "o2_fraction": ("o2_fraction", 2e-5, 0.0, "fraction"),
    "fire": ("fire_reactivity", 1e-4, 0.0, "x"),
    "vent_potential": ("vent_potential", 1e-4, 0.0, "0-1"),
    "bottom_temp": ("bottom_temp_c", 0.01, -10.0, "degC"),
    "sediment": ("sediment_m", 0.5, 0.0, "m"),
    "plant_height": ("plant_height_x", 2e-3, 0.0, "x"),
    "sst": ("sst", 0.01, -100.0, "degC"),
    "productivity": ("productivity", 1e-4, 0.0, "0-1"),
    "current_speed": ("current_speed", 1e-4, 0.0, "m/s"),
    "upwelling": ("upwelling", 0.05, -1600.0, "m/yr"),
}
CATEGORICAL = {"plates": "plate", "age_class": "age_class", "holdridge": "holdridge",
               "seasonality": "seasonality", "landform": "landform", "lithology": "lithology",
               "ground": "ground", "ice": "ice",
               "seabed_type": "seabed_type", "seabed_mineral": "seabed_mineral", "deposits": "deposit_main"}
DETAIL_M = np.array([60, 40, 150, 400, 80, 150, 300, 300, 300, 60, 30, 120], dtype=np.float64)  # by landform
CHUNK = 128
COAST_DETAIL_M = 250.0


def _by_rows(fn, v, dtype, tail=()):
    """fn applied to CHUNK-row slices of v (fn elementwise): the same values as fn(v), without fn's full-size
    float64 scratch (a raster ramp builds five H×W×3 float64 temporaries: ~4 GB at 8192 px)."""
    v = np.asarray(v)
    if v.ndim < 2 or v.shape[0] <= CHUNK:
        return fn(v)
    out = np.empty(v.shape + tuple(tail), dtype)
    for y0 in range(0, v.shape[0], CHUNK):
        out[y0:y0 + CHUNK] = fn(v[y0:y0 + CHUNK])
    return out


def encode(v, scale, offset):
    return _by_rows(lambda x: np.clip(np.round((np.asarray(x, dtype=np.float64) - offset) / scale), 0, 65535)
                    .astype(np.uint16), v, np.uint16)


def decode(raw, scale, offset):
    return np.asarray(raw, dtype=np.float64) * scale + offset


def _rows_xyz(rows, W, H):
    lat = 90.0 - (rows + 0.5) / H * 180.0
    lon = (np.arange(W) + 0.5) / W * 360.0 - 180.0
    LA, LO = np.meshgrid(lat, lon, indexing="ij")
    return latlon_to_xyz(LA.ravel(), LO.ravel())


def _make_render_jit():
    """Compiled ramp and sampling loops (numba optional; WORLDGEN_NO_JIT=1 turns it off): per pixel the same float
    steps in the same order as the numpy statements they replace, so the same bytes, without the temporaries."""
    import os
    if os.environ.get("WORLDGEN_NO_JIT"):
        return None
    try:
        import numba
    except ImportError:
        return None

    @numba.njit(cache=True, nogil=True)
    def ramp(x, lo, span, stops, out):          # x: float64 (n,), out: uint8 (n, c)
        m = stops.shape[0]
        for j in range(x.shape[0]):
            t = (x[j] - lo) / span
            t = 0.0 if t < 0.0 else (1.0 if t > 1.0 else t)
            t = t * (m - 1)
            i = min(np.int64(t), m - 2)
            f = t - i
            for c in range(stops.shape[1]):
                out[j, c] = np.uint8(np.int64(stops[i, c] * (1 - f) + stops[i + 1, c] * f))

    @numba.njit(cache=True, nogil=True)
    def sample(v, idx, wts, out):               # out[p] = Σ_k v[idx[p,k]] * w[p,k], k left to right (as np.sum, k < 8)
        for p in range(idx.shape[0]):
            s = v[idx[p, 0]] * np.float64(wts[p, 0])
            for k in range(1, idx.shape[1]):
                s += v[idx[p, k]] * np.float64(wts[p, k])
            out[p] = s
    return ramp, sample


_render_jit = _make_render_jit()


def pixel_neighbours(g, W, H, k=3):
    from .graph import workers
    tree = cKDTree(g.xyz)
    idx = np.empty((H, W, k), np.int32)
    wts = np.empty((H, W, k), np.float32)
    for y0 in range(0, H, CHUNK):
        rows = np.arange(y0, min(H, y0 + CHUNK))
        d, i = tree.query(_rows_xyz(rows, W, H), k=k, workers=workers())   # per point: the same answer
        w = 1.0 / np.maximum(d, 1e-9) ** 2
        w /= w.sum(axis=1, keepdims=True)
        idx[rows] = i.reshape(len(rows), W, k)
        wts[rows] = w.reshape(len(rows), W, k)
    return idx, wts


def sample_cont(v, idx, wts):
    out = np.empty(idx.shape[:2])
    if (_render_jit is not None and v.dtype == np.float64 and v.ndim == 1 and wts.dtype == np.float32
            and 1 <= idx.shape[-1] < 8 and idx.shape == wts.shape):
        k = idx.shape[-1]
        _render_jit[1](np.ascontiguousarray(v), np.ascontiguousarray(idx).reshape(-1, k),
                       np.ascontiguousarray(wts).reshape(-1, k), out.reshape(-1))
        return out
    for y0 in range(0, idx.shape[0], CHUNK):
        s = slice(y0, y0 + CHUNK)
        out[s] = np.sum(v[idx[s]] * wts[s], axis=-1)
    return out


def sample_cat(v, idx):
    return v[idx[..., 0]]


def pixel_land(ocean_k, z):
    """Pixel land mask: unanimous neighbour cells decide; at the coast (mixed) the sub-cell height does."""
    all_sea = ocean_k.all(axis=-1)
    all_land = ~ocean_k.any(axis=-1)
    return all_land | (~all_sea & ~all_land & (z > 0))


def hillshade(z, radius_km, az=315.0, alt=45.0, exag=4.0, low_memory=False):
    """low_memory: the same values, CHUNK rows at a time (one halo row each side keeps the central differences)."""
    if low_memory:
        H = z.shape[0]
        out = np.empty(z.shape)
        for y0 in range(0, H, CHUNK):
            a, b = max(y0 - 1, 0), min(y0 + CHUNK + 1, H)
            part = _hillshade_rows(z[a:b], a, H, radius_km, az, alt, exag)
            out[y0:min(y0 + CHUNK, H)] = part[y0 - a: y0 - a + min(CHUNK, H - y0)]
        return out
    return _hillshade_rows(z, 0, z.shape[0], radius_km, az, alt, exag)


def _hillshade_rows(z, row0, H, radius_km, az, alt, exag):
    """Hillshade of rows row0.. of an H-row raster (z holds those rows; edges of z use one-sided differences)."""
    W = z.shape[1]
    lat = 90.0 - (np.arange(row0, row0 + z.shape[0]) + 0.5) / H * 180.0
    dy = np.pi * radius_km * 1000.0 / H
    dx = 2 * np.pi * radius_km * 1000.0 * np.maximum(np.cos(np.radians(lat)), 0.01) / W
    gy, gx = np.gradient(z)
    dzdx = gx / dx[:, None] * exag
    dzdn = -gy / dy * exag
    norm = np.sqrt(dzdx**2 + dzdn**2 + 1.0)
    a, b = np.radians(az), np.radians(alt)
    L = (np.sin(a) * np.cos(b), np.cos(a) * np.cos(b), np.sin(b))
    return np.clip((-dzdx * L[0] - dzdn * L[1] + L[2]) / norm, 0.0, 1.0)


def _ramp(v, lo, hi, stops):
    stops = np.asarray(stops, dtype=np.float64)

    def part(x):
        if (_render_jit is not None and isinstance(x, np.ndarray) and x.dtype == np.float64 and stops.ndim == 2
                and len(stops) >= 2 and not np.isnan(x).any()):
            out = np.empty(x.shape + (stops.shape[1],), np.uint8)
            _render_jit[0](np.ascontiguousarray(x).reshape(-1), float(lo), float(hi - lo), stops,
                           out.reshape(-1, stops.shape[1]))
            return out
        t = np.clip((x - lo) / (hi - lo), 0, 1) * (len(stops) - 1)
        i = np.minimum(t.astype(np.int64), len(stops) - 2)
        f = (t - i)[..., None]
        return (stops[i] * (1 - f) + stops[i + 1] * f).astype(np.uint8)
    return _by_rows(part, v, np.uint8, (stops.shape[-1],))


def holdridge_palette():
    ramp = np.array([[216, 200, 160], [208, 196, 140], [200, 200, 120], [152, 168, 96], [106, 150, 80],
                     [70, 125, 68], [50, 105, 62], [37, 90, 56]], dtype=np.float64)
    tint = {"polar": ((232, 236, 239), 0.9), "subpolar": ((170, 176, 160), 0.55), "boreal": ((70, 100, 80), 0.35),
            "tropical": ((20, 90, 40), 0.15)}
    cols = []
    for r in REGIONS:
        k = len(ZONES[r])
        for j in range(k):
            c = ramp[int(round((j / max(k - 1, 1)) * (len(ramp) - 1)))]
            if r in tint:
                c = c * (1 - tint[r][1]) + np.array(tint[r][0]) * tint[r][1]
            cols.append(c)
    return np.array(cols, dtype=np.uint8)


def category_palette(n):
    return np.array([[int(255 * c) for c in colorsys.hsv_to_rgb((i * 0.618034) % 1.0, 0.55, 0.9)]
                     for i in range(max(n, 1))], dtype=np.uint8)


def _region_palette(cols):
    """38 zone colours from per-region (dry, wet) colour pairs."""
    out = []
    for r in REGIONS:
        dry, wet = (np.array(c, dtype=np.float64) for c in cols[r])
        k = len(ZONES[r])
        out += [dry + (wet - dry) * (j / max(k - 1, 1)) for j in range(k)]
    return np.array(out, dtype=np.uint8)


STYLES = {  # alien palettes (config [render] style); None = Earth-like
    "tidal-lock": {
        "zones": {"polar": ((34, 38, 46), (34, 38, 46)), "subpolar": ((70, 72, 78), (96, 104, 112)),
                  "boreal": ((120, 96, 64), (150, 110, 60)), "cool temperate": ((170, 120, 60), (190, 140, 70)),
                  "warm temperate": ((110, 70, 44), (130, 86, 50)), "subtropical": ((46, 36, 34), (60, 44, 38)),
                  "tropical": ((22, 20, 22), (30, 26, 28))},
        "ocean": [[4, 12, 16], [10, 34, 40], [40, 80, 84]], "lake": (60, 96, 104), "river": (70, 110, 118),
        "ground": {6: (92, 86, 80)}, "sheet": (200, 222, 240), "sea_ice": (170, 200, 226), "sea_ice_seasonal": (120, 150, 170)},
    "salt-mirror": {
        "zones": {"polar": ((236, 240, 244), (236, 240, 244)), "subpolar": ((228, 228, 234), (214, 214, 228)),
                  "boreal": ((238, 236, 228), (208, 208, 224)), "cool temperate": ((242, 238, 226), (200, 204, 222)),
                  "warm temperate": ((240, 230, 204), (196, 204, 220)), "subtropical": ((238, 224, 186), (192, 206, 218)),
                  "tropical": ((234, 214, 166), (186, 206, 216))},
        "ocean": [[70, 120, 130], [150, 195, 200], [215, 235, 235]], "lake": (150, 196, 204), "river": (150, 196, 204),
        "ground": {6: (252, 250, 244)}, "sheet": (248, 250, 252), "sea_ice": (236, 244, 246), "sea_ice_seasonal": (220, 236, 238)},
}


def style_palette(style):
    return None if style is None else STYLES[style]


def relief_rgb(z, hs, zone, ground, ice, lake, land=None, vary=None, style=None, low_memory=False):
    """vary: optional (brightness factor, tint) per pixel for open land (not lakes or ice): tint > 0 drier/yellower,
    < 0 lusher (deeper green). style: a STYLES key (alien palette) or None. low_memory: the same pixels, made
    CHUNK rows at a time (no full-size float64 scratch)."""
    if low_memory:
        out = np.empty(np.shape(z) + (3,), np.uint8)
        for y0 in range(0, np.shape(z)[0], CHUNK):
            s = slice(y0, y0 + CHUNK)
            out[s] = relief_rgb(z[s], hs[s], zone[s], ground[s], ice[s], lake[s], None if land is None else land[s],
                                None if vary is None else tuple(np.asarray(v)[s] for v in vary), style)
        return out
    land = z > 0 if land is None else land
    st = style_palette(style)
    pal = holdridge_palette() if st is None else _region_palette(st["zones"])
    rgb = pal[np.clip(zone, 0, 37)].astype(np.float64)
    ocean = _ramp(z, -6500.0 if st is None else -1500.0, 0.0,
                  [[11, 43, 90], [30, 90, 150], [143, 198, 224]] if st is None else st["ocean"]).astype(np.float64)
    rgb = np.where(land[..., None], rgb, ocean)
    grounds = ((1, (111, 143, 106)), (2, (120, 128, 100)), (5, (63, 111, 74)), (6, (239, 233, 220))) if st is None \
        else tuple(st["ground"].items())
    for code, col in grounds:
        rgb = np.where((land & (ground == code))[..., None], col, rgb)
    if vary is not None:
        bright, tint = (np.asarray(v, dtype=np.float64)[..., None] for v in vary)
        shift = np.where(tint > 0, tint * np.array([1.0, 0.6, -0.8]), -tint * np.array([-0.9, -0.2, -0.6]))
        if st is not None:
            shift = np.abs(tint) * np.array([0.4, 0.4, 0.4]) * np.sign(tint)
        rgb = np.where(land[..., None], rgb * bright + shift, rgb)
    rgb = np.where((lake & land)[..., None], (79, 143, 192) if st is None else st["lake"], rgb)
    rgb = np.where(np.isin(ice, [1, 2])[..., None], (244, 248, 251) if st is None else st["sheet"], rgb)
    rgb = np.where((ice == 4)[..., None], (225, 235, 242) if st is None else st["sea_ice"], rgb)
    rgb = np.where((ice == 3)[..., None], 0.5 * rgb + 0.5 * np.array([220, 232, 240] if st is None else st["sea_ice_seasonal"]), rgb)
    shade = np.where(land, 0.55 + 0.45 * hs, 0.85 + 0.15 * hs)
    return np.clip(rgb * shade[..., None], 0, 255).astype(np.uint8)


RIVER_RGB = (58, 112, 176)


def draw_rivers(rgb, g, recv, river, strahler, min_order=2, colour=RIVER_RGB):
    """Draw river segments (cell centre → receiver) of order ≥ min_order; width grows with order."""
    H, W = rgb.shape[:2]
    im = Image.fromarray(rgb)
    draw = ImageDraw.Draw(im)
    x = (np.asarray(g.lon) + 180.0) / 360.0 * W - 0.5
    y = (90.0 - np.asarray(g.lat)) / 180.0 * H - 0.5
    cells = np.flatnonzero(river & (recv != np.arange(g.n)) & (strahler >= min_order))
    for i in cells[np.argsort(strahler[cells])]:
        j = recv[i]
        if abs(x[i] - x[j]) > W / 2:
            continue
        width = max(1, int(round(int(strahler[i]) * W / 8192)))
        draw.line([(x[i], y[i]), (x[j], y[j])], fill=tuple(colour), width=width)
    return np.array(im)


def draw_currents(rgb, g, current, ocean, per_row=60, colour=(255, 255, 255)):
    """Arrows along the surface current on a lattice ≈ `per_row` across the map; length ∝ speed (1 m/s ≈ one
    lattice step), sea only; currents under 2 cm/s get none."""
    H, W = rgb.shape[:2]
    step = W / per_row
    current = np.asarray(current, dtype=np.float64)
    ocean = np.asarray(ocean, bool)
    e, n = east_north(g.xyz)
    tree = cKDTree(g.xyz)
    ys, xs = np.meshgrid(np.arange(step / 2, H, step), np.arange(step / 2, W, step), indexing="ij")
    lat, lon = 90.0 - (ys.ravel() + 0.5) / H * 180.0, (xs.ravel() + 0.5) / W * 360.0 - 180.0
    _, cell = tree.query(latlon_to_xyz(lat, lon))
    ue, un = np.sum(current[cell] * e[cell], axis=1), np.sum(current[cell] * n[cell], axis=1)
    sp = np.hypot(ue, un)
    im = Image.fromarray(rgb)
    draw = ImageDraw.Draw(im)
    for x, y, a, b, s, ok in zip(xs.ravel(), ys.ravel(), ue, un, sp, ocean[cell]):
        if not ok or s < 0.02:
            continue
        L = min(s, 1.5) * 0.9 * step
        dx, dy = a / s * L, -b / s * L
        x0, y0, x1, y1 = x - dx / 2, y - dy / 2, x + dx / 2, y + dy / 2
        draw.line([(x0, y0), (x1, y1)], fill=tuple(colour), width=1)
        for ang in (2.6, -2.6):                                        # arrowhead: two barbs ±150°
            c, sn = np.cos(ang), np.sin(ang)
            draw.line([(x1, y1), (x1 + 0.35 * (c * dx - sn * dy), y1 + 0.35 * (sn * dx + c * dy))],
                      fill=tuple(colour), width=1)
    return np.array(im)


RAMPS = {  # key: (unit, lo, hi, colour stops) — shared by previews, viewer textures and legends
    "elevation": ("m", -6000.0, 6000.0, [[8, 30, 70], [30, 90, 150], [140, 200, 225], [60, 120, 60],
                                          [150, 160, 90], [140, 110, 80], [235, 235, 235]]),
    "T_mean": ("degC", -40.0, 40.0, [[40, 60, 160], [240, 240, 240], [180, 30, 30]]),
    "T_range": ("degC", 0.0, 60.0, [[68, 1, 84], [33, 145, 140], [253, 231, 37]]),
    "P": ("mm/yr", 0.0, 4000.0, [[150, 110, 60], [230, 220, 150], [60, 150, 70], [30, 70, 170]]),
    "po2": ("bar", 0.08, 0.32, [[60, 40, 90], [240, 240, 240], [200, 90, 20]]),
    "gravity": ("g", 0.3, 1.8, [[20, 120, 180], [240, 240, 240], [120, 40, 40]]),
    "pressure": ("bar", 0.3, 2.5, [[40, 60, 120], [240, 240, 240], [150, 60, 30]]),
    "o2_fraction": ("fraction", 0.10, 0.40, [[60, 40, 90], [240, 240, 240], [200, 90, 20]]),
    "fire": ("x", 0.0, 1.5, [[40, 80, 160], [240, 240, 240], [220, 120, 20]]),
    "vent_potential": ("0-1", 0.0, 1.0, [[10, 20, 50], [60, 90, 160], [250, 190, 60], [255, 250, 220]]),
    "bottom_temp": ("degC", -2.0, 30.0, [[30, 40, 120], [80, 160, 200], [240, 200, 120]]),
    "sediment": ("m", 0.0, 3000.0, [[40, 30, 60], [140, 110, 80], [240, 220, 170]]),
    "plant_height": ("x", 0.5, 3.5, [[150, 120, 60], [240, 240, 240], [30, 110, 50]]),
    "sst": ("degC", -2.0, 32.0, [[30, 40, 120], [60, 150, 200], [240, 240, 200], [220, 90, 40]]),
    "productivity": ("0-1", 0.0, 1.0, [[10, 20, 60], [20, 110, 120], [120, 200, 90], [240, 240, 120]]),
    "current_speed": ("m/s", 0.0, 1.0, [[10, 20, 50], [40, 90, 170], [120, 200, 230], [255, 255, 255]]),
    "upwelling": ("m/yr", -200.0, 200.0, [[40, 60, 160], [240, 240, 240], [30, 140, 80]]),
}
VIEWER_LAYERS = [  # id, name, source (continuous raster name | categorical name | "relief")
    ("relief", "Relief", "relief"), ("biomes", "Biomes (Holdridge)", "holdridge"),
    ("elevation", "Elevation", "elevation"), ("temperature", "Mean temperature", "T_mean"),
    ("rainfall", "Rainfall", "P_ann"), ("seasonality", "Seasonality", "seasonality"),
    ("landform", "Landform", "landform"), ("ground", "Ground", "ground"), ("ice", "Ice", "ice"), ("deposits", "Mineral deposits", "deposits"),
    ("plates", "Plates", "plates"), ("o2", "O₂ partial pressure", "po2"), ("gravity", "Gravity", "gravity"),
    ("pressure", "Air pressure", "pressure"), ("fire", "Fire reactivity", "fire"),
    ("seabed", "Sea-floor type", "seabed_type"), ("minerals", "Sea-floor minerals", "seabed_mineral"),
    ("bottom_temp", "Bottom temperature", "bottom_temp"), ("sediment", "Sediment", "sediment"),
    ("vent_potential", "Vent potential", "vent_potential"),
    ("currents", "Ocean currents", "current_speed"), ("sst", "Sea-surface temperature", "sst"),
    ("productivity", "Sea productivity", "productivity"),
]


def _ramp_key(name):
    return "P" if name.startswith("P_") else name


def _colorize(name, v):
    _, lo, hi, stops = RAMPS[_ramp_key(name)]
    return _ramp(v, lo, hi, stops)


def _preview(img: Image.Image, width: int, nearest: bool) -> Image.Image:
    return img.resize((width, width // 2), Image.NEAREST if nearest else Image.LANCZOS)


def contact_sheet(tiles, tile_w):
    cols = 3
    th = tile_w // 2 + 14
    rows = (len(tiles) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * tile_w, rows * th), (20, 20, 24))
    draw = ImageDraw.Draw(sheet)
    for k, (name, im) in enumerate(tiles):
        x, y = (k % cols) * tile_w, (k // cols) * th
        sheet.paste(im.resize((tile_w, tile_w // 2)), (x, y + 14))
        draw.text((x + 4, y + 1), name, fill=(235, 235, 235))
    return sheet


class _Writer:
    """Saves files on a background thread (PNG/zlib encoding releases the GIL) while the next layer is computed;
    at most `depth` waiting, so memory stays bounded. The files are the same bytes as saved in line."""
    def __init__(self, threads: int = 2, depth: int = 4):
        from concurrent.futures import ThreadPoolExecutor
        self.ex, self.pending, self.depth = ThreadPoolExecutor(threads), [], depth

    def __call__(self, fn, *args, **kw):
        while len(self.pending) >= self.depth:
            self.pending.pop(0).result()
        self.pending.append(self.ex.submit(fn, *args, **kw))

    def close(self):
        try:
            for f in self.pending:
                f.result()
        finally:
            self.ex.shutdown(wait=True, cancel_futures=True)


def run(ctx) -> dict:
    writer = _Writer()
    try:
        return _run(ctx, writer)
    finally:
        writer.close()


def _run(ctx, save) -> dict:
    g, cfg, d = ctx.grid, ctx.cfg, ctx.data
    W = int(cfg["build"]["raster_width"])
    if ctx.res < int(cfg["build"]["res_final"]):
        W = int(cfg["build"].get("dev_raster_width", W))
    H = W // 2
    pw = int(cfg["build"]["preview_width"])
    out = ctx.out_dir or ctx.root / "out" / f"r{ctx.res}"
    rdir = out / "raster"
    pdir = ctx.preview_dir or ctx.root / "previews" / f"r{ctx.res}"
    rdir.mkdir(parents=True, exist_ok=True)
    pdir.mkdir(parents=True, exist_ok=True)
    idx, wts = pixel_neighbours(g, W, H)
    legends = {**LEGENDS, "age_class": AGE_NAMES, "ice": ICE_NAMES, "plates": [p["id"] for p in ctx.tect["plate"]],
               "seabed_type": SEABED_NAMES, "seabed_mineral": MINERAL_NAMES}
    meta = {"width": W, "height": H, "projection": "equirectangular", "radius_km": g.radius_km,
            "units": cfg.get("units", {}), "continuous": {}, "categorical": {}}
    tiles = []
    cats = {name: sample_cat(np.asarray(d[key]), idx) for name, key in CATEGORICAL.items()}

    z = sample_cont(np.asarray(d["z_surface_m"], dtype=np.float64), idx, wts)
    amp = DETAIL_M[np.clip(cats["landform"], 0, len(DETAIL_M) - 1)]
    for y0 in range(0, H, CHUNK):
        rows = np.arange(y0, min(H, y0 + CHUNK))
        z[rows] += amp[rows] * fbm(_rows_xyz(rows, W, H), ctx.seed + 61, 5, 64.0).reshape(len(rows), W)
    lake = sample_cat(np.asarray(d["lake"]), idx)
    ocean_cells = np.asarray(d["ocean"]) if "ocean" in d else np.asarray(d["z_surface_m"]) <= 0
    ocean_k = ocean_cells[idx]
    coast = ocean_k.any(axis=-1) & ~ocean_k.all(axis=-1)
    for y0 in range(0, H, CHUNK):                      # extra fine detail where the coastline runs
        rows = np.arange(y0, min(H, y0 + CHUNK))
        cz = COAST_DETAIL_M * fbm(_rows_xyz(rows, W, H), ctx.seed + 67, 5, 256.0).reshape(len(rows), W)
        z[rows] += np.where(coast[rows], cz, 0.0)
    land = pixel_land(ocean_k, z)
    if "lake_level_m" in d:   # the water surface for drawing: lake cells at their level (beds stay in `elevation`)
        lev = np.asarray(d["lake_level_m"], dtype=np.float64)
        dz = sample_cont(np.where(np.isfinite(lev), lev - np.asarray(d["z_surface_m"], dtype=np.float64), 0.0), idx, wts)
        save(Image.fromarray(encode(z + dz, *CONTINUOUS["elevation"][1:3])).save, rdir / "surface.png")
        meta["continuous"]["surface"] = {"file": "surface.png", "scale": CONTINUOUS["elevation"][1],
                                         "offset": CONTINUOUS["elevation"][2], "unit": "m"}
    hs = hillshade(z, g.radius_km, low_memory=ctx.low_memory)
    style = cfg.get("render", {}).get("style")
    rgb = relief_rgb(z, hs, cats["holdridge"], cats["ground"], cats["ice"], lake, land, style=style,
                     low_memory=ctx.low_memory)
    rgb = draw_rivers(rgb, g, np.asarray(d["recv"]), np.asarray(d["river"]), np.asarray(d["strahler"]),
                      colour=RIVER_RGB if style is None else STYLES[style]["river"])
    relief = Image.fromarray(rgb)
    projections.write_all(rgb, pdir, g, ocean_cells, pw, globe_size=max(pw // 2, 64))
    save(relief.copy().save, rdir / "relief.png")
    tiles.append(("relief", _preview(relief, pw, False)))

    vdir = out / "viewer"
    by_src = {}                                        # viewer layers are written as soon as their colours exist
    for vid, vname, src in VIEWER_LAYERS:
        by_src.setdefault(src, []).append((vid, vname))
    ventries = {}

    def emit(src, rgb_, legend):
        for vid, vname in by_src.get(src, []):
            ventries[vid] = viewer_export.write_layer(vdir, {"id": vid, "name": vname, "rgb": rgb_, "legend": legend})

    emit("relief", rgb, None)
    for name, (key, scale, offset, unit) in CONTINUOUS.items():
        v = z if name == "elevation" else sample_cont(np.asarray(d[key], dtype=np.float64), idx, wts)
        save(Image.fromarray(encode(v, scale, offset)).save, rdir / f"{name}.png")
        meta["continuous"][name] = {"file": f"{name}.png", "scale": scale, "offset": offset, "unit": unit}
        crgb = _colorize(name, v)
        if name == "current_speed" and "current" in d:
            crgb = draw_currents(crgb, g, d["current"], ocean_cells)
        if name != "elevation":
            tiles.append((name, _preview(Image.fromarray(crgb), pw, False)))
        if name in by_src:
            unit_, lo, hi, stops = RAMPS[_ramp_key(name)]
            emit(name, crgb, viewer_export.continuous_legend(unit_, lo, hi, stops))
        del crgb
    for name, v in cats.items():
        save(Image.fromarray(v.astype(np.uint8), "L").save, rdir / f"{name}.png")
        meta["categorical"][name] = {"file": f"{name}.png", "legend": legends[name]}
        pal = holdridge_palette() if name == "holdridge" else category_palette(len(legends[name]))
        crgb = pal[np.clip(v, 0, len(pal) - 1)]
        tiles.append((name, _preview(Image.fromarray(crgb), pw, True)))
        emit(name, crgb, viewer_export.categorical_legend(legends[name], pal))
        del crgb
    (out / "fields.json").write_text(json.dumps(meta, indent=1))

    per_cell = {k: v for k, v in d.items() if isinstance(v, np.ndarray) and v.shape[:1] == (g.n,)}
    save(lambda: np.savez_compressed(out / "cells.npz", **per_cell))
    (out / "cells_meta.json").write_text(json.dumps(
        {"res": ctx.res, "radius_km": g.radius_km, "n_cells": g.n, "units": cfg.get("units", {}),
         "planet": {**cfg["planet"], **({"sun_lock": cfg["climate"].get("lock_at", [0.0, 0.0])}
                                        if cfg.get("climate", {}).get("lock") else {})}, "name": cfg.get("render", {}).get("name", "World"), "style": style,
         "legends": legends, "fields": sorted(per_cell),
         "plateaus": [PL.features(p, ctx.seed, g.radius_km) for p in ctx.tect.get("plateau", [])]}, indent=1))

    for name, im in tiles:
        im.save(pdir / f"{name}.png")
    contact_sheet(tiles, max(pw // 3, 64)).save(pdir / "contact_sheet.png")
    viewer_export.write_index(vdir, [ventries[vid] for vid, _, _ in VIEWER_LAYERS])
    geo.write_all(ctx, out, land, lake)
    return {}