// Copyright (C) 2018-2022 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include "legacy/transformations/convert_opset1_to_legacy/convert_sequences_to_sequences_ie.hpp"

#include <memory>

#include <ngraph/opsets/opset5.hpp>
#include <ngraph/rt_info.hpp>
#include <ngraph/pattern/op/wrap_type.hpp>

#include <legacy/ngraph_ops/lstm_sequence_ie.hpp>
#include <legacy/ngraph_ops/gru_sequence_ie.hpp>
#include <legacy/ngraph_ops/rnn_sequence_ie.hpp>

NGRAPH_RTTI_DEFINITION(ngraph::pass::ConvertLSTMSequenceMatcher, "ConvertLSTMSequenceMatcher", 0);
NGRAPH_RTTI_DEFINITION(ngraph::pass::ConvertGRUSequenceMatcher, "ConvertGRUSequenceMatcher", 0);
NGRAPH_RTTI_DEFINITION(ngraph::pass::ConvertRNNSequenceMatcher, "ConvertRNNSequenceMatcher", 0);

namespace {
    int64_t get_seq_axis(const std::shared_ptr<ngraph::Node>& sequence_node) {
        // Optimization.
        // Plug-ins support seq_axis attribute (value 1 or 0) for Seq ops, but according to the spec we don't
        // support this attribute and should insert Transpose layer before and after Seq op in TI to Sequences
        // transformation. Additional Transpose layers affect the performance, so we try to detect pattern
        // Transpose(axis_order={1,0,2}) -> Seq -> Transpose(axis_order={2,1,0,3}
        // and replace unnecessary Transpose ops with SeqIE (seq_axis = 0) to transfer value
        // of the attribute to plug-ins.
        // todo: specify seq_axis attribute for Sequence ops.
        int64_t seq_axis = 1; // default
        const auto& target_inputs = sequence_node->output(0).get_target_inputs();
        if (target_inputs.size() == 1) {
            const auto& transpose_before = std::dynamic_pointer_cast<ngraph::opset5::Transpose>(sequence_node->input_value(0).get_node_shared_ptr());
            const auto& transpose_after = std::dynamic_pointer_cast<ngraph::opset5::Transpose>(target_inputs.begin()->get_node()->shared_from_this());
            if (transpose_after != nullptr && transpose_before != nullptr) {
                auto order_before = std::dynamic_pointer_cast<ngraph::opset5::Constant>(
                        transpose_before->input_value(1).get_node_shared_ptr());
                auto order_after = std::dynamic_pointer_cast<ngraph::opset5::Constant>(
                        transpose_after->input_value(1).get_node_shared_ptr());
                if (order_before != nullptr && order_after != nullptr) {
                    auto order_before_values = order_before->cast_vector<int64_t>();
                    auto order_after_values = order_after->cast_vector<int64_t>();
                    std::vector<int64_t> order_ref_before = {1, 0, 2};
                    std::vector<int64_t> order_ref_after = {2, 1, 0, 3};
                    if (order_before_values == order_ref_before && order_after_values == order_ref_after) {
                        seq_axis = 0;
                    }
                }
            }
        }
        return seq_axis;
    }
} // namespace

