#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Generate the shipped sample of bevdet-p150: ``code/tt_bevdet/samples/synthetic_6cam*`` (a synthetic six-camera
frame made by this script, Apache-2.0 like the code), its calibration preset ``calib/synthetic_6cam.json``, the
stored CPU-reference output the container smoke compares with, and a sparse copy of the reference head maps for the
host tests. CPU only; run it in the research venv (torch + onnx)::

    cd bundles/bevdet-p150
    PY=/home/ubuntu/experiments/tt-models/tools/research-venv/bin/python
    OMP_NUM_THREADS=4 PYTHONPATH=$PWD/code $PY code/scripts/make_sample.py optimise --out /tmp/sample   # ~20-40 min
    OMP_NUM_THREADS=4 PYTHONPATH=$PWD/code $PY code/scripts/make_sample.py finalise --crop /tmp/sample/crop_f32.npy

Why a synthetic frame. The redistributable real data of this model (PandaSet, CC BY 4.0) may not ship until the user
approves PandaSet-derived samples (PLAN.md 6.3, research/DATASETS.md N5), nuScenes is non-commercial, and plain
synthetic images give the network nothing to detect (two or three boxes near the 0.2 threshold), which makes a weak
and fragile smoke reference. So the sample is a **test pattern**: six rendered road views of a generic rig whose
network-visible pixels were optimised against the fp32 CPU reference (gradient ascent through the image encoder,
BEVPool and the BEV decoder, fresh-stream semantics) until the network reports eight objects at chosen places, with
large margins (each target cell confident, its class the clear argmax and its logit above every BEV neighbour; the rest
of the map pushed below the 0.1 decode threshold). The images look like a road with a noisy texture; they show no
real objects. What they give is a stable reference: every published box sits far from the 0.1 / 0.2 thresholds, so
bf16 numerics cannot flip it. (Forcing every neighbour of a target below the decode threshold as well, to rule out
the occasional extra box of a small object beside its target, collapses the targets: the network's heatmap is too
smooth for one-cell peaks.)

- **Rig** (:data:`RIG`, round numbers of our own, nuScenes-like layout): six 1600x900 pinhole cameras at 1.55 m,
  yaws 0, +-55, +-110, 180 deg, fx 1260 px (800 px for the back camera), principal point (799.5, 480); BEVPool
  tables M 355,907 / N 12,498 (capacities 370,000 / 14,000).
- **Objects** (:data:`TARGETS`): three cars, a truck, a bus, a pedestrian, a bicycle and a traffic cone (UNKNOWN in
  Autoware), 4-24 m around the ego vehicle, at rest.
