//*****************************************************************************
// Copyright 2017-2020 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//*****************************************************************************

#include <cstdint>
#include <cstdlib>
#include <vector>

#include "default_opset.hpp"
#include "ngraph/builder/autobroadcast.hpp"
#include "ngraph/builder/reshape.hpp"
#include "ngraph/builder/split.hpp"
#include "ngraph/check.hpp"
#include "ngraph/enum_names.hpp"
#include "recurrent.hpp"

namespace ngraph
{
    namespace onnx_import
    {
        namespace recurrent
        {
            OpInputMap::OpInputMap(const onnx_import::Node& node, std::size_t gates_count)
            {
                const auto& ng_inputs = node.get_ng_inputs();

                m_map[OpInput::X] = ng_inputs.at(0);
                m_map[OpInput::W] = ng_inputs.at(1);
                m_map[OpInput::R] = ng_inputs.at(2);

                const auto el_type = ng_inputs.at(0)->get_output_element_type(0);

                const auto x_pshape = m_map[OpInput::X]->get_output_partial_shape(0);
                const auto w_pshape = m_map[OpInput::W]->get_output_partial_shape(0);
                const auto r_pshape = m_map[OpInput::R]->get_output_partial_shape(0);
                NGRAPH_CHECK(x_pshape.rank().is_static() && x_pshape[0].is_static() &&
                                 x_pshape[1].is_static(),
                             "RecurrentSequence input X must have static \"seq_length\" and "
                             "\"batch_size\" dimensions.");
                NGRAPH_CHECK(w_pshape.rank().is_static() && w_pshape[0].is_static(),
                             "RecurrentSequence input W must have static \"num_directions\" "
                             "(outermost) dimension.");
                NGRAPH_CHECK(r_pshape.rank().is_static() && r_pshape[2].is_static(),
                             "RecurrentSequence input R must have static \"hidden_size\" "
                             "(innermost) dimension.");

                const std::size_t hidden_size = m_map[OpInput::R]->get_shape().back();
                const std::size_t batch_size = m_map[OpInput::X]->get_shape().at(1);
                const std::size_t num_directions = m_map[OpInput::W]->get_shape().front();

                if (ng_inputs.size() > 3 && !ng_inputs.at(3)->is_null())
                {
                    auto bias = ng_inputs.at(3);
                    auto split_bias = builder::opset1::split(bias, 2, 1);
                    m_map[OpInput::B] = split_bias.at(0) + split_bias.at(1);
                }
                else
                {
                    m_map[OpInput::B] = std::make_shared<default_opset::Constant>(
                        el_type, Shape{num_directions, gates_count * hidden_size}, 0.f);
                }
                if (ng_inputs.size() > 4 && !ng_inputs.at(4)->is_null())
                {
                    m_map[OpInput::SEQ_LENGTHS] = ng_inputs.at(4);
                }
                else
                {
                    m_map[OpInput::SEQ_LENGTHS] = std::make_shared<default_opset::Constant>(
                        element::i32, Shape{batch_size}, m_map[OpInput::X]->get_shape().at(0));
                }
                // The initial value of the hidden.
                if (ng_inputs.size() > 5 && !ng_inputs.at(5)->is_null())
                {
                    m_map[OpInput::INIT_H] = ng_inputs.at(5);
                }
                else
                {
                    m_map[OpInput::INIT_H] = std::make_shared<default_opset::Constant>(
                        el_type, Shape{num_directions, batch_size, hidden_size}, 0.f);
                }
            }

            OpInputMap::OpInputMap(container_type&& map)
                : m_map(std::move(map))
            {
            }

            std::shared_ptr<ngraph::Node>& OpInputMap::at(const OpInput& key)
            {
                return m_map.at(key);
            }

            const std::shared_ptr<ngraph::Node>& OpInputMap::at(const OpInput& key) const
            {
                return m_map.at(key);
            }

            // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ATTRIBUTES PARSING ~~~~~~~~~~~~~~~~~~~~~~~~~~~~

