# SPDX-License-Identifier: Apache-2.0
"""The three TensorRT plugins of the deployed graph as plain torch fp32 (no mmcv, no CUDA), plus literal float64 ports
of the CUDA loops that unit-test them. Ported from the verified research reference
(``research/bevformer/scripts/bevformer_ref.py``; checks in ``research/bevformer/test_ref_ops.log``).

``AW/`` = autoware_universe ``perception/autoware_tensorrt_bevformer`` @ 9ceaccf:

- :func:`msda` == ``MultiScaleDeformableAttnTRT`` (``AW/TensorRT/plugin/multi_scale_deformable_attn/
  multiScaleDeformableAttnKernel.cu:637-707``, ``...Plugin.cpp:95-105,208-290``): softmax over levels x points INSIDE
  the op; point ``p`` uses reference ``p % Pg``; ``loc_px = ref * (W, H) + offset`` (offsets in pixels of the level);
  bilinear at ``loc_px - 0.5`` with zeros outside (== ``grid_sample(align_corners=False, padding_mode="zeros")``).
- :func:`mdcn` == ``ModulatedDeformableConv2dTRT`` (``AW/TensorRT/plugin/modulated_deformable_conv2d/
  modulatedDeformableConv2dKernel.cu:291-349``): mmcv DCNv2, offsets interleaved ``(dy_k, dx_k)`` per tap
  ``k = i * 3 + j``, sample at ``(h - 1 + i + dy, w - 1 + j + dx)``, bilinear, zeros outside, times the sigmoid mask,
  then the GEMM (+ the bias the port gets from the folded BatchNorm).
- :func:`rotate` == ``RotateTRT`` nearest (``AW/TensorRT/plugin/rotate/rotateKernel.cu:166-243``): the inverse
  affine grid of torchvision ``rotate`` about ``center`` with ``::round`` (half away from zero) and zero fill.

Importing this module imports torch (the reference runs where torch is: tt-env and the research venv); never ttnn.
"""
from __future__ import annotations

import math
from typing import Sequence, Tuple

import torch
import torch.nn.functional as F

__all__ = ["msda", "msda_naive", "mdcn", "mdcn_naive", "rotate", "round_half_away"]


def msda(value: torch.Tensor, spatial_shapes: Sequence[Tuple[int, int]], reference_points: torch.Tensor,
         sampling_offsets: torch.Tensor, attention_weights: torch.Tensor, q_chunk: int = 4096) -> torch.Tensor:
    """``MultiScaleDeformableAttnTRT`` semantics.

    value             [N, S, H, C]    (S = sum h_l * w_l)
    spatial_shapes    [(h, w), ...]   python ints (ONNX: an INT64 constant [L, 2])
    reference_points  [N, Q, 1, Pg*2] normalised (x, y); sampling point p uses reference p % Pg
    sampling_offsets  [N, Q, H, L*P*2] in PIXELS of the level, (x, y)
    attention_weights [N, Q, H, L*P]  logits: softmax over L*P inside the op
    returns           [N, Q, H, C]
    """
    N, S, H, C = value.shape
    L = len(spatial_shapes)
    Q = reference_points.shape[1]
    Pg = reference_points.shape[-1] // 2
    P = attention_weights.shape[-1] // L
    w = torch.softmax(attention_weights, dim=-1).view(N, Q, H, L, P)
    off = sampling_offsets.view(N, Q, H, L, P, 2)
    ref = reference_points.view(N, Q, 1, Pg, 2)
    ref_p = ref[:, :, :, torch.arange(P) % Pg, :]                         # [N, Q, 1, P, 2]
    out = None
    start = 0
    for lvl, (h, wd) in enumerate(spatial_shapes):
        v = value[:, start:start + h * wd].permute(0, 2, 3, 1).reshape(N * H, C, h, wd)
        start += h * wd
        norm = torch.tensor([wd, h], dtype=value.dtype)
        parts = []
        for q0 in range(0, Q, q_chunk):
            q1 = min(Q, q0 + q_chunk)
            loc = ref_p[:, q0:q1] + off[:, q0:q1, :, lvl] / norm          # [N, q, H, P, 2]
            grid = (2.0 * loc - 1.0).permute(0, 2, 1, 3, 4).reshape(N * H, q1 - q0, P, 2)
            s = F.grid_sample(v, grid, mode="bilinear", padding_mode="zeros", align_corners=False)
            wl = w[:, q0:q1, :, lvl].permute(0, 2, 1, 3).reshape(N * H, 1, q1 - q0, P)
            parts.append((s * wl).sum(-1))                                # [N*H, C, q]
        lvl_out = parts[0] if len(parts) == 1 else torch.cat(parts, dim=2)
        out = lvl_out if out is None else out + lvl_out
    return out.view(N, H, C, Q).permute(0, 3, 1, 2).contiguous()


