"""CPU build of vis4d_cuda_ops, standing in for the published package.

Map-Det3D imports ``iou_box3d`` and ``ms_deform_attn_forward`` from
``vis4d_cuda_ops``. Installing that package on a Space does not work: it has no
wheels, and pip compiles it inside an isolated build environment that resolves
its own, newest torch, while the Space caps the runtime torch several releases
lower. The two ABIs disagree and importing the result dies on
``undefined symbol: c10::NotImplementedError``.

Since ``/home/user/app`` comes first on ``sys.path``, this package answers those
imports instead. It compiles only the CPU translation units of the upstream
sources, against the torch the Space actually runs, so no ABI can drift:

* ``iou_box3d`` runs the upstream CPU kernel, on the handful of boxes the track
  graph compares per frame.
* ``ms_deform_attn_forward`` runs the grid-sample implementation that ships in
  ``mapdet3d.op.layer.ms_deform_attn`` and is used whenever the fused kernel is
  unavailable. It is the same computation and still runs on the GPU, only
  without the fused kernel.

The upstream sources live under ``src/`` unmodified; see ``LICENSE``.
"""

from __future__ import annotations

import os

from torch import Tensor
from torch.utils.cpp_extension import load

_SRC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "src")

_CPU_SOURCES = [
    os.path.join(_SRC, "vision.cpp"),
    os.path.join(_SRC, "iou_box3d", "iou_box3d_cpu.cpp"),
    os.path.join(_SRC, "box_iou_rotated", "box_iou_rotated_cpu.cpp"),
    os.path.join(_SRC, "nms_rotated", "nms_rotated_cpu.cpp"),
]

# Compiled on first import and cached in ~/.cache/torch_extensions afterwards.
_ops = load(
    name="vis4d_cpu_ops",
    sources=_CPU_SOURCES,
    extra_include_paths=[_SRC],
    extra_cflags=["-O2"],
    with_cuda=False,
    verbose=False,
)

box_iou_rotated = _ops.box_iou_rotated
nms_rotated = _ops.nms_rotated
deform_conv_forward = _ops.deform_conv_forward
deform_conv_backward_input = _ops.deform_conv_backward_input
deform_conv_backward_filter = _ops.deform_conv_backward_filter
modulated_deform_conv_forward = _ops.modulated_deform_conv_forward
modulated_deform_conv_backward = _ops.modulated_deform_conv_backward


def iou_box3d(
    boxes1: Tensor, boxes2: Tensor
) -> tuple[Tensor, Tensor]:  # pragma: no cover
    """Intersection volume and IoU of two sets of 3D box corners."""
    device = boxes1.device
    volume, iou = _ops.iou_box3d(boxes1.cpu(), boxes2.cpu())

    return volume.to(device), iou.to(device)


def ms_deform_attn_forward(
    value: Tensor,
    spatial_shapes: Tensor,
    level_start_index: Tensor,
    sampling_loc: Tensor,
    attn_weight: Tensor,
    im2col_step: int,
) -> Tensor:  # pragma: no cover
    """Multi-scale deformable attention, via grid sampling.

    ``level_start_index`` and ``im2col_step`` only exist to drive the fused
    kernel; the grid-sample path splits ``value`` by ``spatial_shapes`` itself.
    """
    # Imported here because mapdet3d imports this module while it is still
    # setting up its own.
    from mapdet3d.op.layer.ms_deform_attn import ms_deformable_attention_cpu

    return ms_deformable_attention_cpu(
        value, spatial_shapes, sampling_loc, attn_weight
    )


def ms_deform_attn_backward(*args: object, **kwargs: object) -> None:
    """Not built: the demo only ever runs under ``torch.no_grad()``."""
    raise NotImplementedError(
        "vis4d_cuda_ops was built without the deformable attention backward "
        "kernel. Install the upstream package to train."
    )


__all__ = [
    "box_iou_rotated",
    "deform_conv_backward_filter",
    "deform_conv_backward_input",
    "deform_conv_forward",
    "iou_box3d",
    "modulated_deform_conv_backward",
    "modulated_deform_conv_forward",
    "ms_deform_attn_backward",
    "ms_deform_attn_forward",
    "nms_rotated",
]