            OpAttributes::OpAttributes(const Node& node)
                : m_hidden_size{node.get_attribute_value<std::int64_t>("hidden_size")}
                , m_clip_threshold{node.get_attribute_value<float>("clip", 0.f)}
                // Recurrent Operators which have more activation functions should override
                // this value in constructor of respective Attributes' struct.
                , m_activations{node.get_attribute_value<std::vector<std::string>>("activations",
                                                                                   {"tanh"})}
                // Default values for activation functions are same as for corresponding
                // ONNX operator.
                , m_activations_alpha{node.get_attribute_value<std::vector<float>>(
                      "activation_alpha", std::vector<float>{})}
                , m_activations_beta{node.get_attribute_value<std::vector<float>>(
                      "activation_beta", std::vector<float>{})}
            {
                m_clip_threshold = std::abs(m_clip_threshold);
                std::string direction =
                    ngraph::to_lower(node.get_attribute_value<std::string>("direction", "forward"));
                m_direction = ngraph::as_enum<ngraph::op::RecurrentSequenceDirection>(direction);
            }

            // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Sequence Computations ~~~~~~~~~~~~~~~~~~~~~~~~~~~~

            RecurrentSequence::RecurrentSequence(OpInputMap& args,
                                                 ngraph::op::RecurrentSequenceDirection direction)
                : m_args(args)
                , m_direction(direction)
            {
            }

            NodeVector RecurrentSequence::run_sequence(const RecurrentCellFunction& kernel)
            {
                NodeVector results;
                if (m_direction == ngraph::op::RecurrentSequenceDirection::FORWARD ||
                    m_direction == ngraph::op::RecurrentSequenceDirection::REVERSE)
                {
                    results = recurrent_sequence_pass(
                        kernel, m_direction == ngraph::op::RecurrentSequenceDirection::REVERSE);
                }
                else if (m_direction == ngraph::op::RecurrentSequenceDirection::BIDIRECTIONAL)
                {
                    NodeVector fwd_results{recurrent_sequence_pass(kernel)};
                    NodeVector rev_results{recurrent_sequence_pass(kernel, true)};

                    // Stack together respective outputs from both forward and reverse passess.
                    std::shared_ptr<ngraph::Node> Y{std::make_shared<default_opset::Concat>(
                        NodeVector{fwd_results.at(0), rev_results.at(0)}, 1)};
                    results.push_back(Y);

                    std::shared_ptr<ngraph::Node> Y_h{std::make_shared<default_opset::Concat>(
                        NodeVector{fwd_results.at(1), rev_results.at(1)}, 0)};
                    results.push_back(Y_h);
                }
                else
                {
                    throw ngraph_error(
                        "RecurrentSequence: unhandled direction mode during decomposition.");
                }
                return results;
            }

