//*****************************************************************************
// Copyright 2017-2019 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 <cstddef>
#include <iterator>
#include <memory>
#include <vector>

#include "exceptions.hpp"
#include "matmul.hpp"
#include "ngraph/coordinate.hpp"
#include "ngraph/log.hpp"
#include "ngraph/op/concat.hpp"
#include "ngraph/op/dot.hpp"
#include "ngraph/op/experimental/quantized_dot.hpp"
#include "ngraph/op/reshape.hpp"
#include "ngraph/op/slice.hpp"
#include "ngraph/op/util/broadcasting.hpp"
#include "ngraph/op/util/reshape.hpp"
#include "ngraph/shape.hpp"
#include "ngraph/util.hpp"
#include "utils/reshape.hpp"

/// \brief      Slice the sub matrix from the input tensor.
///
/// \param[in]  node  The input tensor. Must be at most of rank 3.
/// \param[in]  idx   The index on the first axis, at which to slice sub-matrix.
///
/// \return     The node representing sub matrix.
///
static std::shared_ptr<ngraph::Node> get_sub_matrix(const std::shared_ptr<ngraph::Node>& node,
                                                    std::size_t idx)
{
    const ngraph::Shape& shape{node->get_shape()};
    if (shape.size() < 3)
    {
        return node;
    }
    // Below bounds defines the sub_matrix through ranges for each input node axis.
    ngraph::Coordinate lower_bounds(shape.size());
    ngraph::Coordinate upper_bounds = shape;
    // We assume `node` tensor is of rank equal 3, thus we slice the sub-matrix lying in the last
    // two dimensions at index `idx` of first axis.
    lower_bounds.at(0) = idx;
    upper_bounds.at(0) = idx + 1;

    auto sub_matrix = std::shared_ptr<ngraph::Node>{
        std::make_shared<ngraph::op::Slice>(node, lower_bounds, upper_bounds)};
    // Remove first single entry dim.
    return ngraph::onnx_import::reshape::squeeze(sub_matrix);
}

namespace ngraph
{
    namespace onnx_import
    {
        namespace op
        {
            namespace set_1
            {
                NodeVector make_matmul_op(const Node& node, bool quantized)
                {
                    const NodeVector& ng_inputs{node.get_ng_inputs()};
                    auto left = std::shared_ptr<ngraph::Node>{ng_inputs.at(0)};
                    auto right = std::shared_ptr<ngraph::Node>{};
                    auto scale = std::shared_ptr<ngraph::Node>{};
                    if (quantized)
                    {
                        NGRAPH_WARN
                            << "[" << node.get_name()
                            << "] Zero point different from 0 is not supported. Assuming Zero "
                               "point is 0";
                        right = ng_inputs.at(3);
                        scale = ng_inputs.at(6);
                    }
                    else
                    {
                        right = ng_inputs.at(1);
                    }

                    std::size_t left_rank{left->get_shape().size()};
                    std::size_t right_rank{right->get_shape().size()};

                    if (left_rank == 0 || right_rank == 0)
                    {
                        NGRAPH_WARN
                            << (node) << " "
                            << "ONNX standard doesn't allow scalar operands, however nGraph "
                               "accepts them. Consider use of element-wise multiplication instead "
                               "to conform with ONNX standard.";
                    }

                    // First (easy) case that is already internally handled by Ngraph Dot operator.
                    // Multiply two tensors where both of them has rank lower equal 2.
                    if (left_rank <= 2 && right_rank <= 2)
                    {
                        if (quantized)
                        {
                            right = std::make_shared<ngraph::op::Reshape>(
                                right,
                                AxisVector{1, 0},
                                Shape(right->get_shape().rbegin(), right->get_shape().rend()));

                            return {std::make_shared<ngraph::op::QuantizedDot>(left, right, scale)};
                        }
                        else
                        {
                            return {std::make_shared<ngraph::op::Dot>(left, right)};
                        }
                    }

                    // Second case:
                    // Multiply two tensors where at least one of them is rank greater equal 3.

                    // Broadcast input arguments only if both of them are not vectors.
                    if (left_rank > 1 && right_rank > 1)
                    {
                        const NodeVector& broadcasted_nodes =
                            ngraph::op::numpy_style_broadcast_for_matmul_operation(left, right);

                        left = broadcasted_nodes.at(0);
                        right = broadcasted_nodes.at(1);
                    }
                    const auto& left_shape = left->get_shape();
                    const auto& right_shape = right->get_shape();

                    // Collapse both tensors _stack of matrices_ axes (all except the last two).
                    // This will make easier further dot product calculations.
                    if (left_shape.size() > 3)
                    {
                        left = reshape::collapse(left, 0, left_shape.size() - 3);
                    }
                    if (right_shape.size() > 3)
                    {
                        right = reshape::collapse(right, 0, right_shape.size() - 3);
                    }

                    // Perform multiple small dot products
                    std::size_t groups = left->get_shape().at(0);
                    // If we haven't broadcast earlier this means that one of the inputs is a vector,
                    // thus the number of groups is defined by the shape of the bigger tensor.
                    if (right->get_shape().size() > left->get_shape().size())
                    {
                        groups = right->get_shape().at(0);
                    }
                    NodeVector small_dots(groups);

                    for (std::size_t g = 0; g < groups; ++g)
                    {
                        const auto& sliced_left = get_sub_matrix(left, g);
                        auto sliced_right = get_sub_matrix(right, g);

                        auto sub_dot = std::shared_ptr<ngraph::Node>{};

                        if (quantized)
                        {
                            sliced_right = std::make_shared<ngraph::op::Reshape>(
                                sliced_right,
                                AxisVector{1, 0},
                                Shape(sliced_right->get_shape().rbegin(),
                                      sliced_right->get_shape().rend()));
                            sub_dot = std::make_shared<ngraph::op::QuantizedDot>(
                                sliced_left, sliced_right, scale);
                        }
                        else
                        {
                            sub_dot = std::make_shared<ngraph::op::Dot>(sliced_left, sliced_right);
                        }

                        // Expand sub_dot result with single empty outermost axis, in order to
                        // later concatenate sub_dots at this axis.
                        small_dots.at(g) = reshape::expand_dims(sub_dot);
                    }

                    // Concatenate sub_dots on groups axis.
                    auto result = std::make_shared<ngraph::op::Concat>(small_dots, 0);

                    if (left_shape.size() <= 3 && right_shape.size() <= 3)
                    {
                        return {result};
                    }
                    // Expand result _stack of matrices_ axes to get expected result shape.
                    else
                    {
                        const Shape& shape{result->get_shape()};
                        Shape result_shape(std::next(std::begin(shape)), std::end(shape));
                        result_shape.insert(
                            std::begin(result_shape),
                            std::begin(left_shape),
                            std::next(std::begin(left_shape), left_shape.size() - 2));
                        return {std::make_shared<ngraph::op::Reshape>(
                            result, ngraph::get_default_order(shape.size()), result_shape)};
                    }
                }

                NodeVector matmul(const Node& node) { return make_matmul_op(node, false); }
            } // namespace set_1

        } //namespace op

    } // namespace onnx_import

} // namespace ngraph
