# SPDX-License-Identifier: Apache-2.0
"""Python API: StreamPETR (Autoware camera_streampetr) on one Tenstorrent Blackhole p150.

    from tt_streampetr import StreamPETR, load_sample

    with StreamPETR.from_pretrained(device_id=0) as model:  # weights -> HF cache, device open, traces captured
        out = model(**load_sample("code/tt_streampetr/samples/synthetic_5cam.json"))   # 5 cameras + calibration
        # a sequence: one dict per frame with the same stream id (the temporal memory carries over), e.g.
        # model(images=..., calibration=..., stream={"id": "rig", "timestamp_s": t, "T_world_from_ego": T})
    print(out.to_dict())  # the same JSON as POST /predict

The contract shared by every bundle of the Autoware collection is the vendored ``ttaw.api_base.ModelBase``
(BUNDLE_CONVENTIONS.md section 8): ``from_pretrained`` resolves the pinned weights before it claims the chip, opens it
(ETH dispatch, 12x10), builds the graph and captures every trace variant in ``warmup_variants``, so the first call is
as fast as the later ones; calls are serialised by a lock (one chip, batch 1) and fill ``timing_ms``; ``close()`` is
idempotent, also runs at interpreter exit, and closes the chip only if the model opened it. The HTTP server
(``tt_streampetr.server.app``) calls this class, so ``/predict`` and ``model(...)`` agree bit for bit.

Inputs per call (``_prepare``; host code in ``tt_streampetr.host``, exact Autoware semantics):

- ``images``: the five cameras ``CAM_FRONT``, ``CAM_FRONT_LEFT``, ``CAM_BACK_LEFT``, ``CAM_FRONT_RIGHT``,
  ``CAM_BACK_RIGHT`` (any order, any size; resized / cropped / normalised like the node);
- ``calibration``: per camera ``intrinsics`` (``CameraInfo.P`` or K of the image as sent; optional full 3x4 ``P``) and
  ``T_ref_from_camera`` (camera optical frame -> base_link), or ``{"preset": name}``;
- ``stream``: ``{"id", "reset", "timestamp_s", "T_world_from_ego"}`` drives the temporal object memory: a new id or
  ``reset`` starts like a fresh node (zero memory), ``timestamp_s`` (the CAM_FRONT stamp) gives the memory ages and
  ``T_world_from_ego`` (map <- base_link) the ego-motion compensation (a stream frame without it is skipped, as the
  node skips a cycle whose TF lookup fails: ``meta.skipped = "no_ego_pose"``). Without ``stream``: one frame from zero
  memory. Per-camera ``timestamp_s`` feed the node's 0.15 s camera-sync check (a failure resets the stream and
  returns no objects, ``meta.skipped``).

Device graph (``tt_streampetr.tt``, PORT_LOG.md 1.1): ONE trace ``frame`` = VoVNet-99-eSE + CPFPN, the head (tokens,
temporal alignment, 6 decoder layers, shared branches), the top-256 selection and the memory update, with the
memory (embedding bf16, geometry fp32 in the previous ego frame) kept in persistent device states. The host runs
the node's pre-processing (C14), the position embedding and the spatial-alignment coefficients once per
calibration, the stream clock / ego poses / ages, and the post-processing (all-layer decode, IoU-BEV NMS).
Load-time knobs: :data:`KNOBS` (``STREAMPETR_*``). Importing this module has no side effects (torch / ttnn are
imported in the hooks).
"""
from __future__ import annotations

from collections import OrderedDict
from typing import Any, Dict, Mapping, Optional

import numpy as np

from . import io as tio
from .device import DEVICE_DEFAULTS
from .reference.config import CAMERA_ORDER, CLASS_NAMES, NORMALIZATION_PRESETS
from .ttaw.api_base import ModelBase
from .ttaw.knobs import Knob, Knobs
from .ttaw.outputs import Detections3D

__all__ = ["StreamPETR", "Output", "KNOBS"]

# The result class of this model (ttaw.outputs: Detections3D, Detections2D, Segmentation3D, Mask2D, Trajectory).
Output = Detections3D

