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

Built from :mod:`.weights` (the ONNX initializers addressed by node; BN folded where the TT port folds it: the two
PFN BatchNorm1d and the deblock2 ConvTranspose BatchNorm), it reproduces ``transfusion.onnx`` (S:transfusion:155-208)
from the decorated pillar features on (the in-graph decoration is host pre-processing in this port, computed in
float32 exactly as the graph does by ``host.preprocess.decorate_pillars``):

- ``pfn``: ``h = relu(x @ W0' + b0')`` (N, 20, 32); ``m = max over ALL 20 slots`` (padded slots included: their rows
  are ``relu(b0')``, S:transfusion:157); ``y = max over the slots of relu(concat[h, m] @ W1' + b1')`` -> (N, 64);
- ``scatter``: ``canvas[0, c, iy, ix] = pfn[p, c]`` on a zero (1, 64, 768, 768) canvas (``ScatterElements``);
- ``backbone``: SECOND, 3 blocks of 3x3 convs + ReLU (4 / 6 / 6 layers, first stride 2);
- ``neck``: SECONDFPN: conv 2x2 s2, conv 1x1, ConvTranspose 2x2 s2 (BN folded), each + ReLU, concat (1, 384, 192, 192);
- ``head`` dense part: shared conv 3x3 384 -> 128 (``lidar_feat``, no activation), heatmap head (3x3 + ReLU, 3x3 -> 5
  logits), ``heat = sigmoid``, local max ``heat * (heat == maxpool3x3_s1_p0(heat) with a zero border)`` (border cells
  always suppressed, every class);
- selection: the top ``K = 500`` of the flattened (5 x 36864) local-max heatmap, as ONNX ``TopK`` defines it (largest
  first, **ties broken by the lower index**: a stable descending sort). ``topk_indices=`` replaces the selection by
  given indices (teacher forcing, PLAN.md 2.5: the device top-k breaks ties in no defined order, probe P8);
- query init: ``class = idx // 36864``, ``pos = idx % 36864``; ``query = lidar_feat[:, pos] + class_table[class]``
  (the Conv1d of a one-hot is the table row ``W[:, c] + b``); ``query_pos = (pos % 192 + 0.5, pos // 192 + 0.5)``,
  computed arithmetically in float32 (bit-equal to the graph's ``bev_pos`` gather, checked at load);
- 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)`` (36,864 keys), FFN 128 -> 256 -> 128, LayerNorm eps
  1e-5 after each; ``Q`` is divided by ``sqrt(16) = 4`` before ``Q K^T`` (the graph's ``Div``);
- prediction heads (Conv1d 128 -> 64 + ReLU, Conv1d 64 -> C for center 2, height 1, dim 3, rot 2, vel 2, heatmap 5);
- outputs: ``bbox_pred0 = [center + query_pos (float32), height, dim, vel]`` (1, 8, K), ``dir_cls_pred0 = rot``
  (1, 2, K), ``cls_score0 = query_heat * sigmoid(heads.heatmap)`` (1, 5, K) with ``query_heat`` the local-max
  heatmap of all 5 classes at the query's cell (no one-hot: mmdeploy.patch).

The TT-friendly equivalent forms the port uses (PFN split, K-split shared conv, constant cross-attention K / V tables,
merged prediction heads, the masked local-max chain, 32-slot pillar tiles, head-dim padding) are in :mod:`.rewrites`
and proven equal to these modules by ``tests/test_reference_rewrites_host.py``.

Taps (``ttaw.golden.TapRegistry``; stable names, the keys of the device gates; batch axis kept as the ONNX tensors
have it): :data:`MODULE_TAPS` plus the opt-in per-layer / large taps :data:`EXTRA_TAPS`. Values are float32 numpy
(``head.topk_indices`` 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 .weights import TransFusionWeights

__all__ = ["TransFusionTorch", "MODULE_TAPS", "EXTRA_TAPS", "SELECTION_TAPS", "OUTPUT_NAMES", "stable_topk"]

OUTPUT_NAMES = ("cls_score0", "bbox_pred0", "dir_cls_pred0")
# module-output taps (golden files and device gates)
MODULE_TAPS = ("pfn.input", "pfn", "canvas", "backbone.block0", "backbone.block1", "backbone.block2",
               "neck.deblock0", "neck.deblock1", "neck.deblock2", "head.shared", "head.heatmap", "head.heat_nms",
               "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
# opt-in taps: per layer, the 56 MB neck concat and the input-independent key position embedding
EXTRA_TAPS = ("pfn.layer0", "neck", "head.heatmap.hidden", "decoder.key_pos_embed") + tuple(
    f"backbone.block{b}.conv{j}" for b, n in ((0, 4), (1, 6), (2, 6)) for j in range(n))
# 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


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 TransFusionTorch:
    """The fp32 reference network. ``threads`` caps torch's intra-op threads (shared host: <= 4)."""

    def __init__(self, weights: TransFusionWeights, *, 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
        p = weights.pfn
        self.w0, self.b0, self.w1, self.b1 = t(p.w0), t(p.b0), t(p.w1), t(p.b1)
        self._layers: Dict[str, tuple] = {}
        for layer in weights.conv_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)                             # (36864, 2)
        self.K = int(h.num_proposals)
        self.grid_w, self.grid_h = weights.grid
        self.feat_w = int(round(np.sqrt(h.bev_pos.shape[0])))
        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)}

    # ------------------------------------------------------------------------------------------- modules
    def pfn(self, features: Any, taps: TapRegistry = NULL_TAPS):
        """(N, 20, 11) decorated pillar features (padded slots zero) -> (N, 64) pillar features."""
        torch = _torch()
        x = _tap(taps, "pfn.input", torch.as_tensor(np.asarray(features, dtype=np.float32), device=self.device))
        h = _tap(taps, "pfn.layer0", torch.relu(torch.matmul(x, self.w0) + self.b0))       # (N, 20, 32)
        m = h.amax(dim=1, keepdim=True)                                                    # all 20 slots
        y = torch.relu(torch.matmul(torch.cat([h, m.expand_as(h)], dim=2), self.w1) + self.b1)
        return _tap(taps, "pfn", y.amax(dim=1))

    def scatter(self, pillar_features: Any, coords: np.ndarray, taps: TapRegistry = NULL_TAPS):
        """``ScatterElements`` into the zero canvas: (N, 64) -> (1, 64, 768, 768) with ``canvas[c, iy, ix]``.
        ``coords``: (N, >=2) with (.., iy, ix) in the last two columns (``(0, iy, ix)`` or ONNX ``coors``)."""
        torch = _torch()
        pf = torch.as_tensor(pillar_features, device=self.device)
        canvas = torch.zeros((1, pf.shape[1], self.grid_h, self.grid_w), dtype=torch.float32, device=self.device)
        c = torch.as_tensor(np.asarray(coords)[:, -2:].astype(np.int64), device=self.device)
        canvas[0, :, c[:, 0], c[:, 1]] = pf.T
        return _tap(taps, "canvas", canvas)

    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, canvas, taps: TapRegistry = NULL_TAPS):
        outs, x = [], canvas
        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, 192, 192), heat_nms (1, 5, 192, 192))."""
        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 = torch.sigmoid(logits)
        k = int(self.weights.head.nms_kernel)
        pad = k // 2
        local_max = torch.zeros_like(heat)
        local_max[:, :, pad:-pad, pad:-pad] = F.max_pool2d(heat, k, 1, 0)   # pad 0 + zero border (all classes)
        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))."""
        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.feat_w
        qpos = torch.stack([(pos % w).to(torch.float32) + 0.5,
                            torch.div(pos, w, rounding_mode="floor").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)

    def _posembed(self, which: str, pos):
        """(P, 2) -> (P, 128): Conv1d 2 -> 128 + ReLU, Conv1d 128 -> 128 (as 1x1 convs over the 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)`` (36864, 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                        # (36864, 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, 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 (S:transfusion:241)
        cls = qheat * torch.sigmoid(out["heatmap"])
        box = torch.cat([center, out["height"], out["dim"], out["vel"]], dim=1)
        return {"cls_score0": _tap(taps, "cls_score0", cls), "bbox_pred0": _tap(taps, "bbox_pred0", box),
                "dir_cls_pred0": _tap(taps, "dir_cls_pred0", out["rot"])}

    # --------------------------------------------------------------------------------------------- whole
    def forward(self, features: np.ndarray, coords: np.ndarray, num_pillars: int, *,
                topk_indices: Optional[Any] = None, taps: TapRegistry = NULL_TAPS) -> Dict[str, np.ndarray]:
        """Decorated pillars (the first ``num_pillars`` rows of ``features`` / ``coords`` are used) -> the three
        graph outputs as float32 numpy: ``cls_score0`` (1, 5, K), ``bbox_pred0`` (1, 8, K), ``dir_cls_pred0``
        (1, 2, K). ``topk_indices`` (K,) or (1, K): teacher-forced query selection."""
        torch = _torch()
        n = int(num_pillars)
        with torch.no_grad():
            pf = self.pfn(np.asarray(features)[:n], taps)
            canvas = self.scatter(pf, np.asarray(coords)[:n], taps)
            neck = self.neck(self.backbone(canvas, taps), taps)
            lidar, heat_nms = self.dense_head(neck, taps)
            _, idx = self.select(heat_nms, topk_indices, taps)
            query, qpos, qheat = self.query_init(lidar, heat_nms, idx, taps)
            x = self.decoder(query, qpos, lidar, taps)
            out = self.prediction_heads(x, qpos, qheat, taps)
        return {k: v.cpu().numpy().astype(np.float32) for k, v in out.items()}

    def describe(self) -> Dict[str, Any]:
        w = self.weights
        d = w.head.decoder
        return {"onnx_sha256": w.sha256, "grid": list(w.grid), "num_proposals": self.K,
                "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}}


def layer_names(weights: TransFusionWeights) -> list:
    """Every per-layer conv tap name (backbone convs, deblocks, head convs)."""
    return [layer.name for layer in weights.conv_layers()]