            NodeVector
                RecurrentSequence::recurrent_sequence_pass(const RecurrentCellFunction& kernel,
                                                           bool is_reverse)
            {
                OutputVector h_list;

                // back-up nodes which we may later modify.
                std::shared_ptr<ngraph::Node> orig_W = m_args.at(OpInput::W);
                std::shared_ptr<ngraph::Node> orig_R = m_args.at(OpInput::R);
                std::shared_ptr<ngraph::Node> orig_B = m_args.at(OpInput::B);

                std::shared_ptr<ngraph::Node> X = m_args.at(OpInput::X);
                std::shared_ptr<ngraph::Node> H_t =
                    prepare_input(m_args.at(OpInput::INIT_H), is_reverse);
                std::shared_ptr<ngraph::Node> W = prepare_input(m_args.at(OpInput::W), is_reverse);
                std::shared_ptr<ngraph::Node> R = prepare_input(m_args.at(OpInput::R), is_reverse);
                std::shared_ptr<ngraph::Node> B = prepare_input(m_args.at(OpInput::B), is_reverse);
                std::shared_ptr<ngraph::Node> seq_lengths = m_args.at(OpInput::SEQ_LENGTHS);

                m_args.at(OpInput::W) = W;
                m_args.at(OpInput::R) = R;
                m_args.at(OpInput::B) = B;

                if (is_reverse)
                {
                    X = std::make_shared<default_opset::ReverseSequence>(
                        X, seq_lengths, 1 /*batch_axis*/, 0 /*seq_axis*/);
                }

                OutputVector in_seq_steps =
                    as_output_vector(builder::opset1::split(X, X->get_shape().at(0)));

                for (auto& in_x : in_seq_steps)
                {
                    // remove first empty dim, after above split.
                    in_x = builder::opset1::squeeze(in_x);
                }

                int32_t time_step{1};
                for (const auto& in_x : in_seq_steps)
                {
                    Output<ngraph::Node> H = kernel(m_args, in_x, H_t);

                    // Expand tensors with empty outermost dim, so we can later concatenate
                    // them.
                    // Mask hidden state tensor in order to handle mixed sequence lengths.
                    // This results in zeroing out values in batches with sequence shorter
                    // than current time_step.
                    h_list.push_back(
                        get_masked_node(builder::opset1::expand_dims(H), time_step, 1));

                    // Here we make sure that only appropriate batches (with respect to its sequence
                    // length) are updated. Those batches which has shorter sequences preserve
                    // the last value.
                    H_t = get_masked_node(H, time_step, 0, H_t);
                    time_step++;
                }

                // Get back original nodes.
                m_args.at(OpInput::W) = orig_W;
                m_args.at(OpInput::R) = orig_R;
                m_args.at(OpInput::B) = orig_B;

                // The tensor that concats all the intermediate output values of the hidden.
                // It has shape [seq_length, batch_size, hidden_size]
                std::shared_ptr<ngraph::Node> Y{std::make_shared<default_opset::Concat>(h_list, 0)};

                // Get back the original order of the output data.
                if (is_reverse)
                {
                    Y = std::make_shared<default_opset::ReverseSequence>(
                        Y, seq_lengths, 1 /*batch_axis*/, 0 /*seq_axis*/);
                }

                // Expand Y so that it has expected shape:
                // [seq_length, num_directions, batch_size, hidden_size]
                Y = builder::opset1::expand_dims(Y, 1);

                // Expand H_t so that it has expected shape:
                // [num_directions, batch_size, hidden_size]
                auto Y_h = builder::opset1::expand_dims(H_t);

                return {Y, Y_h};
            }

            std::shared_ptr<ngraph::Node>
                RecurrentSequence::get_masked_node(const Output<ngraph::Node>& data,
                                                   int32_t time_step,
                                                   size_t batch_axis,
                                                   const Output<ngraph::Node>& default_value) const
            {
                Output<ngraph::Node> mask_value = default_value;
                // Create zero mask value node.
                if (!mask_value.get_node_shared_ptr())
                {
                    mask_value = std::make_shared<default_opset::Constant>(
                        data.get_element_type(), data.get_shape(), 0.f);
                }

                // Create predicate nodes. The condition is whether current time step value
                // is greater than sequence length for respective batch inputs.
                std::shared_ptr<ngraph::Node> curr_time_step_node =
                    std::make_shared<default_opset::Constant>(
                        element::i32, data.get_shape(), time_step);

                Output<ngraph::Node> batch_seq_length =
                    builder::opset1::legacy_broadcast_for_binary_operation(
                        curr_time_step_node, m_args.at(OpInput::SEQ_LENGTHS), batch_axis);

                // Create mask node deciding whether or not to mask batch data.
                std::shared_ptr<ngraph::Node> mask_condition =
                    std::make_shared<default_opset::Greater>(curr_time_step_node, batch_seq_length);

                // Select values depnding on mask_condition.
                // Select(<condition>, <true_value>, <false_value>)
                return std::make_shared<default_opset::Select>(mask_condition, mask_value, data);
            }

            std::shared_ptr<ngraph::Node>
                RecurrentSequence::prepare_input(Output<ngraph::Node> node, bool is_reverse) const
            {
                // In bidirectional mode inputs are stacked together, so we must split them.
                std::shared_ptr<ngraph::Node> tmp = node.get_node_shared_ptr();
                if (m_direction == ngraph::op::RecurrentSequenceDirection::BIDIRECTIONAL)
                {
                    tmp = builder::opset1::split(node, 2).at(is_reverse ? 1 : 0);
                }
                // Since we work in forward pass mode, we can squeeze `num_directions` axis from
                // input.
                return builder::opset1::squeeze(tmp);
            }

        } // recurrent
    }     // onnx_import
} // ngraph