# Load-time switches (read once at build; env STREAMPETR_<NAME>). Each optimization adds its knob here (A/B switch).
KNOBS = Knobs("STREAMPETR", [
    Knob("NORMALIZATION", "autoware_main", "input normalisation preset (PLAN D7; SPEC 3.4): autoware_main = universe "
                                           "main RGB planes + RGB stats; autoware_0.52 = an rgb8 camera on 0.52.1; "
                                           "awml_training = BGR planes + BGR stats", choices=NORMALIZATION_PRESETS),
    Knob("USE_TEMPORAL", True, "carry the object memory between frames (the node's use_temporal; False wipes it)"),
    Knob("MAX_STREAMS", 1, "temporal streams kept (PLAN D16; another stream id evicts the least recent one)"),
    Knob("AUTOWARE_COMPAT", True, "reproduce the node's age quirks (default, PORT_LOG section 6 Q1): a sync-failure "
                                  "wipe keeps the stale ages, the 600 s clock-origin reset shifts them (SPEC 6 / risk "
                                  "11); False zeroes / rebases the ages instead (the SPEC's clean variant)"),
    Knob("PREPROCESS_FMA", False, "model nvcc's fmaf contraction in the resize kernel (ttaw.image_triangle; <= 3e-5 "
                                  "difference)", experiment=True),
    Knob("PREPROCESS_WORKERS", 4, "host threads of the five-camera resize / normalise"),
    Knob("HOST_PRE", "numba", "host pre-processing (OPT_PLAN item 1): numba = host.fast.preprocess_rows_bf16 (the C14 "
                              "triangle resize, crop, normalisation and the bf16 cast of the upload rows fused in one "
                              "pass per camera, written into a persistent ttnn host buffer; bit-identical upload); "
                              "numpy = ttaw.image_triangle float32 NCHW + img_rows + bf16 conversion (the release "
                              "path; also the automatic fallback without numba)", choices=("numba", "numpy")),
    Knob("HOST_POST", "numba", "host post-processing (OPT_PLAN item 4): numba = the compiled IoU-BEV NMS "
                               "(host.fast.iou_bev_nms, the same keep decisions); numpy = ttaw.nms.iou_bev_nms (also "
                               "the fallback without numba)", choices=("numba", "numpy")),
    Knob("STEM_IN", "kin", "device input layout (OPT_PLAN item 2): kin = the image as 3,840-byte bf16 rows "
                           "[5*480, 640*3] and per camera a K-IN generic_op (tt/kin.py) writing stem conv 0's L1 "
                           "shards (conv on the L1 path); rows = [5*480*640, 3] 6-byte pixel pages and the stem's "
                           "DRAM-sliced conv (the release path)", choices=("kin", "rows")),
    Knob("CONV_L1", "s3r45n", "3x3 conv placement (OPT_PLAN item 3, tt/placement.py): the OSA 3x3 convs of the listed "
                           "stages (n = the CPFPN 3x3) on ttnn's L1 path (the op moves its DRAM input into L1 shards; "
                           "no TILE reshape pair), DRAM slicing as the first-call fallback. s3r45n (default) = stages "
                           "4-5 + neck + stage 3 except each block's first conv (whose 256 / 512-channel input fits "
                           "L1 only with a changed K split: s345* fails the frozen sample_stream.centre_err gate, "
                           "OPT_REPORT R2.1); s3r45n / s45n are bit-identical to none = every 3x3 conv DRAM-sliced "
                           "(the release path)", choices=("none", "s5", "s45", "s45n", "s3r45n", "s345", "s345n", "s2345n")),
    Knob("CONV_L1_ABH", "256,128,64,32", "with CONV_L1: activation block heights (comma list, multiples of 32) retried "
                                         "in order when a conv's L1 path does not fit with ttnn's own blocks"),
    Knob("ESE", "kernel", "eSE gate application (OPT_PLAN items 5 / 5b): kernel = one ESG generic_op per OSA block "
                           "(tt/ese.py: the gate tiles built in L1 per camera run, the stock FPU multiply and the "
                           "identity add fused; bit-identical); bcast = one row-broadcast multiply of the "
                           "[1, N, HW, C] view by the per-camera gate rows in the stages whose H*W is a multiple of 32 "
                           "(2, 3; exact: the same products); matmul = the gate broadcast to every row by the "
                           "transposed camera-indicator matmul, then a multiply (the release path)",
         choices=("matmul", "bcast", "kernel")),
    Knob("ESE_MEAN", "kernel", "eSE per-camera sums (OPT_PLAN item 5b): kernel = ESM partial sums on the whole grid "
                                "(tt/ese.py) + a fixed-order fp32 reduce (ESR; the fp32 sums are regrouped, "
                                "closer to the exact sums than the matmul's; gated, the 8-frame stream outputs stayed "
                                "bit-identical); matmul = the camera-indicator matmul (the release path)",
         choices=("matmul", "kernel")),
    Knob("OSA_CAT", "kmm", "OSA concatenation [x, y1..y5] (OPT_PLAN item 6b, tt/kcat.py): kernel = one KCAT "
                              "data-movement generic_op (output row tiles over the whole grid; a copy: bit-identical); "
                              "kmm = (OPT round 9, tt/kmm.py) stages 2-3 read the parts straight from the concat "
                              "1x1 matmul's in0 reader (the stock 1D matmul program ported to a generic_op, same "
                              "compute: bit-identical; no concat tensor), KCAT in stages 4-5; "
                              "concat = ttnn.concat (the release path)", choices=("concat", "kernel", "kmm")),
    Knob("CONV_CFG", "adb,wdb", "data-movement Conv2dConfig flags of every backbone 3x3 conv (OPT round 4, "
                                "tt/backbone.conv_flags): comma list of adb (activation double buffer, always "
                                "applied), wdb (weights double buffer), split (forced split reader), reuse "
                                "(activation reuse; NOT exact: it changes the stride-1 DRAM-sliced convs' outputs); "
                                "wdb / split / reuse are tried on the first call and dropped where a conv rejects "
                                "them (never by shrinking activation blocks). adb,wdb is bit-identical; none = "
                                "ttnn's defaults (the release path)"),
    Knob("CONV_CHAIN", "kcat45", "L1 conv chaining in the OSA blocks (OPT_PLAN item 19, OPT round 7, "
                                 "tt/backbone.OSABlock._convs; kcat45 default, bit-identical): "
                               "chain = a conv on ttnn's L1 path whose successor is on the L1 path too keeps its "
                               "sharded output as the successor's input (no InterleavedToSharded; the concat part is "
                               "its DRAM copy); kcat = chain, and every L1-path conv output stays sharded as its "
                               "own concat part (the KCAT-SH reader reads it from L1: no ShardedToInterleaved "
                               "either; needs OSA_CAT=kernel, else DRAM copies); none = every L1 conv returns DRAM "
                               "interleaved (the round-6 path); kcat45 = kcat in stages 4-5, chain in stage 3 (whose "
                               "five outputs do not fit L1 together with the conv buffers)",
         choices=("none", "chain", "kcat45", "kcat")),
    Knob("RESHAPE", "kernel", "input re-tiling of the DRAM-sliced 3x3 convs (OPT_PLAN item 20, OPT round 8, "
                              "tt/krt.py): kernel = the flat [1, 1, N*H*W, C] map unflattened to [N, H, W, C] by one "
                              "KRT data-movement generic_op (16-row half-tile copies, whole grid) so ttnn.conv2d's own "
                              "ReshapeView is a no-op (stage-3 block-first convs, W = 80; a copy: bit-identical); "
                              "stock = ttnn.conv2d's ReshapeView (the round-7 path)", choices=("stock", "kernel")),
    Knob("KMM_BUF", "auto", "KMM circular-buffer depth (OPT round 9, R9.2, tt/kmm.buffers; OSA_CAT=kmm only): auto = "
                            "deeper in0 / in1 buffers (stage 2 12 / 12 blocks, stage 3 4 / 4; data movement only: "
                            "bit-identical); 2 = the stock factory's double buffers (R9.1)", choices=("2", "auto")),
    Knob("OSA_MM", "rowblock", "OSA concat 1x1 matmul config (OPT_PLAN item 6a, tt/layers.row_block_config): rowblock = "
                           "ttnn's 1D config with the output in row blocks (stage 2: 25 x 8 -> 5 x 8 fp32 L1 partials; "
                           "the same K order: bit-identical); auto = ttnn's own config (the release path)",
         choices=("auto", "rowblock")),
    Knob("KV_EPI", True, "split-key K / V operands (OPT_PLAN item 7, tt/kvx.py): True = single-width K / V "
                          "projections and one KVX generic_op per K / V pair writing the per-layer [hi | lo] / [v | v] "
                          "head tensors (replaces the 5-program key split over the doubled width, the head split and "
                          "the per-layer slices; bit-identical); False = the doubled projections + stock programs "
                          "(the release path)"),
    Knob("SELF_KV", "buffer", "self-attention keys (OPT_PLAN item 9, with KV_EPI): buffer = KVX writes the memory and "
                              "query K / V rows straight into per-layer [memory; queries] buffers (bit-identical); "
                              "concat = untilize, ROW_MAJOR concat and tilize per layer (the release path)",
         choices=("concat", "buffer")),
    Knob("S26_KERNEL", True, "exact-product slicing of the temporal chain's sliced maps (OPT_PLAN item 8, tt/ksl.py): "
                              "True = the x slices and their pair-block concatenation as one KSL generic_op per map "
                              "input (the three maps on the sines share one; bit-identical); False = 5 S - 1 stock "
                              "elementwise programs + ttnn.concat per map (the release path)"),
    Knob("SEL", "mk", "top-256 selection + memory update (OPT rounds 5 / 10): rank = (tt/sel.py) exact ranks from fp32 "
                        "comparisons on the whole grid, a one-hot matrix and one-hot matmul gathers (the fp32 rows "
                        "as three bf16 terms; the same set and order as the stable sort: bit-identical); sort = "
                        "ttnn.sort on one core + ttnn.embedding / ttnn.gather (the release path); mk = MK-SEL "
                        "(OPT round 10, tt/mksel.py): the selection + the whole memory update (gathers, ego identity, "
                        "ring shift, state writes) as one generic_op on 72 cores (bit-identical; the taps path keeps rank)",
         choices=("sort", "rank", "mk")),
    Knob("PRECISION", "HiFi4+fp32+l1acc:w=fp32", "precision of every device conv / linear (ttaw.precision spec; "
                                                 "PORT_LOG S10: fp32 weights and biases, HiFi4, fp32 accumulation)"),
    Knob("ATTN_EXP", "exact", "SDPA softmax exponential (OPT round 5): approx = the SFPU's approximate exp "
                              "(exp_approx_mode + math_approx_mode) in every SDPA; exp = the approximate exp only; "
                              "cross / self = the approximate exp in those SDPAs only. The exact exp makes the SDPAs "
                              "SFPU-bound (cross 643 -> 207 us approximate) but the numerics move: gated (approx "
                              "fails sample_stream); exact = the release path (PORT_LOG r3)",
         choices=("exact", "approx", "exp", "cross", "self")),
    Knob("XATTN", "x8", "SDPA implementation (OPT round 6, item 11b, tt/spx.py): sdpa = ttnn's SDPA; spx = the SPX "
                          "generic_op (the stock SDPA program factory ported, stock reader / writer, a copy of the "
                          "compute kernel) with the stock exponential (bit-identical); x8 = SPX whose softmax "
                          "exponential runs the stock exp_fp32_accurate(x * scale) operations as one hand-scheduled "
                          "TTI pass (26 SFPU instructions per row instead of a multiply pass + ~37; bit-identical). "
                          "Needs ATTN_EXP=exact and ATTN_FP32_ACC (else sdpa)",
         choices=("sdpa", "spx", "x8")),
    Knob("ATTN_FP32_ACC", True, "SDPA with fp32 accumulation (probe P7: half the max error)"),
    Knob("HEAD_ACT", "float32", "dtype of the head's token / query / residual activations (PORT_LOG S25; SDPA "
                                "operands stay bf16)", choices=("bfloat16", "float32")),
    Knob("EXACT_MAP", "gather", "fp32 maps of the temporal alignment (PORT_LOG S18): column gather (exact bf16 "
                                "embedding of split terms) + fp32 elementwise, or the bf16-split matmul",
         choices=("gather", "matmul")),
    Knob("KEY_SPLIT", True, "SDPA keys as hi + lo bf16 terms in a doubled head dim (PORT_LOG S25: the bf16 rounding "
                            "of the keys dominates the head's error otherwise)"),
    Knob("TEMPORAL_EXACT", True, "the temporal positional chain (query embedding MLP, MLN, time embedding, their "
                                 "LayerNorms) as exact-product sliced maps (PORT_LOG S26: its matrix-engine error "
                                 "dominates the self-attention error otherwise)"),
])


