#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""The shipped sample of meteor-p150: a synthetic eight-camera frame rendered here (data generated by this repository,
Apache-2.0; no third-party pixels or calibration). CPU only, research venv (torch, OpenCV)::

    cd bundles/meteor-p150; PY=/home/ubuntu/experiments/tt-models/tools/research-venv/bin/python
    OMP_NUM_THREADS=6 PYTHONPATH=$PWD/code $PY code/scripts/make_synthetic_sample.py render     # ~3 min
    OMP_NUM_THREADS=6 PYTHONPATH=$PWD/code $PY code/scripts/make_synthetic_sample.py optimise   # ~15 s per step
    OMP_NUM_THREADS=6 PYTHONPATH=$PWD/code $PY code/scripts/make_synthetic_sample.py finalise   # PNGs + reference

(``--work`` holds the intermediate arrays, default ``logs/synthetic_sample/``; ``optimise`` resumes from its
``images_f32.npy``.) The shipped sample came from ``render``, ``optimise --steps 80`` (stopped after 28 steps to drop
the pedestrian beyond 20 m from the targets), ``optimise --steps 60 --lr 0.015`` and ``finalise``.

A generic METEOR-like rig (``calib/synthetic_8cam.json``: the eight slots in METEOR's order, a 98 deg wide camera
front and back, four corner cameras pitched 25 deg down, two 30.4 deg narrow cameras; fy = 0.87 fx like METEOR's
vertically squashed training images; mounted 1.9 m above the road) looks at a procedural street scene that is
ray-cast per pixel (2x2 supersampled, INTER_AREA down to 768x432): a straight two-lane road with lane lines and a
stop line, kerbs, pavements, building facades, box-shaped vehicles and pedestrians. All eight cameras are present.
Plain renders give the network too little to detect (the strongest 3D heatmap cell was 0.22 < the 0.35 vehicle
threshold), so ``optimise`` then changes the object pixels only (gradient ascent through the fp32 CPU reference: image
encoder, lift, BEV detector, box refiner; the image rounded to uint8 by a straight-through estimator) until every
target object within range (:data:`DROP_FAR`) clears :data:`GOAL` at its cell and nothing else clears :data:`CEIL`:
a stable smoke reference whose published boxes sit far from the thresholds.

Writes:

- ``code/tt_meteor/calib/synthetic_8cam.json``: the rig as a calibration preset (served in ``/info``);
- ``code/tt_meteor/samples/synthetic_8cam/<CAM>.png``: the eight 768x432 RGB images (lossless, so every decoder
  gives the same pixels);
- ``code/tt_meteor/samples/synthetic_8cam.json``: the request manifest (``model(**load_sample(path))``);
- ``code/tt_meteor/samples/synthetic_8cam.reference.json``: the ``/predict`` body of the fp32 CPU reference
  (``tt_meteor.reference.pipeline.MeteorReference``, PCC 1.0 vs ONNX Runtime) on that request as a fresh stream,
  the container smoke's stored reference (``--no-reference`` skips it: 40-60 s of CPU);
- ``code/tt_meteor/tests/goldens/synthetic_8cam_inputs.json``: sha256 of the graph feed (``imgs``, ``K``,
  ``T_cam_ego``, ``v0``) the request produces, so a host test notices any change of the decoding / resize path.
