//*****************************************************************************
// Copyright 2017-2021 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 "gtest/gtest.h"
#include "ngraph/builder/reshape.hpp"
#include "ngraph/ngraph.hpp"
#include "ngraph/opsets/opset5.hpp"
#include "util/type_prop.hpp"

using namespace std;
using namespace ngraph;

TEST(type_prop, tensor_iterator_lstm)
{
    // That which we iterate over
    const size_t N = 32; // Batch size
    const size_t L = 10; // Sequence length
    const size_t I = 8;  // Input size
    const size_t H = 32; // Hidden size
    auto SENT = make_shared<op::Parameter>(element::f32, Shape{N, L, I});

    auto H_init = make_shared<op::Parameter>(element::f32, Shape{N, 1, H});
    auto C_init = make_shared<op::Parameter>(element::f32, Shape{N, 1, H});

    auto W = make_shared<op::Parameter>(element::f32, Shape{4 * H, I});
    auto R = make_shared<op::Parameter>(element::f32, Shape{4 * H, H});
    auto H_t = make_shared<op::Parameter>(element::f32, Shape{N, 1, H});
    auto C_t = make_shared<op::Parameter>(element::f32, Shape{N, 1, H});

    // Body
    auto X = make_shared<op::Parameter>(element::f32, Shape{N, 1, I});
    auto W_body = make_shared<op::Parameter>(element::f32, Shape{4 * H, I});
    auto R_body = make_shared<op::Parameter>(element::f32, Shape{4 * H, H});
    auto LSTM_cell = make_shared<opset5::LSTMCell>(builder::opset1::reshape(X, Shape{N, I}),
                                                   builder::opset1::reshape(H_t, Shape{N, H}),
                                                   builder::opset1::reshape(C_t, Shape{N, H}),
                                                   W_body,
                                                   R_body,
                                                   H);
    auto H_o = builder::opset1::reshape(LSTM_cell->output(0), Shape{N, 1, H});
    auto C_o = builder::opset1::reshape(LSTM_cell->output(1), Shape{N, 1, H});
    auto body = make_shared<ngraph::Function>(OutputVector{H_o, C_o},
                                              ParameterVector{X, H_t, C_t, W_body, R_body});

    auto tensor_iterator = make_shared<op::TensorIterator>();
    tensor_iterator->set_body(body);
    // start=0, stride=1, part_size=1, end=39, axis=1
    tensor_iterator->set_sliced_input(X, SENT, 0, 1, 1, -1, 1);
    // H_t is Hinit on the first iteration, Ho after that
    tensor_iterator->set_merged_input(H_t, H_init, H_o);
    tensor_iterator->set_merged_input(C_t, C_init, C_o);
    tensor_iterator->set_invariant_input(W_body, W);
    tensor_iterator->set_invariant_input(R_body, R);

    // Output 0 is last Ho, result 0 of body
    auto out0 = tensor_iterator->get_iter_value(H_o, -1);
    // Output 1 is last Co, result 1 of body
    auto out1 = tensor_iterator->get_iter_value(C_o, -1);

    auto results = ResultVector{make_shared<op::Result>(out0), make_shared<op::Result>(out1)};
    auto f = make_shared<Function>(results, ParameterVector{SENT, H_init, C_init, W, R});
}

TEST(type_prop, tensor_iterator_2_slice_inputs_part_size_2)
{
    // That which we iterate over
    auto X = make_shared<op::Parameter>(element::f32, Shape{32, 40, 10});
    auto Y = make_shared<op::Parameter>(element::f32, Shape{32, 40, 10});
    auto M = make_shared<op::Parameter>(element::f32, Shape{32, 2, 10});

    // Set up the cell body, a function from (Xi, Yi) -> (Zo)
    // Body parameters
    auto Xi = make_shared<op::Parameter>(element::f32, Shape{32, 2, 10});
    auto Yi = make_shared<op::Parameter>(element::f32, Shape{32, 2, 10});
    auto M_body = make_shared<op::Parameter>(element::f32, Shape{32, 2, 10});

    // Body
    auto Zo = std::make_shared<op::v1::Multiply>(std::make_shared<op::v1::Add>(Xi, Yi), M_body);
    auto body = make_shared<ngraph::Function>(OutputVector{Zo}, ParameterVector{Xi, Yi, M_body});

    auto tensor_iterator = make_shared<op::TensorIterator>();
    tensor_iterator->set_body(body);
    // The Xi are the elements of Xseq
    // start=0, stride=2, part_size=2, end=39, axis=1
    tensor_iterator->set_sliced_input(Xi, X, 0, 2, 2, 39, 1);
    // The Yi are the elements of Yseq
    // start=0, stride=2, part_size=2, end=-1, axis=1
    tensor_iterator->set_sliced_input(Yi, Y, 0, 2, 2, -1, 1);
    tensor_iterator->set_invariant_input(M_body, M);

    // Output 0 is last Zo
    auto out0 = tensor_iterator->get_iter_value(Zo, -1);
    // Output 1 is concat of Zos
    // start=0, stride=2, part_size=2, end=39, axis=1
    auto out1 = tensor_iterator->get_concatenated_slices(Zo, 0, 2, 2, 39, 1);
    auto result0 = make_shared<op::Result>(out0);
    auto result1 = make_shared<op::Result>(out1);
    Shape out0_shape{32, 2, 10};
    Shape out1_shape{32, 40, 10};

    auto results = ResultVector{result0, result1};
    auto f = make_shared<Function>(results, ParameterVector{X, Y, M});
    EXPECT_EQ(result0->get_output_shape(0), out0_shape);
    EXPECT_EQ(result1->get_output_shape(0), out1_shape);
}

