# SPDX-License-Identifier: Apache-2.0
"""fp32 PyTorch re-implementation of the deployed BEVFusion-L network, module by module, with taps.

Built from :mod:`.weights` (the ONNX initializers addressed by node; every BatchNorm folded in float64 where the TT
port folds it), it reproduces ``bevfusion_lidar.onnx`` (S:bevfusion:205-264) from the voxel features on (the VFE mean
is host pre-processing in this port, computed in float32 exactly as the graph does by ``host.preprocess.vfe_mean``):

- ``sparse`` (the 21 ``autoware::ImplicitGemm`` layers): every layer in the **gather form** the device uses (C13 /
  C28): ``y[o] = sum_k W'[:, k, :] @ x[nmap[o, k]] + b'`` over the neighbour maps of ``host.preprocess.SparsePlan``
  (``-1``: no input), ReLU, the basic blocks' residual before their second ReLU; BN folded (``W' = W * s``,
  ``b' = beta - mean * s``). The plan's maps are spconv-faithful (13-query + symmetric write); the independent
  oracle is the research reference's scatter form (``tests/test_reference_cpu.py``);
- ``to-dense``: ``dense[0, x, y, z, :] = conv_out[o]`` on a zero ``[1, 180, 180, 2, 128]`` map, ``Transpose(0, 4, 3,
  1, 2)``, ``Reshape [1, 256, 180, 180]``: channel ``c * 2 + z``, rows x, columns y;
- ``backbone``: SECOND, 2 blocks of 6 3x3 convs + ReLU (first stride 2 in block 1); ``neck``: SECONDFPN deblock0 1x1
  + ReLU, deblock1 ConvTranspose 2x2 s2 + ReLU, concat (1, 512, 180, 180);
- ``head`` dense part: shared conv 3x3 512 -> 128 (``lidar_feat``, no activation), heatmap head (3x3 + ReLU, 3x3 -> 5
  logits), ``heat = sigmoid``, local max ``heat * (heat == local_max)`` with ``local_max`` = MaxPool 3x3 s1 p0 of
  classes 0-3 placed in a zero map (their border cells never survive) and class 4 = the heat itself (passes);
- selection: the top ``K = 500`` of the flattened (5 x 32400) local-max heatmap as ONNX ``TopK`` defines it (largest
  first, **ties to the lower index**: a stable descending sort); ``topk_indices=`` replaces it by given indices
  (teacher forcing, PLAN.md 0.3 item 1: the device top-k breaks ties in no defined order, probe P8);
- query init: ``class = idx // 32400``, ``pos = idx % 32400``; ``query = lidar_feat[:, pos] + class_table[class]``
  (the Conv1d of a one-hot is the table row ``W[:, c] + b``); ``query_pos = (pos // 180 + 0.5, pos % 180 + 0.5)``
  computed arithmetically in float32 (bit-equal to the graph's ``bev_pos`` gather, checked at load);
  ``query_heat = heat_nms[:, :, pos]`` (all 5 classes);
- decoder (one layer, post-norm, 8 heads x 16): self-attention on ``x + pe(query_pos)`` (Q = K = V), cross-attention
  of ``x + pe(query_pos)`` over ``lidar_feat^T + pe(bev_pos)`` (32,400 keys), FFN 128 -> 256 -> 128, LayerNorm eps
  1e-5 after each; ``Q`` divided by ``sqrt(16) = 4`` before ``Q K^T`` (the graph's ``Div``);
- prediction heads (Conv1d 128 -> 64 + BN folded + ReLU, Conv1d 64 -> C: center 2, height 1, dim 3, rot 2, vel 2,
  heatmap 5);
- outputs: ``bbox_pred = [center + query_pos, height, dim, rot, vel]`` (10, K), ``score = sigmoid(heatmap)[label] *
  query_heat[label]`` (K,) (the graph's one-hot product + ReduceMax, equal because both factors are >= 0),
  ``label_pred = class`` (K,) int64.

Taps (``ttaw.golden.TapRegistry``; stable names, the keys of the device gates; dense taps keep the ONNX tensors'
batch axis and layout; sparse taps are ``(N_level, C)`` rows in the plan's level order): :data:`MODULE_TAPS` plus the
opt-in per-layer / large taps :data:`EXTRA_TAPS`. Values are float32 numpy (``head.topk_indices`` and ``label_pred``
int64). torch is imported lazily; the module runs where ttnn is absent.
"""
from __future__ import annotations