"""
from __future__ import annotations

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

import numpy as np

BUNDLE = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(BUNDLE / "code"))
PKG = BUNDLE / "code" / "tt_meteor"
NAME = "synthetic_8cam"
W, H, SS = 768, 432, 2                  # network input size, supersampling factor
EGO_SPEED = 8.0                         # m/s
SEED = 20261010

# slot -> (x, y, z of the camera in base_link, yaw deg CCW from +x, pitch deg down, horizontal FOV deg)
RIG = {
    "CAM_FRONT_WIDE": (2.10, 0.00, 1.90, 0.0, 2.0, 98.0),
    "CAM_FRONT_LEFT": (1.90, 0.85, 1.90, 60.0, 25.0, 98.0),
    "CAM_FRONT_RIGHT": (1.90, -0.85, 1.90, -60.0, 25.0, 98.0),
    "CAM_BACK_WIDE": (-0.90, 0.00, 1.90, 180.0, 2.0, 98.0),
    "CAM_BACK_LEFT": (-0.70, 0.85, 1.90, 120.0, 25.0, 98.0),
    "CAM_BACK_RIGHT": (-0.70, -0.85, 1.90, -120.0, 25.0, 98.0),
    "CAM_FRONT_NARROW": (2.10, 0.00, 1.88, 0.0, 1.0, 30.4),
    "CAM_BACK_NARROW": (-0.90, 0.00, 1.88, 180.0, 2.0, 30.4),
}
FY_OVER_FX = 0.87                       # METEOR's training images are squashed vertically (2880x1860 -> 768x432)

# the street (base_link: x forward, y left, z up, origin on the road below the rear axle; left-hand traffic, ego in
# the left lane): lane edges and lines (y, half width, kind)
ROAD_L, ROAD_R = 1.75, -5.25            # carriageway edges (ego lane 1.75 .. -1.75, oncoming lane -1.75 .. -5.25)
KERB = 0.15                             # kerb height
WALK_L, WALK_R = 5.0, -8.5              # pavement outer edges = building facades
# boxes: (x, y, yaw deg, length, width, height, rgb, kind)
OBJECTS = [
    (11.0, 0.0, 0.0, 4.4, 1.8, 1.5, (180, 30, 35), "car"),          # car ahead in the ego lane
    (19.0, -3.5, 180.0, 4.6, 1.85, 1.55, (40, 70, 160), "car"),     # oncoming cars
    (29.0, -3.4, 180.0, 4.5, 1.8, 1.5, (150, 150, 155), "car"),
    (-7.0, -3.5, 180.0, 4.4, 1.8, 1.5, (200, 170, 40), "car"),
    (42.0, 0.1, 0.0, 8.5, 2.4, 3.0, (225, 225, 220), "truck"),      # truck further ahead
    (-13.0, 0.0, 0.0, 4.3, 1.75, 1.45, (220, 220, 225), "car"),     # car behind
    (6.0, 3.2, 0.0, 4.2, 1.75, 1.45, (60, 60, 65), "car"),          # parked car on the left (on the pavement edge)
    (-4.0, -6.8, 90.0, 0.5, 0.6, 1.7, (40, 40, 90), "ped"),         # pedestrian on the right pavement
    (9.0, 4.0, 0.0, 0.5, 0.6, 1.75, (120, 40, 40), "ped"),          # pedestrian on the left pavement
    (22.0, -7.2, 0.0, 0.5, 0.6, 1.65, (30, 90, 60), "ped"),
]
STOP_LINE_X = 34.0                      # stop line across the ego lane
CROSSWALK_X = (36.0, 40.0)              # zebra across the road


def object_parts():
    """Every object as oriented boxes: (centre x, y, z, yaw deg, half extents (l, w, h), rgb, material)."""
    out = []
    for (bx, by, yaw, L, Wd, Hh, col, kind) in OBJECTS:
        z0 = KERB if (by > ROAD_L or by < ROAD_R) else 0.0                 # on the pavement or on the road
        c, s = math.cos(math.radians(yaw)), math.sin(math.radians(yaw))

        def at(dx, dy, dz, half, rgb, mat, yaw_=yaw):
            out.append((bx + c * dx - s * dy, by + s * dx + c * dy, z0 + dz, yaw_, half, rgb, mat))

        if kind == "car":
            hb = 0.55 * Hh                                                   # lower body 0.25 .. 0.25 + hb
            at(0, 0, 0.25 + hb / 2, (L / 2, Wd / 2, hb / 2), col, "paint")
            at(-0.15 * L, 0, 0.25 + hb + (Hh - 0.25 - hb) / 2, (0.28 * L, 0.44 * Wd, (Hh - 0.25 - hb) / 2),
               (55, 65, 80), "glass")                                        # cabin
            at(-0.15 * L, 0, Hh - 0.02, (0.26 * L, 0.42 * Wd, 0.03), col, "paint")   # roof
            for sx in (0.32, -0.32):
                for sy in (1, -1):
                    at(sx * L, sy * (Wd / 2 - 0.12), 0.32, (0.32, 0.13, 0.32), (22, 22, 24), "rubber")
            for sy in (1, -1):                                               # lamps
                at(L / 2 - 0.02, sy * 0.33 * Wd, 0.25 + 0.7 * hb, (0.04, 0.16, 0.06), (240, 240, 215), "lamp")
                at(-L / 2 + 0.02, sy * 0.33 * Wd, 0.25 + 0.7 * hb, (0.04, 0.16, 0.06), (210, 25, 25), "lamp")
        elif kind == "truck":
            cab = 1.9
            at(L / 2 - cab / 2, 0, 0.35 + (Hh - 0.6 - 0.35) / 2, (cab / 2, Wd / 2, (Hh - 0.6 - 0.35) / 2), (40, 90, 160), "paint")
            at(L / 2 - 0.05, 0, Hh - 1.3, (0.06, 0.45 * Wd, 0.35), (55, 65, 80), "glass")
            at(-cab / 2, 0, 0.6 + (Hh - 0.6) / 2, ((L - cab) / 2 - 0.05, Wd / 2, (Hh - 0.6) / 2), col, "paint")
            for sx in (0.38, 0.0, -0.33):
                for sy in (1, -1):
                    at(sx * L, sy * (Wd / 2 - 0.15), 0.45, (0.45, 0.15, 0.45), (22, 22, 24), "rubber")
        else:                                                                # pedestrian
            at(0, 0, 0.42, (0.12, 0.2, 0.42), (40, 40, 55), "cloth")
            at(0, 0, 0.84 + (Hh - 1.07) / 2, (0.14, 0.24, (Hh - 1.07) / 2), col, "cloth")
            at(0, 0, Hh - 0.115, (0.1, 0.09, 0.115), (215, 170, 140), "skin")
    return out


def shadow(px: np.ndarray, py: np.ndarray) -> np.ndarray:
    """Soft contact shadows of the objects on the ground: a multiplicative factor in (0, 1]."""
    f = np.ones(px.shape, np.float32)
    for (bx, by, yaw, L, Wd, Hh, col, kind) in OBJECTS:
        c, s = math.cos(math.radians(yaw)), math.sin(math.radians(yaw))
        dx, dy = px - bx, py - by
        lx, ly = np.abs(c * dx + s * dy) - L / 2, np.abs(-s * dx + c * dy) - Wd / 2
        out = np.hypot(np.maximum(lx, 0), np.maximum(ly, 0)) + np.minimum(np.maximum(lx, ly), 0)
        f *= (1.0 - 0.75 * np.clip(1.0 - (out + 0.15) / 0.6, 0, 1)).astype(np.float32)
    return f


def rotation(yaw_deg: float, pitch_deg: float) -> np.ndarray:
    """camera optical frame (x right, y down, z forward) -> base_link rotation."""
    p, y = math.radians(pitch_deg), math.radians(yaw_deg)
    f = np.array([math.cos(p) * math.cos(y), math.cos(p) * math.sin(y), -math.sin(p)])
    r = np.array([math.sin(y), -math.cos(y), 0.0])
    d = np.cross(f, r)
    return np.stack([r, d, f], axis=1)


def intrinsics(hfov_deg: float, w: int = W, h: int = H) -> np.ndarray:
    fx = (w / 2.0) / math.tan(math.radians(hfov_deg) / 2.0)
    return np.array([[fx, 0.0, w / 2.0], [0.0, FY_OVER_FX * fx, h / 2.0], [0.0, 0.0, 1.0]])


def value_noise(x: np.ndarray, y: np.ndarray, cell: float, seed: int) -> np.ndarray:
    """Smooth hash noise in [0, 1) on a ground-plane lattice of ``cell`` metres."""
    gx, gy = np.floor(x / cell), np.floor(y / cell)
    fx, fy = x / cell - gx, y / cell - gy

    def h(ix, iy):
        v = (ix.astype(np.int64) * 73856093) ^ (iy.astype(np.int64) * 19349663) ^ seed
        v = (v * 2654435761) & 0xFFFFFFFF
        return (v % 10007) / 10007.0

    sx, sy = fx * fx * (3 - 2 * fx), fy * fy * (3 - 2 * fy)
    a, b, c, d = h(gx, gy), h(gx + 1, gy), h(gx, gy + 1), h(gx + 1, gy + 1)
    return (a * (1 - sx) + b * sx) * (1 - sy) + (c * (1 - sx) + d * sx) * sy


def ground_colour(px: np.ndarray, py: np.ndarray, dist: np.ndarray) -> np.ndarray:
    n = 0.6 * value_noise(px, py, 0.35, SEED) + 0.4 * value_noise(px, py, 2.0, SEED + 1)
    rgb = np.empty(px.shape + (3,), np.float32)
    road = (py <= ROAD_L) & (py >= ROAD_R)
    asphalt = 78 + 34 * n
    rgb[...] = np.stack([asphalt, asphalt, asphalt + 4], -1)[...]
    walk = ~road
    tile = 150 + 30 * n + 12 * ((np.floor(px / 0.6) + np.floor(py / 0.6)) % 2)
    rgb[walk] = np.stack([tile, tile - 6, tile - 14], -1)[walk]
    # markings on the carriageway: edge lines (solid), centre line (dashed 5 m / 5 m), stop line, zebra
    white = np.zeros(px.shape, bool)
    white |= road & (np.abs(py - (ROAD_L - 0.2)) < 0.075)
    white |= road & (np.abs(py - (ROAD_R + 0.2)) < 0.075)
    white |= (np.abs(py + 1.75) < 0.075) & (np.mod(px, 10.0) < 5.0)
    white |= (py <= 1.6) & (py >= -1.75) & (np.abs(px - STOP_LINE_X) < 0.225)
    zebra = road & (px >= CROSSWALK_X[0]) & (px <= CROSSWALK_X[1]) & (np.mod(py - ROAD_R, 0.9) < 0.45)
    white |= zebra
    wv = 215 + 25 * n
    rgb[white] = np.stack([wv, wv, wv - 5], -1)[white]
    rgb *= shadow(px, py)[..., None]
    # haze with distance
    fog = np.clip(dist / 120.0, 0, 0.6)[..., None]
    return rgb * (1 - fog) + np.array([170, 180, 195], np.float32) * fog


def render_camera(R: np.ndarray, t: np.ndarray, K: np.ndarray):
    """Ray-cast the street for one camera; returns (uint8 (H*SS, W*SS, 3), bool object mask (H*SS, W*SS))."""
    w, h = W * SS, H * SS
    Ks = K.copy()
    Ks[:2] *= SS
    u, v = np.meshgrid(np.arange(w) + 0.5, np.arange(h) + 0.5)
    rays_c = np.stack([(u - Ks[0, 2]) / Ks[0, 0], (v - Ks[1, 2]) / Ks[1, 1], np.ones_like(u)], -1)
    d = rays_c @ R.T
    d /= np.linalg.norm(d, axis=-1, keepdims=True)
    best = np.full(d.shape[:2], np.inf)
    rgb = np.zeros(d.shape, np.float32)
    obj = np.zeros(d.shape[:2], bool)
    # sky
    up = np.clip(d[..., 2], -1, 1)
    sky = np.stack([120 + 60 * (1 - up), 160 + 50 * (1 - up), 225 + 20 * (1 - up)], -1)
    rgb[...] = sky
    # ground (z = 0; pavements at kerb height)
    with np.errstate(divide="ignore", invalid="ignore"):
        for z0, keep in ((0.0, lambda py: (py <= ROAD_L) & (py >= ROAD_R)),
                         (KERB, lambda py: (py > ROAD_L) | (py < ROAD_R))):
            tg = (z0 - t[2]) / d[..., 2]
            ok = (tg > 0) & np.isfinite(tg)
            px, py = t[0] + tg * d[..., 0], t[1] + tg * d[..., 1]
            ok &= keep(py) & (py <= WALK_L) & (py >= WALK_R)
            ok &= tg < best
            col = ground_colour(px, py, tg)
            rgb[ok], best[ok] = col[ok], tg[ok]
        # kerb faces (vertical planes y = ROAD_L / ROAD_R, 0 .. KERB)
        for yk in (ROAD_L, ROAD_R):
            tk = (yk - t[1]) / d[..., 1]
            zk = t[2] + tk * d[..., 2]
            ok = (tk > 0) & (zk >= 0) & (zk <= KERB) & (tk < best)
            rgb[ok], best[ok] = np.array([185, 180, 170], np.float32), tk[ok]
        # facades (y = WALK_L / WALK_R), 6-14 m tall blocks with windows
        for yw, sd in ((WALK_L, 3), (WALK_R, 4)):
            tw = (yw - t[1]) / d[..., 1]
            xw, zw = t[0] + tw * d[..., 0], t[2] + tw * d[..., 2]
            block = np.floor(xw / 12.0)
            hb = 6 + 8 * value_noise(block * 12.0, np.zeros_like(block), 12.0, SEED + sd)
            ok = (tw > 0) & (zw >= KERB) & (zw <= hb) & (tw < best)
            base = 0.55 + 0.45 * value_noise(block * 12.0, np.full_like(block, 5.0), 12.0, SEED + 10 + sd)
            wall = np.stack([200 * base, 170 * base, 140 * base], -1)
            win = (np.mod(xw, 3.0) > 0.9) & (np.mod(zw - 0.8, 3.2) < 1.6) & (zw > 3.0)
            wall[win] = np.array([70, 90, 115], np.float32)
            door = (zw < 2.4) & (np.mod(xw, 12.0) > 5.0) & (np.mod(xw, 12.0) < 6.6)
            wall[door] = np.array([60, 45, 35], np.float32)
            fog = np.clip(tw / 120.0, 0, 0.6)[..., None]
            wall = wall * (1 - fog) + np.array([170, 180, 195], np.float32) * fog
            rgb[ok], best[ok] = wall[ok], tw[ok]
    # objects: every part is an oriented box (slab test in the part frame)
    light = np.array([0.4, 0.3, 0.87])
    light /= np.linalg.norm(light)
    for (cx, cy, cz, yaw, half, col, mat) in object_parts():
        c, s = math.cos(math.radians(yaw)), math.sin(math.radians(yaw))
        Rb = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1.0]])
        o = Rb.T @ (t - np.array([cx, cy, cz]))
        dl = d @ Rb
        half = np.asarray(half, np.float64)
        with np.errstate(divide="ignore", invalid="ignore"):
            t1 = (-half - o) / dl
            t2 = (half - o) / dl
        tmin = np.nanmax(np.minimum(t1, t2), axis=-1)
        tmax = np.nanmin(np.maximum(t1, t2), axis=-1)
        hit = (tmax >= tmin) & (tmin > 0) & (tmin < best)
        if not hit.any():
            continue
        p = o + tmin[..., None] * dl                                     # hit point, part frame
        ax = np.argmax(np.abs(p) / half, axis=-1)
        nrm_l = np.zeros_like(p)
        np.put_along_axis(nrm_l, ax[..., None], np.sign(np.take_along_axis(p, ax[..., None], -1)), -1)
        nrm = nrm_l @ Rb.T
        lam = np.clip(nrm @ light, 0, 1)
        base = np.broadcast_to(np.array(col, np.float32), p.shape).copy()
        if mat == "glass":                                                # sky reflection on upward-tilted views
            refl = np.clip(0.35 + 0.65 * np.abs(nrm[..., 2]), 0, 1)
            shade = 0.5 + 0.3 * refl
        elif mat == "paint":
            shade = 0.35 + 0.65 * lam + 0.08 * (value_noise(p[..., 0], p[..., 2], 0.4, SEED + 7) - 0.5)
        else:
            shade = 0.4 + 0.6 * lam
        cshade = base * shade[..., None]
        rgb[hit], best[hit] = cshade[hit], tmin[hit]
        obj |= hit
    obj &= np.isfinite(best)
    return np.clip(rgb, 0, 255).astype(np.uint8), obj


def build_rig():
    cams = {}
    for name, (x, y, z, yaw, pitch, fov) in RIG.items():
        R = rotation(yaw, pitch)
        T = np.eye(4)
        T[:3, :3], T[:3, 3] = R, (x, y, z)
        cams[name] = {"intrinsics": intrinsics(fov).round(4).tolist(), "T_ref_from_camera": T.round(9).tolist(),
                      "image_size": [W, H]}
    return cams


def hm_cell(x: float, y: float):
    """base_link (x, y) -> (row, col) of the 400x250 detection grid (rows x = 79.8 .. -79.8, cols y = 49.8 .. -49.8)."""
    return int(round((79.8 - x) / 0.4)), int(round((49.8 - y) / 0.4))


# what the optimisation asks of the 3D heads: every OBJECTS entry is a target (vehicle: car / truck, VRU: ped)
GOAL = {0: 0.70, 1: 0.50}               # sigmoid at the target cell (thresholds: vehicle 0.35, VRU 0.15)
CEIL = {0: 0.20, 1: 0.07}               # sigmoid ceiling everywhere else (outside 1.6 m of a target of that class)
DROP_FAR = {0: 40.0, 1: 20.0}           # targets beyond this range are scenery (pushed below the ceiling instead)


def targets():
    out = []
    for (bx, by, yaw, L, Wd, Hh, col, kind) in OBJECTS:
        cls = 0 if kind in ("car", "truck") else 1
        if math.hypot(bx, by) <= DROP_FAR[cls]:
            out.append((cls, bx, by))
    return out


def differentiable_hm(net, x, K, T):
    """The 3D heatmap logits [1, 2, 400, 250] of the fp32 reference from float images x [8, 3, 432, 768] in [0, 1]
    (the graph's /255 input normalisation; the same modules as ``MeteorNet.forward``, with gradients)."""
    f = net.image_encoder(x)
    seg = net.seg2d_head(f)
    _, dprob = net.depth_head(f)
    ctx = net.context(f, seg)
    raw = net.lift(ctx, dprob, K, T)
    _, hm_pre, reg_pre = net.det_stem(raw)
    hm, _ = net.box_refiner(hm_pre, reg_pre)
    return hm


def stage_render(work: Path) -> None:
    import cv2

    t0 = time.time()
    cams = build_rig()
    imgs, masks = [], []
    for name in RIG:
        c = cams[name]
        T = np.asarray(c["T_ref_from_camera"])
        img, obj = render_camera(T[:3, :3], T[:3, 3], np.asarray(c["intrinsics"]))
        imgs.append(cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA))
        m = cv2.resize(obj.astype(np.uint8) * 255, (W, H), interpolation=cv2.INTER_AREA) > 0
        masks.append(cv2.dilate(m.astype(np.uint8), np.ones((9, 9), np.uint8)) > 0)
    work.mkdir(parents=True, exist_ok=True)
    np.save(work / "base_u8.npy", np.stack(imgs))                      # [8, H, W, 3] RGB
    np.save(work / "mask.npy", np.stack(masks))                        # [8, H, W]
    print(f"render: {time.time() - t0:.0f} s, object pixels {np.stack(masks).mean():.3%}", flush=True)


