# SPDX-License-Identifier: Apache-2.0

import numpy as np
from .....proto import onnx_proto
from ....common._registration import register_converter
from .Reshape import apply_reshape


def extract_rnn_activation_info(activation):
    activation_type = activation.WhichOneof("NonlinearityType")
    alpha = None
    beta = None

    activation_map = {
        "linear": "Affine",
        "ReLU": "Relu",
        "leakyReLU": "LeakyRelu",
        "thresholdedReLU": "ThresholdedRelu",
        "PReLU": "PRelu",
        "tanh": "Tanh",
        "scaledTanh": "ScaledTanh",
        "sigmoid": "Sigmoid",
        "sigmoidHard": "HardSigmoid",
        "ELU": "Elu",
        "softsign": "Softsign",
        "softplus": "Softplus",
        "parametricSoftplus": "ParametricSoftplus",
    }

    if activation_type not in activation_map:
        raise ValueError("Unsupported activation function: {}".format(activation_type))

    # Notice that if we see a default vaulue (i.e., 0 for float), we may replace it with
    # the real default parameter for the specified activation function if necessary.
    if activation_type == "leakyReLU":
        alpha = activation.leakyReLU.alpha
        if alpha == 0:
            alpha = 0.3
    elif activation_type == "PReLU":
        raise RuntimeError(
            "Unsupported activation function: {}".format(activation_type)
        )
    elif activation_type == "ELU":
        alpha = activation.ELU.alpha
    elif activation_type == "thresholdedReLU":
        alpha = activation.thresholdedReLU.alpha
        if alpha == 0:
            alpha = 1.0
    elif activation_type == "scaledTanh":
        alpha = activation.scaledTanh.alpha
        beta = activation.scaledTanh.beta
    elif activation_type == "linear":
        alpha = activation.linear.alpha
        beta = activation.linear.beta
        if alpha == 0:
            alpha = 1.0
    elif activation_type == "sigmoidHard":
        alpha = activation.sigmoidHard.alpha
        beta = activation.sigmoidHard.beta
        if alpha == 0:
            alpha = 0.2
        if beta == 0:
            beta = 0.5
    elif activation_type == "parametricSoftplus":
        raise RuntimeError(
            "Unsupported activation function: {}".format(activation_type)
        )

    return activation_map[activation_type], alpha, beta