from typing import Any, Dict, Optional

import numpy as np

from ..ttaw.golden import NULL_TAPS, TapRegistry
from .config import BEV_SHAPE, SPARSE_LEVELS
from .weights import ENCODER_PLAN, BEVFusionWeights

__all__ = ["BEVFusionTorch", "MODULE_TAPS", "EXTRA_TAPS", "SELECTION_TAPS", "SPARSE_TAPS", "DENSE_TAPS",
           "OUTPUT_NAMES", "stable_topk"]

OUTPUT_NAMES = ("bbox_pred", "score", "label_pred")
SPARSE_TAPS = ("vfe",) + tuple(name for name, *_ in ENCODER_PLAN if not name.endswith(".conv1")) + ("sparse.bev",)
DENSE_TAPS = ("backbone.block0", "backbone.block1", "neck.deblock0", "neck.deblock1", "head.shared", "head.heatmap",
              "head.heat", "head.heat_nms")
# taps downstream of the data-dependent top-k (compared set-based or teacher-forced, PLAN.md 0.3 item 1)
SELECTION_TAPS = ("head.topk_values", "head.topk_indices", "query.feat", "query.pos", "query.heat",
                  "decoder.query_pos_embed", "decoder.self_attn", "decoder.cross_attn", "decoder.out",
                  "heads.center", "heads.height", "heads.dim", "heads.rot", "heads.vel", "heads.heatmap") + OUTPUT_NAMES
MODULE_TAPS = SPARSE_TAPS + DENSE_TAPS + SELECTION_TAPS
# opt-in taps: per layer, the 66 MB neck concat and the input-independent key position embedding
EXTRA_TAPS = (("neck", "head.heatmap.hidden", "decoder.key_pos_embed")
              + tuple(name for name, *_ in ENCODER_PLAN if name.endswith(".conv1"))
              + tuple(f"backbone.block{b}.conv{j}" for b in range(2) for j in range(6)))


def _torch():
    import torch

    return torch


def _tap(taps: TapRegistry, name: str, value):
    """Record a torch tensor as a numpy copy (float32, or int64 for integer tensors) and return it unchanged; never
    through ``ttaw.tensors.to_numpy`` (which would import ttnn)."""
    if taps.enabled and taps.wants(name):
        a = value.detach().cpu().numpy()
        taps.tap(name, a.astype(np.int64 if a.dtype.kind in "iu" else np.float32, copy=True))
    return value


def stable_topk(scores, k: int):
    """ONNX ``TopK(largest=1, sorted=1)`` on the last axis: values descending, equal values in ascending index order
    (a stable descending sort). Returns ``(values, indices)``."""
    torch = _torch()
    order = torch.sort(scores, dim=-1, descending=True, stable=True)
    return order.values[..., :k], order.indices[..., :k]


