# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# Modified from:
# https://github.com/pytorch/pytorch/blob/9ade03959392e5a90b74261012de1d806cab2253/torch/onnx/symbolic_opset9.py
import warnings

import torch
import torch.onnx.symbolic_helper as sym_help
from torch.onnx.symbolic_helper import _unimplemented
from torch.onnx.symbolic_opset9 import unused

from juxtapose.mmdeploy.core import FUNCTION_REWRITER


@FUNCTION_REWRITER.register_rewriter(
    func_name="torch.onnx.symbolic_opset9._generic_rnn", backend="ncnn"
)
def generic_rnn__ncnn(
    g,
    variant,
    input,
    initial_states,
    all_weights,
    has_biases,
    num_layers,
    dropout,
    train,
    bidirectional,
    batch_first=None,
    batch_sizes=None,
):
    """rewrite of _generic_rnn for ncnn.

    `g.op` will add some nodes for h0 and c0 in LSTM. which is not supported in
    ncnn. So we add a custom domain to avoid it.
    """
    warnings.warn(
        "Exporting a model to ONNX with a batch_size other than 1, "
        + "with a variable length with "
        + variant
        + " can cause an error "
        + "when running the ONNX model with a different batch size. "
        + "Make sure to save the model with a batch size of 1, "
        + "or define the initial states (h0/c0) as inputs of the model. "
    )

    onnxActivations = [
        "Relu",
        "Tanh",
        "Sigmoid",
        "Affine",
        "LeakyRelu",
        "ThresholdedRelu",
        "ScaledTanh",
        "HardSigmoid",
        "Elu",
        "Softsign",
        "Softplus",
    ]
    variantToOnnxActivationMap = dict(
        zip([act_fun.lower() for act_fun in onnxActivations], onnxActivations)
    )
    weights_per_layer = 4 if has_biases else 2
    # this means that projections are used inside LSTM,
    # so need to tell user that it's not supported
    if variant == "LSTM" and len(all_weights) != num_layers * weights_per_layer * (
        1 + bidirectional
    ):
        return _unimplemented("LSTM", "LSTMs with projections")
    assert len(all_weights) == num_layers * weights_per_layer * (1 + bidirectional)
    layer_weights = [
        all_weights[i : i + weights_per_layer]
        for i in range(0, len(all_weights), weights_per_layer)
    ]
    if batch_first:
        # batch, seq, feat -> seq, batch, feat
        input = g.op("Transpose", input, perm_i=[1, 0, 2])
    if dropout and train:
        return _unimplemented("RNN/GRU/LSTM", "dropout in training mode")

    if variant.startswith("RNN"):
        nonlinearity = variantToOnnxActivationMap[variant[4:].lower()]
        variant = "RNN"

    w_hh = all_weights[1]
    hidden_size = sym_help._get_tensor_dim_size(w_hh, 1)
    if hidden_size is None:
        return _unimplemented("RNN/GRU/LSTM", "unknown hidden size")

    unidirectional = not bidirectional

    prev_output = input

    h_outs = []
    if variant == "RNN" or variant == "GRU":
        h0 = initial_states
    elif variant == "LSTM":
        h0, c0 = initial_states
        c_outs = []

    sequence_lens = unused(g) if batch_sizes is None else batch_sizes

    if variant == "GRU":
        # pytorch is reset, input, hidden
        # onnx is    input, reset, hidden
        reform_permutation = [(1, 2), (0, 1), (2, 3)]
    elif variant == "LSTM":
        # pytorch is input, forget, cell, output.
        # onnx is    input, output, forget, cell.
        reform_permutation = [(0, 1), (3, 4), (1, 3)]

    def reform_weights(g, w, n, intervals):
        slices = [
            sym_help._slice_helper(g, w, axes=[0], starts=[x * n], ends=[y * n])
            for x, y in intervals
        ]
        return g.op("Concat", *slices, axis_i=0)

    def transform_weights_no_bias(layer_index):
        weights = layer_weights[layer_index]
        if variant == "RNN":
            weight_ih, weight_hh = weights
        elif variant == "GRU" or variant == "LSTM":
            weight_ih, weight_hh = [
                reform_weights(g, w, hidden_size, reform_permutation) for w in weights
            ]
        return tuple(
            sym_help._unsqueeze_helper(g, x, [0]) for x in (weight_ih, weight_hh)
        )

    def transform_weights(layer_index):
        weights = layer_weights[layer_index]
        if variant == "RNN":
            weight_ih, weight_hh, bias_ih, bias_hh = weights
        elif variant == "GRU" or variant == "LSTM":
            weight_ih, weight_hh, bias_ih, bias_hh = [
                reform_weights(g, w, hidden_size, reform_permutation) for w in weights
            ]
        bias_concat = g.op("Concat", bias_ih, bias_hh, axis_i=0)
        return tuple(
            sym_help._unsqueeze_helper(g, x, [0])
            for x in (weight_ih, weight_hh, bias_concat)
        )

    def retrieve_state(x, start, end):
        return (
            x
            if num_layers == 1
            else sym_help._slice_helper(g, x, axes=[0], starts=[start], ends=[end])
        )

    for i in range(num_layers):
        if unidirectional:
            if weights_per_layer == 4:
                weight_ih, weight_hh, bias_concat = transform_weights(i)
            else:
                weight_ih, weight_hh = transform_weights_no_bias(i)
                bias_concat = unused(g)

            state_indices = i, i + 1
        else:
            if weights_per_layer == 4:
                weight_ih_f, weight_hh_f, bias_f = transform_weights(2 * i)
                weight_ih_b, weight_hh_b, bias_b = transform_weights(2 * i + 1)
                bias_concat = g.op("Concat", bias_f, bias_b, axis_i=0)
            else:
                weight_ih_f, weight_hh_f = transform_weights_no_bias(2 * i)
                weight_ih_b, weight_hh_b = transform_weights_no_bias(2 * i + 1)
                bias_concat = unused(g)

            weight_ih = g.op("Concat", weight_ih_f, weight_ih_b, axis_i=0)
            weight_hh = g.op("Concat", weight_hh_f, weight_hh_b, axis_i=0)

            state_indices = 2 * i, 2 * i + 2

        inputs = [prev_output, weight_ih, weight_hh, bias_concat, sequence_lens]

        inputs.append(retrieve_state(h0, *state_indices))
        if variant == "LSTM":
            inputs.append(retrieve_state(c0, *state_indices))

        extra_kwargs = {} if unidirectional else {"direction_s": "bidirectional"}
        if variant == "RNN":
            if bidirectional:
                activation = [nonlinearity, nonlinearity]
            else:
                activation = [nonlinearity]

            prev_output, h_out = g.op(
                "RNN",
                *inputs,
                outputs=2,
                hidden_size_i=hidden_size,
                activations_s=activation,
                **extra_kwargs,
            )
        elif variant == "GRU":
            prev_output, h_out = g.op(
                "GRU",
                *inputs,
                outputs=2,
                hidden_size_i=hidden_size,
                linear_before_reset_i=1,
                **extra_kwargs,
            )
        elif variant == "LSTM":
            # g.op will add some node to h0 and c0,
            # which is not necessary for us
            prev_output, h_out, c_out = g.op(
                "ncnn::LSTM",
                *inputs,
                outputs=3,
                hidden_size_i=hidden_size,
                **extra_kwargs,
            )
        if bidirectional:
            # The ONNX RNN/GRU/LSTM produce an output of dimensions
            #   seq_len, num_directions, batch, hidden_size
            # We have to convert to match pytorch's expected
            #   seq_len, batch, num_directions * hidden_size
            # by first moving num_directions before hidden_size with
            # Transpose, and then combining it with hidden_size
            # with Reshape.
            prev_output = g.op("Transpose", prev_output, perm_i=[0, 2, 1, 3])
            prev_output = g.op(
                "Reshape",
                prev_output,
                g.op("Constant", value_t=torch.LongTensor([0, 0, -1])),
            )
        else:
            prev_output = sym_help._squeeze_helper(g, prev_output, [1])

        h_outs.append(h_out)
        if variant == "LSTM":
            c_outs.append(c_out)
    if batch_first:
        # seq, batch, num_directions * hidden_size -> batch, seq,
        # num_directions * hidden_size
        prev_output = g.op("Transpose", prev_output, perm_i=[1, 0, 2])
    h_outs = h_out if num_layers == 1 else g.op("Concat", *h_outs, axis_i=0)
    if variant == "RNN" or variant == "GRU":
        return prev_output, h_outs
    elif variant == "LSTM":
        c_outs = c_out if num_layers == 1 else g.op("Concat", *c_outs, axis_i=0)
        return prev_output, h_outs, c_outs