def _load_net(weights_dir, threads):
    import torch

    from tt_meteor.reference.pipeline import MeteorReference
    from tt_meteor.reference.weights import MeteorWeights

    torch.set_num_threads(threads)
    return MeteorReference(weights=MeteorWeights(weights_dir), threads=threads)   # plain tensors: no weight grads


def _calib_tensors():
    import torch

    from tt_meteor.host.preprocess import invert_extrinsics

    cams = build_rig()
    K = np.stack([np.asarray(cams[n]["intrinsics"], np.float32) for n in RIG])[None]
    T = np.stack([invert_extrinsics(cams[n]["T_ref_from_camera"]) for n in RIG])[None]
    return torch.from_numpy(K), torch.from_numpy(T)


def stage_optimise(work: Path, weights_dir, threads: int, steps: int, lr: float) -> None:
    """Gradient ascent on the object pixels (and only those) until every target clears its goal on the uint8-rounded
    images and nothing else clears its ceiling."""
    import torch

    ref = _load_net(weights_dir, threads)
    K, T = _calib_tensors()
    base = torch.from_numpy(np.load(work / "base_u8.npy").astype(np.float32) / 255.0).permute(0, 3, 1, 2)
    mask = torch.from_numpy(np.load(work / "mask.npy")).float()[:, None]
    start = work / "images_f32.npy"
    x0 = torch.from_numpy(np.load(start)) if start.is_file() else base.clone()
    delta = (x0 - base).clone().requires_grad_(True)
    opt = torch.optim.Adam([delta], lr=lr)
    tg = targets()
    logit = lambda p: math.log(p / (1 - p))
    keep = torch.ones(1, 2, 400, 250, dtype=torch.bool)
    rr, cc = torch.meshgrid(torch.arange(400), torch.arange(250), indexing="ij")
    for cls, bx, by in tg:
        r, c = hm_cell(bx, by)
        keep[0, cls] &= (rr - r) ** 2 + (cc - c) ** 2 > 16              # 1.6 m around a target of its class
    ceil = torch.tensor([logit(CEIL[0]), logit(CEIL[1])]).reshape(1, 2, 1, 1)
    log = open(work / "optimise.log", "a")
    for step in range(steps):
        t0 = time.time()
        x = (base + delta * mask).clamp(0, 1)
        xq = x + ((x * 255).round() / 255 - x).detach()                  # straight-through uint8 rounding
        hm = differentiable_hm(ref.net, xq, K, T)
        z = []
        loss = torch.zeros(())
        for cls, bx, by in tg:
            r, c = hm_cell(bx, by)
            zt = hm[0, cls, r - 1:r + 2, c - 1:c + 2].max()             # the 3x3 peak at the target
            z.append(float(torch.sigmoid(zt)))
            loss = loss + torch.nn.functional.softplus(logit(GOAL[cls]) + 0.3 - zt) * 2.0
        over = torch.relu(hm - ceil)[keep]
        loss = loss + (over ** 2).sum() * 0.5
        opt.zero_grad()
        loss.backward()
        opt.step()
        with torch.no_grad():
            delta.clamp_(-0.6, 0.6)
        n_over = int((over > 0).sum())
        ok = all(p >= GOAL[cls] for p, (cls, _, _) in zip(z, tg)) and n_over == 0
        line = (f"step {step}: loss {float(loss):.4f}, targets {[round(p, 3) for p in z]}, cells over the ceiling "
                f"{n_over}, max other {float(torch.sigmoid(hm.detach()[keep]).max()):.3f}, {time.time() - t0:.0f} s")
        print(line, flush=True)
        log.write(line + "\n")
        log.flush()
        np.save(work / "images_f32.npy", x.detach().numpy())
        if ok:
            print("all targets met", flush=True)
            break