def msda_naive(value, spatial_shapes, reference_points, sampling_offsets, attention_weights) -> torch.Tensor:
    """Literal float64 port of the CUDA kernel loop (unit tests on small sizes)."""
    value = value.double()
    reference_points = reference_points.double()
    sampling_offsets = sampling_offsets.double()
    attention_weights = attention_weights.double()
    N, S, H, C = value.shape
    L = len(spatial_shapes)
    Q = reference_points.shape[1]
    Pg = reference_points.shape[-1] // 2
    P = attention_weights.shape[-1] // L
    out = torch.zeros(N, Q, H, C, dtype=torch.float64)
    for b in range(N):
        for q in range(Q):
            for hd in range(H):
                aw = attention_weights[b, q, hd]
                mx = aw.max()
                ssum = 0.0
                acc = torch.zeros(C, dtype=torch.float64)
                widx = 0
                start = 0
                for (h, wd) in spatial_shapes:
                    vl = value[b, start:start + h * wd, hd].view(h, wd, C)
                    for p in range(P):
                        pg = p % Pg
                        lw = reference_points[b, q, 0, 2 * pg] * wd + sampling_offsets[b, q, hd, 2 * widx]
                        lh = reference_points[b, q, 0, 2 * pg + 1] * h + sampling_offsets[b, q, hd, 2 * widx + 1]
                        wt = torch.exp(aw[widx] - mx)
                        ssum = ssum + wt
                        him, wim = lh - 0.5, lw - 0.5
                        if -1 < him < h and -1 < wim < wd:
                            hl, wl = int(math.floor(him)), int(math.floor(wim))
                            hh, wh = hl + 1, wl + 1
                            dh, dw = him - hl, wim - wl
                            val = torch.zeros(C, dtype=torch.float64)
                            if hl >= 0 and wl >= 0:
                                val += (1 - dh) * (1 - dw) * vl[hl, wl]
                            if hl >= 0 and wh <= wd - 1:
                                val += (1 - dh) * dw * vl[hl, wh]
                            if hh <= h - 1 and wl >= 0:
                                val += dh * (1 - dw) * vl[hh, wl]
                            if hh <= h - 1 and wh <= wd - 1:
                                val += dh * dw * vl[hh, wh]
                            acc += val * wt
                        widx += 1
                    start += h * wd
                out[b, q, hd] = acc / ssum
    return out


def mdcn(x: torch.Tensor, offset: torch.Tensor, mask: torch.Tensor, weight: torch.Tensor,
         bias: torch.Tensor = None, stride: int = 1, padding: int = 1, dilation: int = 1) -> torch.Tensor:
    """``ModulatedDeformableConv2dTRT`` (groups = deform_groups = 1). ``offset`` [N, 2K, Ho, Wo] interleaved
    (dy_k, dx_k), ``mask`` [N, K, Ho, Wo] (post-sigmoid); nine ``grid_sample`` calls + a GEMM per tap."""
    N, C, H, W = x.shape
    Cout, Cin, kh, kw = weight.shape
    Ho = (H + 2 * padding - dilation * (kh - 1) - 1) // stride + 1
    Wo = (W + 2 * padding - dilation * (kw - 1) - 1) // stride + 1
    by = (torch.arange(Ho, dtype=x.dtype) * stride - padding).view(1, Ho, 1)
    bx = (torch.arange(Wo, dtype=x.dtype) * stride - padding).view(1, 1, Wo)
    out = None
    for i in range(kh):
        for j in range(kw):
            k = i * kw + j
            py = by + i * dilation + offset[:, 2 * k]
            px = bx + j * dilation + offset[:, 2 * k + 1]
            grid = torch.stack([(2.0 * px + 1.0) / W - 1.0, (2.0 * py + 1.0) / H - 1.0], dim=-1)
            s = F.grid_sample(x, grid, mode="bilinear", padding_mode="zeros", align_corners=False)
            s = s * mask[:, k:k + 1]
            contrib = torch.einsum("nchw,oc->nohw", s, weight[:, :, i, j])
            out = contrib if out is None else out + contrib
    if bias is not None:
        out = out + bias.view(1, -1, 1, 1)
    return out


