# Copyright (c) OpenMMLab. All rights reserved.
import torch
from torch.autograd import Function

from mmdeploy.core.optimizers import mark
from mmdeploy.core.rewriters import FUNCTION_REWRITER


class MultiLevelRoiAlign(Function):
    """Create MMCVMultiLevelRoiAlign op.

    This class is used to create a MultiLevelRoiAlign in ONNX for the TensorRT
    backend.
    """

    def __init__(self) -> None:
        super().__init__()

    @staticmethod
    def symbolic(g, *args):
        """Symbolic function for creating onnx op."""
        aligned = args[-1]
        featmap_strides = args[-2]
        finest_scale = args[-3]
        roi_scale_factor = args[-4]
        sampling_ratio = args[-5]
        output_size = args[-6]
        inputs = args[:len(featmap_strides)]
        rois = args[len(featmap_strides)]
        return g.op(
            'mmdeploy::MMCVMultiLevelRoiAlign',
            rois,
            *inputs,
            output_height_i=output_size[1],
            output_width_i=output_size[0],
            sampling_ratio_i=sampling_ratio,
            roi_scale_factor_f=roi_scale_factor,
            finest_scale_i=finest_scale,
            featmap_strides_f=featmap_strides,
            aligned_i=aligned)

    @staticmethod
    def forward(g, *args):
        """Run forward."""
        # aligned = args[-1]
        featmap_strides = args[-2]
        # finest_scale = args[-3]
        # roi_scale_factor = args[-4]
        # sampling_ratio = args[-5]
        output_size = args[-6]
        inputs = args[:len(featmap_strides)]
        rois = args[len(featmap_strides)]

        num_proposals = rois.shape[0]
        channel = inputs[0].shape[1]

        return rois.new_zeros(
            (num_proposals, channel, output_size[1], output_size[0]))


@FUNCTION_REWRITER.register_rewriter(
    'mmdet.models.roi_heads.roi_extractors.'
    'single_level_roi_extractor.SingleRoIExtractor.forward',
    backend='tensorrt')
@mark('roi_extractor', inputs=['feats', 'rois'], outputs=['bbox_feats'])
def single_roi_extractor__forward__tensorrt(ctx,
                                            self,
                                            feats,
                                            rois,
                                            roi_scale_factor=None):
    """Rewrite `forward` of `SingleRoIExtractor` for TensorRT backend.

    This function uses MMCVMultiLevelRoiAlign op for TensorRT deployment.
    """
    featmap_strides = self.featmap_strides
    finest_scale = self.finest_scale

    roi_layer = self.roi_layers[0]
    out_size = roi_layer.output_size
    sampling_ratio = roi_layer.sampling_ratio
    aligned = roi_layer.aligned
    if roi_scale_factor is None:
        roi_scale_factor = 1.0

    featmap_strides = [float(s) for s in featmap_strides]
    return MultiLevelRoiAlign.apply(*feats, rois, out_size, sampling_ratio,
                                    roi_scale_factor, finest_scale,
                                    featmap_strides, aligned)


@FUNCTION_REWRITER.register_rewriter(
    func_name='mmdet.models.roi_heads.SingleRoIExtractor.forward')
