from typing import Any, cast
import random
random.seed(0)
import numpy as np
np.random.seed(0)
import tensorflow as tf
from onnx import numpy_helper
import onnx2tf.gs as gs
from onnx2tf.utils.common_functions import (
    print_node_info,
    make_tf_node_info,
    get_constant_or_variable,
)


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

    Parameters
    ----------
    graph_node: gs.Node
        graph_surgeon Node

    tf_layers_dict: dict
        optype, shape, dtype, tensorflow graph
    """
    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,
    }

    # Generation of TF OP
    # https://github.com/onnx/onnx-tensorflow/blob/main/onnx_tf/handlers/backend/constant_of_shape.py
    shape = get_constant_or_variable(
        graph_node.inputs[0],
        before_op_output_shape_trans=False,
        is_bias=True,
    )
    if isinstance(shape, gs.Variable) \
        and graph_node.inputs[0].name in tf_layers_dict:
        shape = tf_layers_dict[graph_node.inputs[0].name]['tf_node']
    if not tf.is_tensor(shape):
        shape = tf.convert_to_tensor(cast(Any, shape))

    # make sure the shape dtype is either int32 or int64
    if shape.dtype not in [tf.int64, tf.int32]:
        shape = tf.cast(cast(Any, shape), tf.int64)

    # the default value is 0, float32
    constant_tensor = None
    if "value" in graph_node.attrs:
        attr_value = graph_node.attrs['value']
        if hasattr(attr_value, 'values'):
            value = attr_value.values
        elif isinstance(attr_value, np.ndarray):
            value = attr_value
        else:
            # onnx2tf.gs may keep ConstantOfShape attrs['value'] as TensorProto
            value = numpy_helper.to_array(attr_value)
        value = np.asarray(value)
        constant_tensor = value.flatten()[0] if value.size > 0 else 0.0
    else:
        constant_tensor = 0.0

    cons = tf.fill(
        dims=shape,
        value=constant_tensor,
        name=graph_node.name,
    )

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

    # Generation of Debug Info
    tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
        make_tf_node_info(
            node_info={
                'tf_op_type': tf.constant,
                'tf_inputs': {
                    'value': constant_tensor,
                },
                'tf_outputs': {
                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
                },
            }
        )
