from typing import Any, cast
import sys
import random
random.seed(0)
import numpy as np
np.random.seed(0)
import tensorflow as tf
import tf_keras
import onnx2tf.gs as gs
from onnx2tf.utils.common_functions import (
    replace_parameter,
    get_constant_or_variable,
    print_node_info,
    inverted_operation_enable_disable,
    make_tf_node_info,
    get_replacement_parameter,
    pre_process_transpose,
    post_process_transpose,
    transpose_with_flexing_deterrence,
)
from onnx2tf.utils.logging import *

from tensorflow.python.framework import ops
from tensorflow.python.ops import gen_image_ops
from tensorflow.python.util import dispatch


class NMSLayer(tf_keras.layers.Layer):
    def __init__(self, switch_nms_version='v4'):
        super(NMSLayer, self).__init__()
        self.switch_nms_version = switch_nms_version

    @dispatch.add_dispatch_support
    def non_max_suppression(
        self,
        boxes,
        scores,
        max_output_size,
        iou_threshold=0.5,
        score_threshold=float('-inf'),
        pad_to_max_output_size=False,
        name=None,
    ):
        with ops.name_scope(name, 'non_max_suppression'):
            selected_indices = tf.constant([], dtype=tf.int32)
            num_valid = tf.constant(0, dtype=tf.int32)
            if self.switch_nms_version == 'v4':
                selected_indices, num_valid = gen_image_ops.non_max_suppression_v4(
                    boxes=boxes,
                    scores=scores,
                    max_output_size=max_output_size \
                        if not isinstance(max_output_size, np.ndarray) \
                            else tf.convert_to_tensor(
                                value=max_output_size,
                                name='max_output_size'
                            ),
                    iou_threshold=iou_threshold \
                        if not isinstance(iou_threshold, np.ndarray) \
                            else tf.convert_to_tensor(
                                value=iou_threshold,
                                name='iou_threshold',
                            ),
                    score_threshold=score_threshold \
                        if not isinstance(score_threshold, np.ndarray) \
                            else tf.convert_to_tensor(
                                value=score_threshold,
                                name='score_threshold',
                            ),
                    pad_to_max_output_size=pad_to_max_output_size,
                )

            elif self.switch_nms_version == 'v5':
                selected_indices, selected_scores, num_valid = gen_image_ops.non_max_suppression_v5(
                    boxes=boxes,
                    scores=scores,
                    max_output_size=max_output_size \
                        if not isinstance(max_output_size, np.ndarray) \
                            else tf.convert_to_tensor(
                                value=max_output_size,
                                name='max_output_size'
                            ),
                    iou_threshold=iou_threshold \
                        if not isinstance(iou_threshold, np.ndarray) \
                            else tf.convert_to_tensor(
                                value=iou_threshold,
                                name='iou_threshold',
                            ),
                    score_threshold=score_threshold \
                        if not isinstance(score_threshold, np.ndarray) \
                            else tf.convert_to_tensor(
                                value=score_threshold,
                                name='score_threshold',
                            ),
                    soft_nms_sigma=0.0,
                    pad_to_max_output_size=pad_to_max_output_size,
                )
            if pad_to_max_output_size:
                return selected_indices

            else:
                valid_count = tf.cast(num_valid, tf.int32)
                return tf.gather(selected_indices, tf.range(valid_count))

    def call(
        self,
        boxes,
        scores,
        max_output_size,
        iou_threshold=0.5,
        score_threshold=float('-inf'),
        pad_to_max_output_size=False,
        name=None,
    ):
        return self.non_max_suppression(
            boxes,
            scores,
            max_output_size,
            iou_threshold,
            score_threshold,
            pad_to_max_output_size,
            name,
        )


