from typing import Any
import re
import sys
import random
random.seed(0)
import numpy as np
np.random.seed(0)
import importlib
import tensorflow as tf
from tensorflow.python.keras.backend import switch
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.enums import NUMPY_DTYPES_TO_TF_DTYPES
from onnx2tf.utils.logging import *


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

    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,
    )
    input_tensor = tf_layers_dict[graph_node_input.name]['tf_node'] \
        if isinstance(graph_node_input, gs.Variable) else graph_node_input
    graph_node_outputs = [] + graph_node.outputs

    # Then branch
    subgraph_kwargs = dict(kwargs)
    subgraph_kwargs['suppress_log'] = True
    then_branch_graph: gs.Graph = graph_node.attrs['then_branch']
    then_branch_graph_outputs = then_branch_graph.outputs
    for then_branch_graph_node in then_branch_graph.nodes:
        optype = then_branch_graph_node.op
        try:
            op = importlib.import_module(f'onnx2tf.ops.{optype}')
        except ModuleNotFoundError as ex:
            error(
                f'{optype} OP is not yet implemented.'
            )
            sys.exit(1)
        # substitution because saved_model does not allow colons
        then_branch_graph_node.name = then_branch_graph_node.name.replace(':','__')
        # Substitution because saved_model does not allow leading slashes in op names
        if kwargs['output_signaturedefs']:
            then_branch_graph_node.name = re.sub('^/', 'wa/', then_branch_graph_node.name)
        op.make_node(
            graph_node=then_branch_graph_node,
            tf_layers_dict=tf_layers_dict,
            **subgraph_kwargs,
        )
    # Then branch - Resister constant
    for output in then_branch_graph_outputs:
        if output.name not in tf_layers_dict and isinstance(output, gs.Constant):
            tf_layers_dict[output.name] = {
                'optype': 'Constant',
                'shape': output.values.shape,
                'dtype': output.values.dtype,
            }
            tf_layers_dict[output.name]['tf_node'] = \
                tf.constant(
                    output.values,
                    dtype=NUMPY_DTYPES_TO_TF_DTYPES[output.values.dtype] \
                        if isinstance(output.values.dtype, np.dtype) else output.values.dtype,
                )
    then_branch_ops = []
    for then_branch_graph_output in then_branch_graph_outputs:
        then_branch_ops.append(
            tf_layers_dict[then_branch_graph_output.name]['tf_node']
        )

    # Else branch
    else_branch_graph: gs.Graph = graph_node.attrs['else_branch']
    else_branch_graph_outputs = else_branch_graph.outputs
    for else_branch_graph_node in else_branch_graph.nodes:
        optype = else_branch_graph_node.op
        try:
            op = importlib.import_module(f'onnx2tf.ops.{optype}')
        except ModuleNotFoundError as ex:
            error(
                f'{optype} OP is not yet implemented.'
            )
            sys.exit(1)
        # substitution because saved_model does not allow colons
        else_branch_graph_node.name = else_branch_graph_node.name.replace(':','__')
        # Substitution because saved_model does not allow leading slashes in op names
        if kwargs['output_signaturedefs']:
            else_branch_graph_node.name = re.sub('^/', 'wa/', else_branch_graph_node.name)
        op.make_node(
            graph_node=else_branch_graph_node,
            tf_layers_dict=tf_layers_dict,
            **subgraph_kwargs,
        )
    # Else branch - Resister constant
    for output in else_branch_graph_outputs:
        if output.name not in tf_layers_dict and isinstance(output, gs.Constant):
            tf_layers_dict[output.name] = {
                'optype': 'Constant',
                'shape': output.values.shape,
                'dtype': output.values.dtype,
            }
            tf_layers_dict[output.name]['tf_node'] = \
                tf.constant(
                    output.values,
                    dtype=NUMPY_DTYPES_TO_TF_DTYPES[output.values.dtype] \
                        if isinstance(output.values.dtype, np.dtype) else output.values.dtype,
                )
    else_branch_ops = []
    for else_branch_graph_output in else_branch_graph_outputs:
        else_branch_ops.append(
            tf_layers_dict[else_branch_graph_output.name]['tf_node']
        )

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

    # Generation of TF OP
    for graph_node_output in graph_node_outputs:
        tf_layers_dict[graph_node_output.name] = {
            'optype': graph_node.op,
            'shape': graph_node_output.shape,
            'dtype': graph_node_output.dtype,
        }

    if_cond_outputs = []
    for then_op, else_op in zip(then_branch_ops, else_branch_ops):
        if_cond_outputs.append(
            switch(
                condition=input_tensor,
                then_expression=then_op,
                else_expression=else_op,
            )
        )

    for graph_node_output, if_cond_output in zip(graph_node_outputs, if_cond_outputs):
        tf_layers_dict[graph_node_output.name]['tf_node'] = if_cond_output
        # 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_output.name,
            **kwargs,
        )

    # Generation of Debug Info
    tf_outputs = {f"output{idx}": value for idx, value in enumerate(if_cond_outputs)}
    for graph_node_output in graph_node_outputs:
        tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
            make_tf_node_info(
                node_info={
                    'tf_op_type': tf.cond,
                    'tf_inputs': {
                        'input': input_tensor,
                        'true_fn': then_branch_ops,
                        'false_fn': else_branch_ops,
                    },
                    'tf_outputs': {
                        'output': tf_outputs,
                    },
                }
            )
