from typing import Any
import sys
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.logging import *


def _build_col2im_kernel(
    *,
    k_h,
    k_w,
    dilation_h,
    dilation_w,
    dtype,
):
    k_h = tf.cast(k_h, tf.int32)
    k_w = tf.cast(k_w, tf.int32)
    eff_k_h = tf.add(
        tf.multiply(tf.subtract(k_h, 1), dilation_h),
        1,
    )
    eff_k_w = tf.add(
        tf.multiply(tf.subtract(k_w, 1), dilation_w),
        1,
    )

    ky = tf.reshape(tf.repeat(tf.range(k_h), k_w), tf.stack([k_h, k_w]))
    kx = tf.reshape(tf.tile(tf.range(k_w), [k_h]), tf.stack([k_h, k_w]))

    positions = ky * dilation_h * eff_k_w + kx * dilation_w
    positions = tf.reshape(positions, [-1])
    one_hot = tf.one_hot(positions, depth=eff_k_h * eff_k_w, dtype=dtype)
    kernel = tf.reshape(one_hot, tf.stack([tf.multiply(k_h, k_w), eff_k_h, eff_k_w]))
    kernel = tf.transpose(kernel, [1, 2, 0])
    kernel = tf.expand_dims(kernel, axis=2)
    return kernel, eff_k_h, eff_k_w