ngraph::pass::ConvertLSTMSequenceMatcher::ConvertLSTMSequenceMatcher() {
    auto lstm_sequence_ngraph = ngraph::pattern::wrap_type<ngraph::opset5::LSTMSequence>();

    ngraph::matcher_pass_callback callback = [](pattern::Matcher &m) {
        auto lstm_sequence = std::dynamic_pointer_cast<ngraph::opset5::LSTMSequence>(m.get_match_root());
        if (!lstm_sequence) {
            return false;
        }

        const auto& W = lstm_sequence->input_value(4);
        const auto& R = lstm_sequence->input_value(5);

        // Bidirectional cases are not supported
        if (lstm_sequence->get_direction() == ngraph::op::RecurrentSequenceDirection::BIDIRECTIONAL)
            return false;

        // Detect pattern: Transpose_before -> Seq -> Transpose_after
        auto seq_axis = get_seq_axis(lstm_sequence);
        ngraph::Output<ngraph::Node> in_0 = lstm_sequence->input(0).get_source_output();
        if (seq_axis == 0) {
            // input(0) to Transpose_before
            in_0 = lstm_sequence->get_input_source_output(0).get_node_shared_ptr()->get_input_source_output(0);
        }
        // for forward/reverse cases we can squeeze num_direction dimension
        auto axis_1 = ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {1});
        auto in_1 = std::make_shared<ngraph::opset5::Squeeze>(lstm_sequence->input_value(1), axis_1);
        auto in_2 = std::make_shared<ngraph::opset5::Squeeze>(lstm_sequence->input_value(2), axis_1);
        auto concat = std::make_shared<ngraph::opset5::Concat>(ngraph::OutputVector{W, R}, 2);
        auto axis_2 = ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0});
        auto in_3 = std::make_shared<ngraph::opset5::Squeeze>(concat->output(0), axis_2);
        auto in_4 = std::make_shared<ngraph::opset5::Squeeze>(lstm_sequence->input_value(6), axis_2);
        auto lstm_sequence_ie = std::make_shared<ngraph::op::LSTMSequenceIE>(
                in_0,  // X
                in_1,  // initial_hidden_state
                in_2,  // initial_cell_state
                lstm_sequence->input_value(3),
                in_3,  // WR
                in_4,  // B
                lstm_sequence->get_hidden_size(),
                lstm_sequence->get_direction(),
                lstm_sequence->get_activations(),
                lstm_sequence->get_activations_alpha(),
                lstm_sequence->get_activations_beta(),
                lstm_sequence->get_clip(),
                seq_axis);

        auto unsqueeze_axis = ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {1});
        auto unsqueeze_1 = std::make_shared<ngraph::opset5::Unsqueeze>(lstm_sequence_ie->output(0), unsqueeze_axis);
        auto unsqueeze_2 = std::make_shared<ngraph::opset5::Unsqueeze>(lstm_sequence_ie->output(1), unsqueeze_axis);
        auto unsqueeze_3 = std::make_shared<ngraph::opset5::Unsqueeze>(lstm_sequence_ie->output(2), unsqueeze_axis);

        ngraph::copy_runtime_info(lstm_sequence, {concat, lstm_sequence_ie, in_1, in_2, in_3, in_4, unsqueeze_1,
                                                  unsqueeze_2, unsqueeze_3});
        unsqueeze_1->set_friendly_name(lstm_sequence->get_friendly_name()+".0");
        unsqueeze_2->set_friendly_name(lstm_sequence->get_friendly_name()+".1");
        unsqueeze_3->set_friendly_name(lstm_sequence->get_friendly_name()+".2");
        if (seq_axis == 1) {
            ngraph::replace_node(lstm_sequence, {unsqueeze_1->output(0), unsqueeze_2->output(0), unsqueeze_3->output(0)});
        } else {
            const auto &lstm_target_inputs = lstm_sequence->output(0).get_target_inputs();
            if (lstm_target_inputs.empty())
                return false;
            auto transpose_after = lstm_target_inputs.begin()->get_node()->shared_from_this();
            unsqueeze_1->set_friendly_name(transpose_after->get_friendly_name());
            ngraph::replace_node(transpose_after, unsqueeze_1);
            ngraph::replace_node(lstm_sequence, {lstm_sequence_ie->output(0), unsqueeze_2->output(0), unsqueeze_3->output(0)});
        }
        return true;
    };

    auto m = std::make_shared<ngraph::pattern::Matcher>(lstm_sequence_ngraph, "ConvertLSTMSequenceToLSTMSequenceIE");
    this->register_matcher(m, callback);
}