def stage_finalise(work: Path, weights_dir, threads: int, reference: bool) -> None:
    import cv2

    t0 = time.time()
    cams = build_rig()
    preset = {"frame_id": "base_link",
              "_source": "generated by code/scripts/make_synthetic_sample.py (a generic METEOR-like 8-camera rig of "
                         "round numbers; no third-party data), Apache-2.0",
              "_rig": {k: dict(zip(("x", "y", "z", "yaw_deg", "pitch_down_deg", "hfov_deg"), v)) for k, v in RIG.items()},
              "cameras": cams}
    (PKG / "calib" / f"{NAME}.json").write_text(json.dumps(preset, indent=1) + "\n")
    src = work / "images_f32.npy"
    if src.is_file():
        imgs = np.round(np.load(src) * 255).clip(0, 255).astype(np.uint8).transpose(0, 2, 3, 1)
    else:
        imgs = np.load(work / "base_u8.npy")
    out = PKG / "samples" / NAME
    out.mkdir(parents=True, exist_ok=True)
    for name, img in zip(RIG, imgs):
        cv2.imwrite(str(out / f"{name}.png"), np.ascontiguousarray(img[:, :, ::-1]), [cv2.IMWRITE_PNG_COMPRESSION, 9])
    req = {"images": {name: f"{NAME}/{name}.png" for name in RIG},
           "calibration": {"preset": NAME},
           "ego_speed": EGO_SPEED,
           "stream": {"id": NAME, "reset": True, "T_world_from_ego": {"x": 0.0, "y": 0.0, "z": 0.0, "roll": 0.0,
                                                                        "pitch": 0.0, "yaw": 0.0}},
           "_source": "synthetic street scene rendered and optimised by code/scripts/make_synthetic_sample.py (data "
                      "generated by this repository, Apache-2.0); calibration preset calib/synthetic_8cam.json; all "
                      "eight cameras present"}
    (PKG / "samples" / f"{NAME}.json").write_text(json.dumps(req, indent=1) + "\n")
    size = sum(p.stat().st_size for p in out.iterdir()) / 1e6
    print(f"wrote samples/{NAME}/ ({size:.2f} MB), samples/{NAME}.json, calib/{NAME}.json", flush=True)

    from tt_meteor.api import load_sample
    from tt_meteor.host.inputs import prepare_request

    kw = load_sample(PKG / "samples" / f"{NAME}.json")
    frame = prepare_request(kw["images"], kw["calibration"], kw["ego_speed"], kw.get("stream"))
    assert np.array_equal(frame.imgs[0].transpose(0, 2, 3, 1), imgs), "the PNGs do not read back bit-exactly"
    feed_sha = {k: hashlib.sha256(np.ascontiguousarray(v).tobytes()).hexdigest() for k, v in frame.feed().items()}
    (PKG / "tests" / "goldens" / f"{NAME}_inputs.json").write_text(json.dumps(
        {"_doc": f"sha256 of the graph feed of samples/{NAME}.json (code/scripts/make_synthetic_sample.py)",
         "input_sha256": feed_sha, "present": frame.present.astype(int).tolist()}, indent=1) + "\n")
    if not reference:
        return
    ref = _load_net(weights_dir, threads)
    t1 = time.time()
    body = ref(**kw).to_dict()
    body["timing_ms"] = {}
    (PKG / "samples" / f"{NAME}.reference.json").write_text(json.dumps(body) + "\n")
    labels = {}
    for d in body["detections"]:
        labels[d["label"]] = labels.get(d["label"], 0) + 1
    n2d = sum(len(v) for v in body["detections_2d"].values())
    print(f"reference ({time.time() - t1:.0f} s): {body['num_detections']} 3D boxes {labels} scores "
          f"{[round(d['score'], 3) for d in body['detections']]}, {n2d} 2D boxes, mode {body['plan']['mode']} "
          f"(probs {[round(v, 3) for v in body['plan']['mode_probs']]}), TL {body['traffic_light']['state']}, "
          f"path end {[round(v, 2) for v in body['trajectory'][-1]]}; total {time.time() - t0:.0f} s", flush=True)


def default_weights_dir():
    snaps = Path.home() / ".cache/huggingface/hub/models--AutowareFoundation--meteor/snapshots"
    for cand in sorted(snaps.glob("*")):
        if (cand / "meteor_v157c3Z.onnx").is_file():
            return cand
    return Path("/home/ubuntu/experiments/tt-models/assets/meteor/hf_meteor")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("stage", choices=("render", "optimise", "finalise", "all"))
    ap.add_argument("--work", type=Path, default=BUNDLE / "logs" / "synthetic_sample")
    ap.add_argument("--threads", type=int, default=4)
    ap.add_argument("--steps", type=int, default=60)
    ap.add_argument("--lr", type=float, default=0.02)
    ap.add_argument("--no-reference", action="store_true")
    ap.add_argument("--weights-dir", type=Path, default=None)
    a = ap.parse_args()
    wd = a.weights_dir or default_weights_dir()
    if a.stage in ("render", "all"):
        stage_render(a.work)
    if a.stage in ("optimise", "all"):
        stage_optimise(a.work, wd, a.threads, a.steps, a.lr)
    if a.stage in ("finalise", "all"):
        stage_finalise(a.work, wd, a.threads, not a.no_reference)


if __name__ == "__main__":
    main()