def convert_simple_rnn(scope, operator, container):
    # The RNN inputs are feature vector, X, and initial hidden state, h_init. Let C_in and C_out denote
    # the input feature length and the output dimension, respectively. Assume that S is the sequence length
    # (i.e., S-axis is the time axis of a sequence). In CorML, the shapes of X and h_init are [S, C_in] and
    # [1, C_out] respectively. In ONNX, the two shapes become [S, N, C_in] (X) and [D, N, C_out]
    # (h_init), where N is the batch size. To simulate CoreML RNN under ONNX, we need to introduce some extra
    # operators. Note that N=1 and D=1 always hold in ONNX if we are considering RNN from CoreML because there
    # is no batch size in CoreML and CoreML RNN is always uni-directional.
    #
    # Below we provide a visualization of our conversion for CoreML RNN.
    #
    # Symbols:
    #
    #  X: input features of CoreML RNN
    #  h_init: initial RNN state of CoreML
    #  Y: CoreML RNN's output. It can be [S, C_out] (if sequence_output is on ) or [1, C_out] (if sequence_output is off)
    #  Y_h: CoreML RNN's last hidden state
    #
    #  X': input features of ONNX RNN
    #  h_init': initial RNN state of ONNX
    #  Y': ONNX RNN's output
    #  Y_h': ONNX RNN's last hidden state
    #
    # Computational graph of CoreML RNN (sequence_output is on):
    #
    # X [S, C_in] ---> CoreML RNN ---> Y [S, C_out]
    #                    ^     |
    #                    |     |
    # h_init [1, C_out] -'     '---> Y_h [1, C_out]
    #
    # Computational graph we use for represent CoreML RNN into ONNX (sequence_output is on):
    #
    # X [S, C_in] --> Reshape --> X' [S, 1, C_in] -----> ONNX RNN --> Y' [S, 1, C_out] --> Reshape --> Y [S, C_out]
    #                                                        ^ |
    #                                                        | |
    # h_init [1, C_out] --> Reshape -> h_init' [1, 1, C_out]-' '--> Y_h' [1, 1, C_out] ---> Reshape --> Y_h [1, C_out]
    #
    # Computational graph of CoreML RNN (sequence_output is off):
    #
    # X [S, C_in] ---> CoreML RNN ---> Y [1, C_out]
    #                    ^     |
    #                    |     |
    # h_init [1, C_cou] -'     '---> Y_h [1, C_out] Note that in this case, Y=Y_h.
    #
    # Computational graph we use to represent CoreML RNN into ONNX (sequence_output is off):
    #
    # X [S, C_in] --> Reshape --> X' [S, 1, C_in] -----> ONNX RNN --> Y' [S, 1, C_out] Here Y' is useless.
    #                                                        ^ |
    #                                                        | |
    # h_init [1, C_out] --> Reshape -> h_init' [1, 1, C_out]-' '--> Y_h' [1, 1, C_out] ---> Reshape --> Y [1, C_out]
    #                                                                                                   |
    #                                                                                                   v
    #                                                                                                 Identity
    #                                                                                                   |
    #                                                                                                   v
    #                                                                                               Y_h [1, C_out]

    params = operator.raw_operator.simpleRecurrent
    input_size = params.inputVectorSize
    hidden_size = params.outputVectorSize

    X_name = operator.inputs[0].full_name
    X_reshape_name = scope.get_unique_variable_name("X")
    apply_reshape(
        scope, X_name, X_reshape_name, container, desired_shape=[-1, 1, input_size]
    )

    rnn_op_name = scope.get_unique_operator_name("RNN")
    rnn_attrs = {"name": rnn_op_name}
    rnn_inputs = [X_reshape_name]

    # Load RNN's weight matrix and add it into RNN's input list
    rnn_w_name = scope.get_unique_variable_name(rnn_op_name + "_W")
    container.add_initializer(
        rnn_w_name,
        onnx_proto.TensorProto.FLOAT,
        [1, hidden_size, input_size],
        params.weightMatrix.floatValue,
    )
    rnn_inputs.append(rnn_w_name)

    # Load RNN's recursion matrix and add it into RNN's input list
    rnn_r_name = scope.get_unique_variable_name(rnn_op_name + "_R")
    container.add_initializer(
        rnn_r_name,
        onnx_proto.TensorProto.FLOAT,
        [1, hidden_size, hidden_size],
        params.recursionMatrix.floatValue,
    )
    rnn_inputs.append(rnn_r_name)

    if params.hasBiasVector:
        # Load RNN's bias vector and add it into RNN's input list
        rnn_b_name = scope.get_unique_variable_name(rnn_op_name + "_B")
        rnn_b_content = np.concatenate(
            [params.biasVector.floatValue, np.zeros(hidden_size)]
        ).flatten()
        container.add_initializer(
            rnn_b_name,
            onnx_proto.TensorProto.FLOAT,
            [1, 2 * hidden_size],
            rnn_b_content,
        )
        rnn_inputs.append(rnn_b_name)
    else:
        # Input names are position-sensitive, so for optional
        # but missing inputs, we need to provide an empty string.
        rnn_inputs.append("")

    # The input, sequence_lens, in ONNX is alwasy optional for
    # this conversion, so here is always an empty string.
    rnn_inputs.append("")

    # If initial hidden state is provided, we add it into RNN's
    # input list after adjusting its shape.
    if len(operator.inputs) == 2:
        rnn_h_init_reshape_name = scope.get_unique_variable_name(
            rnn_op_name + "_h_init"
        )
        apply_reshape(
            scope,
            operator.inputs[1].full_name,
            rnn_h_init_reshape_name,
            container,
            desired_shape=[1, 1, hidden_size],
        )

        rnn_inputs.append(rnn_h_init_reshape_name)
        # Add a zero initializer to initial hidden state so that
        # this variable becomes optional
        container.add_initializer(
            operator.inputs[1].full_name,
            onnx_proto.TensorProto.FLOAT,
            operator.inputs[1].type.shape,
            np.zeros(shape=operator.inputs[1].type.shape).flatten(),
        )
    else:
        # Input names are position-sensitive, so for optional
        # but missing inputs, we need to provide an empty string.
        rnn_inputs.append("")

    # Add RNN's information of activation function
    activation, alpha, beta = extract_rnn_activation_info(params.activation)
    rnn_attrs["activations"] = [activation.encode("utf-8")]
    if alpha is not None:
        rnn_attrs["activation_alpha"] = [alpha]
    if beta is not None:
        rnn_attrs["activation_beta"] = [beta]

    # Set up other attributes
    rnn_attrs["direction"] = "reverse" if params.reverseInput else "forward"
    rnn_attrs["hidden_size"] = hidden_size

    # Set up version-dependent attributes
    if container.target_opset < 7:
        rnn_attrs["output_sequence"] = params.sequenceOutput
        op_version = 1
    else:
        op_version = 7

    # We use the collected information to build ONNX's RNN.
    # ONNX RNN's outputs will be saved onto two intermediate
    # tensors and we will adjust them subsequently to mimic
    # Keras output format.
    rnn_y_name = scope.get_unique_variable_name(rnn_op_name + "_Y")
    rnn_h_name = scope.get_unique_variable_name(rnn_op_name + "_Y_h")
    container.add_node(
        "RNN", rnn_inputs, [rnn_y_name, rnn_h_name], op_version=op_version, **rnn_attrs
    )

    # Set up outputs' of RNN
    if params.sequenceOutput:
        # Connect ONNX's output and CoreML's output via a reshape operator
        apply_reshape(
            scope,
            rnn_y_name,
            operator.outputs[0].full_name,
            container,
            desired_shape=[-1, hidden_size],
        )

        # Handel the second RNN output (aka last hidden state), which is optional.
        if len(operator.outputs) == 2:
            # Connect ONNX's output and CoreML's output via a reshape operator
            apply_reshape(
                scope,
                rnn_h_name,
                operator.outputs[1].full_name,
                container,
                desired_shape=[1, hidden_size],
            )
    else:
        # According to CoreML, its two outputs are always identical, so
        # we just need to compute one of them and produce
        # the other one using an identity operator. Note that the
        # first ONNX RNN output is undefined in this case.

        # Reshape last hidden state's ONNX format to its CoreML format
        apply_reshape(
            scope,
            rnn_h_name,
            operator.outputs[0].full_name,
            container,
            desired_shape=[1, hidden_size],
        )

        if len(operator.outputs) == 2:
            # Copy the first output to the second output
            container.add_node(
                "Identity",
                operator.outputs[0].full_name,
                operator.outputs[1].full_name,
                name=scope.get_unique_operator_name("Identity"),
            )


register_converter("simpleRecurrent", convert_simple_rnn)