@mark('roi_extractor', inputs=['feats', 'rois'], outputs=['bbox_feats'])
def single_roi_extractor__forward(ctx,
                                  self,
                                  feats,
                                  rois,
                                  roi_scale_factor=None):
    """Rewrite `forward` of SingleRoIExtractor for default backend.

    Rewrite this function to enable exporting to onnx even though the input
    image contains no targets. Note that, `ScatterND` of onnx may conflict with
    `Reshape` if a tensor have a dim size of 0. Thus, we have to cat zeros to
    the dim 0 of `roi_feats` and recover back after all roi align finished.

    Besides, this function adds mark for roi_extractor forward and remove
    unnecessary code of origin forward function.
    """
    out_size = self.roi_layers[0].output_size
    num_levels = len(feats)
    roi_feats = feats[0].new_zeros(rois.shape[0], self.out_channels, *out_size)
    if num_levels == 1:
        assert len(rois) > 0, 'The number of rois should be positive'
        return self.roi_layers[0](feats[0], rois)

    target_lvls = self.map_roi_levels(rois, num_levels)

    if roi_scale_factor is not None:
        rois = self.roi_rescale(rois, roi_scale_factor)

    # concat len num_levels * 2 of zero tensors to dim 0 of roi_feats
    roi_feats = torch.cat(
        (roi_feats.new_zeros(num_levels * 2,
                             *roi_feats.shape[-3:]), roi_feats))
    for i in range(num_levels):
        mask = target_lvls == i
        inds = mask.nonzero(as_tuple=False).squeeze(1)

        # concat len 2 zero tensors to dim 0 of roi_feats
        rois_i = torch.cat((rois.new_zeros(2, 5), rois[inds]))

        roi_feats_t = self.roi_layers[i](feats[i], rois_i)

        # correspondingly change the inds
        inds = torch.cat([
            torch.tensor([2 * i, 2 * i + 1],
                         device=inds.device,
                         dtype=inds.dtype), inds + num_levels * 2
        ])
        roi_feats[inds] = roi_feats_t

    # slice and recover tensors
    roi_feats = roi_feats[num_levels * (2):]
    return roi_feats


class SingleRoIExtractorOpenVINO(Function):
    """This class adds support for ExperimentalDetectronROIFeatureExtractor
    when exporting to OpenVINO.

    The `forward` method returns the original output, which is calculated in
    advance and added to the SingleRoIExtractorOpenVINO class. In addition, the
    list of arguments is changed here to be more suitable for
    ExperimentalDetectronROIFeatureExtractor.
    """

    def __init__(self) -> None:
        super().__init__()

    @staticmethod
    def forward(g, output_size, featmap_strides, sample_num, rois, *feats):
        return SingleRoIExtractorOpenVINO.origin_output

    @staticmethod
    def symbolic(g, output_size, featmap_strides, sample_num, rois, *feats):
        from torch.onnx.symbolic_helper import _slice_helper
        rois = _slice_helper(g, rois, axes=[1], starts=[1], ends=[5])
        domain = 'org.openvinotoolkit'
        op_name = 'ExperimentalDetectronROIFeatureExtractor'
        roi_feats = g.op(
            f'{domain}::{op_name}',
            rois,
            *feats,
            output_size_i=output_size,
            pyramid_scales_i=featmap_strides,
            sampling_ratio_i=sample_num,
            image_id_i=0,
            distribute_rois_between_levels_i=1,
            preserve_rois_order_i=0,
            aligned_i=1,
            outputs=1)
        return roi_feats


@FUNCTION_REWRITER.register_rewriter(
    'mmdet.models.roi_heads.roi_extractors.'
    'single_level_roi_extractor.SingleRoIExtractor.forward',
    backend='openvino')
def single_roi_extractor__forward__openvino(ctx,
                                            self,
                                            feats,
                                            rois,
                                            roi_scale_factor=None):
    """Replaces SingleRoIExtractor with SingleRoIExtractorOpenVINO when
    exporting to OpenVINO.

    This function uses ExperimentalDetectronROIFeatureExtractor for OpenVINO.
    """

    # Adding original output to SingleRoIExtractorOpenVINO.
    state = torch._C._get_tracing_state()
    origin_output = ctx.origin_func(self, feats, rois, roi_scale_factor)
    setattr(SingleRoIExtractorOpenVINO, 'origin_output', origin_output)
    torch._C._set_tracing_state(state)

    output_size = self.roi_layers[0].output_size[0]
    featmap_strides = self.featmap_strides
    sample_num = self.roi_layers[0].sampling_ratio

    args = (output_size, featmap_strides, sample_num, rois, *feats)
    result = SingleRoIExtractorOpenVINO.apply(*args)
    return result