- **Images**: rows 318-899 of each camera (what the network's crop reads) carry the optimised pattern, upsampled
  from the 704x256 network grid by nearest neighbour, so the Preprocess gather of the 1600x900 PNG returns exactly
  the optimised uint8 crop; rows above keep the rendered sky. PNG, lossless (no decoder dependence).
- **Reference**: the ``/predict`` body of ``tt_bevdet.reference.BEVDetReference`` for the sample sent alone on a
  fresh stream with ``params: {"score_threshold": 0.1}`` (every vendor box: the smoke's agreement rule needs them),
  ``timing_ms`` emptied.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import sys
import time
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np

CODE = Path(__file__).resolve().parents[1]
if str(CODE) not in sys.path:
    sys.path.insert(0, str(CODE))

from tt_bevdet.reference.config import CAMERA_ORDER, CLASS_NAMES, default_config  # noqa: E402

PKG = CODE / "tt_bevdet"
NAME = "synthetic_6cam"
SRC_W, SRC_H = 1600, 900
CX, CY = 799.5, 480.0
# camera optical frame (x right, y down, z forward) -> base_link (x forward, y left, z up)
R_OPT = np.array([[0.0, 0.0, 1.0], [-1.0, 0.0, 0.0], [0.0, -1.0, 0.0]])
# name -> (x, y, z, yaw_deg, fx): camera position in base_link (metres), yaw about +z, pinhole focal length (px)
RIG: Dict[str, Tuple[float, float, float, float, float]] = {
    "CAM_FRONT_LEFT": (1.50, 0.50, 1.55, 55.0, 1260.0),
    "CAM_FRONT": (1.70, 0.00, 1.55, 0.0, 1260.0),
    "CAM_FRONT_RIGHT": (1.50, -0.50, 1.55, -55.0, 1260.0),
    "CAM_BACK_LEFT": (1.05, 0.50, 1.55, 110.0, 1260.0),
    "CAM_BACK": (0.05, 0.00, 1.55, 180.0, 800.0),
    "CAM_BACK_RIGHT": (1.05, -0.50, 1.55, -110.0, 1260.0),
}
# class, x, y, yaw (rad), (length, width, height); bottom on the road (base_link z = 0), at rest
TARGETS: List[Tuple[str, float, float, float, Tuple[float, float, float]]] = [
    ("car", 12.4, 0.4, 0.0, (4.6, 1.9, 1.6)),
    ("truck", 23.6, -3.6, 0.0, (7.6, 2.5, 3.2)),
    ("pedestrian", 8.8, 4.4, 1.57, (0.8, 0.7, 1.75)),
    ("car", 6.0, 9.2, 0.0, (4.5, 1.9, 1.5)),
    ("bicycle", 4.4, -7.6, 0.0, (1.7, 0.6, 1.3)),
    ("traffic_cone", 7.6, -3.6, 0.0, (0.4, 0.4, 0.7)),
    ("bus", -17.2, 0.4, 0.0, (11.5, 2.9, 3.4)),
    ("car", -8.4, -6.0, 3.14, (4.6, 1.9, 1.6)),
]
ONNX_DEFAULT = Path("/home/ubuntu/experiments/tt-models/assets/bevdet/hf_tensorrt_bevdet/bevdet_one_lt_d.onnx")


# --------------------------------------------------------------------------------------------- rig and render

def camera_T(x: float, y: float, z: float, yaw_deg: float) -> np.ndarray:
    """4x4 T_base_link_from_camera_optical for a level camera at (x, y, z) looking along yaw."""
    a = math.radians(yaw_deg)
    rz = np.array([[math.cos(a), -math.sin(a), 0.0], [math.sin(a), math.cos(a), 0.0], [0.0, 0.0, 1.0]])
    T = np.eye(4)
    T[:3, :3] = rz @ R_OPT
    T[:3, 3] = (x, y, z)
    return T


def calibration() -> Dict:
    """The calibration envelope of the rig (``calib/synthetic_6cam.json``)."""
    cams = {}
    for name in CAMERA_ORDER:
        x, y, z, yaw, f = RIG[name]
        T = camera_T(x, y, z, yaw)
        cams[name] = {"intrinsics": [[f, 0.0, CX], [0.0, f, CY], [0.0, 0.0, 1.0]],
                      "T_ref_from_camera": [[round(float(v), 12) for v in row] for row in T]}
    return {"frame_id": "base_link", "cameras": cams}


def render_road(K, T, seed: int = 0, noise: float = 1.5) -> np.ndarray:
    """A camera's view of a flat world: sky gradient, a 14 m wide asphalt road along base_link x with lane markings
    (solid edges, dashed separators every 3.5 m), sand-coloured shoulders, haze toward the horizon. (900, 1600, 3)
    uint8 RGB."""
    rng = np.random.default_rng(seed)
    u, v = np.meshgrid(np.arange(SRC_W, dtype=np.float64), np.arange(SRC_H, dtype=np.float64), indexing="xy")
    d = np.stack([u, v, np.ones_like(u)], -1) @ np.linalg.inv(np.asarray(K, np.float64)).T
    R, c = np.asarray(T)[:3, :3], np.asarray(T)[:3, 3]
    db = d @ R.T
    el = np.arctan2(db[..., 2], np.hypot(db[..., 0], db[..., 1]))
    sky = np.clip(el / 0.6, 0, 1)[..., None]
    img = (1 - sky) * np.array([205.0, 210.0, 214.0]) + sky * np.array([120.0, 160.0, 205.0])
    ground = db[..., 2] < -1e-6
    with np.errstate(divide="ignore", invalid="ignore"):
        t = np.where(ground, -c[2] / np.where(ground, db[..., 2], -1.0), np.inf)
        px, py = c[0] + t * db[..., 0], c[1] + t * db[..., 1]
    dist = np.hypot(px, py)
    g = ground & (dist < 80.0)
    road = g & (np.abs(py) < 7.0)
    side = g & ~road
    tex = rng.normal(0, 4.0, size=(SRC_H, SRC_W, 1))
    img[road] = np.array([92.0, 94.0, 96.0]) + tex[road]
    img[side] = np.array([140.0, 132.0, 118.0]) + tex[side]
    for yl, dashed in ((6.8, False), (-6.8, False), (0.0, True), (3.5, True), (-3.5, True)):
        on = g & (np.abs(py - yl) < 0.08)
        if dashed:
            on &= np.mod(np.where(g, px, 0.0), 9.0) < 3.0
        img[on] = np.array([225.0, 225.0, 220.0])
    haze = np.clip((np.where(g, dist, 0.0) - 30.0) / 50.0, 0, 1)[..., None] * g[..., None]
    img = img * (1 - 0.5 * haze) + 0.5 * haze * np.array([200.0, 205.0, 210.0])
    img += rng.normal(0, noise, size=img.shape)
    return np.clip(np.rint(img), 0, 255).astype(np.uint8)


def renders(cal: Dict) -> Dict[str, np.ndarray]:
    return {n: render_road(cal["cameras"][n]["intrinsics"], np.array(cal["cameras"][n]["T_ref_from_camera"]), seed=i)
            for i, n in enumerate(CAMERA_ORDER)}


def box_corners(x: float, y: float, yaw: float, l: float, w: float, h: float) -> np.ndarray:
    c, s = math.cos(yaw), math.sin(yaw)
    return np.array([[x + c * dx * l - s * dy * w, y + s * dx * l + c * dy * w, dz] for dz in (0.0, h)
                     for dx, dy in ((0.5, 0.5), (0.5, -0.5), (-0.5, -0.5), (-0.5, 0.5))])


def target_masks(cal: Dict, rows: np.ndarray, cols: np.ndarray, grow: float = 0.35) -> np.ndarray:
    """[6, 256, 704] float32: the network-grid pixels covered by each target's projected 3D box (dilated)."""
    masks = np.zeros((len(CAMERA_ORDER), len(rows), len(cols)), np.float32)
    for ci, name in enumerate(CAMERA_ORDER):
        K = np.asarray(cal["cameras"][name]["intrinsics"], np.float64)
        Tinv = np.linalg.inv(np.asarray(cal["cameras"][name]["T_ref_from_camera"], np.float64))
        for _, x, y, yaw, (l, w, h) in TARGETS:
            pc = box_corners(x, y, yaw, l, w, h) @ Tinv[:3, :3].T + Tinv[:3, 3]
            if (pc[:, 2] < 0.5).any():
                continue
            uv = pc @ K.T
            uv = uv[:, :2] / uv[:, 2:3]
            (u0, v0), (u1, v1) = uv.min(0), uv.max(0)
            du, dv = (u1 - u0) * grow, (v1 - v0) * grow
            if u1 + du < 0 or u0 - du > SRC_W - 1 or v1 + dv < 0 or v0 - dv > SRC_H - 1:
                continue
            rr = (rows >= v0 - dv) & (rows <= v1 + dv)
            cc = (cols >= u0 - du) & (cols <= u1 + du)
            masks[ci][np.ix_(rr, cc)] = 1.0
    return masks


def cell_of(x: float, y: float, cfg) -> Tuple[int, int, float, float]:
    """BEV cell (row, col) of a base_link point and its sub-cell offsets (the reg head's targets)."""
    gx, gy = cfg.grid_x, cfg.grid_y
    fx = (x - float(gx[0])) / float(gx[2])
    fy = (y - float(gy[0])) / float(gy[2])
    return int(math.floor(fy)), int(math.floor(fx)), fx - math.floor(fx), fy - math.floor(fy)


# ------------------------------------------------------------------------------------------------ optimise

def optimise(a) -> int:
    """Gradient ascent on the six uint8-range network crops (see the module docstring)."""
    import torch
    import torch.nn.functional as F

    from tt_bevdet.host.calib import bevpool_tables, cam_params
    from tt_bevdet.host.inputs import prepare_frame
    from tt_bevdet.host.preprocess import gather_index
    from tt_bevdet.reference.model import BEVDetNet
    from tt_bevdet.reference.plugins import bev_pool
    from tt_bevdet.reference.weights import BEVDetWeights
    from tt_bevdet.ttaw.golden import NULL_TAPS

    torch.manual_seed(0)
    torch.set_num_threads(a.threads)
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    cfg = default_config()
    cal = calibration()
    rows, cols = gather_index(cfg)
    imgs = renders(cal)
    frame = prepare_frame([{"camera": n, "image": imgs[n]} for n in CAMERA_ORDER], cal, cfg)
    tables = bevpool_tables(cfg, frame.calib)
    print("tables", tables.stats(), flush=True)
    cp = torch.as_tensor(cam_params(cfg, frame.calib))
    net = BEVDetNet(BEVDetWeights(a.onnx), cfg)
    mean = torch.as_tensor(np.asarray(cfg.mean, np.float32)).view(1, 3, 1, 1)
    std = torch.as_tensor(np.asarray(cfg.std, np.float32)).view(1, 3, 1, 1)
    crop0 = torch.as_tensor(frame.crop_u8.astype(np.float32))            # [6, 3, 256, 704] B, G, R, 0..255
    x = (torch.as_tensor(np.load(a.resume).astype(np.float32)) if a.resume else crop0.clone()).requires_grad_(True)
    masks = torch.as_tensor(target_masks(cal, rows, cols))[:, None]
    opt = torch.optim.Adam([x], lr=a.lr * 255.0 / 8.0)
    ny, nx = cfg.bev_hw
    tg = []
    near = torch.zeros(ny, nx, dtype=torch.bool)
    for cls, tx, ty, yaw, size in TARGETS:
        r, c, fx, fy = cell_of(tx, ty, cfg)
        tg.append((CLASS_NAMES.index(cls), r, c, fx, fy, yaw, size))
        near[max(r - 2, 0):r + 3, max(c - 2, 0):c + 3] = True

    def heads_of(xc):
        c4, c5 = net.backbone((xc - mean) / std, NULL_TAPS)
        depth, feat = net.depth_net(net.neck(c4, c5, NULL_TAPS), cp, taps=NULL_TAPS)
        curr = bev_pool(depth, feat, tables, cfg)                            # [1, 80, 128, 128]
        bev_in = torch.cat([curr, torch.zeros(1, cfg.bev_in_channels - curr.shape[1], ny, nx)], 1)  # fresh stream
        return net.center_head(net.bev_encoder(bev_in, NULL_TAPS), NULL_TAPS)

    hist = []
    for it in range(a.iters):
        t0 = time.time()
        opt.zero_grad()
        hd = heads_of(x)
        hm = hd["heatmap_0"][0]
        loss_t, loss_r, scores = 0.0, 0.0, []
        for k, r, c, fx, fy, yaw, (l, w, h) in tg:
            lg = hm[k, r, c]
            others = torch.cat([hm[:k, r, c], hm[k + 1:, r, c]])
            win = hm[:, r - 1:r + 2, c - 1:c + 2].reshape(10, 9)            # the 3x3 neighbourhood minus the centre
            nb = torch.cat([win[:, :4], win[:, 5:]], 1)
            # confident target cell, its class the clear argmax, a peak above every neighbour of any class
            loss_t = loss_t + F.softplus(4.0 - lg) + F.softplus(others.max() - lg + 2.5) \
                + F.softplus(nb.max() - lg + 1.0)
            loss_r = loss_r + ((hd["reg_0"][0, :, r, c] - torch.tensor([fx, fy])) ** 2).sum() \
                + (hd["height_0"][0, 0, r, c] - h / 2.0) ** 2 \
                + ((hd["dim_0"][0, :, r, c] - torch.log(torch.tensor([l, w, h]))) ** 2).sum() \
                + ((hd["rot_0"][0, :, r, c] - torch.tensor([math.sin(yaw), math.cos(yaw)])) ** 2).sum() \
                + (hd["vel_0"][0, :, r, c] ** 2).sum()
            scores.append(round(float(torch.sigmoid(lg.detach())), 3))
        bg = hm[:, ~near]
        loss_bg = F.softplus(bg + 4.0).mean() * 4.0 + F.softplus(bg + 2.5).max()
        tv = ((x[:, :, 1:] - x[:, :, :-1]).abs().mean() + (x[..., 1:] - x[..., :-1]).abs().mean()) / 255.0
        dev = (((x - crop0) / 255.0) ** 2 * (1.0 - masks)).mean()
        loss = loss_t + 0.5 * loss_r + loss_bg + a.tv * tv + 20.0 * dev
        loss.backward()
        with torch.no_grad():
            x.grad.mul_(masks + a.outside * (1.0 - masks))
            opt.step()
            x.clamp_(0.0, 255.0)
        hist.append({"it": it, "loss": round(float(loss.detach()), 4), "targets": round(float(loss_t.detach()), 4),
                     "bg": round(float(loss_bg.detach()), 4), "scores": scores, "s": round(time.time() - t0, 1)})
        print(json.dumps(hist[-1]), flush=True)
        if it % 10 == 9:
            np.save(out / "crop_f32.npy", x.detach().numpy())
    np.save(out / "crop_f32.npy", x.detach().numpy())
    args = {k: (str(v) if isinstance(v, Path) else v) for k, v in vars(a).items()}
    (out / "history.json").write_text(json.dumps({"args": args, "history": hist}, indent=0))
    return 0


# ------------------------------------------------------------------------------------------------ finalise

def embed(crop_u8: np.ndarray, base: np.ndarray, rows: np.ndarray, cols: np.ndarray) -> np.ndarray:
    """A 1600x900 RGB image whose Preprocess gather (rows x cols, the node's nearest crop) is exactly ``crop_u8``
    ([3, 256, 704] B, G, R planes): rows from the first gathered row down are the crop upsampled by nearest
    neighbour, the rows above keep ``base``."""
    rgb = crop_u8[::-1].transpose(1, 2, 0)
    ri = np.abs(np.arange(SRC_H)[:, None] - rows[None, :]).argmin(1)
    ci = np.abs(np.arange(SRC_W)[:, None] - cols[None, :]).argmin(1)
    out = base.copy()
    top = int(rows[0])
    out[top:] = rgb[ri[top:]][:, ci]
    if not (out[rows][:, cols] == rgb).all():
        raise AssertionError("the gather of the embedded image is not the crop")
    return out


def sha256(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def finalise(a) -> int:
    from PIL import Image

    from tt_bevdet.host.inputs import load_sample, prepare_frame
    from tt_bevdet.host.preprocess import gather_index
    from tt_bevdet.reference.config import HEAD_NAMES
    from tt_bevdet.reference.pipeline import BEVDetReference
    from tt_bevdet.ttaw.golden import TapRegistry

    cfg = default_config()
    cal = calibration()
    rows, cols = gather_index(cfg)
    crop = np.clip(np.rint(np.load(a.crop)), 0, 255).astype(np.uint8)
    base = renders(cal)
    samples, calib_dir, goldens = PKG / "samples", PKG / "calib", PKG / "tests" / "goldens"
    img_dir = samples / NAME
    img_dir.mkdir(parents=True, exist_ok=True)
    files = {}
    for k, name in enumerate(CAMERA_ORDER):
        im = embed(crop[k], base[name], rows, cols)
        path = img_dir / f"{name}.png"
        Image.fromarray(im).save(path, format="PNG", optimize=True)
        files[name] = f"{NAME}/{name}.png"
    preset = dict(calibration())
    preset = {"_doc": ("Calibration preset of the shipped synthetic sample (samples/synthetic_6cam): a generic six-camera "
                       "rig with round numbers of our own (code/scripts/make_sample.py RIG): 1600x900 pinhole cameras "
                       "1.55 m above the road, yaws 0 / +-55 / +-110 / 180 deg, fx 1260 px (back camera 800 px), "
                       "principal point (799.5, 480). intrinsics = K of the image as sent; T_ref_from_camera = camera "
                       "optical frame -> base_link. Generated data, Apache-2.0."), **preset}
    (calib_dir / f"{NAME}.json").write_text(json.dumps(preset, indent=1) + "\n")
    manifest = {
        "_doc": ("bevdet-p150 sample: a synthetic six-camera test pattern generated by code/scripts/make_sample.py "
                 "(rendered road views whose network-visible pixels were optimised so that the fp32 reference "
                 "reports the eight objects below; no real scene). model(**tt_bevdet.load_sample(path)). "
                 "Apache-2.0."),
        "images": files,
        "calibration": {"preset": NAME},
        "targets": [{"label": c, "center": [x, y, round(s[2] / 2, 3)], "size": list(s), "yaw": yaw}
                    for c, x, y, yaw, s in TARGETS],
        "crop_u8_sha256": sha256(np.ascontiguousarray(crop).tobytes()),
    }
    man_path = samples / f"{NAME}.json"
    man_path.write_text(json.dumps(manifest, indent=1) + "\n")
    # the decoded PNGs must give back the optimised crop exactly (lossless; the node's gather)
    kw = load_sample(man_path)
    frame = prepare_frame(kw["images"], kw["calibration"], cfg)
    if sha256(np.ascontiguousarray(frame.crop_u8).tobytes()) != manifest["crop_u8_sha256"]:
        raise AssertionError("the PNG round trip changed the crop")
    ref = BEVDetReference(a.onnx, threads=a.threads)
    taps = TapRegistry(include=["head.*_0"])
    det = ref(**kw, stream={"id": f"smoke-{NAME}"}, score_threshold=0.1, taps=taps)
    body = dict(det.to_dict())
    body["timing_ms"] = {}
    body["meta"]["request"] = {"stream": "fresh", "params": {"score_threshold": 0.1}}
    (samples / f"{NAME}.reference.json").write_text(json.dumps(body, indent=1) + "\n")
    # sparse head maps: every channel at the cells the vendor decode keeps (score > 0.1); the others below it
    heads = {k: np.asarray(taps[f"head.{k}"], np.float32) for k in HEAD_NAMES}
    hm = heads["heatmap_0"][0]
    keep = 1.0 / (1.0 + np.exp(-hm.max(0))) > 0.1
    rr, cc = np.nonzero(keep)
    sparse = {"cells": np.stack([rr, cc], 1).astype(np.int32)}
    for k, v in heads.items():
        sparse[k] = v[0][:, rr, cc].T.astype(np.float32)                    # [K, channels]
    goldens.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(goldens / f"{NAME}_heads.npz", **sparse)
    pub = [d for d in body["detections"] if d["score"] >= 0.2]
    print(f"{NAME}: {len(body['detections'])} boxes > 0.1, {len(pub)} >= 0.2, {int(keep.sum())} decoded cells")
    for d in pub:
        print("  ", d["label"], d["score"], d["center"], d["size"], d["yaw"])
    sizes = {p.name: p.stat().st_size for p in sorted(img_dir.glob("*.png"))}
    print("png bytes", sizes, sum(sizes.values()))
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    o = sub.add_parser("optimise", help="gradient ascent on the network crops -> <out>/crop_f32.npy")
    o.add_argument("--out", required=True)
    o.add_argument("--iters", type=int, default=150)
    o.add_argument("--lr", type=float, default=0.06)
    o.add_argument("--outside", type=float, default=0.15, help="gradient scale outside the target masks")
    o.add_argument("--tv", type=float, default=2.0, help="total-variation weight")
    o.add_argument("--resume", help="start from a crop_f32.npy")
    f = sub.add_parser("finalise", help="crop -> sample images, preset, manifest, CPU reference, sparse heads")
    f.add_argument("--crop", required=True)
    for p in (o, f):
        p.add_argument("--onnx", type=Path, default=ONNX_DEFAULT)
        p.add_argument("--threads", type=int, default=4)
    a = ap.parse_args()
    return optimise(a) if a.cmd == "optimise" else finalise(a)


if __name__ == "__main__":
    sys.exit(main())