# SPDX-License-Identifier: Apache-2.0


"""
tf2onnx.rewriter.gru_tf2_rewriter - Rewrites GRU pattern used by tf2.
"""

from tf2onnx.graph_builder import GraphBuilder
from tf2onnx.graph_matcher import GraphMatcher
from tf2onnx.rewriter.gru_rewriter import GRUUnitRewriter
from tf2onnx.rewriter.rnn_utils import keras_gru_pattern, make_grucell_pattern
from tf2onnx.rewriter.unit_rnn_rewriter_base import UnitRnnContext
from tf2onnx.tf_loader import find_function

# pylint: disable=invalid-name,unused-argument,missing-docstring, unused-variable


def rewrite_gru_tf2(g, ops):
    pattern1 = make_grucell_pattern("Identity")
    pattern2 = keras_gru_pattern

    for pattern in [pattern1, pattern2]:
        matcher = GraphMatcher(pattern, allow_reorder=True)
        match_results = list(matcher.match_ops(ops))
        for match_result in match_results:
            activation_op = match_result.get_op("optional_activation")
            activations = ["Sigmoid", activation_op.type]
            if activation_op.type not in ["Relu", "Tanh", "Sigmoid"]:
                continue

            if pattern is pattern1:
                concat = match_result.get_op("cell_inputs")
                if len(concat.inputs) != 3:
                    continue
                get_item = concat.inputs[0]
                init_state = concat.inputs[1]
            else:
                get_item = match_result.get_op("gru_input")
                init_state = match_result.get_op("state")
            if not get_item.type == "TensorListGetItem":
                continue
            x_e = get_item.inputs[0]
            if not x_e.is_graph_input():
                continue
            x_idx = g.input_names.index(x_e.output[0])
            if not init_state.is_graph_input():
                continue
            init_state_idx = g.input_names.index(init_state.output[0])

            cell_output = match_result.get_op("cell_output")
            final_consumers = g.find_output_consumers(cell_output.output[0])
            select_ops = [n for n in final_consumers if n.type == "Select"]
            def has_tensor_list_consumer(n):
                return any(c.type == "TensorListSetItem" for c in g.find_output_consumers(n.output[0]))
            select_ops = [n for n in select_ops if has_tensor_list_consumer(n)]
            if len(select_ops) == 1:
                greater_eq = select_ops[0].inputs[0]
                if greater_eq.type != "GreaterEqual":
                    continue
                seq_len = greater_eq.inputs[1]
                if not seq_len.is_graph_input():
                    continue
                seq_len_idx = g.input_names.index(seq_len.output[0])
                final_consumers = g.find_output_consumers(select_ops[0].output[0])
            else:
                seq_len_idx = None

            tensor_set_items = [n for n in final_consumers if n.type == "TensorListSetItem"]
            if len(tensor_set_items) != 1:
                continue

            if not tensor_set_items[0].inputs[0].is_graph_input():
                continue
            out_idx = g.input_names.index(tensor_set_items[0].input[0])

            hk = match_result.get_op("hidden_kernel")
            while hk.type == "Identity":
                hk = hk.inputs[0]
            if not hk.is_graph_input():
                continue
            hk_idx = g.input_names.index(hk.output[0])

            hb = match_result.get_op("hidden_bias")
            if not hb.is_graph_input():
                continue
            hb_idx = g.input_names.index(hb.output[0])

            gk = match_result.get_op("gate_kernel")
            while gk.type == "Identity":
                gk = gk.inputs[0]
            if not gk.is_graph_input():
                continue
            gk_idx = g.input_names.index(gk.output[0])

            gb = match_result.get_op("gate_bias")
            if not gb.is_graph_input():
                continue
            gb_idx = g.input_names.index(gb.output[0])

            bias_add = match_result.get_op("bias_add")
            if bias_add is not None and bias_add.data_format != "NHWC":
                continue

            g.gru_rewriter_context = {
                "x_idx": x_idx,
                "out_idx": out_idx,
                "initial_state_idx": init_state_idx,
                "hidden_kernel_idx": hk_idx,
                "hidden_bias_idx": hb_idx,
                "gate_kernel_idx": gk_idx,
                "gate_bias_idx": gb_idx,
                "seq_len_idx": seq_len_idx,
                "activations": activations,
                "from_keras": pattern is pattern2,
                "linear_before_reset": 1 if pattern is pattern2 else 0,
            }

    for op in ops:
        if op.is_while():
            body_graph = find_function(op.get_attr_str("body"))
            if body_graph.gru_rewriter_context is None:
                continue
            body_context = body_graph.gru_rewriter_context
            hk = op.input[body_context["hidden_kernel_idx"]]
            hb = op.input[body_context["hidden_bias_idx"]]
            gk = op.input[body_context["gate_kernel_idx"]]
            gb = op.input[body_context["gate_bias_idx"]]
            if not all(g.is_const(w) for w in [hk, hb, gk, gb]):
                continue
            hk_const = g.get_tensor_value(hk, as_list=False)
            hb_const = g.get_tensor_value(hb, as_list=False)
            gk_const = g.get_tensor_value(gk, as_list=False)
            gb_const = g.get_tensor_value(gb, as_list=False)

            initial_state_sq = op.input[body_context["initial_state_idx"]]
            initial_state = GraphBuilder(g).make_unsqueeze({"data": initial_state_sq, "axes": [0]})

            context = UnitRnnContext()
            context.from_keras = body_context["from_keras"]
            context.weights.update({
                "hidden_kernel": hk_const,
                "hidden_bias": hb_const,
                "gate_kernel": gk_const,
                "gate_bias": gb_const
            })
            context.attributes["activations"] = body_context["activations"]
            context.attributes["linear_before_reset"] = body_context["linear_before_reset"]
            tensor_array_inp = op.inputs[body_context["x_idx"]]
            if not tensor_array_inp.type == "TensorListFromTensor":
                continue

            final_consumers = g.find_output_consumers(op.output[body_context["out_idx"]])
            output_ys = [n.output[0] for n in final_consumers if n.type == "TensorListStack"]

            context.onnx_input_ids["X"] = tensor_array_inp.input[0]
            if body_context["seq_len_idx"] is None:
                context.onnx_input_ids["sequence_lens"] = ""
            else:
                context.onnx_input_ids["sequence_lens"] = op.input[body_context["seq_len_idx"]]
            context.onnx_input_ids["initial_state"] = initial_state

            gru_rewriter = GRUUnitRewriter(g)
            gru_rewriter.process_weights_and_bias(context)
            gru_node = gru_rewriter.create_rnn_node(context)
            squeeze_output = GraphBuilder(g).make_squeeze({"data": gru_node.output[0], "axes": [1]})
            for output in output_ys:
                g.replace_all_inputs(output, squeeze_output)

            f_state_squeeze = GraphBuilder(g).make_squeeze({"data": gru_node.output[1], "axes": [0]})
            g.replace_all_inputs(op.output[body_context["initial_state_idx"]], f_state_squeeze)

    return g.get_nodes()
