from typing import Any, List, cast
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,
    get_weights_constant_or_variable,
    get_padding_as_op,
    print_node_info,
    inverted_operation_enable_disable,
    make_tf_node_info,
    transpose_with_flexing_deterrence,
    _is_output_nhwc_assumed,
)
from onnx2tf.utils.logging import *

INF_INDEX_VALUE: int = 4294967296


def _dequantize_tensor(
    *,
    base,
    zero_point,
    scale,
):
    # Do computation in float32
    base = tf.cast(base, tf.float32)
    zero_point = tf.cast(zero_point, tf.float32)
    scale = tf.cast(scale, tf.float32)
    return tf.multiply(
        tf.subtract(base, zero_point),
        scale,
    )


def _reshape_per_output_channel(
    *,
    value,
    weights,
):
    value_rank = len(value.shape)
    weights_rank = len(weights.shape)
    if value_rank == 0:
        return value
    if value_rank == 1 and weights_rank is not None:
        shape = [1] * weights_rank
        shape[-1] = -1
        return tf.reshape(value, shape)
    return value


def _dequantize_weights(
    *,
    base,
    zero_point,
    scale,
):
    # Do computation in float32
    casted_base = tf.cast(base, tf.float32)
    casted_zero_point = tf.cast(zero_point, tf.float32)
    casted_scale = tf.cast(scale, tf.float32)
    casted_zero_point = _reshape_per_output_channel(
        value=casted_zero_point,
        weights=casted_base,
    )
    casted_scale = _reshape_per_output_channel(
        value=casted_scale,
        weights=casted_base,
    )
    return tf.multiply(
        tf.subtract(casted_base, casted_zero_point),
        casted_scale,
    )


