from typing import Any, List, Optional, cast
import math
import random
random.seed(0)
import numpy as np
np.random.seed(0)
import tensorflow as tf
import onnx2tf.gs as gs
from tensorflow.python.keras.layers import (
    AveragePooling1D,
    AveragePooling2D,
    AveragePooling3D,
)
from onnx2tf.utils.logging import *
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,
    calc_tf_pooling_pads,
    calc_extra_padding_with_ceil,
    transpose_with_flexing_deterrence,
)

INF_INDEX_VALUE: int = 4294967296


def summarize_multiplier(arr):
    """
    summarize consecutive numbers in average multiplier
    Parameters
    ----------
    arr: List
        average multiplier

    Returns
    -------
    summary: List
        summarized average multiplier
    """
    summary = []
    for sub_arr in arr:
        sub_summary = []
        i = 0
        while i < len(sub_arr):
            start_index = i
            value = sub_arr[i]
            while i < len(sub_arr) - 1 and sub_arr[i] == sub_arr[i+1]:
                i += 1
            end_index = i
            sub_summary.append((start_index, end_index, value))
            i += 1
        summary.append(sub_summary)
    return summary


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

    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 = get_constant_or_variable(
        graph_node.inputs[0],
        before_op_output_shape_trans,
    )
    graph_node_output: gs.Variable = graph_node.outputs[0]
    shape = graph_node_output.shape
    dtype = graph_node_output.dtype

    input_tensor = tf_layers_dict[graph_node_input.name]['tf_node'] \
        if isinstance(graph_node_input, gs.Variable) else graph_node_input
    input_tensor_shape = input_tensor.shape

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

    # 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 the NWC or NHWC or NDHWC format.
    # However, if all dimensions of CW or CHW or CDHW have the same value,
    # the forced transposition process is skipped because it may destroy the structure of the model.
    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:

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

    auto_pad = graph_node.attrs.get('auto_pad', 'NOTSET')
    ceil_mode = bool(graph_node.attrs.get('ceil_mode', 0))
    count_include_pad = bool(graph_node.attrs.get('count_include_pad', 0))
    kernel_shape = graph_node.attrs['kernel_shape']
    spatial_size = len(kernel_shape)
    dilations = graph_node.attrs.get('dilations', [1] * spatial_size)
    pads = graph_node.attrs.get('pads', [0] * spatial_size * 2)
    strides = graph_node.attrs.get('strides', [1] * spatial_size)

    input_tensor_shape = input_tensor.shape.as_list()
    is_known_shape = None not in input_tensor_shape[1:]
    extra_pads = [0] * spatial_size
    average_multiplier = None

    # default tensorflow action is 'SAME_UPPER' mode (extra padding in the end for odd numbers)
    # explicit pad layer is added for tensorflow incompatible cases
    tf_pad_mode = 'VALID'
    is_explicit_padding = False
    tf_pads = calc_tf_pooling_pads(
        input_shape=input_tensor_shape,
        kernel=kernel_shape,
        strides=strides,
        input_tensor=input_tensor
    )

    build_tf_pads_from_begin_end_fn: Optional[Any] = None
    calc_extra_padding_with_ceil_dynamic_fn: Optional[Any] = None
    if not is_known_shape:
        def build_tf_pads_from_begin_end(pads_begin, pads_end):
            pads_begin = tf.cast(pads_begin, tf.int32)
            pads_end = tf.cast(pads_end, tf.int32)
            spatial_pads = tf.stack([pads_begin, pads_end], axis=1)
            return tf.concat(
                [
                    tf.zeros((1, 2), dtype=tf.int32),
                    spatial_pads,
                    tf.zeros((1, 2), dtype=tf.int32),
                ],
                axis=0,
            )

        def calc_extra_padding_with_ceil_dynamic(input_tensor, pads, kernel_shape, dilations, strides):
            input_shape = tf.shape(input_tensor)
            input_rank = tf.size(input_shape)
            input_spatial = tf.slice(
                input_shape,
                begin=[1],
                size=[tf.subtract(input_rank, 2)],
            )
            pads_begin = tf.constant(pads[:len(pads) // 2], dtype=tf.int32)
            pads_end = tf.constant(pads[len(pads) // 2:], dtype=tf.int32)
            pads_along_axis = tf.add(pads_begin, pads_end)
            k = tf.constant(kernel_shape, dtype=tf.int32)
            d = tf.constant(dilations, dtype=tf.int32)
            s = tf.constant(strides, dtype=tf.int32)

            numerator = tf.subtract(
                tf.subtract(
                    tf.add(input_spatial, pads_along_axis),
                    tf.multiply(d, tf.subtract(k, 1)),
                ),
                1,
            )
            output_spatial = tf.cast(
                tf.math.ceil(
                    tf.add(
                        tf.math.divide(
                            tf.cast(numerator, tf.float32),
                            tf.cast(s, tf.float32),
                        ),
                        1.0,
                    )
                ),
                tf.int32,
            )
            last_stride_starts = tf.multiply(tf.subtract(output_spatial, 1), s)
            last_stride_validity = tf.less(
                last_stride_starts,
                tf.add(input_spatial, pads_begin),
            )

            extra_pads = tf.where(
                last_stride_validity,
                tf.subtract(
                    tf.add(
                        tf.add(
                            last_stride_starts,
                            tf.multiply(tf.subtract(k, 1), d),
                        ),
                        1,
                    ),
                    tf.add(input_spatial, pads_along_axis),
                ),
                tf.zeros_like(input_spatial),
            )
            return extra_pads

        build_tf_pads_from_begin_end_fn = build_tf_pads_from_begin_end
        calc_extra_padding_with_ceil_dynamic_fn = calc_extra_padding_with_ceil_dynamic

        def compute_output_spatial_shape_from_tensor(input_tensor, pads, kernel_shape, dilations, strides, ceil_mode=False):
            input_shape = tf.shape(input_tensor)  # Get dynamic shape
            input_rank = tf.size(input_shape)
            input_spatial = tf.slice(
                input_shape,
                begin=[1],
                size=[tf.subtract(input_rank, 2)],
            )  # Extract spatial dimensions only (NHWC format)

            pad_begin = pads[:len(pads) // 2]
            pad_end = pads[len(pads) // 2:]

            round_func = tf.math.ceil if ceil_mode else tf.math.floor

            output_spatial = []
            input_spatial_elems = cast(List[Any], tf.unstack(input_spatial))
            for i, pb, pe, k, d, s in zip(input_spatial_elems, pad_begin, pad_end, kernel_shape, dilations, strides):
                i = tf.cast(i, tf.float32)
                pb = tf.constant(pb, dtype=tf.float32)
                pe = tf.constant(pe, dtype=tf.float32)
                k = tf.constant(k, dtype=tf.float32)
                d = tf.constant(d, dtype=tf.float32)
                s = tf.constant(s, dtype=tf.float32)

                numerator = tf.subtract(
                    tf.subtract(
                        tf.add(tf.add(i, pb), pe),
                        tf.multiply(d, tf.subtract(k, 1)),
                    ),
                    1,
                )
                raw_output = tf.add(tf.math.divide(numerator, s), 1)
                output_dim = tf.cast(round_func(raw_output), tf.int32)
                output_spatial.append(output_dim)

            return output_spatial

        output_spatial_shape = compute_output_spatial_shape_from_tensor(
            input_tensor=input_tensor,
            pads=pads,
            kernel_shape=kernel_shape,
            dilations=dilations,
            strides=strides,
            ceil_mode=ceil_mode
        )
    else:
        func = math.ceil if ceil_mode else math.floor
        output_spatial_shape = [
            func((i + pb + pe - d * (k - 1) - 1) / s + 1)
            for i, pb, pe, k, d, s in zip(input_tensor_shape[1:-1], pads[:len(pads) // 2], pads[len(pads) // 2:], kernel_shape, dilations, strides)
        ]

    # onnx padding value is ignored if auto_pad is not 'NOTSET'
    if auto_pad == 'NOTSET':

        # check if onnx padding is same with tensorflow padding mode 'SAME'
        # this is to avoid flex operations since tflite builtin pooling operator does not support manual padding
        if is_known_shape and pads != [0] * spatial_size * 2 and tf_pads == pads:
            auto_pad = 'SAME_UPPER'
            tf_pad_mode = 'SAME'

        else:
            auto_pad = 'VALID'
            is_explicit_padding = True

            # extra padding to end side (right, bottom) may be needed when ceil_mode is True
            # this extra padding should not be counted as padding when count_include_pad is True
            if not is_known_shape:
                pads_begin = tf.constant(pads[:len(pads) // 2], dtype=tf.int32)
                pads_end = tf.constant(pads[len(pads) // 2:], dtype=tf.int32)
                if ceil_mode:
                    assert calc_extra_padding_with_ceil_dynamic_fn is not None
                    extra_pads = calc_extra_padding_with_ceil_dynamic_fn(
                        input_tensor=input_tensor,
                        pads=pads,
                        kernel_shape=kernel_shape,
                        dilations=dilations,
                        strides=strides,
                    )
                    pads_end = pads_end + extra_pads
                assert build_tf_pads_from_begin_end_fn is not None
                tf_pads = build_tf_pads_from_begin_end_fn(pads_begin, pads_end)
            else:
                if ceil_mode:
                    extra_pads = \
                        calc_extra_padding_with_ceil(
                            input_shape=input_tensor_shape[1:-1],
                            kernel=kernel_shape,
                            pads=pads,
                            dilations=dilations,
                            strides=strides,
                        )
                    pads = pads[:len(pads) // 2] + [p + e for p, e in zip(pads[len(pads) // 2:], extra_pads)]

                tf_pads = pads

    elif auto_pad == 'SAME_UPPER':
        tf_pad_mode = 'SAME'

    elif auto_pad == 'SAME_LOWER':
        is_explicit_padding = True

    elif auto_pad == 'VALID':
        tf_pads = [0] * spatial_size * 2

    else:
        error_msg = Color.RED(f'ERROR:') + ' ' + \
                    f'Wrong auto_pad parameter in AveragePool: {auto_pad}.'
        raise ValueError(error_msg)

    # count nonzero elements in kernel each strides for the case count_include_pad is False
    non_zero_counts = []

    if not is_known_shape:
        def compute_non_zero_counts_loop(input_tensor, output_spatial_shape, kernel_shape, dilations, strides, pads):
            
            counts_list = []
            
            for dim in range(len(kernel_shape)):
                k = kernel_shape[dim]
                counts = [k] * (output_spatial_shape[dim].numpy() if tf.is_tensor(output_spatial_shape[dim]) and hasattr(output_spatial_shape[dim], 'numpy') else output_spatial_shape[dim])
                counts_list.append(counts)
                
            return counts_list

        non_zero_counts = compute_non_zero_counts_loop(
            input_tensor=input_tensor,
            output_spatial_shape=output_spatial_shape,
            kernel_shape=kernel_shape,
            dilations=dilations,
            strides=strides,
            pads=pads
        )
    else:
        for input_spatial_shape, output_size, kernel, dilation, stride, pads_begin, pads_end \
                in zip(input_tensor_shape[1:-1], output_spatial_shape, kernel_shape,
                       dilations, strides, pads[:len(pads) // 2], pads[len(pads) // 2:]):
            sample_target = np.concatenate([
                np.zeros(pads_begin),
                np.ones(input_spatial_shape),
                np.zeros(pads_end)]
            )
            sample_kernel = np.zeros((kernel - 1) * dilation + 1)
            sample_kernel[::dilation] = 1

            counts = []
            for i in range(output_size):
                start = i * stride
                stride_target = sample_target[start:start+len(sample_kernel)]
                # pad target to match size
                stride_target = np.concatenate([stride_target, np.zeros(len(sample_kernel) - len(stride_target))])
                counts.extend(np.convolve(stride_target, sample_kernel, mode='valid'))

            non_zero_counts.append(counts)

    if not is_known_shape:
        need_multiplier = False  # Default to False for dynamic tensors to avoid errors
    else:
        need_multiplier = len(set([i for sublist in non_zero_counts for i in sublist])) != 1

    # default tensorflow option for count_include_pad is True and cannot control
    # average value should be compensated in cases below
    # 1. when extra padding layer is added and count_include_pad is False
    # 2. when extra padding layer is not added and count_include_pad is True
    # 3. when last stride has extra padding due to ceil_mode and count_include_pad is True
    if is_explicit_padding:
        warn(
            f'Tensorflow incompatible padding detected. ' \
            f'Extra pad layer is inserted automatically. '
        )
        if not is_known_shape:
            padded_tensor = tf.pad(
                tensor=input_tensor,
                paddings=cast(Any, tf_pads),
                mode='CONSTANT',
            )
            if input_tensor_shape is not None and len(input_tensor_shape) == spatial_size + 2:
                # Preserve known batch/channel dims since dynamic paddings erase shape info.
                padded_tensor = tf.ensure_shape(
                    padded_tensor,
                    [input_tensor_shape[0]] + [None] * spatial_size + [input_tensor_shape[-1]],
                )
        else:
            if auto_pad == 'SAME_LOWER':
                # switch the order of pads
                tf_pads_list = cast(List[int], tf_pads)
                tf_pads = [i for tup in zip(tf_pads_list[len(tf_pads_list) // 2:], tf_pads_list[:len(tf_pads_list) // 2]) for i in tup]

            if not count_include_pad and need_multiplier:
                average_multiplier = []
                for k, non_zero_count in zip(kernel_shape, non_zero_counts):
                    multiplier = [k / n if n != 0 else 1 for n in non_zero_count]
                    average_multiplier.append(multiplier)

            # convert to tensorflow padding format
            tf_pads_list = cast(List[int], tf_pads)
            tf_pads = \
                [[0, 0]] + \
                [list(i) for i in zip(tf_pads_list[:len(tf_pads_list) // 2], tf_pads_list[len(tf_pads_list) // 2:])] + \
                [[0, 0]]

            if spatial_size == 1 and kernel_shape[0] > input_tensor_shape[1]:
                padded_tensor = input_tensor
            else:
                padded_tensor = tf.pad(
                    tensor=input_tensor,
                    paddings=cast(Any, tf_pads),
                    mode='CONSTANT',
                )

    else:
        padded_tensor = input_tensor

        if count_include_pad and need_multiplier:
            average_multiplier = []
            for k, non_zero_count in zip(kernel_shape, non_zero_counts):
                multiplier = [n / k for n in non_zero_count]
                average_multiplier.append(multiplier)

    if count_include_pad and is_known_shape and extra_pads != [0] * spatial_size:
        # extra padding in last stride should not be included in averaging
        if average_multiplier is None:
            average_multiplier = []
            extra_pads_list = cast(List[int], extra_pads)
            for k, non_zero_count, extra_pad in zip(kernel_shape, non_zero_counts, extra_pads_list):
                multiplier = [1 for _ in non_zero_count]
                multiplier[-1] = k / (k - extra_pad)
                average_multiplier.append(multiplier)
        else:
            for i in range(len(kernel_shape)):
                k = kernel_shape[i]
                extra_pad = extra_pads[i]
                average_multiplier[i][-1] = k / (k - extra_pad)

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

    # Generation of TF OP
    tf_op_type = None
    if len(kernel_shape) == 1:
        dim1 = padded_tensor.shape[1]
        if isinstance(dim1, int) and kernel_shape[0] > dim1:
            pool_dim = int(dim1)
            pooled_tensor = AveragePooling1D(
                pool_size=pool_dim,
                strides=pool_dim,
                padding=tf_pad_mode.upper(),
            )(padded_tensor)
        else:
            pooled_tensor = AveragePooling1D(
                pool_size=kernel_shape,
                strides=strides,
                padding=tf_pad_mode.upper(),
            )(padded_tensor)
        tf_op_type = AveragePooling1D

    elif len(kernel_shape) == 2:
        pooled_tensor = AveragePooling2D(
            pool_size=kernel_shape,
            strides=strides,
            padding=tf_pad_mode.upper(),
        )(padded_tensor)
        tf_op_type = AveragePooling2D

    elif len(kernel_shape) == 3:
        pooled_tensor = AveragePooling3D(
            pool_size=kernel_shape,
            strides=strides,
            padding=tf_pad_mode.upper(),
        )(padded_tensor)
        tf_op_type = AveragePooling3D

    else:
        error_msg = f'' +\
            Color.RED(f'ERROR:') + ' ' +\
            f'AveragePool supports only 1D, 2D, and 3D. ' +\
            f'opname: {graph_node.name} Type: AveragePool{len(kernel_shape)}D'
        print(error_msg)
        raise AssertionError(error_msg)
    pooled_tensor = tf.convert_to_tensor(cast(Any, pooled_tensor))

    # Dynamic shape compensation for count_include_pad=False with explicit padding.
    # Use pooled mask to compute valid element counts per window.
    if not is_known_shape and is_explicit_padding and not count_include_pad:
        mask = tf.ones_like(input_tensor, dtype=pooled_tensor.dtype)
        if tf_pads is not None:
            if tf.is_tensor(tf_pads):
                mask = tf.pad(
                    tensor=mask,
                    paddings=cast(Any, tf_pads),
                    mode='CONSTANT',
                )
            elif tf_pads != [0] * spatial_size * 2:
                mask = tf.pad(
                    tensor=mask,
                    paddings=cast(Any, tf_pads),
                    mode='CONSTANT',
                )
        if len(kernel_shape) == 1:
            mask_pooled = AveragePooling1D(
                pool_size=kernel_shape,
                strides=strides,
                padding=tf_pad_mode.upper(),
            )(mask)
        elif len(kernel_shape) == 2:
            mask_pooled = AveragePooling2D(
                pool_size=kernel_shape,
                strides=strides,
                padding=tf_pad_mode.upper(),
            )(mask)
        else:
            mask_pooled = AveragePooling3D(
                pool_size=kernel_shape,
                strides=strides,
                padding=tf_pad_mode.upper(),
            )(mask)
        mask_pooled = tf.convert_to_tensor(cast(Any, mask_pooled))
        kernel_volume = float(np.prod(kernel_shape))
        count_valid = tf.multiply(
            mask_pooled,
            tf.cast(kernel_volume, dtype=mask_pooled.dtype),
        )
        multiplier = tf.math.divide_no_nan(
            tf.cast(kernel_volume, dtype=mask_pooled.dtype),
            count_valid,
        )
        pooled_tensor = tf.multiply(pooled_tensor, multiplier)

    # tensorflow average pooling needs extra process to get same output with onnx
    # https://github.com/PINTO0309/onnx2tf/issues/124
    if average_multiplier is not None:
        warn(
            f'Tensorflow incompatible action detected. ' \
            f'Some additional layers are inserted to reproduce same output. ' \
            f'Please refer to the following link for more information: ' \
            f'https://github.com/PINTO0309/onnx2tf/issues/124'
        )

        average_multiplier = summarize_multiplier(average_multiplier)

        for i, multiplier in enumerate(average_multiplier, start=1):
            slice_list = [slice(None) for _ in range(spatial_size * 2)]
            multiplied_slices = []

            for m in multiplier:
                start, stop, value = m
                slice_list[i] = slice(start, stop + 1)
                sliced_tensor = cast(Any, pooled_tensor)[tuple(slice_list)]
                multiplied_slices.append(tf.multiply(sliced_tensor, value))

            pooled_tensor = tf.concat(multiplied_slices, axis=i)

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

    # 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_op_type,
                'tf_inputs': {
                    'x': input_tensor,
                    'pool_size': kernel_shape,
                    'strides': strides,
                    'padding': tf_pads if tf_pad_mode != 'same' else tf_pad_mode,
                },
                'tf_outputs': {
                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
                },
            }
        )