class BEVFusionTorch:
    """The fp32 reference network. ``threads`` caps torch's intra-op threads (shared host: <= 4)."""

    def __init__(self, weights: BEVFusionWeights, *, threads: Optional[int] = None, device: str = "cpu"):
        torch = _torch()
        if threads is not None:
            torch.set_num_threads(int(threads))
        self.weights = weights
        self.device = device
        t = self._t
        self._sparse = [(layer, t(layer.weight.reshape(layer.out_channels, -1, layer.in_channels)), t(layer.bias))
                        for layer in weights.encoder]
        self._layers: Dict[str, tuple] = {}
        for layer in weights.dense_layers():
            self._layers[layer.name] = (layer, t(layer.weight), t(layer.bias))
        h = weights.head
        d = h.decoder
        for layer in (*d.self_posembed, *d.cross_posembed) + tuple(x for _, a, b in h.prediction_heads for x in (a, b)):
            self._layers[layer.name] = (layer, t(layer.weight), t(layer.bias))
        self.class_table = t(h.class_table)                     # (5, 128)
        self.bev_pos = t(h.bev_pos)                             # (32400, 2)
        self.K = int(h.num_proposals)
        self.bev_h, self.bev_w = BEV_SHAPE
        self.pooled = tuple(h.pooled_classes)
        self._attn = {"self_attn": self._attn_tensors(d.self_attn), "cross_attn": self._attn_tensors(d.cross_attn)}
        self._norms = [(t(g), t(b), float(eps)) for g, b, eps in d.norms]
        self._ffn = (t(d.ffn_w1), t(d.ffn_b1), t(d.ffn_w2), t(d.ffn_b2))
        self._kpe = None                                        # cross_posembed(bev_pos): input-independent, cached

    def _t(self, a: Any):
        torch = _torch()
        return torch.from_numpy(np.array(a, dtype=np.float32, copy=True)).to(self.device)

    def _attn_tensors(self, a) -> dict:
        t = self._t
        return {"wq": t(a.wq), "bq": t(a.bq), "wk": t(a.wk), "bk": t(a.bk), "wv": t(a.wv), "bv": t(a.bv),
                "wo": t(a.wo), "bo": t(a.bo), "heads": a.num_heads, "hd": a.head_dim, "div": float(a.q_divisor)}

    # ---------------------------------------------------------------------------------------- sparse encoder
    def sparse_conv(self, x, nmap: np.ndarray, w, b):
        """Gather-form sparse conv: ``x`` (N_in, Cin) torch, ``nmap`` (N_out, K) int (-1 = none), ``w`` (Cout, K,
        Cin) -> (N_out, Cout) (bias added, no activation). Taps accumulate in tap order, fp32."""
        torch = _torch()
        n_in, cin = x.shape
        xp = torch.cat([x, torch.zeros((1, cin), dtype=x.dtype, device=x.device)], 0)    # explicit zero row
        idx = torch.from_numpy(np.where(nmap >= 0, nmap, n_in).astype(np.int64)).to(x.device)
        out = torch.zeros((idx.shape[0], w.shape[0]), dtype=x.dtype, device=x.device)
        for k in range(idx.shape[1]):
            col = idx[:, k]
            if bool((col < n_in).any()):
                out += xp[col] @ w[:, k, :].T
        return out + b

    def sparse_encoder(self, vfe: np.ndarray, plan: Any, taps: TapRegistry = NULL_TAPS):
        """(V, 4) voxel features + ``host.preprocess.SparsePlan`` -> the BEV (1, 256, 180, 180)."""
        torch = _torch()
        x = _tap(taps, "vfe", torch.as_tensor(np.asarray(vfe, dtype=np.float32), device=self.device))
        if x.shape[0] != plan.levels["l1"].num_active:
            raise ValueError(f"{x.shape[0]} voxel features for {plan.levels['l1'].num_active} active voxels")
        outs: Dict[str, Any] = {}
        for layer, w, b in self._sparse:
            y = self.sparse_conv(x, plan.maps[layer.map], w, b)
            if layer.residual is not None:
                y = y + outs[layer.residual]
            x = _tap(taps, layer.name, torch.relu(y))
            outs[layer.name] = x
        lvl = plan.levels["out"]
        X, Y, Z = lvl.spatial_shape
        if (X, Y) != BEV_SHAPE or lvl.spatial_shape != SPARSE_LEVELS["out"]:
            raise ValueError(f"conv_out grid {lvl.spatial_shape}")
        dense = torch.zeros((1, X, Y, Z, x.shape[1]), dtype=x.dtype, device=x.device)
        oc = torch.from_numpy(lvl.coords).to(x.device)
        dense[0, oc[:, 0], oc[:, 1], oc[:, 2]] = x
        bev = dense.permute(0, 4, 3, 1, 2).reshape(1, x.shape[1] * Z, X, Y)
        return _tap(taps, "sparse.bev", bev)

    # ----------------------------------------------------------------------------------------- dense layers
    def layer(self, name: str, x):
        torch = _torch()
        F = torch.nn.functional
        layer, w, b = self._layers[name]
        if layer.kind == "conv":
            y = F.conv2d(x, w, b, stride=layer.stride, padding=layer.padding)
        elif layer.kind == "conv_transpose":
            y = F.conv_transpose2d(x, w, b, stride=layer.stride, padding=layer.padding)
        else:
            y = F.conv1d(x, w, b, stride=layer.stride, padding=layer.padding)
        return torch.relu(y) if layer.relu else y

    def backbone(self, bev, taps: TapRegistry = NULL_TAPS):
        outs, x = [], bev
        for i, block in enumerate(self.weights.blocks):
            for layer in block:
                x = _tap(taps, layer.name, self.layer(layer.name, x))
            outs.append(_tap(taps, f"backbone.block{i}", x))
        return outs

    def neck(self, blocks, taps: TapRegistry = NULL_TAPS):
        torch = _torch()
        ups = [_tap(taps, layer.name, self.layer(layer.name, x)) for layer, x in zip(self.weights.deblocks, blocks)]
        return _tap(taps, "neck", torch.cat(ups, dim=1))

    def dense_head(self, neck, taps: TapRegistry = NULL_TAPS):
        """-> (lidar_feat (1, 128, 180, 180), heat_nms (1, 5, 180, 180))."""
        torch = _torch()
        F = torch.nn.functional
        lidar = _tap(taps, "head.shared", self.layer("head.shared", neck))
        hid = _tap(taps, "head.heatmap.hidden", self.layer("head.heatmap.hidden", lidar))
        logits = _tap(taps, "head.heatmap", self.layer("head.heatmap", hid))
        heat = _tap(taps, "head.heat", torch.sigmoid(logits))
        k = int(self.weights.head.nms_kernel)
        pad = k // 2
        pooled = list(self.pooled)
        local_max = torch.zeros_like(heat)
        local_max[:, pooled, pad:-pad, pad:-pad] = F.max_pool2d(heat[:, pooled], k, 1, 0)
        rest = [c for c in range(heat.shape[1]) if c not in self.pooled]
        local_max[:, rest] = heat[:, rest]
        heat_nms = _tap(taps, "head.heat_nms", heat * (heat == local_max).to(heat.dtype))
        return lidar, heat_nms

    def select(self, heat_nms, topk_indices: Optional[Any] = None, taps: TapRegistry = NULL_TAPS):
        """Top-K of the flattened local-max heatmap -> (values (1, K), indices (1, K) int64); given ``topk_indices``
        (teacher forcing) are used as they are and their heat values gathered."""
        torch = _torch()
        flat = heat_nms.reshape(1, -1)
        if topk_indices is None:
            values, idx = stable_topk(flat, self.K)
        else:
            idx = torch.as_tensor(np.asarray(topk_indices, dtype=np.int64).reshape(1, -1), device=self.device)
            if idx.shape[1] != self.K or int(idx.min()) < 0 or int(idx.max()) >= flat.shape[1]:
                raise ValueError(f"topk_indices must be {self.K} indices in [0, {flat.shape[1]})")
            values = flat.gather(1, idx)
        return _tap(taps, "head.topk_values", values), _tap(taps, "head.topk_indices", idx)

    def query_init(self, lidar, heat_nms, idx, taps: TapRegistry = NULL_TAPS):
        """-> (query (1, 128, K), query_pos (1, K, 2), query_heat (1, 5, K), class (K,))."""
        torch = _torch()
        n_cells = lidar.shape[2] * lidar.shape[3]
        cls_id = torch.div(idx, n_cells, rounding_mode="floor")
        pos = idx - cls_id * n_cells
        flat = lidar.reshape(1, lidar.shape[1], -1)
        q = flat.gather(2, pos[:, None, :].expand(-1, flat.shape[1], -1)) + self.class_table[cls_id[0]].T[None]
        w = self.bev_w
        qpos = torch.stack([torch.div(pos, w, rounding_mode="floor").to(torch.float32) + 0.5,
                            (pos % w).to(torch.float32) + 0.5], dim=-1)
        heat_flat = heat_nms.reshape(1, heat_nms.shape[1], -1)
        qheat = heat_flat.gather(2, pos[:, None, :].expand(-1, heat_flat.shape[1], -1))
        return (_tap(taps, "query.feat", q), _tap(taps, "query.pos", qpos), _tap(taps, "query.heat", qheat),
                cls_id[0])

    def _posembed(self, which: str, pos):
        """(P, 2) -> (P, 128): Conv1d 2 -> 128 + ReLU (BN folded), Conv1d 128 -> 128 (1x1 convs over P positions)."""
        x = pos.T[None]                                                        # (1, 2, P)
        x = self.layer(f"decoder.{which}.0", x)
        return self.layer(f"decoder.{which}.1", x)[0].T                         # (P, 128)

    def key_pos_embed(self, taps: TapRegistry = NULL_TAPS):
        """``cross_posembed(bev_pos)`` (32400, 128): input-independent, computed once."""
        torch = _torch()
        if self._kpe is None:
            with torch.no_grad():
                self._kpe = self._posembed("cross_posembed", self.bev_pos)
        return _tap(taps, "decoder.key_pos_embed", self._kpe[None])[0]

    def _mha(self, which: str, q_in, kv_in, v_in=None):
        """Multi-head attention as exported: projections, per-head ``softmax((q / div) @ k^T) @ v``, out proj."""
        torch = _torch()
        a = self._attn[which]
        v_in = kv_in if v_in is None else v_in
        q = q_in @ a["wq"] + a["bq"]
        k = kv_in @ a["wk"] + a["bk"]
        v = v_in @ a["wv"] + a["bv"]
        H, D = a["heads"], a["hd"]
        lq, lk = q.shape[0], k.shape[0]
        q = q.reshape(lq, H, D).transpose(0, 1) / a["div"]
        k = k.reshape(lk, H, D).permute(1, 2, 0)
        v = v.reshape(lk, H, D).transpose(0, 1)
        o = torch.softmax(q @ k, dim=-1) @ v                                   # (H, Lq, D)
        o = o.transpose(0, 1).reshape(lq, H * D)
        return o @ a["wo"] + a["bo"]

    def _ln(self, i: int, x):
        torch = _torch()
        g, b, eps = self._norms[i]
        return torch.nn.functional.layer_norm(x, (x.shape[-1],), g, b, eps)

    def decoder(self, query, qpos, lidar, taps: TapRegistry = NULL_TAPS):
        """One post-norm decoder layer -> (K, 128)."""
        torch = _torch()
        qpe = _tap(taps, "decoder.query_pos_embed", self._posembed("self_posembed", qpos[0])[None])[0]
        kpe = self.key_pos_embed(taps)
        x = query[0].T                                                         # (K, 128)
        key = lidar.reshape(lidar.shape[1], -1).T + kpe                        # (32400, 128)
        s = x + qpe
        x = _tap(taps, "decoder.self_attn", self._ln(0, x + self._mha("self_attn", s, s))[None])[0]
        x = _tap(taps, "decoder.cross_attn", self._ln(1, x + self._mha("cross_attn", x + qpe, key))[None])[0]
        w1, b1, w2, b2 = self._ffn
        ff = torch.relu(x @ w1 + b1) @ w2 + b2
        return _tap(taps, "decoder.out", self._ln(2, x + ff)[None])[0]

    def prediction_heads(self, x, qpos, qheat, cls_id, taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]:
        torch = _torch()
        y = x.T[None]                                                          # (1, 128, K)
        out = {}
        for name, hid, fin in self.weights.head.prediction_heads:
            out[name] = _tap(taps, fin.name, self.layer(fin.name, self.layer(hid.name, y)))
        center = out["center"] + qpos.transpose(1, 2)                          # float32
        box = torch.cat([center, out["height"], out["dim"], out["rot"], out["vel"]], dim=1)[0]
        cls = (torch.sigmoid(out["heatmap"]) * qheat)[0]                       # (5, K)
        score = cls.gather(0, cls_id[None])[0]
        return {"bbox_pred": _tap(taps, "bbox_pred", box), "score": _tap(taps, "score", score),
                "label_pred": _tap(taps, "label_pred", cls_id.to(torch.int64))}

    # --------------------------------------------------------------------------------------------- whole
    def dense_forward(self, bev: Any, *, topk_indices: Optional[Any] = None,
                      taps: TapRegistry = NULL_TAPS) -> Dict[str, np.ndarray]:
        """The plugin-free part (the dense sub-graph of the ONNX): BEV (1, 256, 180, 180) -> outputs."""
        torch = _torch()
        with torch.no_grad():
            bev = torch.as_tensor(np.asarray(bev, dtype=np.float32), device=self.device) \
                if not isinstance(bev, torch.Tensor) else bev
            neck = self.neck(self.backbone(bev, taps), taps)
            lidar, heat_nms = self.dense_head(neck, taps)
            _, idx = self.select(heat_nms, topk_indices, taps)
            query, qpos, qheat, cls_id = self.query_init(lidar, heat_nms, idx, taps)
            x = self.decoder(query, qpos, lidar, taps)
            out = self.prediction_heads(x, qpos, qheat, cls_id, taps)
        return {"bbox_pred": out["bbox_pred"].cpu().numpy().astype(np.float32),
                "score": out["score"].cpu().numpy().astype(np.float32),
                "label_pred": out["label_pred"].cpu().numpy().astype(np.int64)}

    def forward(self, vfe: np.ndarray, plan: Any, *, topk_indices: Optional[Any] = None,
                taps: TapRegistry = NULL_TAPS) -> Dict[str, np.ndarray]:
        """Voxel features + sparse plan -> the three graph outputs: ``bbox_pred`` (10, K) float32, ``score`` (K,)
        float32, ``label_pred`` (K,) int64. ``topk_indices`` (K,) or (1, K): teacher-forced query selection."""
        torch = _torch()
        with torch.no_grad():
            bev = self.sparse_encoder(vfe, plan, taps)
        return self.dense_forward(bev, topk_indices=topk_indices, taps=taps)

    def describe(self) -> Dict[str, Any]:
        w = self.weights
        d = w.head.decoder
        return {"onnx_sha256": w.sha256, "num_proposals": self.K, "bev": list(BEV_SHAPE),
                "encoder": [(lay.name, lay.map, lay.in_channels, lay.out_channels, lay.kernel_volume)
                            for lay in w.encoder],
                "blocks": [[(lay.in_channels, lay.out_channels, lay.stride) for lay in b] for b in w.blocks],
                "deblocks": [(lay.kind, lay.in_channels, lay.out_channels, lay.kernel, lay.stride) for lay in w.deblocks],
                "decoder": {"heads": d.self_attn.num_heads, "head_dim": d.self_attn.head_dim,
                            "ffn": int(d.ffn_w1.shape[1])},
                "prediction_heads": {n: fin.out_channels for n, _, fin in w.head.prediction_heads},
                "params_after_folds": w.num_params()}