# Copyright (c) 2020  PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
from paddle2onnx.utils import logging
from paddle2onnx.constant import dtypes
from paddle2onnx.op_mapper import OpMapper as op_mapper


@op_mapper(
    ['multiclass_nms', 'multiclass_nms2', 'matrix_nms', 'multiclass_nms3'])
class MultiClassNMS():
    support_opset_verision_range = (10, 12)
    """
    Convert the paddle multiclass_nms to onnx op.
    This op is get the select boxes from origin boxes.
    """

    @classmethod
    def opset_10(cls, graph, node, **kw):
        logging.warning("Operator:{} only supports input[batch_size] == 1.".
                        format(node.type))
        scores = node.input('Scores', 0)
        bboxes = node.input('BBoxes', 0)
        num_class = node.input_shape('Scores', 0)[1]
        if len(node.input_shape('Scores', 0)) == 2:
            # inputs: scores & bboxes is lod tensor
            scores = graph.make_node('Transpose', inputs=[scores], perm=[1, 0])
            scores = graph.make_node('Unsqueeze', inputs=[scores], axes=[0])
            scores_list = graph.make_node(
                'Split',
                inputs=scores,
                outputs=num_class,
                axis=1,
                split=[1] * num_class)
            bboxes = graph.make_node('Transpose', inputs=bboxes, perm=[1, 0, 2])
            bboxes_list = graph.make_node(
                'Split',
                inputs=bboxes,
                outputs=num_class,
                axis=0,
                split=[1] * num_class)
            bbox_ids = []
            for i in range(num_class):
                bbox_id = cls.nms(graph,
                                  node,
                                  scores_list[i],
                                  bboxes_list[i],
                                  class_id=i)
                bbox_ids.append(bbox_id)
            bbox_ids = graph.make_node('Concat', inputs=bbox_ids, axis=0)
            const_shape = graph.make_node(
                'Constant', dtype=dtypes.ONNX.INT64, value=[1, -1, 4])
            bboxes = graph.make_node('Reshape', inputs=[bboxes, const_shape])
            cls.keep_top_k(
                graph, node, bbox_ids, scores, bboxes, is_lod_input=True)
        else:
            bbox_ids = cls.nms(graph, node, scores, bboxes)
            cls.keep_top_k(graph, node, bbox_ids, scores, bboxes)

    @classmethod
    def nms(cls, graph, node, scores, bboxes, class_id=None):
        normalized = node.attr('normalized')
        nms_top_k = node.attr('nms_top_k')
        if node.type == 'matrix_nms':
            iou_threshold = 0.5
            logging.warning(
                "Operator:{} is not supported completely, so we use traditional"
                " NMS (nms_theshold={}) to instead it, which introduce some difference.".
                format(node.type, str(iou_threshold)))
        else:
            iou_threshold = node.attr('nms_threshold')
        if nms_top_k == -1:
            nms_top_k = 100000

        #convert the paddle attribute to onnx tensor
        score_threshold = graph.make_node(
            'Constant',
            dtype=dtypes.ONNX.FLOAT,
            value=[float(node.attr('score_threshold'))])
        iou_threshold = graph.make_node(
            'Constant', dtype=dtypes.ONNX.FLOAT, value=[float(iou_threshold)])
        nms_top_k = graph.make_node(
            'Constant', dtype=dtypes.ONNX.INT64, value=[np.int64(nms_top_k)])

        # the paddle data format is x1,y1,x2,y2
        kwargs = {'center_point_box': 0}

        if normalized:
            select_bbox_indices = graph.make_node(
                'NonMaxSuppression',
                inputs=[
                    bboxes, scores, nms_top_k, iou_threshold, score_threshold
                ])
        elif not normalized:
            value_one = graph.make_node(
                'Constant', dims=[1], dtype=dtypes.ONNX.FLOAT, value=1.0)
            new_bboxes = graph.make_node(
                'Split', inputs=[bboxes], outputs=4, axis=2,
                split=[1, 1, 1, 1])
            new_xmax = graph.make_node('Add', inputs=[new_bboxes[2], value_one])
            new_ymax = graph.make_node('Add', inputs=[new_bboxes[3], value_one])
            new_bboxes = graph.make_node(
                'Concat',
                inputs=[new_bboxes[0], new_bboxes[1], new_xmax, new_ymax],
                axis=2)
            select_bbox_indices = graph.make_node(
                'NonMaxSuppression',
                inputs=[
                    new_bboxes, scores, nms_top_k, iou_threshold,
                    score_threshold
                ])

        if class_id is not None and class_id != 0:
            class_id = graph.make_node(
                'Constant', dtype=dtypes.ONNX.INT64, value=[0, class_id, 0])
            class_id = graph.make_node('Unsqueeze', inputs=[class_id], axes=[0])
            select_bbox_indices = graph.make_node(
                'Add', inputs=[select_bbox_indices, class_id])

        return select_bbox_indices

    @classmethod
    def keep_top_k(cls,
                   graph,
                   node,
                   select_bbox_indices,
                   scores,
                   bboxes,
                   is_lod_input=False):
        # step 1 nodes select the nms class
        # create some const value to use
        background = node.attr('background_label')
        const_values = []
        for value in [0, 1, 2, -1]:
            const_value = graph.make_node(
                'Constant', dtype=dtypes.ONNX.INT64, value=[value])
            const_values.append(const_value)

        # In this code block, we will deocde the raw score data, reshape N * C * M to 1 * N*C*M
        # and the same time, decode the select indices to 1 * D, gather the select_indices
        class_id = graph.make_node(
            'Gather', inputs=[select_bbox_indices, const_values[1]], axis=1)

        squeezed_class_id = graph.make_node(
            'Squeeze', inputs=[class_id], axes=[1])

        bbox_id = graph.make_node(
            'Gather', inputs=[select_bbox_indices, const_values[2]], axis=1)

        if background == 0:
            nonzero = graph.make_node('NonZero', inputs=[squeezed_class_id])
        else:
            filter_cls_id = graph.make_node(
                'Constant', dtype=dtypes.ONNX.INT32, value=[background])
            cast = graph.make_node(
                'Cast', inputs=[squeezed_class_id], to=dtypes.ONNX.INT32)
            filter_index = graph.make_node('Sub', inputs=[cast, filter_cls_id])
            nonzero = graph.make_node('NonZero', inputs=[filter_index])

        class_id = graph.make_node('Gather', inputs=[class_id, nonzero], axis=0)

        bbox_id = graph.make_node('Gather', inputs=[bbox_id, nonzero], axis=0)

        # get the shape of scores
        shape_scores = graph.make_node('Shape', inputs=scores)

        # gather the index: 2 shape of scores
        class_num = graph.make_node(
            'Gather', inputs=[shape_scores, const_values[2]], axis=0)

        # reshape scores N * C * M to (N*C*M) * 1
        scores = graph.make_node('Reshape', inputs=[scores, const_values[-1]])

        # mul class * M
        mul_classnum_boxnum = graph.make_node(
            'Mul', inputs=[class_id, class_num])

        # add class * M * index
        add_class_indices = graph.make_node(
            'Add', inputs=[mul_classnum_boxnum, bbox_id])

        # Squeeze the indices to 1 dim
        score_indices = graph.make_node(
            'Squeeze', inputs=[add_class_indices], axes=[0, 2])

        # gather the data from flatten scores
        scores = graph.make_node(
            'Gather', inputs=[scores, score_indices], axis=0)

        keep_top_k = node.attr('keep_top_k')
        keep_top_k = graph.make_node(
            'Constant',
            dtype=dtypes.ONNX.INT64,
            dims=[1, 1],
            value=[node.attr('keep_top_k')])

        # get min(topK, num_select)
        shape_select_num = graph.make_node('Shape', inputs=[scores])
        const_zero = graph.make_node(
            'Constant', dtype=dtypes.ONNX.INT64, value=[0])
        gather_select_num = graph.make_node(
            'Gather', inputs=[shape_select_num, const_zero], axis=0)
        unsqueeze_select_num = graph.make_node(
            'Unsqueeze', inputs=[gather_select_num], axes=[0])
        concat_topK_select_num = graph.make_node(
            'Concat', inputs=[unsqueeze_select_num, keep_top_k], axis=0)
        cast_concat_topK_select_num = graph.make_node(
            'Cast', inputs=[concat_topK_select_num], to=6)
        keep_top_k = graph.make_node(
            'ReduceMin', inputs=[cast_concat_topK_select_num], keepdims=0)
        # unsqueeze the indices to 1D tensor
        keep_top_k = graph.make_node('Unsqueeze', inputs=[keep_top_k], axes=[0])
        # cast the indices to INT64
        keep_top_k = graph.make_node('Cast', inputs=[keep_top_k], to=7)

        # select topk scores  indices
        keep_topk_scores, keep_topk_indices = graph.make_node(
            'TopK', inputs=[scores, keep_top_k], outputs=2)

        # gather topk label, scores, boxes
        gather_topk_scores = graph.make_node(
            'Gather', inputs=[scores, keep_topk_indices], axis=0)

        gather_topk_class = graph.make_node(
            'Gather', inputs=[class_id, keep_topk_indices], axis=1)

        # gather the boxes need to gather the boxes id, then get boxes
        if is_lod_input:
            gather_topk_boxes_id = graph.make_node(
                'Gather', [add_class_indices, keep_topk_indices], axis=1)
        else:
            gather_topk_boxes_id = graph.make_node(
                'Gather', [bbox_id, keep_topk_indices], axis=1)

        # squeeze the gather_topk_boxes_id to 1 dim
        squeeze_topk_boxes_id = graph.make_node(
            'Squeeze', inputs=[gather_topk_boxes_id], axes=[0, 2])

        gather_select_boxes = graph.make_node(
            'Gather', inputs=[bboxes, squeeze_topk_boxes_id], axis=1)

        # concat the final result
        # before concat need to cast the class to float
        cast_topk_class = graph.make_node(
            'Cast', inputs=[gather_topk_class], to=1)

        unsqueeze_topk_scores = graph.make_node(
            'Unsqueeze', inputs=[gather_topk_scores], axes=[0, 2])

        inputs_concat_final_results = [
            cast_topk_class, unsqueeze_topk_scores, gather_select_boxes
        ]

        sort_by_socre_results = graph.make_node(
            'Concat', inputs=inputs_concat_final_results, axis=2)

        # sort by class_id
        squeeze_cast_topk_class = graph.make_node(
            'Squeeze', inputs=[cast_topk_class], axes=[0, 2])

        neg_squeeze_cast_topk_class = graph.make_node(
            'Neg', inputs=[squeeze_cast_topk_class])

        data, indices = graph.make_node(
            'TopK', inputs=[neg_squeeze_cast_topk_class, keep_top_k], outputs=2)

        concat_final_results = graph.make_node(
            'Gather', inputs=[sort_by_socre_results, indices], axis=1)

        concat_final_results = graph.make_node(
            'Squeeze',
            inputs=[concat_final_results],
            outputs=[node.output('Out', 0)],
            axes=[0])

        if node.type in ['multiclass_nms2', 'matrix_nms', 'multiclass_nms3']:
            final_indices = graph.make_node(
                'Squeeze',
                inputs=[bbox_id],
                outputs=node.output('Index'),
                axes=[0])
            if node.type in ['matrix_nms', 'multiclass_nms3']:
                select_bboxes_shape = graph.make_node(
                    'Shape', inputs=[concat_final_results])
                select_bboxes_shape1 = graph.make_node(
                    'Cast', inputs=[select_bboxes_shape], to=dtypes.ONNX.INT32)
                indices = graph.make_node(
                    'Constant', dtype=dtypes.ONNX.INT64, value=[0])
                rois_num = None
                if 'NmsRoisNum' in node.outputs:
                    rois_num = node.output('NmsRoisNum')
                elif 'RoisNum' in node.outputs:
                    rois_num = node.output('RoisNum')
                if rois_num is not None:
                    graph.make_node(
                        "Gather",
                        inputs=[select_bboxes_shape1, indices],
                        outputs=rois_num)