class StreamPETR(ModelBase):
    """StreamPETR (Autoware camera_streampetr) on one Blackhole p150. Create it with :meth:`from_pretrained`."""

    MODEL_NAME = "streampetr-p150"
    ENV_PREFIX = "STREAMPETR"  # prefix of the environment knobs (SERVING.md section 3.4)
    DEFAULT_REPO = "AutowareFoundation/camera_streampetr"
    DEFAULT_TAG = "v1.0"  # the Autoware ansible artifacts pin
    # the commit DEFAULT_TAG points to (pinned: tags can move)
    DEFAULT_REVISION = "90f37d64186686ebc458509d69d3d4ca16d185c7"
    ALLOW_PATTERNS = ["simplify_extract_img_feat.onnx", "simplify_position_embedding.onnx",
                      "simplify_pts_head_memory.onnx", "ml_package_camera_streampetr.param.yaml"]
    VARIANTS = ["default"]  # load-time: selects weights files and trace shapes
    DEFAULT_VARIANT = "default"
    INPUT_KIND = "multicam"  # lidar | camera | multicam | lidar+multicam | planner
    CAMERA_ORDER = CAMERA_ORDER  # the network's camera slots (HF ml_package yaml:10)
    POINT_FIELDS = tio.DEFAULT_POINT_FIELDS  # unused (camera model); kept for the shared /info schema
    # class names in network output order (HF ml_package yaml:4); they are also the Autoware labels the node publishes
    LABELS = list(CLASS_NAMES)
    # Per-request knobs: name -> (type, min, max, default). Host-side post-processing only (S:5, S:9.1).
    RUNTIME_PARAMS = {
        "score_threshold": (float, 0.0, 1.0, None),            # None: the per-class node thresholds .36/.39/.38/.41/.43
        "iou_nms_threshold": (float, 0.0, 1.0, 0.5),           # IouBevNms iou_threshold (0 = off)
        "iou_nms_search_distance": (float, 0.0, 1000.0, 10.0),  # IouBevNms search_distance_2d [m]
        "circle_nms_distance": (float, 0.0, 100.0, 0.0),       # circle NMS radius [m] (0 = off, the node's default)
        "decode_layers": (int, 1, 6, 6),                        # last decoder layers decoded (6 = all, the node)
        "awml_decode": (bool, None, None, False),               # AWML's evaluation decode (last layer, top-300)
        "max_detections": (int, 1, 5400, 500),
    }
    EXTRA_INPUTS = ()
    DEVICE_DEFAULTS = DEVICE_DEFAULTS  # validated open parameters (device.py)

    # ---- port-specific hooks (called by ModelBase; keep host work out of _forward) ---------------------------
    def _build(self) -> None:
        """Weights (``reference.weights.StreamPETRWeights``: the three ONNX files by consuming node), the host state
        and the device graph (``tt.model.TtStreamPETR``, the ``frame`` variant: backbone + head + in-trace memory
        update). The position embedding and the spatial-alignment coefficients run on the host once per
        calibration (PORT_LOG S5). No capture here (``_warm_one``)."""
        from .reference.model import Head, PositionEmbedding
        from .reference.weights import StreamPETRWeights, find_weights_dir
        from .tt.model import TtStreamPETR

        self._build_host()
        wd = find_weights_dir(self.weights_path)
        if wd is None:
            raise FileNotFoundError(f"{self.weights_path}: the three camera_streampetr v1.0 ONNX files are needed")
        self.net_weights = StreamPETRWeights(wd)
        self.ref_head = Head(self.net_weights, self.cfg)
        self.ref_pe = PositionEmbedding(self.net_weights, self.cfg)
        self.tt = TtStreamPETR(self.device, self.net_weights, self.cfg, variants=("frame",),
                               precision=str(self.knobs.PRECISION), attn_fp32_acc=bool(self.knobs.ATTN_FP32_ACC),
                               head_act=str(self.knobs.HEAD_ACT), key_split=bool(self.knobs.KEY_SPLIT),
                               exact_map=str(self.knobs.EXACT_MAP),
                               temporal_exact=bool(self.knobs.TEMPORAL_EXACT),
                               state_banks=int(self.knobs.MAX_STREAMS),
                               head=self.ref_head, stem_in=str(self.knobs.STEM_IN),
                               conv_l1=str(self.knobs.CONV_L1), conv_l1_abh=str(self.knobs.CONV_L1_ABH),
                               ese=str(self.knobs.ESE), osa_mm=str(self.knobs.OSA_MM),
                               kv_epi=bool(self.knobs.KV_EPI), self_kv=str(self.knobs.SELF_KV),
                               s26_kernel=bool(self.knobs.S26_KERNEL), ese_mean=str(self.knobs.ESE_MEAN),
                               osa_cat=str(self.knobs.OSA_CAT), conv_cfg=str(self.knobs.CONV_CFG),
                               sel=str(self.knobs.SEL), attn_exp=str(self.knobs.ATTN_EXP),
                               xattn=str(self.knobs.XATTN), conv_chain=str(self.knobs.CONV_CHAIN),
                               reshape=str(self.knobs.RESHAPE), kmm_buf=str(self.knobs.KMM_BUF))
        self.runner = self.tt.runner
        self._img_stage = None
        if self._fast_pre():
            from .host import fast as hfast
            from .ttaw.tensors import HostStaging

            stage = HostStaging(self.tt.img_shape, "bfloat16")
            self._img_stage = stage if stage.zero_copy else None
            hfast.warm()
        self.calib_fn = self._calibration_constants
        self._active_calib: Optional[str] = None       # rig key whose PE / MLN coefficients are on the device
        self._device_owner: Optional[str] = None       # stream id whose memory the device states hold
        self._banks: "OrderedDict[str, int]" = OrderedDict()   # stream id -> state bank (PLAN D16)

    def _warm_one(self, variant: Any) -> None:
        """Capture the ``frame`` trace (warm-up runs of the eager graph first; ``TraceRunner.capture``)."""
        self.tt.capture()

    def _calibration_constants(self, rig: Any) -> Dict[str, Any]:
        """Host fp32 PE graph (``reference.model.PositionEmbedding``) and the spatial-alignment MLN coefficients of
        its ``cone`` (``Head.spatial_alignment_coeffs``) -> the device params of one calibration."""
        import torch

        with torch.no_grad():
            pos, cone = self.ref_pe(rig.img_metas_pad, rig.intrinsics, rig.img2lidar)
            gamma, beta = self.ref_head.spatial_alignment_coeffs(cone)
        params = self.tt.calibration_params(pos.numpy(), gamma.numpy(), beta.numpy())
        return {"key": rig.key(), "params": params}

    def _build_host(self) -> None:
        """The host half of the model: configuration, load-time knobs, the per-calibration cache and the per-stream
        temporal states (no device, no weights; the host tests call it directly)."""
        from .reference.config import default_config

        params = {k.upper(): v for k, v in getattr(self, "compile_params", {}).items() if k.upper() in KNOBS.knobs}
        self.knobs = KNOBS.read(**params)
        self.cfg = default_config(normalization=str(self.knobs.NORMALIZATION),
                                  use_temporal=bool(self.knobs.USE_TEMPORAL))
        self._streams: "OrderedDict[str, Any]" = OrderedDict()
        self._calib: "OrderedDict[str, Dict[str, Any]]" = OrderedDict()      # rig key -> PE / cone / MLN coeffs
        self.calib_fn = None          # rig -> {"pos_embed", "cone", "sa_gamma", "sa_beta"} (set by _build)
        self._img_bits: Optional[np.ndarray] = None   # HOST_PRE=numba without a zero-copy staging buffer
        self._img_stage = getattr(self, "_img_stage", None)

    def _fast_pre(self) -> bool:
        """HOST_PRE=numba applies (numba importable; else the numpy path)."""
        from .host import fast as hfast

        return str(self.knobs.HOST_PRE) == "numba" and hfast.available()

    def _img_target(self) -> np.ndarray:
        """Where HOST_PRE=numba writes the bf16 upload rows: the persistent ttnn host buffer (zero copy) when the
        device model exists, else a reused numpy buffer. Calls are serialized by the model lock, and each frame's
        upload completes inside its ``_forward``, so one buffer suffices."""
        if self._img_stage is not None:
            return self._img_stage.view
        n, (h, w) = self.cfg.num_cams, tuple(self.cfg.input_hw)
        if self._img_bits is None:
            self._img_bits = np.empty((n * h * w, 3), np.uint16)
        return self._img_bits

    def _img_input(self, prepared: Dict[str, Any]) -> Any:
        """The value of the ``img`` trace input: the staging host tensor, a ttnn host tensor of the bf16 bits, or
        (HOST_PRE=numpy) the float32 NHWC rows."""
        if "img_bits" not in prepared:
            from .tt.model import img_rows

            return self.tt.img_input(img_rows(prepared["img"]))
        if self._img_stage is not None and prepared["img_bits"] is self._img_stage.view:
            return self._img_stage.tensor
        import torch
        import ttnn

        bits = np.ascontiguousarray(prepared["img_bits"]).reshape(self.tt.img_shape)
        t = torch.from_numpy(bits.view(np.int16)).view(torch.bfloat16)
        return ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT)

    def _stream_state(self, sid: str, reset: bool) -> Any:
        """The host temporal state of ``sid`` (``host.temporal.StreamState``); a new id or ``reset`` starts fresh;
        beyond ``MAX_STREAMS`` the least recently used stream is dropped (PLAN D16, first publish: one stream)."""
        from .host.temporal import StreamState

        if reset or sid not in self._streams:
            self._streams[sid] = StreamState(self.cfg, autoware_compat=bool(self.knobs.AUTOWARE_COMPAT),
                                             use_temporal=bool(self.knobs.USE_TEMPORAL))
        self._streams.move_to_end(sid)
        while len(self._streams) > int(self.knobs.MAX_STREAMS):
            self._streams.popitem(last=False)
        return self._streams[sid]

    def _calibration(self, rig: Any) -> Optional[Dict[str, Any]]:
        """Once-per-calibration constants (PE, cone, spatial-alignment gamma / beta), cached by the rig's hash."""
        if self.calib_fn is None:
            return None
        key = rig.key()
        if key not in self._calib:
            self._calib[key] = self.calib_fn(rig)
            while len(self._calib) > 8:
                self._calib.popitem(last=False)
        self._calib.move_to_end(key)
        return self._calib[key]

    def _prepare(self, points: Any = None, *, images: Any = None, calibration: Optional[Mapping[str, Any]] = None,
                 stream: Optional[Mapping[str, Any]] = None, **_: Any) -> Dict[str, Any]:
        """Five images + calibration + stream -> the host values of one frame: the network input ``img``
        [1, 5, 3, 480, 640] float32 (preset ``NORMALIZATION``), the rig calibration (+ its cached PE constants once
        the device model is built) and the stream's frame plan (ts, E / E^-1, the ages, fresh / skipped)."""
        from .host.inputs import prepare_frame
        from .host.preprocess import preprocess_images
        from .host.temporal import StreamState

        if points is not None:
            raise tio.InputError("this model takes five camera images, not a point cloud")
        frame = prepare_frame(images, calibration, stream, self.cfg)
        st = frame.stream
        if st.stateless:
            sid, state = None, StreamState(self.cfg, autoware_compat=bool(self.knobs.AUTOWARE_COMPAT),
                                           use_temporal=bool(self.knobs.USE_TEMPORAL))
        else:
            sid, state = st.id, self._stream_state(st.id, st.reset)
        anchor = frame.anchor_stamp if frame.anchor_stamp is not None else 0.0
        T = np.eye(4) if st.stateless else st.T_world_from_ego     # None: no ego pose -> skipped (no_ego_pose)
        plan = state.begin(anchor, T, frame.camera_stamps)
        out: Dict[str, Any] = {"frame": frame, "plan": plan, "state": state, "stream_id": sid}
        if plan.skipped is None:
            if self._fast_pre():
                from .host import fast as hfast

                out["img_bits"] = hfast.preprocess_rows_bf16(
                    frame.images, preset=self.cfg.normalization, dst_hw=tuple(self.cfg.input_hw),
                    fma=bool(self.knobs.PREPROCESS_FMA), threads=int(self.knobs.PREPROCESS_WORKERS),
                    out=self._img_target())
            else:
                out["img"] = preprocess_images(frame.images, self.cfg, fma=bool(self.knobs.PREPROCESS_FMA),
                                               workers=int(self.knobs.PREPROCESS_WORKERS))
            out["calib"] = self._calibration(frame.rig)
        return out

    def _select_device_memory(self, prepared: Dict[str, Any]) -> None:
        """Bring the right temporal memory into the device states before the frame (PLAN D16, PORT_LOG S11):
        the memory of the stream whose frame runs (saved to / loaded from its DRAM bank when another stream or a
        stateless frame ran in between), zeros for a fresh stream or a stateless frame, and the node's per-frame
        wipe when ``USE_TEMPORAL`` is off."""
        r, sid, plan, state = self.runner, prepared["stream_id"], prepared["plan"], prepared["state"]
        owner = self._device_owner
        if owner is not None and owner != sid and owner in self._streams and owner in self._banks:
            for name in ("emb", "geo"):
                r.save_state(name, self._banks[owner])
        if sid is not None and sid not in self._banks:
            used = set(self._banks.values())
            free_banks = [b for b in range(int(self.knobs.MAX_STREAMS)) if b not in used]
            if not free_banks:                                   # the least recently used stream lost its memory
                old, bank = self._banks.popitem(last=False)
                free_banks = [bank]
            self._banks[sid] = free_banks[0]
        if sid is not None:
            self._banks.move_to_end(sid)
        for gone in [k for k in self._banks if k not in self._streams]:
            del self._banks[gone]
        if sid is None or plan.fresh:
            self.tt.set_memory()                                  # a fresh node: zero memory
        elif not state.use_temporal:
            ref = state.wiped_device_memory()
            n = self.cfg.pre_memory_length
            self.tt.set_memory(None, np.tile(ref, (n, 1)), np.zeros((n, 4, 4), np.float32),
                               np.zeros((n, 2), np.float32))
        elif owner != sid:
            for name in ("emb", "geo"):
                r.load_state(name, self._banks[sid])
        self._device_owner = sid

    def _forward(self, prepared: Dict[str, Any]) -> Any:
        """The traced device graph: upload ``img`` and the per-frame ages / relative motion (and the calibration
        constants when the calibration changed), replay ``frame`` (backbone + head + memory update), read the packed
        cls / bbox. A skipped frame (camera sync) runs nothing."""
        import time

        from .tt.model import TtStreamPETR
        from .tt.params import age_input

        plan = prepared["plan"]
        if plan.skipped is not None:
            return None
        t0 = time.perf_counter()
        self._select_device_memory(prepared)
        params = self.tt.geometry_params(plan.M)
        calib = prepared["calib"]
        if calib["key"] != self._active_calib:
            params.update(calib["params"])
            self._active_calib = calib["key"]
        inputs = {"img": self._img_input(prepared),
                  "ages": age_input(plan.pre["pre_memory_timestamp"], device_shape=True)}
        t1 = time.perf_counter()
        packed = self.runner("frame", inputs=inputs, params=params)
        t2 = time.perf_counter()
        prepared["state"].end_device(plan)
        out = TtStreamPETR.unpack(packed["cls"], packed["reg"])
        out["timing_ms"] = {"host_in": (t1 - t0) * 1e3, "trace": (t2 - t1) * 1e3}
        return out

    def _postprocess(self, raw: Any, prepared: Dict[str, Any], params: Dict[str, Any]) -> Output:
        """Head outputs (``raw``: ``all_cls_scores`` [1, 5, 5400], ``all_bbox_preds`` [1, 10, 5400]) -> the node's
        decode of all 6 layers, thresholds, NMS -> ``Detections3D``; a skipped frame (camera sync) -> no objects."""
        from .host import postprocess as hpost
        from .reference.pipeline import to_detections

        plan = prepared["plan"]
        meta = {"stream": prepared["stream_id"], "fresh": bool(plan.fresh), "skipped": plan.skipped,
                "ts": float(plan.ts), "frame_index": int(plan.frame_index)}
        if plan.skipped is not None or raw is None:
            empty = {"objects": hpost.decode(np.zeros((5, 0), np.float32), np.zeros((10, 0), np.float32), self.cfg),
                     "num_pre_nms": 0}
            return to_detections(empty, self.cfg, model=self.MODEL_NAME, meta=meta)
        res = hpost.autoware_postprocess(raw["all_cls_scores"], raw["all_bbox_preds"], self.cfg,
                                         hpost.params_from(params, self.cfg),
                                         fast_nms=str(self.knobs.HOST_POST) == "numba")
        out = to_detections(res, self.cfg, model=self.MODEL_NAME, meta=meta)
        out.timing_ms.update({k: round(v, 3) for k, v in (raw.get("timing_ms") or {}).items()})
        return out

    def _release(self) -> None:
        """Release the traces and persistent device tensors; also called when ``from_pretrained`` fails half-way."""
        runner = getattr(self, "runner", None)
        if runner is not None:
            runner.release()

    def extra_info(self) -> Dict[str, Any]:
        """Additions to ``model.info`` and ``/info``: the trace variants, CQs and persistent tensors of the runner,
        the load-time knobs."""
        info: Dict[str, Any] = {}
        runner = getattr(self, "runner", None)
        if runner is not None:
            info["trace"] = runner.describe()
        knobs = getattr(self, "knobs", None)
        if knobs is not None:
            info["knobs"] = knobs.as_dict()
        return info