# ---------------------------------------------
# Copyright (c) OpenMMLab. All rights reserved.
# ---------------------------------------------
#  Modified by Zhiqi Li
# ---------------------------------------------
#  Refactored to use torch.library.custom_op + register_autograd
#  for torch.compile compatibility
# ---------------------------------------------

import torch
from typing import Tuple
from lightop import op as ext_module


# ============ Custom Op Definition (compile-friendly) ============

@torch.library.custom_op("mmcv::ms_deform_attn", mutates_args=())
def ms_deform_attn(
    value: torch.Tensor,
    value_spatial_shapes: torch.Tensor,
    value_level_start_index: torch.Tensor,
    sampling_locations: torch.Tensor,
    attention_weights: torch.Tensor,
    im2col_step: int) -> torch.Tensor:
    """Multi-scale deformable attention forward (compile-friendly)."""
    return ext_module.ms_deform_attn_forward(
        value,
        value_spatial_shapes,
        value_level_start_index,
        sampling_locations,
        attention_weights,
        im2col_step)


@ms_deform_attn.register_fake
def _(
    value: torch.Tensor,
    value_spatial_shapes: torch.Tensor,
    value_level_start_index: torch.Tensor,
    sampling_locations: torch.Tensor,
    attention_weights: torch.Tensor,
    im2col_step: int) -> torch.Tensor:
    bs = value.shape[0]
    num_queries = sampling_locations.shape[1]
    embed_dims = value.shape[3] * sampling_locations.shape[2]
    return torch.empty(bs, num_queries, embed_dims, device=value.device, dtype=value.dtype)


@torch.library.custom_op("mmcv::ms_deform_attn_backward", mutates_args=())
def ms_deform_attn_backward(
    value: torch.Tensor,
    value_spatial_shapes: torch.Tensor,
    value_level_start_index: torch.Tensor,
    sampling_locations: torch.Tensor,
    attention_weights: torch.Tensor,
    grad_output: torch.Tensor,
    im2col_step: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Multi-scale deformable attention backward (compile-friendly)."""
    if not grad_output.is_contiguous():
        grad_output = grad_output.contiguous()

    grad_value = torch.zeros_like(value)
    grad_sampling_loc = torch.zeros_like(sampling_locations)
    grad_attn_weight = torch.zeros_like(attention_weights)

    ext_module.ms_deform_attn_backward(
        value,
        value_spatial_shapes,
        value_level_start_index,
        sampling_locations,
        attention_weights,
        grad_output,
        grad_value,
        grad_sampling_loc,
        grad_attn_weight,
        im2col_step)

    return grad_value, grad_sampling_loc, grad_attn_weight


@ms_deform_attn_backward.register_fake
def _(
    value: torch.Tensor,
    value_spatial_shapes: torch.Tensor,
    value_level_start_index: torch.Tensor,
    sampling_locations: torch.Tensor,
    attention_weights: torch.Tensor,
    grad_output: torch.Tensor,
    im2col_step: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    return (
        torch.empty_like(value),
        torch.empty_like(sampling_locations),
        torch.empty_like(attention_weights),
    )


# ============ Autograd Registration ============

def _ms_deform_attn_setup_context(ctx, inputs, output):
    value, value_spatial_shapes, value_level_start_index, \
        sampling_locations, attention_weights, im2col_step = inputs
    ctx.save_for_backward(value, value_spatial_shapes,
                          value_level_start_index, sampling_locations,
                          attention_weights)
    ctx.im2col_step = im2col_step


def _ms_deform_attn_backward(ctx, grad_output):
    value, value_spatial_shapes, value_level_start_index, \
        sampling_locations, attention_weights = ctx.saved_tensors

    grad_value, grad_sampling_loc, grad_attn_weight = ms_deform_attn_backward(
        value,
        value_spatial_shapes,
        value_level_start_index,
        sampling_locations,
        attention_weights,
        grad_output,
        ctx.im2col_step)

    return grad_value, None, None, \
        grad_sampling_loc, grad_attn_weight, None


torch.library.register_autograd(
    "mmcv::ms_deform_attn",
    _ms_deform_attn_backward,
    setup_context=_ms_deform_attn_setup_context,
)


# ============ Convenience Wrappers (backward-compatible .apply() interface) ============

class MultiScaleDeformableAttnFunction_fp16:
    """FP16 version: casts inputs to float16 under autocast, then calls compile-friendly custom_op."""

    @staticmethod
    def apply(value, value_spatial_shapes, value_level_start_index,
              sampling_locations, attention_weights, im2col_step):
        # Equivalent to original @custom_fwd(cast_inputs=torch.float16)
        if torch.is_autocast_enabled():
            value = value.to(torch.float16)
            sampling_locations = sampling_locations.to(torch.float16)
            attention_weights = attention_weights.to(torch.float16)
        return ms_deform_attn(
            value, value_spatial_shapes, value_level_start_index,
            sampling_locations, attention_weights, im2col_step)


class MultiScaleDeformableAttnFunction_fp32:
    """FP32 version: casts inputs to float32 under autocast, then calls compile-friendly custom_op."""

    @staticmethod
    def apply(value, value_spatial_shapes, value_level_start_index,
              sampling_locations, attention_weights, im2col_step):
        # Equivalent to original @custom_fwd(cast_inputs=torch.float32)
        if torch.is_autocast_enabled():
            value = value.to(torch.float32)
            sampling_locations = sampling_locations.to(torch.float32)
            attention_weights = attention_weights.to(torch.float32)
        return ms_deform_attn(
            value, value_spatial_shapes, value_level_start_index,
            sampling_locations, attention_weights, im2col_step)
