# SPDX-License-Identifier: Apache-2.0
"""Pure-PyTorch fp32 re-implementation of ``meteor_v157c3Z.onnx`` with per-module taps.

A node-by-node transliteration of the deployed graph (2,021 ONNX nodes; ``research/meteor/SPEC.md`` sections 4.2-4.4,
cited ``S:<line>``), written from the graph itself (no upstream code is imported or executed): every convolution runs
with the weights and the attributes of its ONNX node (``reference.weights``), every non-conv constant (attention
pools, LayerNorm, gates, pooling matrices, lift ground points) is read from its consuming node, and the arithmetic
keeps the graph's operation order where it decides a rounding or a discrete choice (lift projection, validity and
depth-bin selection, risk sampling, decoupled decoder, depth_mean). Shape glue (Shape / Gather / Concat / Reshape of
shapes, ~470 nodes) disappears. ``MeteorNet.forward`` returns the 19 graph outputs with the ONNX names, dtypes and
shapes; ``taps=`` records the intermediates below (names shared with ``reference.ort.ORT_TAPS``, the device gates
and the goldens).

Module boundaries (what a TT stage computes) and taps:

==========================  ===========================  ====================================================
tap                         shape                        graph tensor / meaning
==========================  ===========================  ====================================================
``img.input``               [8, 3, 432, 768]             uint8 / 255 (+ ImageNet mean/std with input_norm=imagenet)
``img.stem``                [8, 64, 108, 192]            7x7 s2 + ReLU + maxpool 3x3 s2
``img.layer1..4``           [8, 64..512, ...]            ResNet-34 stages (BasicBlocks, BN folded)
``img.fpn``                 [8, 160, 108, 192]           lat1..4 + bilinear (x2, x4, 14x24 -> 108x192) + fuse 3x3 ReLU
``seg2d.logits``            [8, 21, 108, 192]            SegHeadED raw logits (the PointPainting input, unclamped)
``seg2d.clamped``           [8, 21, 108, 192]            nan_to_num + clamp +-30 (the seg2d argmax input only)
``depth.logits`` / ``.prob``  [8, 64, 108, 192]          depth head 1x1 / softmax over the 64 linear bins
``ctx.painted``             [8, 96, 108, 192]            ctx 1x1 + paint_proj(softmax(seg)[2,3,4,5,6,7,8,13])
``lift.grid``               [8, 100000, 1, 2]            normalised (u, v) of the 400x250 ground points per camera
``lift.valid``              [8, 100000] bool             z > 0.5, inside the image, range < 90 m
``lift.weight``             [8, 1, 100000]               (interpolated bin probability + 0.05) * valid
``lift.bev``                [1, 96, 400, 250]            sum_n ctx_s w / max(sum_n w, 1e-4)
``bev.raw``                 [1, 96, 800, 500]            bilinear x2 of lift.bev
``bev.gate``                [1, 1, 800, 500]             4 softmax(tgate(raw))[0]
``bev.fused``               [1, 96, 800, 500]            raw + tfuse(raw * gate)
``lane.pre``                [1, 9, 800, 500]             lane decoder + lane_branch on channels 4..6
``det.feat``                [1, 128, 400, 250]           det stem (ED) output
``det.hm_pre/.reg_pre``     [1, 2 | 6, 400, 250]         heads before the box refiner
``occ.feat``                [1, 192, 200, 200]           occ stem on raw rows 200:600, cols 50:450
``traj.stem / .feat``       [1, 128 | 256, 400, 250]     traj stem; cat(stem, det.feat) + agent_delta
``traj.agent_delta``        [1, 256, 1, 1]               agent attention pool -> 1x1 conv
``stat.feat``               [1, 48, 400, 250]            delta_stat convs on avgpool2(|raw|)
``risk.feat``               [1, 64, 400, 250]            risk ConvBlocks on fused rows 200:600, cols 125:375
``ego.*``                   [1, 42] (and others)         planner chain: stem, mlp, attn, plus_attn, sem_in, sem,
                                                         plus_sem, kin, plus_kin, risk_sample, after_risk,
                                                         dec_head, dec_wp, after_dec
``bev.fused_mean``          [1, 96]                      global mean of the fused BEV (computed once; S:368-369)
``refiner.seg_res``         [1, 9, 800, 500]             seg refiner residual (before 8 tanh(x / 8))
``refiner.box_res``         [1, 8, 400, 250]             box refiner residual (before 6 tanh(x / 6))
``refiner.e2e_ln / _res``   [1, 139] / [1, 36]           e2e LayerNorm output / 3 tanh(mlp)
==========================  ===========================  ====================================================

The 19 outputs are tapped under their ONNX names (``lane``, ``depth``, ``seg2d``, ``hm`` ... ``depth_mean``).
"""
from __future__ import annotations

from typing import Any, Dict, Optional

import numpy as np

from ..ttaw.golden import NULL_TAPS, TapRegistry
from . import config as C
from .config import MeteorConfig, default_config
from .weights import MeteorWeights

__all__ = ["MeteorNet", "OUTPUT_NAMES", "child"]