def _maybe_align_col2im_input_layout(
    *,
    input_tensor: tf.Tensor,
    k_prod: tf.Tensor,
    out_hw: tf.Tensor,
) -> tf.Tensor:
    """
    Col2Im input is expected as [N, C*K, L] (K = k_h*k_w, L = out_h*out_w).
    Some graphs reach this op as [N, L, C*K] after global transpose propagation.
    When static hints are sufficient, swap axes to restore canonical layout.
    """
    static_shape = input_tensor.shape
    if static_shape is None or static_shape.rank != 3:
        return input_tensor
    d1 = static_shape[1]
    d2 = static_shape[2]
    if d1 is None or d2 is None:
        return input_tensor
    d1_int = d1 if isinstance(d1, int) else getattr(d1, 'value', None)
    d2_int = d2 if isinstance(d2, int) else getattr(d2, 'value', None)
    if d1_int is None or d2_int is None:
        return input_tensor

    k_prod_static = tf.get_static_value(k_prod)
    if k_prod_static is None:
        return input_tensor
    k_prod_int = int(k_prod_static)
    if k_prod_int <= 0:
        return input_tensor

    out_hw_static = tf.get_static_value(out_hw)
    out_hw_int = int(out_hw_static) if out_hw_static is not None else None

    as_nckl_valid = (d1_int % k_prod_int) == 0
    as_nlkc_valid = (d2_int % k_prod_int) == 0
    if out_hw_int is not None:
        as_nckl_valid = as_nckl_valid and d2_int == out_hw_int
        as_nlkc_valid = as_nlkc_valid and d1_int == out_hw_int

    if as_nlkc_valid and not as_nckl_valid:
        return tf.transpose(input_tensor, perm=[0, 2, 1])
    return input_tensor


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

    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=False,
    )
    graph_node_input_3 = get_constant_or_variable(
        graph_node.inputs[2],
        before_op_output_shape_trans=False,
    )
    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_1.name]['tf_node'] \
        if isinstance(graph_node_input_1, gs.Variable) else graph_node_input_1
    input_image_shape = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
        if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
    input_block_shape = tf_layers_dict[graph_node_input_3.name]['tf_node'] \
        if isinstance(graph_node_input_3, gs.Variable) else graph_node_input_3

    dilations = graph_node.attrs.get('dilations', [1, 1])
    pads = graph_node.attrs.get('pads', [0, 0, 0, 0])
    strides = graph_node.attrs.get('strides', [1, 1])

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

    # Pre-process transpose
    before_trans_shape = input_tensor.shape
    input_tensor = pre_process_transpose(
        value_before_transpose=input_tensor,
        param_target='inputs',
        param_name=graph_node.inputs[0].name,
        **kwargs,
    )
    after_trans_shape = input_tensor.shape
    if 'nhwc' in tf_layers_dict[graph_node_output.name].keys() \
        and tf_layers_dict[graph_node_output.name]['nhwc'] == True \
        and before_trans_shape != after_trans_shape:
        tf_layers_dict[graph_node_output.name].pop('nhwc')

    # Generation of TF OP
    original_dtype = input_tensor.dtype
    compute_dtype = original_dtype
    if original_dtype not in (tf.float16, tf.float32, tf.float64, tf.bfloat16):
        if original_dtype in (tf.complex64, tf.complex128):
            error('Col2Im does not support complex types in non-Flex implementation.')
            sys.exit(1)
        compute_dtype = tf.float32
        input_tensor = tf.cast(input_tensor, compute_dtype)

    input_image_shape = tf.cast(input_image_shape, tf.int32)
    input_block_shape = tf.cast(input_block_shape, tf.int32)

    if input_image_shape.shape is not None \
        and input_image_shape.shape.rank is not None \
        and input_image_shape.shape.rank != 1:
        error('Col2Im supports only 2D image_shape input.')
        sys.exit(1)

    if input_block_shape.shape is not None \
        and input_block_shape.shape.rank is not None \
        and input_block_shape.shape.rank != 1:
        error('Col2Im supports only 2D block_shape input.')
        sys.exit(1)

    k_h = tf.gather(input_block_shape, indices=0)
    k_w = tf.gather(input_block_shape, indices=1)
    h_img = tf.gather(input_image_shape, indices=0)
    w_img = tf.gather(input_image_shape, indices=1)

    stride_h, stride_w = strides
    dilation_h, dilation_w = dilations
    pad_top, pad_left, pad_bottom, pad_right = pads

    kernel, eff_k_h, eff_k_w = _build_col2im_kernel(
        k_h=k_h,
        k_w=k_w,
        dilation_h=dilation_h,
        dilation_w=dilation_w,
        dtype=compute_dtype,
    )

    h_pad = h_img + pad_top + pad_bottom
    w_pad = w_img + pad_left + pad_right

    out_h = tf.math.floordiv(h_pad - eff_k_h, stride_h) + 1
    out_w = tf.math.floordiv(w_pad - eff_k_w, stride_w) + 1

    k_prod = k_h * k_w
    out_hw = out_h * out_w
    input_tensor = tf.convert_to_tensor(input_tensor)
    input_tensor = _maybe_align_col2im_input_layout(
        input_tensor=input_tensor,
        k_prod=k_prod,
        out_hw=out_hw,
    )

    input_shape = tf.shape(input_tensor)
    n = tf.gather(input_shape, indices=0)
    ck = tf.gather(input_shape, indices=1)
    c = tf.math.floordiv(ck, k_prod)

    cols = tf.reshape(
        input_tensor,
        tf.stack([n, c, k_prod, out_h, out_w]),
    )
    cols = tf.transpose(cols, [0, 1, 3, 4, 2])
    cols = tf.reshape(cols, tf.stack([n * c, out_h, out_w, k_prod]))

    output_shape = tf.stack([n * c, h_pad, w_pad, 1])
    output = tf.nn.conv2d_transpose(
        cols,
        kernel,
        output_shape=output_shape,
        strides=[1, stride_h, stride_w, 1],
        padding='VALID',
    )

    output = tf.reshape(output, tf.stack([n, c, h_pad, w_pad]))
    output = tf.transpose(output, [0, 2, 3, 1])

    output = tf.slice(
        output,
        tf.stack([0, pad_top, pad_left, 0]),
        tf.stack([-1, h_img, w_img, -1]),
    )

    if output.dtype != original_dtype:
        output = tf.cast(output, original_dtype)

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

    # Post-process transpose
    before_trans_shape = tf_layers_dict[graph_node_output.name]['tf_node'].shape
    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,
    )
    after_trans_shape = tf_layers_dict[graph_node_output.name]['tf_node'].shape
    if 'nhwc' in tf_layers_dict[graph_node_output.name].keys() \
        and tf_layers_dict[graph_node_output.name]['nhwc'] == True \
        and before_trans_shape != after_trans_shape:
        tf_layers_dict[graph_node_output.name].pop('nhwc')

    # Generation of Debug Info
    tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
        make_tf_node_info(
            node_info={
                'tf_op_type': 'Col2Im',
                'tf_inputs': {
                    'input': input_tensor,
                    'image_shape': input_image_shape,
                    'block_shape': input_block_shape,
                    'dilations': dilations,
                    'pads': pads,
                    'strides': strides,
                },
                'tf_outputs': {
                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
                },
            }
        )