ngraph::pass::ConvertGRUSequenceMatcher::ConvertGRUSequenceMatcher() {
    auto gru_sequence_ngraph = ngraph::pattern::wrap_type<ngraph::opset5::GRUSequence>();

    ngraph::matcher_pass_callback callback = [](pattern::Matcher &m) {
        auto gru_sequence = std::dynamic_pointer_cast<ngraph::opset5::GRUSequence>(m.get_match_root());
        if (!gru_sequence) {
            return false;
        }

        auto W = gru_sequence->input_value(3);
        auto R = gru_sequence->input_value(4);

        // Bidirectional cases are not supported
        if (gru_sequence->get_direction() == ngraph::op::RecurrentSequenceDirection::BIDIRECTIONAL)
            return false;

        // Detect pattern: Transpose_before -> Seq -> Transpose_after
        auto seq_axis = get_seq_axis(gru_sequence);
        ngraph::Output<ngraph::Node> in_0 = gru_sequence->input(0).get_source_output();
        if (seq_axis == 0) {
            // input(0) to Transpose_before
            in_0 = gru_sequence->get_input_source_output(0).get_node_shared_ptr()->get_input_source_output(0);
        }
        // for forward/reverse cases we can squeeze num_direction dimension
        auto axis_1 = ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {1});
        auto in_1 = std::make_shared<ngraph::opset5::Squeeze>(gru_sequence->input_value(1), axis_1);
        auto concat = std::make_shared<ngraph::opset5::Concat>(ngraph::OutputVector{W, R}, 2);
        auto axis_2 = ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0});
        auto in_3 = std::make_shared<ngraph::opset5::Squeeze>(concat->output(0), axis_2);
        auto in_4 = std::make_shared<ngraph::opset5::Squeeze>(gru_sequence->input_value(5), axis_2);

        auto gru_sequence_ie = std::make_shared<ngraph::op::GRUSequenceIE>(
                in_0, // X
                in_1,  // initial_hidden_state
                gru_sequence->input_value(2),
                in_3,  // WR
                in_4,  // B
                gru_sequence->get_hidden_size(),
                gru_sequence->get_direction(),
                gru_sequence->get_activations(),
                gru_sequence->get_activations_alpha(),
                gru_sequence->get_activations_beta(),
                gru_sequence->get_clip(),
                gru_sequence->get_linear_before_reset(),
                seq_axis);

        auto unsqueeze_axis = ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {1});
        auto unsqueeze_1 = std::make_shared<ngraph::opset5::Unsqueeze>(gru_sequence_ie->output(0), unsqueeze_axis);
        auto unsqueeze_2 = std::make_shared<ngraph::opset5::Unsqueeze>(gru_sequence_ie->output(1), unsqueeze_axis);

        ngraph::copy_runtime_info(gru_sequence, {concat, gru_sequence_ie, unsqueeze_1, unsqueeze_2, in_1, in_3, in_4});
        unsqueeze_1->set_friendly_name(gru_sequence->get_friendly_name()+".0");
        unsqueeze_2->set_friendly_name(gru_sequence->get_friendly_name()+".1");
        if (seq_axis == 1) {
            ngraph::replace_node(gru_sequence, {unsqueeze_1->output(0), unsqueeze_2->output(0)});
        } else {
            const auto &gru_target_inputs = gru_sequence->output(0).get_target_inputs();
            if (gru_target_inputs.empty())
                return false;
            auto transpose_after = gru_target_inputs.begin()->get_node()->shared_from_this();
            unsqueeze_1->set_friendly_name(transpose_after->get_friendly_name());
            ngraph::replace_node(transpose_after, unsqueeze_1);
            ngraph::replace_node(gru_sequence, {gru_sequence_ie->output(0), unsqueeze_2->output(0)});
        }
        return true;
    };

    auto m = std::make_shared<ngraph::pattern::Matcher>(gru_sequence_ngraph, "ConvertGRUSequenceToGRUSequenceIE");
    this->register_matcher(m, callback);
}