OUTPUT_NAMES = ("lane", "depth", "seg2d", "hm", "reg", "hm2d_s0", "hm2d_s1", "hm2d_s2", "reg2d_s0", "reg2d_s1",
                "reg2d_s2", "ego", "occ", "traj", "stationary", "tl", "risk", "lane_logit", "depth_mean")


def _torch():
    import torch

    return torch


def child(module: str, index: int) -> str:
    """ONNX scope of child ``index`` of a Sequential: ``child("seg_head/d1/d1.3", 0)`` = ``seg_head/d1/d1.3/d1.3.0``."""
    return f"{module}/{module.rsplit('/', 1)[-1]}.{index}"


class MeteorNet:
    """The network as fp32 torch functions over the ONNX weights (no autograd, no nn.Module state)."""

    def __init__(self, weights: MeteorWeights, cfg: Optional[MeteorConfig] = None, *, check_graph: bool = True):
        torch = _torch()
        self.cfg = (cfg or default_config()).check()
        self.weights = weights
        if check_graph:
            weights.check_graph()
        f32 = lambda a: torch.from_numpy(np.array(a, dtype=np.float32))  # own writable copy
        self._conv: Dict[str, Any] = {}
        for module, p in weights.iter_convs():
            pads = tuple(int(v) for v in p.pads)
            if len(pads) != 4 or pads[:2] != pads[2:] or p.group != 1 or tuple(p.dilations) != (1, 1):
                raise ValueError(f"{module}: unsupported conv attributes {p}")
            self._conv[module] = (f32(p.weight), None if p.bias is None else f32(p.bias),
                                  tuple(int(s) for s in p.strides), pads[:2])
        self._gemm: Dict[str, Any] = {}
        for module in weights.gemm_modules():
            g = weights.gemm(module)
            if g.trans_a or not g.trans_b or g.alpha != 1.0 or g.beta != 1.0:
                raise ValueError(f"{module}: unsupported Gemm attributes")
            self._gemm[module] = (f32(g.weight), f32(g.bias))
        self._pool = {name: f32(weights.pool_matrix(name)) for name in
                      ("ego_stem_w", "ego_stem_h", "tl_w", "tl_h", "fused_pool_w", "fused_pool_h", "sem_w", "sem_h",
                       "det_pool_w", "det_pool_h", "fused_mean_w", "fused_mean_h", "e2e_mean_w", "e2e_mean_h")}
        self._mha = {name: weights.mha(name) for name in ("ego_attn", "agent_attn")}
        self._mha_t = {name: (f32(m.in_weight), f32(m.in_bias), f32(m.out_weight), f32(m.out_bias))
                       for name, m in self._mha.items()}
        self._queries = {"ego": f32(weights.queries("ego")), "agent": f32(weights.queries("agent"))}
        g, b, eps = weights.layer_norm()
        self._ln = (f32(g), f32(b), float(eps))
        self._risk_gate = f32(np.asarray(weights.risk_gate(), np.float32).reshape(1))
        self._dec_gate = f32(np.asarray(weights.dec_gate(), np.float32).reshape(1))
        self._ground = f32(weights.lift_ground_points())                     # [1, 4, 100000]
        zc = weights.zero_constants()
        self._hist = f32(zc["hist"])                                         # [1, 96, 800, 500] zeros (baked history)
        self._mask = f32(zc["mask"])                                         # [1, 1, 800, 500] zeros
        self._paint = torch.tensor(weights.paint_classes(), dtype=torch.long)
        self._depth_centres = f32(C.depth_bin_centres(self.cfg.depth_mean_bins)).reshape(1, 1, C.DEPTH_BINS, 1, 1)

    # ---- primitives ------------------------------------------------------------------------------------------------
    def conv(self, x, module: str, relu: bool = False):
        F = _torch().nn.functional
        w, b, stride, pad = self._conv[module]
        y = F.conv2d(x, w, b, stride=stride, padding=pad)
        return F.relu(y) if relu else y

    def block(self, x, module: str):
        """ConvBlock = (3x3 conv, BN (folded), ReLU) x 2: children ``.0`` and ``.3`` (S:318)."""
        return self.conv(self.conv(x, child(module, 0), True), child(module, 3), True)

    def down(self, x, module: str):
        """stride-2 3x3 conv + ReLU (child ``.0``) then a ConvBlock (child ``.3``): the encoder steps of the EDs."""
        return self.block(self.conv(x, child(module, 0), True), child(module, 3))

    def gemm(self, x, module: str):
        w, b = self._gemm[module]
        return _torch().addmm(b, x, w.t())

    @staticmethod
    def resize(x, hw):
        """ONNX Resize linear / half_pixel to a target size = F.interpolate(bilinear, align_corners=False)."""
        F = _torch().nn.functional
        return F.interpolate(x, size=tuple(int(v) for v in hw), mode="bilinear", align_corners=False)

    def pool2(self, x, w_name: str, h_name: str):
        """The exported adaptive pool / mean: (H_mat @ (x @ W_mat)) on the last two dims (S:356-366)."""
        torch = _torch()
        return torch.matmul(self._pool[h_name], torch.matmul(x, self._pool[w_name]))

    @staticmethod
    def nan_to_num(x):
        """The graph's IsNaN / IsInf / Where chain (ONNX nodes 190-213 and 233-255): NaN -> 0, +inf -> 30,
        -inf -> -30; the identity for finite values."""
        torch = _torch()
        x = torch.where(torch.isnan(x), torch.zeros_like(x), x)
        x = torch.where(torch.isinf(x) & (x > 0), torch.full_like(x, 30.0), x)
        return torch.where(torch.isinf(x) & (x < 0), torch.full_like(x, -30.0), x)

    def mha_pool(self, name: str, keys, queries):
        """``nn.MultiheadAttention(q=learned queries, k=v=keys)`` as exported (torch 2.1, need_weights=False):
        q and k|v projections, q / sqrt(head_dim), softmax over the 400 keys, out-projection.
        ``keys`` [L, 1, E] (sequence first), ``queries`` [1, Q, E] -> [1, Q, E]."""
        torch = _torch()
        m = self._mha[name]
        w_in, b_in, w_out, b_out = self._mha_t[name]
        e, h = m.embed_dim, m.num_heads
        hd = e // h
        q_in = queries.transpose(0, 1)                                      # [Q, 1, E]
        q = b_in[:e] + torch.matmul(q_in, w_in[:e].t())                     # /Add_4 = bias + MatMul
        kv = b_in[e:] + torch.matmul(keys, w_in[e:].t())                    # [L, 1, 2E]
        k, v = kv[..., :e], kv[..., e:]
        nq, nk = q.shape[0], k.shape[0]
        qh = q.reshape(nq, h, hd).transpose(0, 1)                           # [H, Q, hd]
        kh = k.reshape(nk, h, hd).transpose(0, 1)                           # [H, L, hd]
        vh = v.reshape(nk, h, hd).transpose(0, 1)
        att = torch.softmax(torch.matmul(qh / m.scale_divisor, kh.transpose(1, 2)), dim=-1)
        o = torch.matmul(att, vh).transpose(0, 1).reshape(nq, e)            # [Q, E]
        o = torch.addmm(b_out, o, w_out.t())                                # out-projection (Gemm)
        return o.reshape(nq, 1, e).transpose(0, 1)                          # [1, Q, E]

    # ---- stages ----------------------------------------------------------------------------------------------------
    def normalize(self, imgs_u8, present=None):
        """uint8 [1, 8, 3, 432, 768] -> [8, 3, 432, 768] float32: /255 (``onnx``) or (x/255 - mean)/std with absent
        cameras zeroed AFTER normalisation (``imagenet``, SPEC section 2 imgs row)."""
        torch = _torch()
        x = imgs_u8.to(torch.float32) / 255.0
        if self.cfg.input_norm == "imagenet":
            mean = torch.tensor(C.IMAGENET_MEAN, dtype=torch.float32).reshape(1, 1, 3, 1, 1)
            std = torch.tensor(C.IMAGENET_STD, dtype=torch.float32).reshape(1, 1, 3, 1, 1)
            x = (x - mean) / std
            if present is not None:
                keep = torch.as_tensor(np.asarray(present, np.float32)).reshape(1, -1, 1, 1, 1)
                x = x * keep
        return x.reshape(C.N_CAMS, 3, C.IMG_H, C.IMG_W)

    def image_encoder(self, x, taps: TapRegistry = NULL_TAPS):
        """ResNet-34 (torchvision layout, BN folded) + FPN -> f [8, 160, 108, 192] (S:301-302)."""
        F = _torch().nn.functional
        y = F.max_pool2d(self.conv(x, "stem/stem.0", True), kernel_size=3, stride=2, padding=1)
        taps.tap("img.stem", y)
        feats = []
        for li, nb in enumerate((3, 4, 6, 3), start=1):
            for bi in range(nb):
                p = f"layer{li}/layer{li}.{bi}"
                ds = f"{p}/downsample/downsample.0"
                idt = self.conv(y, ds) if ds in self._conv else y
                y = F.relu(self.conv(self.conv(y, f"{p}/conv1", True), f"{p}/conv2") + idt)
            feats.append(taps.tap(f"img.layer{li}", y))
        c1, c2, c3, c4 = feats
        hw = (C.FEAT_H, C.FEAT_W)
        s = self.conv(c1, "lat1") + self.resize(self.conv(c2, "lat2"), hw)
        s = s + self.resize(self.conv(c3, "lat3"), hw)
        s = s + self.resize(self.conv(c4, "lat4"), hw)                  # 14x24 -> 108x192: non-integer (S:302)
        return taps.tap("img.fpn", self.conv(s, "fuse/fuse.0", True))

    def seg2d_head(self, f, taps: TapRegistry = NULL_TAPS):
        """SegHeadED -> raw logits [8, 21, 108, 192] (S:303)."""
        d1 = self.down(f, "seg_head/d1")
        d2 = self.down(d1, "seg_head/d2")
        m1 = self.block(d1 + self.resize(self.conv(d2, "seg_head/u1"), d1.shape[-2:]), "seg_head/m1")
        u2 = self.resize(self.conv(m1, "seg_head/u2"), f.shape[-2:])
        y = self.conv(self.conv(self.conv(f, "seg_head/skip") + u2, "seg_head/out/out.0", True), "seg_head/out/out.3")
        return taps.tap("seg2d.logits", y)

    def depth_head(self, f, taps: TapRegistry = NULL_TAPS):
        """4 ConvBlocks 160->128->128->96->64, 1x1 -> 64 bins, softmax over the bins (S:304)."""
        torch = _torch()
        y = f
        for i in range(4):
            y = self.block(y, f"depth_head/depth_head.{i}")
        logits = taps.tap("depth.logits", self.conv(y, "depth_head/depth_head.4"))
        return logits, taps.tap("depth.prob", torch.softmax(logits, dim=1))

    def det2d_heads(self, f):
        """det2d stem + d8 / d16 pyramid, 1x1 heads at strides 4 / 8 / 16 (S:305) -> the six [1, 8, k, h, w] outputs."""
        s = self.block(f, "det2d_stem")
        d8 = self.down(s, "det2d_d8")
        d16 = self.down(d8, "det2d_d16")
        out = {}
        for i, (x, sfx) in enumerate(((s, ""), (d8, "8"), (d16, "16"))):
            for kind, ch in (("hm2d", 10), ("reg2d", 4)):
                y = self.conv(x, f"{kind}_head{sfx}")
                out[f"{kind}_s{i}"] = y.reshape(1, C.N_CAMS, ch, y.shape[-2], y.shape[-1])
        return out

    def tl_head(self, f):
        """cat(f[cam 0], f[cam 6]) -> s2 conv + CB -> s2 conv + CB -> global mean (MatMul consts) -> Linear (S:306)."""
        torch = _torch()
        x = torch.cat([f[0:1], f[6:7]], dim=1)
        x = self.block(self.conv(x, "tl_head/tl_head.0", True), "tl_head/tl_head.3")
        x = self.block(self.conv(x, "tl_head/tl_head.4", True), "tl_head/tl_head.7")
        x = self.pool2(x, "tl_w", "tl_h").flatten(1)
        return self.gemm(x, "tl_fc")

    def context(self, f, seg_logits, taps: TapRegistry = NULL_TAPS):
        """ctx 1x1 + paint_proj(softmax(nan_to_num(raw seg logits))[paint classes]) (S:307; UNCLAMPED logits)."""
        torch = _torch()
        p = torch.softmax(self.nan_to_num(seg_logits), dim=1).index_select(1, self._paint)
        return taps.tap("ctx.painted", self.conv(f, "ctx") + self.conv(p, "paint_proj"))

    def lift(self, ctx, prob, K, T_cam_ego, taps: TapRegistry = NULL_TAPS):
        """Depth-gated IPM lift to the 0.4 m grid, then bilinear x2 to the raw BEV (S:333-351, ONNX nodes 261-421)."""
        torch = _torch()
        F = torch.nn.functional
        n = C.N_CAMS
        T = T_cam_ego.reshape(n, 4, 4).to(torch.float32)
        Km = K.reshape(n, 3, 3).to(torch.float32)
        pc = torch.matmul(T, self._ground)                                  # [8, 4, 100000]
        x, y, z = pc[:, 0], pc[:, 1], pc[:, 2]
        zc = torch.clamp(z, min=C.LIFT_MIN_Z)
        u = Km[:, 0, 0].unsqueeze(-1) * x / zc + Km[:, 0, 2].unsqueeze(-1)
        v = Km[:, 1, 1].unsqueeze(-1) * y / zc + Km[:, 1, 2].unsqueeze(-1)
        dist = torch.sqrt(x * x + y * y + z * z)
        valid = ((z > C.LIFT_MIN_Z) & (u >= 0.0) & (u < float(C.IMG_W)) & (v >= 0.0) & (v < float(C.IMG_H))
                 & (dist < C.LIFT_MAX_DIST))
        taps.tap("lift.valid", valid)
        gu = torch.clamp(u / float(C.IMG_W - 1) * 2.0 - 1.0, -2.0, 2.0)
        gv = torch.clamp(v / float(C.IMG_H - 1) * 2.0 - 1.0, -2.0, 2.0)
        grid = taps.tap("lift.grid", torch.stack([gu, gv], dim=-1).unsqueeze(2))      # [8, 100000, 1, 2]
        ctx_s = F.grid_sample(ctx, grid, mode="bilinear", padding_mode="zeros", align_corners=False).squeeze(3)
        prob_s = F.grid_sample(prob, grid, mode="bilinear", padding_mode="zeros", align_corners=False).squeeze(3)
        b = torch.clamp((dist - C.DEPTH_MIN) / C.DEPTH_STEP, 0.0, float(np.float32(C.LIFT_BIN_CLIP)))
        b0 = torch.floor(b).to(torch.int64)
        fr = (b - b0.to(torch.float32)).unsqueeze(1)                        # [8, 1, 100000]
        g0 = torch.gather(prob_s, 1, b0.unsqueeze(1))
        g1 = torch.gather(prob_s, 1, torch.clamp(b0 + 1, max=C.DEPTH_BINS - 1).unsqueeze(1))
        w = (g0 * (1.0 - fr) + g1 * fr + C.LIFT_W_EPS) * valid.unsqueeze(1).to(torch.float32)
        taps.tap("lift.weight", w)
        wn = w.reshape(1, n, 1, -1)
        num = (ctx_s.reshape(1, n, C.CTX_C, -1) * wn).sum(dim=1)            # [1, 96, 100000]
        den = torch.clamp(wn.sum(dim=1), min=C.LIFT_DEN_EPS)                # [1, 1, 100000]
        bev04 = taps.tap("lift.bev", (num / den).reshape(1, C.CTX_C, C.LIFT_H, C.LIFT_W))
        return taps.tap("bev.raw", self.resize(bev04, (C.BEV_H, C.BEV_W)))

    def temporal_fuse(self, raw, taps: TapRegistry = NULL_TAPS):
        """No-history temporal fuse: fused = raw + tfuse(raw * 4 softmax(tgate(raw))[0]) (S:308)."""
        torch = _torch()
        gate = taps.tap("bev.gate", torch.softmax(self.conv(raw, "tgate_slim"), dim=1)[:, 0:1] * C.TGATE_SCALE)
        t = self.block(self.conv(raw * gate, "tfuse3_slim/tfuse3_slim.0", True), "tfuse3_slim/tfuse3_slim.3")
        return taps.tap("bev.fused", raw + t)

    def lane_decoder(self, raw, taps: TapRegistry = NULL_TAPS):
        """LaneDecED + lane_branch added to channels 4..6 -> lane_pre [1, 9, 800, 500] (S:309)."""
        torch = _torch()
        skip = self.block(raw, "dec/skip")
        d1 = self.down(raw, "dec/d1")
        d2 = self.down(d1, "dec/d2")
        m1 = self.block(d1 + self.resize(self.conv(d2, "dec/u1"), d1.shape[-2:]), "dec/m1")
        u2 = self.resize(self.conv(m1, "dec/u2"), skip.shape[-2:])
        out = self.conv(self.conv(skip + u2, "dec/out/out.0", True), "dec/out/out.3")
        lb = self.conv(self.conv(self.conv(raw, "lane_branch/lane_branch.0", True), "lane_branch/lane_branch.3", True),
                       "lane_branch/lane_branch.6")
        return taps.tap("lane.pre", torch.cat([out[:, :4], out[:, 4:7] + lb, out[:, 7:]], dim=1))

    def det_stem(self, raw, taps: TapRegistry = NULL_TAPS):
        """_DetStemED on the raw BEV -> det_feat [1, 128, 400, 250], hm_pre, reg_pre (S:310)."""
        s = self.down(raw, "det_stem/stem")
        d8 = self.down(s, "det_stem/d8")
        feat = self.conv(s + self.resize(self.conv(d8, "det_stem/u"), s.shape[-2:]), "det_stem/m/m.0", True)
        taps.tap("det.feat", feat)
        hm = taps.tap("det.hm_pre", self.conv(feat, "hm_head"))
        return feat, hm, taps.tap("det.reg_pre", self.conv(feat, "reg_head"))

    def occupancy(self, raw, taps: TapRegistry = NULL_TAPS):
        """occ stem on the raw BEV crop rows 200:600, cols 50:450 -> 1x1 160 -> view [1, 10, 16, 200, 200] (S:311)."""
        (r0, r1), (c0, c1) = C.OCC_CROP
        x = self.down(raw[:, :, r0:r1, c0:c1], "occ_stem")
        taps.tap("occ.feat", x)
        y = self.conv(x, "occ_head")
        return y.reshape(1, len(C.OCC_CLASSES), C.OCC_Z_BINS, y.shape[-2], y.shape[-1])

    def agent_forecast(self, raw, fused, det_feat, reg_pre, taps: TapRegistry = NULL_TAPS):
        """traj_stem on cat(fused, mot = (raw - hist) * mask = 0), agent attention pool, traj head (S:312)."""
        torch = _torch()
        mot = (raw - self._hist) * self._mask                              # the baked-out history: all zeros
        stem = self.block(self.conv(torch.cat([fused, mot], dim=1), "traj_stem/traj_stem.0", True),
                          "traj_stem/traj_stem.3")
        del mot
        taps.tap("traj.stem", stem)
        reg46 = reg_pre[:, 4:6]
        keys = self.pool2(det_feat, "det_pool_w", "det_pool_h").reshape(1, det_feat.shape[1], -1).permute(2, 0, 1)
        att = self.mha_pool("agent_attn", keys, self._queries["agent"])     # [1, 4, 128]
        delta = self.conv(att.mean(dim=1).reshape(1, -1, 1, 1), "agent_delta")
        taps.tap("traj.agent_delta", delta)
        feat = taps.tap("traj.feat", torch.cat([stem, det_feat], dim=1) + delta)
        return self.conv(torch.cat([feat, reg46], dim=1), "traj_head")

    def stationary(self, raw, taps: TapRegistry = NULL_TAPS):
        """delta_stat on avgpool2(|raw - 0|) (S:313)."""
        F = _torch().nn.functional
        x = F.avg_pool2d((raw - self._hist).abs(), kernel_size=2, stride=2)
        x = self.conv(self.conv(x, "delta_stat/delta_stat.1", True), "delta_stat/delta_stat.4", True)
        taps.tap("stat.feat", x)
        return self.conv(x, "delta_stat/delta_stat.7")

    def risk(self, fused, taps: TapRegistry = NULL_TAPS):
        """risk head on the fused BEV crop rows 200:600, cols 125:375 (S:314)."""
        (r0, r1), (c0, c1) = C.RISK_CROP
        x = self.block(self.block(fused[:, :, r0:r1, c0:c1], "risk_head/risk_head.0"), "risk_head/risk_head.1")
        taps.tap("risk.feat", x)
        return self.conv(x, "risk_head/risk_head.2")

    def planner(self, fused, lane_pre, hm_pre, risk, v0, taps: TapRegistry = NULL_TAPS):
        """Ego planner up to the e2e refiner (S:315): ego_stem + MLP, + ego_attn, + sem_ego, + kin, risk-aware
        mode logits, decoupled decoder. Returns ``(ego_after_dec [1, 42], fused_mean [1, 96])``."""
        torch = _torch()
        F = torch.nn.functional
        with taps.scope("ego"):
            e = self.block(self.conv(fused, "ego_stem/ego_stem.0", True), "ego_stem/ego_stem.3")     # s4
            e = self.block(self.conv(e, "ego_stem/ego_stem.4", True), "ego_stem/ego_stem.7")
            e = self.block(self.conv(e, "ego_stem/ego_stem.8", True), "ego_stem/ego_stem.11")
            e = self.conv(e, "ego_stem/ego_stem.12", True)                                           # [1, 256, 25, 16]
            e = taps.tap("stem", self.pool2(e, "ego_stem_w", "ego_stem_h").flatten(1))
            v0 = v0.reshape(1, 1).to(torch.float32)
            h = torch.cat([e, v0], dim=1)
            for i in (0, 2, 4):
                h = F.relu(self.gemm(h, f"ego_mlp/ego_mlp.{i}"))
            ego = taps.tap("mlp", self.gemm(h, "ego_mlp/ego_mlp.6"))
            keys = self.pool2(fused, "fused_pool_w", "fused_pool_h").reshape(1, fused.shape[1], -1).permute(2, 0, 1)
            att = self.mha_pool("ego_attn", keys, self._queries["ego"])                              # [1, 3, 96]
            ego = taps.tap("plus_attn", ego + taps.tap("attn", self.gemm(att.flatten(1), "ego_delta")))
            sem_in = torch.cat([F.avg_pool2d(torch.softmax(lane_pre, dim=1), kernel_size=2, stride=2),
                                torch.sigmoid(hm_pre)], dim=1)
            taps.tap("sem_in", sem_in)
            s = self.conv(self.conv(self.conv(sem_in, "sem_ego/sem_ego.0", True), "sem_ego/sem_ego.3", True),
                          "sem_ego/sem_ego.6", True)
            s = self.pool2(s, "sem_w", "sem_h").flatten(1)
            s = self.gemm(F.relu(self.gemm(s, "sem_ego/sem_ego.11")), "sem_ego/sem_ego.13")
            ego = taps.tap("plus_sem", ego + taps.tap("sem", s))
            kin_in = torch.zeros(1, C.KIN_DIM, dtype=torch.float32)
            kin_in[0, 0] = v0.reshape(()) / C.V0_SCALE
            ego = taps.tap("plus_kin", ego + taps.tap("kin", self.gemm(kin_in, "kin_delta")))
            # risk-aware mode logits: sample sigmoid(risk) (border padding) at each path's 6 waypoints (S:315)
            paths = ego[:, :36].reshape(1, C.EGO_MODES, C.EGO_STEPS, 2)
            gx = -paths[..., 1] / C.RISK_Y_HALF
            gy = (C.RISK_X_HALF - paths[..., 0]) / C.RISK_X_HALF - 1.0
            samp = F.grid_sample(torch.sigmoid(risk), torch.stack([gx, gy], dim=-1), mode="bilinear",
                                 padding_mode="border", align_corners=False)                        # [1, 1, 3, 6]
            samp = taps.tap("risk_sample", samp.mean(dim=3).squeeze(1))                              # [1, 3]
            ego = torch.cat([ego[:, :36], ego[:, 36:39] - self._risk_gate * samp, ego[:, 39:]], dim=1)
            taps.tap("after_risk", ego)
            # decoupled decoder on mean(fused): heading phi and speed v = min(softplus, 25) -> cumulative waypoints
            fmean = self.pool2(fused, "fused_mean_w", "fused_mean_h").flatten(1)
            dh = taps.tap("dec_head", self.gemm(fmean, "dec_head")).reshape(1, C.EGO_MODES, 12)
            phi, vr = dh[..., :6], dh[..., 6:]
            sp = torch.clamp(vr, min=0.0) + torch.log(1.0 + torch.exp(-vr.abs()))
            step = torch.clamp(sp, max=C.DEC_SPEED_MAX) * 0.5
            wp = taps.tap("dec_wp", torch.stack([torch.cumsum(step * torch.cos(phi), dim=2),
                                                 torch.cumsum(step * torch.sin(phi), dim=2)], dim=-1))
            e36 = ego[:, :36].reshape(1, C.EGO_MODES, C.EGO_STEPS, 2)
            e36 = e36 + torch.tanh(self._dec_gate) * (wp - e36)
            ego = taps.tap("after_dec", torch.cat([e36.reshape(1, 36), ego[:, 36:]], dim=1))
        return ego, fmean

    @staticmethod
    def _range_channel(h: int, w: int):
        """The refiners' constant range channel: row i -> (80 - (i + 0.5) * (160 / h)) / 80, float32 in graph order
        (Range + 0.5, Reciprocal(h) * 160, Mul, 80 - ., / 80)."""
        torch = _torch()
        step = (1.0 / torch.tensor(float(h), dtype=torch.float32)) * 160.0
        rows = (80.0 - (torch.arange(h, dtype=torch.float32) + 0.5) * step) / 80.0
        return rows.reshape(1, 1, h, 1).expand(1, 1, h, w)

    def seg_refiner(self, lane_pre, taps: TapRegistry = NULL_TAPS):
        """BEVSegRefiner: lane_logit = clamp(lane_pre, +-20) + 8 tanh(res / 8) (S:316)."""
        torch = _torch()
        lc = torch.clamp(lane_pre, -C.SEG_REFINE_CLAMP, C.SEG_REFINE_CLAMP)
        x = torch.cat([lc, self._range_channel(lc.shape[-2], lc.shape[-1])], dim=1)
        p = "refiner/seg"
        stem = self.conv(x, f"{p}/stem/stem.0", True)
        d1 = self.down(stem, f"{p}/d1")
        d2 = self.down(d1, f"{p}/d2")
        d3 = self.down(d2, f"{p}/d3")
        d4 = self.down(d3, f"{p}/d4")
        m3 = self.block(d3 + self.resize(self.conv(d4, f"{p}/u4"), d3.shape[-2:]), f"{p}/m3")
        m2 = self.block(d2 + self.resize(self.conv(m3, f"{p}/u3"), d2.shape[-2:]), f"{p}/m2")
        m1 = self.block(d1 + self.resize(self.conv(m2, f"{p}/u2"), d1.shape[-2:]), f"{p}/m1")
        o = stem + self.resize(self.conv(m1, f"{p}/u1"), stem.shape[-2:])
        res = taps.tap("refiner.seg_res", self.conv(self.conv(o, f"{p}/out/out.0", True), f"{p}/out/out.3"))
        return lc + torch.tanh(res / C.SEG_REFINE_SCALE) * C.SEG_REFINE_SCALE

    def box_refiner(self, hm_pre, reg_pre, taps: TapRegistry = NULL_TAPS):
        """BEVBoxRefiner: hm, reg = clamp(.) + 6 tanh(res / 6) (S:317)."""
        torch = _torch()
        hc = torch.clamp(hm_pre, -C.BOX_HM_CLAMP, C.BOX_HM_CLAMP)
        rc = torch.clamp(reg_pre, -C.BOX_REG_CLAMP, C.BOX_REG_CLAMP)
        x = torch.cat([hc, rc, self._range_channel(hc.shape[-2], hc.shape[-1])], dim=1)
        p = "refiner/box"
        stem = self.conv(x, f"{p}/stem/stem.0", True)
        d1 = self.down(stem, f"{p}/d1")
        d2 = self.down(d1, f"{p}/d2")
        d3 = self.down(d2, f"{p}/d3")
        m2 = self.block(d2 + self.resize(self.conv(d3, f"{p}/u3"), d2.shape[-2:]), f"{p}/m2")
        m1 = self.block(d1 + self.resize(self.conv(m2, f"{p}/u2"), d1.shape[-2:]), f"{p}/m1")
        o = stem + self.resize(self.conv(m1, f"{p}/u1"), stem.shape[-2:])
        res = taps.tap("refiner.box_res", self.conv(self.conv(o, f"{p}/out/out.0", True), f"{p}/out/out.3"))
        r = torch.tanh(res / C.BOX_REFINE_SCALE) * C.BOX_REFINE_SCALE
        return hc + r[:, :2], rc + r[:, 2:]

    def e2e_refiner(self, ego, v0, fused_mean, taps: TapRegistry = NULL_TAPS):
        """E2ERefiner: ego[:36] += 3 tanh(mlp(LN([ego, v0, mean(fused)]))) (S:318; LN at the logical width 139)."""
        torch = _torch()
        F = torch.nn.functional
        g, b, eps = self._ln
        x = torch.cat([ego, v0.reshape(1, 1).to(torch.float32), fused_mean], dim=1)
        h = taps.tap("refiner.e2e_ln", F.layer_norm(x, (x.shape[-1],), g, b, eps))
        h = F.relu(self.gemm(h, "refiner/e2e/mlp/mlp.0"))
        h = F.relu(self.gemm(h, "refiner/e2e/mlp/mlp.2"))
        r = taps.tap("refiner.e2e_res", torch.tanh(self.gemm(h, "refiner/e2e/mlp/mlp.4")) * C.E2E_SCALE)
        return torch.cat([ego[:, :36] + r, ego[:, 36:]], dim=1)

    def depth_mean(self, depth_logits):
        """avgpool2 of the 64 logits -> softmax -> argmax -> +-2-bin window -> sum p c / max(sum p, 1e-6) -> fp16
        (ONNX nodes 1946-2014; bin centres per ``cfg.depth_mean_bins``)."""
        torch = _torch()
        F = torch.nn.functional
        lp = F.avg_pool2d(depth_logits, kernel_size=2, stride=2)
        p = torch.softmax(lp.reshape(1, C.N_CAMS, C.DEPTH_BINS, lp.shape[-2], lp.shape[-1]), dim=2)
        am = torch.argmax(p, dim=2, keepdim=True)
        r = torch.arange(C.DEPTH_BINS, dtype=torch.int64).reshape(1, 1, C.DEPTH_BINS, 1, 1)
        pm = p * ((r - am).abs() <= 2).to(torch.float32)
        num = (pm * self._depth_centres).sum(dim=2)
        den = torch.clamp(pm.sum(dim=2), min=1e-6)
        return (num / den).to(torch.float16)

    # ---- whole graph ------------------------------------------------------------------------------------------------
    def forward(self, imgs, K, T_cam_ego, v0, *, present=None, taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]:
        """The 19 outputs of the graph (torch tensors, ONNX dtypes) for one frame.

        ``imgs`` uint8 [1, 8, 3, 432, 768] RGB, ``K`` [1, 8, 3, 3] at 768x432, ``T_cam_ego`` [1, 8, 4, 4], ``v0`` [1];
        numpy or torch. ``present`` (8 flags) only matters for ``input_norm="imagenet"``."""
        torch = _torch()
        t = lambda a, dt: a.to(dt) if isinstance(a, torch.Tensor) else torch.from_numpy(np.ascontiguousarray(a)).to(dt)
        imgs = t(imgs, torch.uint8)
        K = t(K, torch.float32)
        T = t(T_cam_ego, torch.float32)
        v0 = t(v0, torch.float32).reshape(1)
        out: Dict[str, Any] = {}
        with torch.inference_mode():
            x = taps.tap("img.input", self.normalize(imgs, present))
            f = self.image_encoder(x, taps)
            del x
            seg = self.seg2d_head(f, taps)
            seg_c = taps.tap("seg2d.clamped", torch.clamp(self.nan_to_num(seg), -C.SEG2D_CLAMP, C.SEG2D_CLAMP))
            dlog, dprob = self.depth_head(f, taps)
            ctx = self.context(f, seg, taps)
            out.update(self.det2d_heads(f))
            out["tl"] = self.tl_head(f)
            del f
            raw = self.lift(ctx, dprob, K, T, taps)
            del ctx, dprob
            fused = self.temporal_fuse(raw, taps)
            lane_pre = self.lane_decoder(raw, taps)
            det_feat, hm_pre, reg_pre = self.det_stem(raw, taps)
            out["occ"] = self.occupancy(raw, taps)
            out["traj"] = self.agent_forecast(raw, fused, det_feat, reg_pre, taps)
            out["stationary"] = self.stationary(raw, taps)
            del raw, det_feat
            out["risk"] = self.risk(fused, taps)
            ego, fmean = self.planner(fused, lane_pre, hm_pre, out["risk"], v0, taps)
            taps.tap("bev.fused_mean", fmean)
            out["lane_logit"] = self.seg_refiner(lane_pre, taps)
            out["hm"], out["reg"] = self.box_refiner(hm_pre, reg_pre, taps)
            e2e_mean = self.pool2(fused, "e2e_mean_w", "e2e_mean_h").flatten(1)
            out["ego"] = self.e2e_refiner(ego, v0, e2e_mean, taps)
            del fused
            out["lane"] = torch.argmax(out["lane_logit"], dim=1).to(torch.uint8)
            dl5 = dlog.reshape(1, C.N_CAMS, C.DEPTH_BINS, C.FEAT_H, C.FEAT_W)
            out["depth"] = torch.argmax(dl5, dim=2).to(torch.uint8)
            out["seg2d"] = torch.argmax(seg_c.reshape(1, C.N_CAMS, len(C.SEG2D_CLASSES), C.FEAT_H, C.FEAT_W),
                                        dim=2).to(torch.uint8)
            out["depth_mean"] = self.depth_mean(dlog)
            for name in OUTPUT_NAMES:
                taps.tap(name, out[name])
        return {name: out[name] for name in OUTPUT_NAMES}

    __call__ = forward