def mdcn_naive(x, offset, mask, weight, stride: int = 1, padding: int = 1, dilation: int = 1) -> torch.Tensor:
    """Literal float64 port of ``modulated_deformable_im2col`` + GEMM (unit tests)."""
    x, offset, mask, weight = x.double(), offset.double(), mask.double(), weight.double()
    N, C, H, W = x.shape
    Cout, Cin, kh, kw = weight.shape
    Ho = (H + 2 * padding - dilation * (kh - 1) - 1) // stride + 1
    Wo = (W + 2 * padding - dilation * (kw - 1) - 1) // stride + 1
    cols = torch.zeros(N, C, kh * kw, Ho, Wo, dtype=torch.float64)
    for b in range(N):
        for ho in range(Ho):
            for wo in range(Wo):
                for i in range(kh):
                    for j in range(kw):
                        k = i * kw + j
                        him = ho * stride - padding + i * dilation + offset[b, 2 * k, ho, wo].item()
                        wim = wo * stride - padding + j * dilation + offset[b, 2 * k + 1, ho, wo].item()
                        val = torch.zeros(C, dtype=torch.float64)
                        if -1 < him < H and -1 < wim < W:
                            hl, wl = int(math.floor(him)), int(math.floor(wim))
                            hh, wh = hl + 1, wl + 1
                            dh, dw = him - hl, wim - wl
                            if hl >= 0 and wl >= 0:
                                val += (1 - dh) * (1 - dw) * x[b, :, hl, wl]
                            if hl >= 0 and wh <= W - 1:
                                val += (1 - dh) * dw * x[b, :, hl, wh]
                            if hh <= H - 1 and wl >= 0:
                                val += dh * (1 - dw) * x[b, :, hh, wl]
                            if hh <= H - 1 and wh <= W - 1:
                                val += dh * dw * x[b, :, hh, wh]
                        cols[b, :, k, ho, wo] = val * mask[b, k, ho, wo]
    return torch.einsum("nckhw,ock->nohw", cols, weight.view(Cout, Cin, kh * kw))


def round_half_away(x: torch.Tensor) -> torch.Tensor:
    """CUDA ``::round``: half away from zero."""
    return torch.sign(x) * torch.floor(torch.abs(x) + 0.5)


def rotate(img: torch.Tensor, angle_deg: torch.Tensor, center: Sequence[float], *,
           plugin_rounding: bool = True) -> torch.Tensor:
    """``RotateTRT`` (nearest) on ``img`` [C, H, W]; ``angle_deg`` (CCW, degrees, the graph's ``can_bus[17]``);
    ``center`` (x, y) in pixels. ``plugin_rounding=False`` is torchvision's ``rotate`` (training: round half to even);
    they differ only on exact ties (2 of 22,500 pixels at 0.37 degrees, research/bevformer/test_ref_ops.log)."""
    C, H, W = img.shape
    ang = -angle_deg.reshape(()) * math.pi / 180.0
    cx = center[0] - 0.5 * W
    cy = center[1] - 0.5 * H
    ca, sa = torch.cos(ang), torch.sin(ang)
    m = [ca, sa, -cx * ca - cy * sa + cx, -sa, ca, cx * sa - cy * ca + cy]
    xs = (torch.arange(W, dtype=img.dtype) - W * 0.5 + 0.5).view(1, W).expand(H, W)
    ys = (torch.arange(H, dtype=img.dtype) - H * 0.5 + 0.5).view(H, 1).expand(H, W)
    gx = (m[0] * xs + m[1] * ys + m[2]) / (0.5 * W)
    gy = (m[3] * xs + m[4] * ys + m[5]) / (0.5 * H)
    if not plugin_rounding:
        grid = torch.stack([gx, gy], -1).unsqueeze(0)
        return F.grid_sample(img.unsqueeze(0), grid, mode="nearest", padding_mode="zeros",
                             align_corners=False).squeeze(0)
    ix = ((gx + 1.0) * W - 1.0) / 2.0
    iy = ((gy + 1.0) * H - 1.0) / 2.0
    ixn = round_half_away(ix)
    iyn = round_half_away(iy)
    valid = (ixn >= 0) & (ixn <= W - 1) & (iyn >= 0) & (iyn <= H - 1)
    flat = (iyn.clamp(0, H - 1) * W + ixn.clamp(0, W - 1)).long().view(-1)
    out = img.reshape(C, H * W)[:, flat].view(C, H, W)
    return out * valid.to(img.dtype).unsqueeze(0)