def _get_qmin_qmax(dtype: tf.dtypes.DType):
    if dtype == tf.uint8:
        return 0.0, 255.0
    if dtype == tf.int8:
        return -128.0, 127.0
    if dtype == tf.uint16:
        return 0.0, 65535.0
    if dtype == tf.int16:
        return -32768.0, 32767.0
    return None, None


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

    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 = \
        before_op_output_shape_trans_1

    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 = get_constant_or_variable(
        graph_node.inputs[2],
        before_op_output_shape_trans,
    )
    kernel_shape = graph_node.attrs.get('kernel_shape', [])
    kernel_size = len(kernel_shape)
    graph_node_input_4 = get_weights_constant_or_variable(
        const_or_var=graph_node.inputs[3],
        kernel_size=kernel_size,
    )
    graph_node_input_5 = get_constant_or_variable(
        graph_node.inputs[4],
        before_op_output_shape_trans,
    )
    graph_node_input_6 = get_constant_or_variable(
        graph_node.inputs[5],
        before_op_output_shape_trans,
    )
    graph_node_input_7 = get_constant_or_variable(
        graph_node.inputs[6],
        before_op_output_shape_trans,
    )
    graph_node_input_8 = get_constant_or_variable(
        graph_node.inputs[7],
        before_op_output_shape_trans,
    )
    graph_node_input_9 = None
    if len(graph_node.inputs) >= 9:
        graph_node_input_9 = get_constant_or_variable(
            graph_node.inputs[8],
            before_op_output_shape_trans,
        )

    graph_node_output: gs.Variable = graph_node.outputs[0]
    try:
        output_tensor_shape = graph_node.o().outputs[0].shape
    except:
        output_tensor_shape = graph_node.outputs[0].shape
    dtype = graph_node_output.dtype

    input_tensor = tf_layers_dict[graph_node_input_1.name]['tf_node'] \
        if isinstance(graph_node_input_1, gs.Variable) else graph_node_input_1
    input_is_dequantized = False
    input_nhwc_assumed = False
    if isinstance(graph_node_input_1, gs.Variable):
        input_is_dequantized = tf_layers_dict.get(graph_node_input_1.name, {}).get('is_dequantized', False)
        input_nhwc_assumed = _is_output_nhwc_assumed(
            graph_node_input=graph_node_input_1,
            tf_layers_dict=tf_layers_dict,
        )
    input_tensor_scale = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
        if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
    input_tensor_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
    input_weights = tf_layers_dict[graph_node_input_4.name]['tf_node'] \
        if isinstance(graph_node_input_4, gs.Variable) else graph_node_input_4
    input_weights_scale = tf_layers_dict[graph_node_input_5.name]['tf_node'] \
        if isinstance(graph_node_input_5, gs.Variable) else graph_node_input_5
    input_weights_zero_point = tf_layers_dict[graph_node_input_6.name]['tf_node'] \
        if isinstance(graph_node_input_6, gs.Variable) else graph_node_input_6
    y_scale = tf_layers_dict[graph_node_input_7.name]['tf_node'] \
        if isinstance(graph_node_input_7, gs.Variable) else graph_node_input_7
    y_zero_point = tf_layers_dict[graph_node_input_8.name]['tf_node'] \
        if isinstance(graph_node_input_8, gs.Variable) else graph_node_input_8
    input_bias = tf_layers_dict[graph_node_input_9.name]['tf_node'] \
        if isinstance(graph_node_input_9, gs.Variable) else graph_node_input_9
    output_quant_dtype = y_zero_point.dtype

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

    # Generation of TF OP

    # Dequantize variables to float32
    if input_is_dequantized:
        input_tensor = tf.cast(input_tensor, tf.float32)
    else:
        input_tensor = _dequantize_tensor(
            base=input_tensor,
            zero_point=input_tensor_zero_point,
            scale=input_tensor_scale,
        )
    input_weights = _dequantize_weights(
        base=input_weights,
        zero_point=input_weights_zero_point,
        scale=input_weights_scale,
    )

    # if bias is defined save it here
    if input_bias is not None:
        input_bias = tf.cast(input_bias, tf.float32)
        input_bias_scale = tf.multiply(
            tf.cast(input_tensor_scale, tf.float32),
            tf.cast(input_weights_scale, tf.float32),
        )
        input_bias = tf.multiply(input_bias, input_bias_scale)

    # Workaround to avoid as many conversion failures as possible
    # for models with useless Transpose immediately before them.
    # If the input geometry of the ONNX and the input geometry of the TF model match,
    # the input geometry on the TF model side is forcibly transposed to NWC/NHWC/NDHWC.
    # However, if all dimensions in C* axes are the same value, skip this process.
    input_tensor_shape = input_tensor.shape
    onnx_input_shape = [
        dim if isinstance(dim, int) else None for dim in graph_node.inputs[0].shape
    ] if graph_node.inputs[0].shape is not None else None
    tf_input_shape = [
        dim if isinstance(dim, int) else None for dim in input_tensor_shape
    ]
    if onnx_input_shape is not None \
        and len(onnx_input_shape) > 1 and len(tf_input_shape) > 1 \
        and onnx_input_shape == tf_input_shape \
        and not input_nhwc_assumed:

        shape_for_judging_skip = [
            dim if dim is not None else INF_INDEX_VALUE for dim in onnx_input_shape[1:]
        ]
        if shape_for_judging_skip.count(shape_for_judging_skip[0]) != len(shape_for_judging_skip):
            if len(onnx_input_shape) == 3:
                # 1D
                input_tensor = transpose_with_flexing_deterrence(
                    input_tensor=input_tensor,
                    perm=[0,2,1],
                    **kwargs,
                )
            elif len(onnx_input_shape) == 4:
                # 2D
                input_tensor = transpose_with_flexing_deterrence(
                    input_tensor=input_tensor,
                    perm=[0,2,3,1],
                    **kwargs,
                )
            elif len(onnx_input_shape) == 5:
                # 3D
                input_tensor = transpose_with_flexing_deterrence(
                    input_tensor=input_tensor,
                    perm=[0,2,3,4,1],
                    **kwargs,
                )

    input_tensor_shape = input_tensor.shape
    input_tensor_rank = len(input_tensor_shape)
    spatial_size = input_tensor_rank - 2
    input_weights_shape = input_weights.shape

    def _to_static_int_dim(dim):
        if isinstance(dim, int):
            return int(dim)
        try:
            if dim is not None:
                return int(dim)
        except Exception:
            pass
        return None

    # Additional layout guard:
    # Some quantized model branches (e.g. split/reshape/transpose chains with
    # unknown ONNX shapes) may reach QLinearConv with NCHW-like tensors even
    # though tf.nn.convolution/depthwise_conv2d expects NHWC.
    #
    # Detect and fix only when the channel mismatch is unambiguous:
    #   input NHWC-channel axis does not match expected in-channels
    #   input NCHW-channel axis matches expected in-channels
    #
    # expected_in_channels:
    #   = (weights HWIO input-channel dim) * group
    # This covers both normal/grouped and depthwise-like cases.
    if input_tensor_rank == 4 and len(input_weights_shape) == 4:
        group = graph_node.attrs.get('group', 1)
        group = _to_static_int_dim(group)
        if group is None or group <= 0:
            group = 1
        tf_in_ch = _to_static_int_dim(input_tensor_shape[-1])
        nchw_in_ch = _to_static_int_dim(input_tensor_shape[1])
        filt_in_ch_per_group = _to_static_int_dim(input_weights_shape[-2])
        expected_in_ch = None
        if filt_in_ch_per_group is not None:
            expected_in_ch = int(filt_in_ch_per_group) * int(group)
        if tf_in_ch is not None and nchw_in_ch is not None and expected_in_ch is not None:
            if tf_in_ch != expected_in_ch and nchw_in_ch == expected_in_ch:
                input_tensor = transpose_with_flexing_deterrence(
                    input_tensor=input_tensor,
                    perm=[0, 2, 3, 1],
                    **kwargs,
                )
                input_tensor_shape = input_tensor.shape
    auto_pad = graph_node.attrs.get('auto_pad', 'NOTSET')
    dilations = graph_node.attrs.get('dilations', [1] * spatial_size)
    group = graph_node.attrs.get('group', 1)
    pads = graph_node.attrs.get('pads', [0, 0] * spatial_size)
    strides = graph_node.attrs.get('strides', [1] * spatial_size)

    """
    Conv1D
    Conv2D
    Conv3D
    DepthwiseConv2D
    SeparableConv2D
    """
    # Check auto_pad nonexistent or NOTSET first
    pad_mode = 'VALID'
    pads_axes_opposite_same = True
    for axis_begin, axis_end in zip(pads[0:spatial_size:1], pads[spatial_size::1]):
        if axis_begin != axis_end:
            pads_axes_opposite_same = False
            break

    if auto_pad == 'NOTSET':
        if pads_axes_opposite_same \
            and input_tensor_rank >=2 \
            and graph_node.inputs[0].shape is not None \
            and output_tensor_shape is not None \
            and graph_node.inputs[0].shape[2:] == output_tensor_shape[2:]:
            pad_mode = "SAME"
        elif pads != [0, 0] * spatial_size:
            input_tensor = get_padding_as_op(
                x=input_tensor,
                pads=pads,
            )
            pad_mode = 'VALID'
        else:
            pad_mode = 'VALID'
    # Then we use auto_pad to setup pad_mode
    elif auto_pad == "SAME_UPPER":
        pad_mode = "SAME"
    elif auto_pad == "VALID":
        pad_mode = "VALID"
    elif auto_pad == "SAME_LOWER":
        error_msg = f'' +\
            Color.RED(f'ERROR:') + ' ' +\
            f'Invalid auto_pad attribute: {auto_pad}'
        print(error_msg)
        assert False, error_msg
    else:
        error_msg = f'' +\
            Color.RED(f'ERROR:') + ' ' +\
            f'Invalid auto_pad attribute: {auto_pad}'
        print(error_msg)
        assert False, error_msg


    depthwise = (input_tensor_rank == 4 and len(input_weights_shape) == 4 and group != 1 and not (None in input_weights_shape))
    if depthwise and input_tensor_shape[-1] != None:
        depthwise = bool(group == input_tensor_shape[-1])

    if depthwise is True:
        depthwise_filter_shape = list(input_weights_shape[0:2]) + [-1, input_weights_shape[3] // group]
        input_weights = tf.reshape(input_weights, depthwise_filter_shape)

    # Conv
    tf_op_type = None
    if not depthwise:
        if group == 1:
            # Conv1D, Conv2D, Conv3D - No Bias
            conv_node = \
                tf.nn.convolution(
                    input=input_tensor,
                    filters=input_weights,
                    strides=strides,
                    padding=pad_mode,
                    dilations=dilations,
                )
            tf_op_type = tf.nn.convolution
        else:
            # SeparableConv
            input_tensor_splits = cast(List[Any], tf.split(input_tensor, num_or_size_splits=group, axis=-1))
            weight_splits = cast(List[Any], tf.split(input_weights, num_or_size_splits=group, axis=-1))
            conv_node = \
                tf.concat(
                    values=[
                        tf.nn.convolution(
                            input=input_tensor_split,
                            filters=weight_split,
                            padding=pad_mode,
                            strides=strides,
                            dilations=dilations,
                        ) for (input_tensor_split, weight_split) in zip(input_tensor_splits, weight_splits)
                    ],
                    axis=-1
                )
            tf_op_type = tf.nn.convolution
    else:
        # DepthwiseConv2D
        strides = [1] + strides + [1]
        conv_node = \
            tf.nn.depthwise_conv2d(
                input=input_tensor,
                filter=input_weights,
                padding=pad_mode,
                strides=strides,
                dilations=dilations,
            )
        tf_op_type = tf.nn.depthwise_conv2d

    # Add bias to the convolution (float)
    if input_bias is not None:
        conv_node = tf.add(
            x=conv_node,
            y=input_bias,
        )

    # quantize then dequantize to float32
    y_scale = tf.cast(y_scale, tf.float32)
    y_zero_point = tf.cast(y_zero_point, tf.float32)
    quantized = tf.add(tf.round(tf.divide(conv_node, y_scale)), y_zero_point)
    qmin, qmax = _get_qmin_qmax(output_quant_dtype)
    if qmin is not None and qmax is not None:
        quantized = tf.clip_by_value(quantized, qmin, qmax)
    dequantized = tf.multiply(tf.subtract(quantized, y_zero_point), y_scale)

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

    # Generation of Debug Info
    tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
        make_tf_node_info(
            node_info={
                'tf_op_type': tf_op_type,
                'tf_inputs': {
                    'input': input_tensor,
                    'weights': input_weights,
                    'bias': input_bias,
                    'strides': strides,
                    'dilations': dilations,
                    'padding': pad_mode,
                    'group': group,
                },
                'tf_outputs': {
                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
                },
            }
        )