ngraph::pass::ConvertRNNSequenceMatcher::ConvertRNNSequenceMatcher() {
    auto rnn_sequence_ngraph = ngraph::pattern::wrap_type<ngraph::opset5::RNNSequence>();

    ngraph::matcher_pass_callback callback = [](pattern::Matcher &m) {
        auto rnn_sequence = std::dynamic_pointer_cast<ngraph::opset5::RNNSequence>(m.get_match_root());
        if (!rnn_sequence) {
            return false;
        }

        // Bidirectional cases are not supported
        if (rnn_sequence->get_direction() == ngraph::op::RecurrentSequenceDirection::BIDIRECTIONAL)
            return false;

        // Detect pattern: Transpose_before -> Seq -> Transpose_after
        auto seq_axis = get_seq_axis(rnn_sequence);
        ngraph::Output<ngraph::Node> in_0 = rnn_sequence->input(0).get_source_output();
        if (seq_axis == 0) {
            // input(0) to Transpose_before
            in_0 = rnn_sequence->get_input_source_output(0).get_node_shared_ptr()->get_input_source_output(0);
        }

        auto W = rnn_sequence->input_value(3);
        auto R = rnn_sequence->input_value(4);

        // for forward/reverse cases we can squeeze num_direction dimension
        auto axis_1 = ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {1});
        auto in_1 = std::make_shared<ngraph::opset5::Squeeze>(rnn_sequence->input_value(1), axis_1);
        auto concat = std::make_shared<ngraph::opset5::Concat>(ngraph::OutputVector{W, R}, 2);
        auto axis_2 = ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0});
        auto in_3 = std::make_shared<ngraph::opset5::Squeeze>(concat->output(0), axis_2);
        auto in_4 = std::make_shared<ngraph::opset5::Squeeze>(rnn_sequence->input_value(5), axis_2);
        auto rnn_sequence_ie = std::make_shared<ngraph::op::RNNSequenceIE>(
                in_0,  // X
                in_1,  // initial_hidden_state
                rnn_sequence->input_value(2),
                in_3,  // WR
                in_4,  // B
                rnn_sequence->get_hidden_size(),
                rnn_sequence->get_direction(),
                rnn_sequence->get_activations(),
                rnn_sequence->get_activations_alpha(),
                rnn_sequence->get_activations_beta(),
                rnn_sequence->get_clip(),
                seq_axis);

        auto unsqueeze_axis = ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {1});
        auto unsqueeze_1 = std::make_shared<ngraph::opset5::Unsqueeze>(rnn_sequence_ie->output(0), unsqueeze_axis);
        auto unsqueeze_2 = std::make_shared<ngraph::opset5::Unsqueeze>(rnn_sequence_ie->output(1), unsqueeze_axis);

        ngraph::copy_runtime_info(rnn_sequence, {concat, rnn_sequence_ie, in_1, in_3, in_4, unsqueeze_1,
                                                 unsqueeze_2});
        unsqueeze_1->set_friendly_name(rnn_sequence->get_friendly_name()+".0");
        unsqueeze_2->set_friendly_name(rnn_sequence->get_friendly_name()+".1");

        if (seq_axis == 1) {
            ngraph::replace_node(rnn_sequence, {unsqueeze_1->output(0), unsqueeze_2->output(0)});
        } else {
            const auto &rnn_target_inputs = rnn_sequence->output(0).get_target_inputs();
            if (rnn_target_inputs.empty())
                return false;
            auto transpose_after = rnn_target_inputs.begin()->get_node()->shared_from_this();
            unsqueeze_1->set_friendly_name(transpose_after->get_friendly_name());
            ngraph::replace_node(transpose_after, unsqueeze_1);
            ngraph::replace_node(rnn_sequence, {rnn_sequence_ie->output(0), unsqueeze_2->output(0)});
        }
        return true;
    };

    auto m = std::make_shared<ngraph::pattern::Matcher>(rnn_sequence_ngraph, "ConvertRNNSequenceToRNNSequenceIE");
    this->register_matcher(m, callback);
}