TEST(type_prop, tensor_iterator_2_slice_inputs_part_size_2_dynamic)
{
    // That which we iterate over
    auto X = make_shared<op::Parameter>(element::f32, Shape{32, 40, 10});
    auto Y = make_shared<op::Parameter>(element::f32, Shape{32, 40, 10});
    auto M = make_shared<op::Parameter>(element::f32, Shape{32, 2, 10});

    // Set up the cell body, a function from (Xi, Yi) -> (Zo)
    // Body parameters
    auto Xi = make_shared<op::Parameter>(element::f32, PartialShape::dynamic());
    auto Yi = make_shared<op::Parameter>(element::f32, PartialShape::dynamic());
    auto M_body = make_shared<op::Parameter>(element::f32, PartialShape::dynamic());

    // Body
    auto Zo = std::make_shared<op::v1::Multiply>(std::make_shared<op::v1::Add>(Xi, Yi), M_body);
    auto body = make_shared<ngraph::Function>(OutputVector{Zo}, ParameterVector{Xi, Yi, M_body});

    auto tensor_iterator = make_shared<op::TensorIterator>();
    tensor_iterator->set_body(body);
    // The Xi are the elements of Xseq
    // start=0, stride=2, part_size=2, end=38, axis=1
    tensor_iterator->set_sliced_input(Xi, X, 0, 2, 2, 38, 1);
    // The Yi are the elements of Yseq
    // start=0, stride=2, part_size=2, end=-2, axis=1
    tensor_iterator->set_sliced_input(Yi, Y, 0, 2, 2, -2, 1);
    tensor_iterator->set_invariant_input(M_body, M);

    // check input descriptors
    for (auto& desc : tensor_iterator->get_input_descriptions())
    {
        auto type_info = desc->get_type_info();
        if (std::strcmp(type_info.name, "InvariantInputDescription") == 0)
        {
            auto input_desc =
                as_type_ptr<ngraph::op::TensorIterator::InvariantInputDescription>(desc);
            EXPECT_NE(input_desc, nullptr);
        }
        else if (std::strcmp(type_info.name, "SliceInputDescription") == 0)
        {
            auto input_desc = as_type_ptr<ngraph::op::TensorIterator::SliceInputDescription>(desc);
            EXPECT_NE(input_desc, nullptr);
        }
        else if (std::strcmp(type_info.name, "MergedInputDescription") == 0)
        {
            auto input_desc = as_type_ptr<ngraph::op::TensorIterator::MergedInputDescription>(desc);
            EXPECT_NE(input_desc, nullptr);
        }
    }

    // Output 0 is last Zo
    auto out0 = tensor_iterator->get_iter_value(Zo, -1);
    // Output 1 is concat of Zos
    // start=0, stride=2, part_size=2, end=38, axis=1
    auto out1 = tensor_iterator->get_concatenated_slices(Zo, 0, 2, 2, 38, 1);

    // check output descriptors
    for (auto& desc : tensor_iterator->get_output_descriptions())
    {
        auto type_info = desc->get_type_info();
        if (std::strcmp(type_info.name, "ConcatOutputDescription") == 0)
        {
            auto output_desc =
                as_type_ptr<ngraph::op::TensorIterator::ConcatOutputDescription>(desc);
            EXPECT_NE(output_desc, nullptr);
        }
        else if (std::strcmp(type_info.name, "BodyOutputDescription") == 0)
        {
            auto output_desc = as_type_ptr<ngraph::op::TensorIterator::BodyOutputDescription>(desc);
            EXPECT_NE(output_desc, nullptr);
        }
    }
    auto result0 = make_shared<op::Result>(out0);
    auto result1 = make_shared<op::Result>(out1);
    Shape out0_shape{32, 2, 10};
    Shape out1_shape{32, 38, 10};

    auto results = ResultVector{result0, result1};
    auto f = make_shared<Function>(results, ParameterVector{X, Y, M});
    EXPECT_EQ(result0->get_output_shape(0), out0_shape);
    EXPECT_EQ(result1->get_output_shape(0), out1_shape);

    EXPECT_EQ(body->get_results()[0]->get_output_shape(0), out0_shape);
}
