/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/

#ifndef THIRD_PARTY_TENSORFLOW_CORE_KERNELS_EIGEN_PATCH_3D_H_
#define THIRD_PARTY_TENSORFLOW_CORE_KERNELS_EIGEN_PATCH_3D_H_

#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"

#if not defined(__CUDACC__)
#include <type_traits>
#endif

namespace Eigen {
namespace internal {

/** Extract3DPatches
 * \ingroup CXX11_NeuralNetworksModule
 *
 * \brief Extracts 3D patches from a multichannel input volume.
 *
 * The input parameter is expected to be a tensor with a rank of 4 or more
 * (channels, depth, height, width, optional others in col-major, and the
 * reverse order in row-major).

 * The return value will be a tensor of 3 more dimension than the input tensor.
 * In col-major, the first 4 dimensions of the result are: channels,
 patch_depth,
 * patch_height, patch_width. The next dimensions will identify the patch
 * position on the 3D grid of extracted patches: z, y, x. The remaining
 * dimensions, if any, will be the same as the 'other' dimensions of the input
 * tensor.
 */

template <typename Input>
EIGEN_ALWAYS_INLINE static const TensorStridingOp<
    const array<typename internal::traits<Input>::Index,
                internal::traits<Input>::NumDimensions + 3>,
    const TensorReshapingOp<
        const DSizes<typename internal::traits<Input>::Index,
                     internal::traits<Input>::NumDimensions + 3>,
        const TensorPatchOp<
            const DSizes<typename internal::traits<Input>::Index,
                         internal::traits<Input>::NumDimensions>,
            const TensorPaddingOp<
                const array<IndexPair<typename internal::traits<Input>::Index>,
                            internal::traits<Input>::NumDimensions>,
                const Input> > > >
Extract3DPatches(
    const Input& input, const DenseIndex patchPlanes,
    const DenseIndex patchRows, const DenseIndex patchCols,
    const DenseIndex stridePlanes, const DenseIndex strideRows,
    const DenseIndex strideCols, const DenseIndex paddingZTop,
    const DenseIndex paddingZBottom, const DenseIndex paddingTop,
    const DenseIndex paddingBottom, const DenseIndex paddingLeft,
    const DenseIndex paddingRight,
    const typename internal::traits<Input>::Scalar padding_value = 0) {
  typedef typename internal::traits<Input>::Index TensorIndex;
  TensorRef<Tensor<typename internal::traits<Input>::Scalar,
                   internal::traits<Input>::NumDimensions,
                   internal::traits<Input>::Layout, TensorIndex> >
      in(input);

  EIGEN_STATIC_ASSERT(internal::traits<Input>::NumDimensions >= 4,
                      YOU_MADE_A_PROGRAMMING_MISTAKE);

  static const bool isColMajor = (internal::traits<Input>::Layout == ColMajor);
  static const int NumDims = internal::traits<Input>::NumDimensions;
  static const int ExtDims = NumDims + 3;

  // Tensor size after patch extraction. We add three dimensions to unpack the
  // linear patch index into a 3D grid over which stride() can work.
  DSizes<TensorIndex, ExtDims> pre_stride_dims;

  if (isColMajor) {
    pre_stride_dims[0] = in.dimension(0);
    pre_stride_dims[1] = patchPlanes;
    pre_stride_dims[2] = patchRows;
    pre_stride_dims[3] = patchCols;
  } else {
    pre_stride_dims[ExtDims - 1] = in.dimension(NumDims - 1);
    pre_stride_dims[ExtDims - 4] = patchCols;
    pre_stride_dims[ExtDims - 3] = patchRows;
    pre_stride_dims[ExtDims - 2] = patchPlanes;
  }

  const TensorIndex inputPlanes =
      isColMajor ? in.dimension(1) : in.dimension(NumDims - 2);
  const TensorIndex inputRows =
      isColMajor ? in.dimension(2) : in.dimension(NumDims - 3);
  const TensorIndex inputCols =
      isColMajor ? in.dimension(3) : in.dimension(NumDims - 4);

  array<IndexPair<TensorIndex>, NumDims> paddings;
  for (int i = 0; i < NumDims; ++i) {
    paddings[i] = IndexPair<TensorIndex>(0, 0);
  }

  paddings[isColMajor ? 1 : (NumDims - 2)] =
      IndexPair<TensorIndex>(paddingZTop, paddingZBottom);
  paddings[isColMajor ? 2 : (NumDims - 3)] =
      IndexPair<TensorIndex>(paddingTop, paddingBottom);
  paddings[isColMajor ? 3 : (NumDims - 4)] =
      IndexPair<TensorIndex>(paddingLeft, paddingRight);

  pre_stride_dims[isColMajor ? 4 : (ExtDims - 5)] =
      inputPlanes + paddingZBottom + paddingZTop - patchPlanes + 1;
  pre_stride_dims[isColMajor ? 5 : (ExtDims - 6)] =
      inputRows + paddingTop + paddingBottom - patchRows + 1;
  pre_stride_dims[isColMajor ? 6 : (ExtDims - 7)] =
      inputCols + paddingLeft + paddingRight - patchCols + 1;

  if (isColMajor) {
    for (int i = 7; i < NumDims + 3; ++i) {
      pre_stride_dims[i] = in.dimension(i - 3);
    }
  } else {
    for (int i = 0; i < NumDims - 4; ++i) {
      pre_stride_dims[i] = in.dimension(i);
    }
  }

  DSizes<TensorIndex, NumDims> patch_dims;
  if (isColMajor) {
    patch_dims[0] = in.dimension(0);
    patch_dims[1] = patchPlanes;
    patch_dims[2] = patchRows;
    patch_dims[3] = patchCols;
    for (int i = 4; i < NumDims; ++i) {
      patch_dims[i] = 1;
    }
  } else {
    patch_dims[NumDims - 1] = in.dimension(NumDims - 1);
    patch_dims[NumDims - 4] = patchCols;
    patch_dims[NumDims - 3] = patchRows;
    patch_dims[NumDims - 2] = patchPlanes;
    for (int i = 0; i < NumDims - 4; i++) {
      patch_dims[i] = 1;
    }
  }

  array<TensorIndex, NumDims + 3> strides;
  if (isColMajor) {
    // No striding within the patches.
    for (int i = 0; i < 4; ++i) {
      strides[i] = 1;
    }
    // Apply striding in the spatial patch grid dimensions only.
    strides[4] = stridePlanes;
    strides[5] = strideRows;
    strides[6] = strideCols;
    // No striding in the remaining dimensions (batches, ...).
    for (int i = 7; i < NumDims + 3; i++) {
      strides[i] = 1;
    }
  } else {
    // No striding within the patches.
    for (int i = 1; i <= 4; ++i) {
      strides[ExtDims - i] = 1;
    }
    // Apply striding in the spatial patch grid dimensions only.
    strides[ExtDims - 7] = strideCols;
    strides[ExtDims - 6] = strideRows;
    strides[ExtDims - 5] = stridePlanes;
    // No striding in the remaining dimensions (batches, ...).
    for (int i = 0; i < NumDims - 4; i++) {
      strides[i] = 1;
    }
  }

  // TODO(mjanusz): Consider getting rid of pad(), and stride() and extend
  // extract_patches to take additional parameters for padding/striding,
  // similarly to etract_image_patches.
  return input.pad(paddings, padding_value)
      .extract_patches(patch_dims)
      .reshape(pre_stride_dims)
      .stride(strides);
}

template <typename Input>
EIGEN_ALWAYS_INLINE static const TensorStridingOp<
    const array<typename internal::traits<Input>::Index,
                internal::traits<Input>::NumDimensions + 3>,
    const TensorReshapingOp<
        const DSizes<typename internal::traits<Input>::Index,
                     internal::traits<Input>::NumDimensions + 3>,
        const TensorPatchOp<
            const DSizes<typename internal::traits<Input>::Index,
                         internal::traits<Input>::NumDimensions>,
            const TensorPaddingOp<
                const array<IndexPair<typename internal::traits<Input>::Index>,
                            internal::traits<Input>::NumDimensions>,
                const Input> > > >
Extract3DPatches(
    const Input& input, const DenseIndex patchPlanes,
    const DenseIndex patchRows, const DenseIndex patchCols,
    const DenseIndex stridePlanes, const DenseIndex strideRows,
    const DenseIndex strideCols, const PaddingType padding_type,
    const typename internal::traits<Input>::Scalar padding_value = 0) {
  typedef typename internal::traits<Input>::Index TensorIndex;
  TensorRef<Tensor<typename internal::traits<Input>::Scalar,
                   internal::traits<Input>::NumDimensions,
                   internal::traits<Input>::Layout, TensorIndex> >
      in(input);

  EIGEN_STATIC_ASSERT(internal::traits<Input>::NumDimensions >= 4,
                      YOU_MADE_A_PROGRAMMING_MISTAKE);

  static const bool isColMajor = (internal::traits<Input>::Layout == ColMajor);
  static const int NumDims = internal::traits<Input>::NumDimensions;

  const TensorIndex inputPlanes =
      isColMajor ? in.dimension(1) : in.dimension(NumDims - 2);
  const TensorIndex inputRows =
      isColMajor ? in.dimension(2) : in.dimension(NumDims - 3);
  const TensorIndex inputCols =
      isColMajor ? in.dimension(3) : in.dimension(NumDims - 4);

  switch (padding_type) {
    case PADDING_VALID:
      // No padding in any dimension.
      return Extract3DPatches(input, patchPlanes, patchRows, patchCols,
                              stridePlanes, strideRows, strideCols, 0, 0, 0, 0,
                              0, 0, padding_value);
    case PADDING_SAME: {
      // The side of the tensor before striding should be just the expected
      // output times the stride.
      const TensorIndex size_z =
          numext::ceil(inputPlanes / static_cast<float>(stridePlanes));
      const TensorIndex size_y =
          numext::ceil(inputRows / static_cast<float>(strideRows));
      const TensorIndex size_x =
          numext::ceil(inputCols / static_cast<float>(strideCols));

      // The size of the patch space is going to be: padded_input_size -
      // patch_size + 1.
      // This has to match the expected size before striding (pre_stride_dims).
      // The deltas below extend the input to the expected size.
      const TensorIndex dz =
          (size_z - 1) * stridePlanes + patchPlanes - inputPlanes;
      const TensorIndex dy = (size_y - 1) * strideRows + patchRows - inputRows;
      const TensorIndex dx = (size_x - 1) * strideCols + patchCols - inputCols;

      return Extract3DPatches(input, patchPlanes, patchRows, patchCols,
                              stridePlanes, strideRows, strideCols, dz / 2,
                              dz - dz / 2, dy / 2, dy - dy / 2, dx / 2,
                              dx - dx / 2, padding_value);
    }
    default:
      eigen_assert(false && "unexpected padding");
      // unreachable code to avoid missing return warning.
      return Extract3DPatches(input, patchPlanes, patchRows, patchCols,
                              stridePlanes, strideRows, strideCols, 0, 0, 0, 0,
                              0, 0, padding_value);
  }
}

// TODO(mjanusz): Switch this to a 'using' alias once CUDA supports C++11.
template <typename Input>
struct Extract3DPatchesType {
  typedef const TensorStridingOp<
      const array<typename internal::traits<Input>::Index,
                  internal::traits<Input>::NumDimensions + 3>,
      const TensorReshapingOp<
          const DSizes<typename internal::traits<Input>::Index,
                       internal::traits<Input>::NumDimensions + 3>,
          const TensorPatchOp<
              const DSizes<typename internal::traits<Input>::Index,
                           internal::traits<Input>::NumDimensions>,
              const TensorPaddingOp<
                  const array<
                      IndexPair<typename internal::traits<Input>::Index>,
                      internal::traits<Input>::NumDimensions>,
                  const Input> > > >
      type;
};

}  // end namespace internal
}  // end namespace Eigen

#endif  // THIRD_PARTY_TENSORFLOW_CORE_KERNELS_EIGEN_PATCH_3D_H_
