# SPDX-License-Identifier: Apache-2.0
"""The three in-network TensorRT plugins of ``bevdet_one_lt_d.onnx`` in torch fp32 (the vendor fp32 path).

Clean-room from research/bevdet/SPEC.md section 4.2 (no vendor code), ported from the verified research reference
(``research/bevdet/scripts/bevdet_common.py``):

- :func:`bev_pool`: ``bevdet::BEVPool``: ``out[c, ranks_bev[k]] += feat[ranks_feat[k] -> (n, c, hw)] *
  depth_flat[ranks_depth[k]]`` over every table entry, zero-initialised (the intended semantics; the TRT plugin
  leaves never-hit cells uninitialised, S:162). One ``index_add_`` in fp32.
- :func:`align_bev`: ``bevdet::AlignBEV``: for history slot n and BEV pixel (h, w): ``ix = a w + b h + c``,
  ``iy = d w + e h + f`` (pixel coordinates), bilinear over the 4 neighbours with zero outside
  (= ``grid_sample(align_corners=True, padding_mode="zeros")``, as in training). The 4 taps are evaluated literally
  in float32 like the CUDA kernel, so an identity transform is an exact copy.
- :func:`gather_bev`: ``bevdet::GatherBEV``: ``cat(curr, s_1 .. s_8)`` with ``s_j = aligned[j - 1]`` if ``flag``
  else ``curr``; channel order [t, t-1, ..., t-8].

Inputs and outputs are torch tensors or numpy arrays (outputs: torch). The Preprocess plugin is host work
(``host.preprocess``).
"""
from __future__ import annotations

from typing import Any

import numpy as np

from ..host.calib import BEVPoolTables
from .config import BEVDetConfig

__all__ = ["bev_pool", "align_bev", "gather_bev"]


def _t(x: Any):
    import torch

    return torch.as_tensor(np.asarray(x, np.float32)) if not isinstance(x, torch.Tensor) else x.to(torch.float32)


def bev_pool(depth: Any, feat: Any, tables: BEVPoolTables, cfg: BEVDetConfig):
    """depth [N, D, H, W] (softmax), feat [N, C, H, W] -> curr [1, C, ny, nx] (fp32, cells never hit are 0)."""
    import torch

    depth_t, feat_t = _t(depth), _t(feat)
    N, C, H, W = feat_t.shape
    ny, nx = cfg.bev_hw
    nz = cfg.bev_xyz[2]
    dflat = depth_t.reshape(-1)
    fflat = feat_t.reshape(N, C, H * W).permute(1, 0, 2).reshape(C, N * H * W)
    rd = torch.from_numpy(tables.ranks_depth.astype(np.int64))
    rf = torch.from_numpy(tables.ranks_feat.astype(np.int64))
    rb = torch.from_numpy(tables.ranks_bev.astype(np.int64))
    contrib = fflat[:, rf] * dflat[rd][None, :]
    out = torch.zeros(C, nx * ny * nz, dtype=torch.float32)
    out.index_add_(1, rb, contrib)
    return out.reshape(1, C, ny, nx)


def align_bev(adj_feats: Any, transforms: Any):
    """adj_feats [1, A, C, H, W], transforms [1, A, 6] -> aligned [1, A, C, H, W] (bilinear, zero outside)."""
    import torch

    adj = _t(adj_feats)
    trs = _t(transforms)
    B, A, C, H, W = adj.shape
    out = torch.zeros_like(adj)
    hh, ww = torch.meshgrid(torch.arange(H, dtype=torch.float32), torch.arange(W, dtype=torch.float32),
                            indexing="ij")
    identity = torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float32)
    for n in range(A):
        if torch.equal(trs[0, n], identity):     # one tap with weight 1, three with weight 0: a copy
            out[0, n] = adj[0, n]
            continue
        a, b, c0, d, e, f = (trs[0, n, k] for k in range(6))
        ix = a * ww + b * hh + c0                                  # float32, like the kernel
        iy = d * ww + e * hh + f
        x0 = torch.floor(ix).double()
        y0 = torch.floor(iy).double()
        x1, y1 = x0 + 1, y0 + 1
        ixd, iyd = ix.double(), iy.double()
        # weights: exact differences (float64), one rounding of the product to float32 (research reference)
        taps = ((y0, x0, (x1 - ixd) * (y1 - iyd)), (y0, x1, (ixd - x0) * (y1 - iyd)),
                (y1, x0, (x1 - ixd) * (iyd - y0)), (y1, x1, (ixd - x0) * (iyd - y0)))
        src = adj[0, n]
        acc = torch.zeros((C, H, W), dtype=torch.float32)
        for yy, xx, wt in taps:
            ok = (yy >= 0) & (yy < H) & (xx >= 0) & (xx < W)
            yi = yy.clamp(0, H - 1).long()
            xi = xx.clamp(0, W - 1).long()
            vals = src[:, yi, xi]
            acc += torch.where(ok[None], vals * wt.float()[None], torch.zeros((), dtype=torch.float32))
        out[0, n] = acc
    return out


def gather_bev(aligned: Any, curr: Any, flag: int):
    """aligned [1, A, C, H, W], curr [1, C, H, W] -> [1, (1 + A) C, H, W]."""
    import torch

    al, cu = _t(aligned), _t(curr)
    A = al.shape[1]
    parts = [cu[0]] + [cu[0] if int(flag) == 0 else al[0, n] for n in range(A)]
    return torch.cat(parts, 0)[None]