from typing import Any
import random
random.seed(0)
import numpy as np
np.random.seed(0)
import tensorflow as tf
from tensorflow.python.keras.layers import PReLU
import onnx2tf.gs as gs
from onnx2tf.utils.common_functions import (
    get_constant_or_variable,
    print_node_info,
    inverted_operation_enable_disable,
    pre_explicit_broadcast,
    explicit_broadcast,
    make_tf_node_info,
    get_replacement_parameter,
    pre_process_transpose,
    post_process_transpose,
)


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

    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,
    )
    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_tensor_shape = input_tensor.shape
    input_tensor_rank = len(input_tensor_shape)
    graph_node_input_2 = get_constant_or_variable(
        graph_node.inputs[1],
        before_op_output_shape_trans \
            if graph_node.inputs[1].shape is not None \
                and input_tensor_rank == len(graph_node.inputs[1].shape) else False,
    )
    slope = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
        if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2

    replace_prelu_to_pseudo_prelu = "prelu" in kwargs['replace_to_pseudo_operators']

    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,
        'nhwc': tf_layers_dict[graph_node_input_1.name]['nhwc'] \
            if isinstance(graph_node_input_1, gs.Variable) \
                and 'nhwc' in tf_layers_dict[graph_node_input_1.name].keys() else False
    }

    input_tensor, slope = \
        pre_explicit_broadcast(
            input_tensor_1=input_tensor,
            input_tensor_2=slope,
        )
    input_tensor, slope = \
        explicit_broadcast(
            const_or_var_1=input_tensor,
            const_or_var_2=slope,
            graph_node=graph_node,
            tf_layers_dict= tf_layers_dict,
        )

    slope = tf.convert_to_tensor(slope)

    # 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')

    # Align slope layout to the runtime input tensor layout.
    # Handles cases like slope=[C,1,1] with input=[N,H,W,C].
    input_shape = input_tensor.shape
    slope_shape = slope.shape if hasattr(slope, 'shape') else None
    if input_shape is not None and slope_shape is not None:
        input_rank = len(input_shape)
        slope_rank = len(slope_shape)
        if input_rank >= 3 and slope_rank >= 1:
            slope_shape_list = list(slope_shape)
            non_one_dims = [
                int(dim) for dim in slope_shape_list
                if isinstance(dim, (int, np.integer)) and dim != 1
            ]
            if slope_rank == 1 and isinstance(slope_shape_list[0], (int, np.integer)):
                non_one_dims = [int(slope_shape_list[0])]
            if len(non_one_dims) == 1:
                slope_channel_dim = non_one_dims[0]
                channel_first_dim = input_shape[1] if input_rank >= 2 else None
                channel_last_dim = input_shape[-1]
                target_shape = None
                if isinstance(channel_last_dim, (int, np.integer)) and int(channel_last_dim) == slope_channel_dim:
                    target_shape = [1 for _ in range(input_rank - 2)] + [slope_channel_dim]
                elif isinstance(channel_first_dim, (int, np.integer)) and int(channel_first_dim) == slope_channel_dim:
                    target_shape = [slope_channel_dim] + [1 for _ in range(input_rank - 2)]
                if target_shape is not None and list(slope_shape) != target_shape:
                    slope = tf.reshape(slope, target_shape)

    # Generation of TF OP
    shared_axes = None
    input_shape = input_tensor.shape
    slope_shape = slope.shape if hasattr(slope, 'shape') else None
    if input_shape is not None and slope_shape is not None:
        input_rank = len(input_shape)
        if len(slope_shape) == input_rank - 1:
            shared_axes = [
                i + 1 for i, dim in enumerate(slope_shape)
                if dim is not None and dim == 1
            ]
        elif len(slope_shape) == 1 and input_rank >= 3:
            slope_dim = slope_shape[0]
            channel_axis = None
            if isinstance(slope_dim, int):
                if input_shape[1] == slope_dim:
                    channel_axis = 1
                elif input_shape[-1] == slope_dim:
                    channel_axis = input_rank - 1
            if channel_axis is not None:
                shared_axes = [ax for ax in range(1, input_rank) if ax != channel_axis]

    if shared_axes is None:
        if slope.shape is not None \
            and len(slope.shape) > 0 \
            and sum([1 if dim is not None and dim == 1 else 0 for dim in slope.shape]) == len(slope.shape):
            shared_axes = [val + 1 for val in range(len(input_tensor.shape) - 1)]
        else:
            input_nhwc = tf_layers_dict.get(graph_node_output.name, {}).get('nhwc', False)
            if input_nhwc:
                shared_axes = [val + 1 for val in range(len(input_tensor.shape) - 2)]
            else:
                shared_axes = [val + 2 for val in range(len(input_tensor.shape) - 2)]

    use_native_prelu = not replace_prelu_to_pseudo_prelu
    if not use_native_prelu:
        pos = tf.nn.relu(input_tensor)
        neg = tf.math.multiply(
            tf.math.subtract(input_tensor, tf.math.abs(input_tensor)),
            tf.math.multiply(slope, 0.5),
        )
        tf_layers_dict[graph_node_output.name]['tf_node'] = tf.math.add(pos, neg)
    else:
        try:
            tf_layers_dict[graph_node_output.name]['tf_node'] = \
                PReLU(
                    weights=slope,
                    shared_axes=shared_axes,
                )(input_tensor)
        except Exception:
            pos = tf.nn.relu(input_tensor)
            neg = tf.math.multiply(
                tf.math.subtract(input_tensor, tf.math.abs(input_tensor)),
                tf.math.multiply(slope, 0.5),
            )
            tf_layers_dict[graph_node_output.name]['tf_node'] = tf.math.add(pos, neg)

    # 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': 'PReLU',
                'tf_inputs': {
                    'x': input_tensor,
                    'slope': slope,
                },
                'tf_outputs': {
                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
                },
            }
        )