@print_node_info
@inverted_operation_enable_disable
@get_replacement_parameter
def make_node(
    *,
    graph_node: gs.Node,
    tf_layers_dict: dict,
    **kwargs: Any,
):
    """NonMaxSuppression

    Parameters
    ----------
    graph_node: gs.Node
        graph_surgeon Node

    tf_layers_dict: dict
        optype, shape, dtype, tensorflow graph
    """
    before_op_output_shape_trans_1 = \
        tf_layers_dict.get(graph_node.inputs[0].name, {}).get('before_op_output_shape_trans', True)
    before_op_output_shape_trans_2 = \
        tf_layers_dict.get(graph_node.inputs[1].name, {}).get('before_op_output_shape_trans', True)
    before_op_output_shape_trans = \
        before_op_output_shape_trans_1 \
        and before_op_output_shape_trans_2

    graph_node_input_1 = get_constant_or_variable(
        graph_node.inputs[0],
        before_op_output_shape_trans,
    )
    graph_node_input_2 = get_constant_or_variable(
        graph_node.inputs[1],
        before_op_output_shape_trans,
    )
    boxes = tf_layers_dict[graph_node_input_1.name]['tf_node'] \
        if isinstance(graph_node_input_1, gs.Variable) else graph_node_input_1
    scores = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
        if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2

    output_nms_with_dynamic_tensor: bool = kwargs['output_nms_with_dynamic_tensor']
    output_nms_with_argmax: bool = kwargs.get('output_nms_with_argmax', False)
    switch_nms_version: str = kwargs['switch_nms_version']

    # Pre-process transpose
    boxes = pre_process_transpose(
        value_before_transpose=boxes,
        param_target='inputs',
        param_name=graph_node.inputs[0].name,
        **kwargs,
    )
    scores = pre_process_transpose(
        value_before_transpose=scores,
        param_target='inputs',
        param_name=graph_node.inputs[1].name,
        **kwargs,
    )
    boxes = tf.convert_to_tensor(boxes)
    scores = tf.convert_to_tensor(scores)

    graph_node_input_3 = None
    if len(graph_node.inputs) >= 3:
        graph_node_input_3 = get_constant_or_variable(
            graph_node.inputs[2],
            before_op_output_shape_trans,
        )
    graph_node_input_4 = None
    if len(graph_node.inputs) >= 4:
        graph_node_input_4 = get_constant_or_variable(
            graph_node.inputs[3],
            before_op_output_shape_trans,
        )
    graph_node_input_5 = None
    if len(graph_node.inputs) >= 5:
        graph_node_input_5 = get_constant_or_variable(
            graph_node.inputs[4],
            before_op_output_shape_trans,
        )
    max_output_boxes_per_class = tf_layers_dict[graph_node_input_3.name]['tf_node'] \
        if isinstance(graph_node_input_3, gs.Variable) else graph_node_input_3
    iou_threshold = tf_layers_dict[graph_node_input_4.name]['tf_node'] \
        if isinstance(graph_node_input_4, gs.Variable) else graph_node_input_4
    score_threshold = tf_layers_dict[graph_node_input_5.name]['tf_node'] \
        if isinstance(graph_node_input_5, gs.Variable) else graph_node_input_5

    if max_output_boxes_per_class is None \
        or (isinstance(max_output_boxes_per_class, str) and max_output_boxes_per_class == ""):
        max_output_boxes_per_class = tf.constant(0, tf.int32)
    max_output_boxes_per_class = tf.cast(
        tf.convert_to_tensor(max_output_boxes_per_class),
        tf.int32,
    )

    boxes_shape = cast(Any, tf.shape(boxes, out_type=tf.int32))
    boxes_dim_1 = tf.gather(boxes_shape, indices=1)
    max_output_boxes_per_class = tf.where(
        condition=tf.less_equal(max_output_boxes_per_class, 0),
        x=boxes_dim_1,
        y=max_output_boxes_per_class,
    )

    max_output_boxes_per_class = tf.squeeze(max_output_boxes_per_class) \
        if len(max_output_boxes_per_class.shape) == 1 else max_output_boxes_per_class

    if iou_threshold is None or (isinstance(iou_threshold, str) and iou_threshold == ""):
        iou_threshold = tf.constant(0, tf.float32)
    iou_threshold = tf.convert_to_tensor(iou_threshold, dtype=tf.float32)
    iou_threshold = tf.squeeze(iou_threshold) \
        if len(iou_threshold.shape) == 1 else iou_threshold

    if score_threshold is None or (isinstance(score_threshold, str) and score_threshold == ""):
        score_threshold = tf.constant(-np.inf)
    score_threshold = tf.convert_to_tensor(score_threshold, dtype=tf.float32)
    score_threshold = tf.squeeze(score_threshold) \
        if len(score_threshold.shape) == 1 else score_threshold

    center_point_box = graph_node.attrs.get('center_point_box', 0)

    graph_node_output: gs.Variable = graph_node.outputs[0]
    shape = graph_node_output.shape
    dtype = graph_node_output.dtype

    # Preserving Graph Structure (Dict)
    tf_layers_dict[graph_node_output.name] = {
        'optype': graph_node.op,
        'shape': shape,
        'dtype': dtype,
    }

    # Param replacement
    boxes = replace_parameter(
        value_before_replacement=boxes,
        param_target='inputs',
        param_name=graph_node.inputs[0].name,
        **kwargs,
    )
    scores = replace_parameter(
        value_before_replacement=scores,
        param_target='inputs',
        param_name=graph_node.inputs[1].name,
        **kwargs,
    )
    if len(graph_node.inputs) >= 3:
        max_output_boxes_per_class = replace_parameter(
            value_before_replacement=max_output_boxes_per_class,
            param_target='inputs',
            param_name=graph_node.inputs[2].name,
            **kwargs,
        )
    if len(graph_node.inputs) >= 4:
        iou_threshold = replace_parameter(
            value_before_replacement=iou_threshold,
            param_target='inputs',
            param_name=graph_node.inputs[3].name,
            **kwargs,
        )
    if len(graph_node.inputs) >= 5:
        score_threshold = replace_parameter(
            value_before_replacement=score_threshold,
            param_target='inputs',
            param_name=graph_node.inputs[4].name,
            **kwargs,
        )
    center_point_box = replace_parameter(
        value_before_replacement=center_point_box,
        param_target='attributes',
        param_name='center_point_box',
        **kwargs,
    )
    boxes = tf.convert_to_tensor(boxes)
    scores = tf.convert_to_tensor(scores)

    # Generation of TF OP
    if center_point_box == 1:
        boxes_t = transpose_with_flexing_deterrence(
            input_tensor=boxes,
            perm=[0, 2, 1],
            **kwargs,
        )
        x_centers = tf.slice(boxes_t, [0, 0, 0], [-1, 1, -1])
        y_centers = tf.slice(boxes_t, [0, 1, 0], [-1, 1, -1])
        widths = tf.slice(boxes_t, [0, 2, 0], [-1, 1, -1])
        heights = tf.slice(boxes_t, [0, 3, 0], [-1, 1, -1])
        y1 = tf.subtract(y_centers, tf.divide(heights, 2))
        x1 = tf.subtract(x_centers, tf.divide(widths, 2))
        y2 = tf.add(y_centers, tf.divide(heights, 2))
        x2 = tf.add(x_centers, tf.divide(widths, 2))
        boxes_t = tf.concat(
            values=cast(Any, [y1, x1, y2, x2]),
            axis=1,
        )
        boxes = transpose_with_flexing_deterrence(
            input_tensor=boxes_t,
            perm=[0, 2, 1],
            **kwargs,
        )
        boxes = tf.convert_to_tensor(boxes)

    # Canonicalize NMS tensor layouts:
    # boxes:  [batch, num_boxes, 4]
    # scores: [batch, num_classes, num_boxes]
    def _static_int(dim):
        return int(dim) if isinstance(dim, (int, np.integer)) else None

    boxes_rank = len(boxes.shape)
    scores_rank = len(scores.shape)
    if boxes_rank == 3 and scores_rank == 3:
        boxes_shape = [_static_int(dim) for dim in list(boxes.shape)]
        scores_shape = [_static_int(dim) for dim in list(scores.shape)]

        # Fix [batch, 4, num_boxes] -> [batch, num_boxes, 4]
        if boxes_shape[2] != 4 and boxes_shape[1] == 4:
            boxes = transpose_with_flexing_deterrence(
                input_tensor=boxes,
                perm=[0, 2, 1],
                **kwargs,
            )
            boxes_shape = [_static_int(dim) for dim in list(boxes.shape)]

        # Fix [batch, num_boxes, num_classes] -> [batch, num_classes, num_boxes]
        num_boxes_dim = boxes_shape[1]
        if num_boxes_dim is not None:
            if scores_shape[2] != num_boxes_dim and scores_shape[1] == num_boxes_dim:
                scores = transpose_with_flexing_deterrence(
                    input_tensor=scores,
                    perm=[0, 2, 1],
                    **kwargs,
                )
                scores = tf.convert_to_tensor(scores)

    scores_argmax_classes = None
    if output_nms_with_argmax:
        scores_rank = len(scores.shape)
        if scores_rank == 3:
            scores_argmax_classes = tf.argmax(
                scores,
                axis=1,
                output_type=tf.int64,
            )
            scores = tf.reduce_max(
                scores,
                axis=1,
                keepdims=True,
            )
        else:
            warn(
                f'--output_nms_with_argmax is enabled, but scores rank is not 3. ' +
                f'Argmax shrink is skipped. graph_node.name: {graph_node.name} scores.shape: {scores.shape}'
            )

    num_batches = boxes.shape[0]

    if num_batches is None:
        error(
            f'It is not possible to specify a dynamic shape ' +
            f'for the batch size of the input tensor in NonMaxSuppression. ' +
            f'Use the --batch_size option to change the batch size to a fixed size. \n' +
            f'graph_node.name: {graph_node.name} boxes.shape: {boxes.shape} scores.shape: {scores.shape}'
        )
        sys.exit(1)

    boxes_dtype_any = getattr(boxes, 'dtype', None)
    boxes_dtype: tf.dtypes.DType = cast(tf.dtypes.DType, boxes_dtype_any) if boxes_dtype_any is not None else tf.float32
    scores_dtype_any = getattr(scores, 'dtype', None)
    scores_dtype: tf.dtypes.DType = cast(tf.dtypes.DType, scores_dtype_any) if scores_dtype_any is not None else tf.float32
    tf_boxes = tf.zeros([0, 4], dtype=boxes_dtype)
    tf_scores = tf.zeros([0], dtype=scores_dtype)
    result = tf.zeros([0, 3], dtype=tf.int64)
    for batch_i in tf.range(num_batches):
        # get boxes in batch_i only
        begin_ = [0 if idx != 0 else batch_i for idx in range(len(boxes.shape))]
        end_ = [0 if idx != 0 else batch_i + 1 for idx in range(len(boxes.shape))]
        begin_mask_ = sum([2**idx if idx != 0 else 0 for idx in range(len(boxes.shape))])
        end_mask_ = begin_mask_
        tf_boxes = \
            tf.squeeze(
                input=tf.strided_slice(
                    input_=boxes,
                    begin=begin_,
                    end=end_,
                    begin_mask=begin_mask_,
                    end_mask=end_mask_,
                ),
                axis=0,
            )
        # get scores of all classes in batch_i only
        begin_ = [0 if idx != 0 else batch_i for idx in range(len(scores.shape))]
        end_ = [0 if idx != 0 else batch_i + 1 for idx in range(len(scores.shape))]
        begin_mask_ = sum([2**idx if idx != 0 else 0 for idx in range(len(scores.shape))])
        end_mask_ = begin_mask_
        batch_i_scores = \
            tf.squeeze(
                input=tf.strided_slice(
                    input_=scores,
                    begin=begin_,
                    end=end_,
                    begin_mask=begin_mask_,
                    end_mask=end_mask_,
                ),
                axis=0,
            )
        batch_i_class_ids = None
        if output_nms_with_argmax and scores_argmax_classes is not None:
            batch_i_class_ids = tf.gather(
                scores_argmax_classes,
                batch_i,
            )
        # get number of classess in batch_i only
        num_classes = batch_i_scores.shape[0]
        for class_j in tf.range(num_classes):
            # get scores in class_j for batch_i only
            begin_ = [0 if idx != 0 else class_j for idx in range(len(batch_i_scores.shape))]
            end_ = [0 if idx != 0 else class_j + 1 for idx in range(len(batch_i_scores.shape))]
            begin_mask_ = sum([2**idx if idx != 0 else 0 for idx in range(len(batch_i_scores.shape))])
            end_mask_ = begin_mask_
            tf_scores = \
                tf.squeeze(
                    input=tf.strided_slice(
                        input_=batch_i_scores,
                        begin=begin_,
                        end=end_,
                        begin_mask=begin_mask_,
                        end_mask=end_mask_,
                    ),
                    axis=0,
                )
            # get the selected boxes indices
            nms = NMSLayer(switch_nms_version=switch_nms_version)
            selected_indices = nms(
                boxes=tf_boxes,
                scores=tf_scores,
                max_output_size=max_output_boxes_per_class,
                iou_threshold=iou_threshold,
                score_threshold=score_threshold,
                pad_to_max_output_size=not output_nms_with_dynamic_tensor,
            )

            # add batch and class information into the indices
            selected_indices_int64 = tf.cast(selected_indices, dtype=tf.int64)
            if output_nms_with_argmax and batch_i_class_ids is not None:
                selected_class_ids = tf.gather(
                    batch_i_class_ids,
                    tf.cast(selected_indices, dtype=tf.int32),
                )
                output = tf.stack(
                    [
                        tf.fill(
                            dims=tf.shape(selected_indices_int64),
                            value=tf.cast(batch_i, dtype=tf.int64),
                        ),
                        tf.cast(selected_class_ids, dtype=tf.int64),
                        selected_indices_int64,
                    ],
                    axis=1,
                )
            else:
                output = tf.transpose([selected_indices_int64])
                paddings = tf.constant([[0, 0], [1, 0]])
                tf_pad = cast(Any, tf.pad)
                output = tf_pad(
                    output,
                    paddings,
                    constant_values=tf.cast(class_j, dtype=tf.int64),
                )
                output = tf_pad(
                    output,
                    paddings,
                    constant_values=tf.cast(batch_i, dtype=tf.int64),
                )
            # tf.function will auto convert "result" from variable to placeholder
            # therefore don't need to use assign here
            result = output if tf.equal(batch_i, 0) and tf.equal(class_j, 0) else tf.concat([result, output], 0)

    tf_layers_dict[graph_node_output.name]['tf_node'] = result

    # Post-process transpose
    tf_layers_dict[graph_node_output.name]['tf_node'] = post_process_transpose(
        value_before_transpose=tf_layers_dict[graph_node_output.name]['tf_node'],
        param_target='outputs',
        param_name=graph_node.outputs[0].name,
        **kwargs,
    )

    # Generation of Debug Info
    tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
        make_tf_node_info(
            node_info={
                'tf_op_type': tf.image.non_max_suppression,
                'tf_inputs': {
                    'boxes': tf_boxes,
                    'scores': tf_scores,
                    'max_output_boxes_per_class': max_output_boxes_per_class,
                    'iou_threshold': iou_threshold,
                    'score_threshold': score_threshold,
                },
                'tf_outputs': {
                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
                },
            }
        )
