from typing import Any, Optional, cast
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,
    convert_axis,
    print_node_info,
    inverted_operation_enable_disable,
    make_tf_node_info,
    process_neg_idx,
    get_replacement_parameter,
    pre_process_transpose,
    post_process_transpose,
)
from onnx2tf.utils.logging import *
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,
):
    """GatherND

    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,
    )
    # Indices must not be layout-transposed.
    graph_node_input_2 = get_constant_or_variable(
        graph_node.inputs[1],
        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
    indices_tensor = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
        if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2

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

    tensor_rank = len(input_tensor.shape)

    batch_dims = graph_node.attrs.get('batch_dims', 0)
    batch_dims = convert_axis(
        axis=batch_dims,
        tensor_rank=tensor_rank,
        before_op_output_shape_trans=before_op_output_shape_trans,
    )

    replace_gathernd_to_pseudo_gathernd = "gathernd" in kwargs['replace_to_pseudo_operators']

    # If params is transposed, adjust indices to match the transposed layout.
    op_rep_params = kwargs.get('op_rep_params', [])
    params_perm = None
    indices_perm_specified = False
    for op_rep_param in op_rep_params:
        if op_rep_param['param_target'] == 'inputs' and op_rep_param['param_name'] == graph_node.inputs[0].name:
            params_perm = op_rep_param.get('pre_process_transpose_perm', None)
        if op_rep_param['param_target'] == 'inputs' and op_rep_param['param_name'] == graph_node.inputs[1].name:
            if op_rep_param.get('pre_process_transpose_perm', None) is not None:
                indices_perm_specified = True
    if params_perm is not None and not indices_perm_specified:
        # Only handle standard layout swaps that keep batch dims at the front.
        if batch_dims <= len(params_perm) \
            and list(params_perm[:batch_dims]) == list(range(batch_dims)):
            perm_tail = [p - batch_dims for p in params_perm if p >= batch_dims]
            try:
                if isinstance(indices_tensor, np.ndarray):
                    if indices_tensor.shape and indices_tensor.shape[-1] == len(perm_tail):
                        indices_tensor = indices_tensor[..., perm_tail]
                else:
                    idx_last = indices_tensor.shape[-1] if indices_tensor.shape is not None else None
                    if idx_last is None or idx_last == len(perm_tail):
                        indices_tensor = tf.gather(indices_tensor, perm_tail, axis=-1)
            except Exception:
                pass

    # 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 = input_tensor \
        if not isinstance(input_tensor, np.ndarray) \
            else tf.convert_to_tensor(input_tensor)
    indices_tensor = indices_tensor \
        if not isinstance(indices_tensor, np.ndarray) \
            else tf.convert_to_tensor(indices_tensor)
    indices_tensor = tf.convert_to_tensor(cast(Any, indices_tensor))

    # Complex GatherND -> Simple GatherND
    simple_gathernd = False
    # Verify if negative numbers need to be converted to positive numbers
    static_indices: Optional[np.ndarray] = None
    if isinstance(indices_tensor, np.ndarray):
        static_indices = np.asarray(indices_tensor)
    elif tf.is_tensor(indices_tensor):
        resolved_indices = tf.get_static_value(indices_tensor)
        if resolved_indices is not None:
            static_indices = np.asarray(resolved_indices)
    if static_indices is not None:
        flatten_indices_tensor = static_indices.flatten()
        simple_gathernd = bool(np.sum(np.where(flatten_indices_tensor < 0, 1, 0)) == 0)
    elif isinstance(indices_tensor, int):
        simple_gathernd = bool(indices_tensor >= 0)
    else:
        simple_gathernd = False

    if not simple_gathernd:
        input_tensor_shape = getattr(input_tensor, 'shape', None)
        indices_tensor_shape = getattr(indices_tensor, 'shape', None)
        if input_tensor_shape is not None \
            and indices_tensor_shape is not None \
            and hasattr(input_tensor.shape, '__len__') \
            and hasattr(indices_tensor_shape, '__len__') \
            and len(input_tensor.shape) > 0 \
            and len(indices_tensor_shape) > 0 \
            and None not in input_tensor.shape \
            and None not in indices_tensor_shape:
            indices_tensor = process_neg_idx(
                data=input_tensor,
                indices=indices_tensor,
                batch_dims=batch_dims,
            )

    if not replace_gathernd_to_pseudo_gathernd:
        tf_layers_dict[graph_node_output.name]['tf_node'] = \
            tf.gather_nd(
                params=input_tensor \
                    if not isinstance(input_tensor, np.ndarray) \
                        else tf.convert_to_tensor(input_tensor),
                indices=indices_tensor \
                    if not isinstance(indices_tensor, np.ndarray) \
                        else tf.convert_to_tensor(indices_tensor),
                batch_dims=batch_dims,
                name=graph_node.name,
            )
    else:
        if batch_dims != 0:
            error(
                f'--replace_gathernd_to_pseudo_gathernd is supported only if batch_dims=0.'+
                f'graph_node.name: {graph_node.name}'
            )
            sys.exit(1)

        params_shape = tf.shape(input_tensor)
        idx_shape = tf.shape(indices_tensor)
        idx_dims = tf.shape(indices_tensor)[-1]
        gather_shape = params_shape[idx_dims:]
        params_flat = tf.reshape(
            input_tensor,
            tf.concat([[-1], gather_shape], axis=0),
        )
        axis_step = tf.math.cumprod(
            params_shape[:idx_dims],
            exclusive=True,
            reverse=True,
        )
        mul = tf.math.multiply(
            indices_tensor,
            tf.cast(
                axis_step,
                dtype= NUMPY_DTYPES_TO_TF_DTYPES[indices_tensor.dtype] \
                    if isinstance(indices_tensor.dtype, np.dtype) else indices_tensor.dtype,
            ),
        )
        indices_flat = tf.reduce_sum(
            mul,
            axis=-1,
        )
        result_flat = tf.gather(
            params_flat,
            indices_flat,
        )
        pseudo_gathernd = None
        indices_rank = indices_tensor.shape.rank
        if indices_rank is not None and indices_rank > 1:
            pseudo_gathernd = tf.reshape(
                result_flat,
                tf.concat([tf.shape(indices_tensor)[:-1], gather_shape], axis=0),
            )
        else:
            pseudo_gathernd = result_flat
        tf_layers_dict[graph_node_output.name]['tf_node'] = pseudo_gathernd

    # 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.gather_nd,
                'tf_inputs': {
                    'params': input_tensor,
                    'indices': indices_tensor,
                    'batch_dims': batch_dims,
                },
                'tf_outputs': {
                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
                },
            }
        )
