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

///////////////////////////////////////////////////////////////////////////////////////////////////
#pragma once
#include "primitive.hpp"

namespace cldnn {
/// @addtogroup cpp_api C++ API
/// @{
/// @addtogroup cpp_topology Network Topology
/// @{
/// @addtogroup cpp_primitives Primitives
/// @{

/// @brief SpaceToBatch operation divides "spatial" dimensions [1, ..., N - 1], N ∈ {4,5,6} of the data input
/// into a grid of blocks of shape block_shape, and interleaves these blocks with the batch dimension (0) such that in the output,
/// the spatial dimensions [1, ..., N - 1], N ∈ {4,5,6} correspond to the position within the grid,
/// and the batch dimension combines both the position within a spatial block and the original batch position.
/// Prior to division into blocks, the spatial dimensions of the input are optionally zero padded according to pads_begin and pads_end.
/// @details The SpaceToBatch operation is similar to the TensorFlow* operation SpaceToBatchND (https://www.tensorflow.org/api_docs/python/tf/space_to_batch_nd)
/// There are 4 inputs of this operation:
/// 1) data - input N-D tensor [batch, D_1, D_2 ... D_{N-1}], N ∈ {4,5,6}. Required.
/// 2) block_shape - input 1-D tensor with shape [N], N ∈ {4,5,6}. Consists of block_sizes each of which specifies the size of the value block to be moved.
/// All values must be >= 1 and required. block_shape[0] is expected to be 1.
/// 3) pads_begin - input 1-D tensor with shape [N], N ∈ {4,5,6}. Specifies the padding for the beginning along each axis of data input.
/// All values must be non-negative and required. pads_begin[0] is expected to be 0.
/// 4) pads_end - input 1-D tensor with shape [N], N ∈ {4,5,6}. Specifies the padding for the ending along each axis of data input.
/// All values must be non-negative and required. pads_end[0] is expected to be 0.
/// 3-4 inputs required that block_shape[i] divides data_shape[i] + pads_begin[i] + pads_end[i]
///
/// The operation is equivalent to the following transformation of the input tensor data of shape [batch, D_1, D_2 ... D_{N - 1}], N ∈ {4,5,6}
/// and block_shape, pads_begin, pads_end of shapes [N] to Y output tensor.
/// Zero-pad the start and end of dimensions [D_0, ..., D_{N - 1}] of the input according to `pads_begin` and `pads_end`
///
/// x' = reshape(x, [batch, (D_1 + P_1) / B_1, B_1, (D_2 + P_2) / B_2, B_2, ..., (D_{N - 1} + P_{N - 1}) / B_{N - 1}, B_{N - 1}]), where B_i = block_shape[i]
///
/// x'' = transpose(x',  [2, 4, ..., (N - 1) + (N - 1), 0, 1, 3, ..., N + (N - 1)])
///
/// y = reshape(x'', [batch * B_1 * ... * B_{N - 1}, (D_1 + P_1) / B_1, (D_2 + P_2) / B_2, ... , (D_{N - 1} + P_{N - 1}) / B_{N - 1}])

struct space_to_batch : public primitive_base<space_to_batch> {
    CLDNN_DECLARE_PRIMITIVE(space_to_batch)

    /// @brief Constructs space_to_batch primitive.
    /// @param id This primitive id.
    /// @param input Input data primitive id.
    /// @param block_shape Array of block sizes.
    /// @param pads_begin Amount to pad for the beginning along each axis of data input.
    /// @param pads_end Amount to pad for the ending along each axis of data input.
    /// @param out_size Size of output tensor.
    space_to_batch(const primitive_id& id,
                   const primitive_id& input,
                   const tensor& block_shape,
                   const tensor& pads_begin,
                   const tensor& pads_end,
                   const tensor& out_size,
                   const primitive_id& ext_prim_id = "",
                   const padding& output_padding = padding())
        : primitive_base(id, {input}, ext_prim_id, output_padding),
          block_shape(block_shape),
          pads_begin(pads_begin),
          pads_end(pads_end),
          out_size(out_size) {}

    tensor block_shape;
    tensor pads_begin;
    tensor pads_end;
    tensor out_size;
};
/// @}
/// @}
/// @}
}  // namespace cldnn
