from typing import Any
import random
random.seed(0)
import numpy as np
np.random.seed(0)
import tensorflow as tf
import onnx2tf.gs as gs
from onnx2tf.utils.common_functions import (
    get_constant_or_variable,
    print_node_info,
    inverted_operation_enable_disable,
    make_tf_node_info,
    get_replacement_parameter,
    pre_process_transpose,
    post_process_transpose,
)
from onnx2tf.utils.enums import NUMPY_DTYPES_TO_TF_DTYPES


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

    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,
    )
    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_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,
    }

    # Generation of TF OP
    input_tensor_1 = tf_layers_dict[graph_node_input_1.name]['tf_node'] \
        if isinstance(graph_node_input_1, gs.Variable) else graph_node_input_1
    input_tensor_2 = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
        if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
    a_zero_point = tf_layers_dict[graph_node_input_3.name]['tf_node'] \
        if isinstance(graph_node_input_3, gs.Variable) else graph_node_input_3
    b_zero_point = tf_layers_dict[graph_node_input_4.name]['tf_node'] \
        if isinstance(graph_node_input_4, gs.Variable) else graph_node_input_4

    if 'matmulinteger' in kwargs['replace_to_pseudo_operators']:
        tf_matmul_dtype = tf.float32
    else:
        tf_matmul_dtype = tf.int32

    casted_input_tensor_1 = tf.cast(input_tensor_1, tf_matmul_dtype)
    casted_input_tensor_2 = tf.cast(input_tensor_2, tf_matmul_dtype)

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

    # apply a_zero_point to A
    if a_zero_point is not None:
        if None not in a_zero_point.shape:
            shape = a_zero_point.shape
            if len(shape) > 0  and shape[0] > 1:
                # reshape a_zero_point before subtract it from A
                a_zero_point = tf.reshape(a_zero_point, [shape[0], 1])
        else:
            @tf.function
            def get_a_zero_point(a_zero_point):
                shape = tf.shape(a_zero_point)
                shape_rank = tf.size(shape)
                first_dim = tf.gather(shape, indices=0)
                if tf.logical_and(
                    tf.greater(shape_rank, 0),
                    tf.greater(first_dim, 1),
                ):
                    # reshape a_zero_point before subtract it from A
                    a_zero_point = tf.reshape(a_zero_point, [first_dim, 1])
                return a_zero_point
            a_zero_point = get_a_zero_point(a_zero_point)

        a_zero_point = tf.cast(a_zero_point, tf_matmul_dtype)
        casted_input_tensor_1 = tf.subtract(casted_input_tensor_1, a_zero_point)

    # apply b_zero_point to B
    if b_zero_point is not None:
        b_zero_point = tf.cast(b_zero_point, tf_matmul_dtype)
        casted_input_tensor_2 = tf.subtract(casted_input_tensor_2, b_zero_point)

    tf_layers_dict[graph_node_output.name]['tf_node'] = \
        tf.matmul(
            a=casted_input_tensor_1,
            b=casted_input_tensor_2,
            output_type=tf_matmul_dtype,
            name=graph_node.name,
        )

    output_dtype = NUMPY_DTYPES_TO_TF_DTYPES[dtype] \
        if isinstance(dtype, np.dtype) else dtype

    if tf_matmul_dtype != output_dtype:
        tf_layers_dict[graph_node_output.name]['tf_node'] = \
            x=tf.round(tf_layers_dict[graph_node_output.name]['tf_node'])
        # MatMulInteger is commonly used in dynamically quantized ONNX models,
        # in which case its output is directly casted to float32, making the
        # cast here unnecessary
        if not (len(graph_node_output.outputs) == 1 and \
                graph_node_output.outputs[0].op == 'Cast' and \
                graph_node_output.outputs[0].attrs['to'] == 1):
            tf_layers_dict[graph_node_output.name]['tf_node'] = \
                tf.cast(
                    x=tf_layers_dict[graph_node_output.name]['tf_node'],
                    dtype=output_dtype,
                )

    # 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.matmul,
                'tf_inputs': {
                    'a': casted_input_tensor_1,
                    'b': casted_input_tensor_2,
                    'output_type': dtype,
                },
                'tf_outputs': {
                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
                },
            }
        )
