#ifndef EIGEN_CXX11_NEURAL_NETWORKS_BACKWARD_CUBOID_CONVOLUTIONS_H
#define EIGEN_CXX11_NEURAL_NETWORKS_BACKWARD_CUBOID_CONVOLUTIONS_H

#include "Patch3d.h"

namespace Eigen {

/** CuboidConvolutionBackwardInput
  * \ingroup CXX11_NeuralNetworks_Module
  *
  * \brief Computes the backprop for the input of a 3D convolution.
  *
  * The output_backward parameter is expected to be a tensor with a rank of 4 or more (channels, depth, height, width, and optionally others)
  * The kernel parameter is expected to be a 5D tensor (filters, channels, kernel_depth, kernel_height, kernel_width)
  * output_backward and kernel have to be in the same layout.
  *
  * The dimensions of the result will be filters, depth, height, width (and others if applicable).
  *
  * It is possible to swap the order of the depth, width and height dimensions provided that the same order is used in the input, the kernel, and the output.
  *
  * All dimension orders above are given for col-major, and should be reversed for row-major.
  */

template <typename OutputBackward, typename Kernel>
EIGEN_ALWAYS_INLINE static const typename internal::conditional<
    internal::traits<OutputBackward>::Layout == ColMajor,
    TensorReshapingOp<
        const DSizes<typename internal::traits<OutputBackward>::Index,
                     internal::traits<OutputBackward>::NumDimensions>,
        const TensorContractionOp<
            const array< IndexPair<typename internal::traits<OutputBackward>::Index>, 2>,
            const TensorReshapingOp<
                const DSizes< typename internal::traits<OutputBackward>::Index, 3>,
                const TensorReverseOp<const array<bool, 5>, const Kernel>
            >,
            const TensorReshapingOp<
                const DSizes< typename internal::traits<OutputBackward>::Index, 3>,
                const TensorVolumePatchOp<Dynamic, Dynamic, Dynamic, const OutputBackward>
            >
        >
    >,
    TensorReshapingOp<
        const DSizes<typename internal::traits<OutputBackward>::Index,
                     internal::traits<OutputBackward>::NumDimensions>,
        const TensorContractionOp<
            const array< IndexPair<typename internal::traits<OutputBackward>::Index>, 2>,
            const TensorReshapingOp<
                const DSizes< typename internal::traits<OutputBackward>::Index, 3>,
                const TensorVolumePatchOp<Dynamic, Dynamic, Dynamic, const OutputBackward>
            >,
            const TensorReshapingOp<
                const DSizes<typename internal::traits<OutputBackward>::Index, 3>,
                const TensorReverseOp<const array<bool, 5>, const Kernel>
            >
        >
    >
>::type
CuboidConvolutionBackwardInput(
    const Kernel& kernel, const OutputBackward& output_backward,
    typename internal::traits<OutputBackward>::Index inputPlanes,
    typename internal::traits<OutputBackward>::Index inputRows,
    typename internal::traits<OutputBackward>::Index inputCols,
    const DenseIndex stridePlanes = 1, const DenseIndex strideRows = 1,
    const DenseIndex strideCols = 1) {
  typedef typename internal::traits<OutputBackward>::Index TensorIndex;
  const TensorRef<const Tensor<typename internal::traits<Kernel>::Scalar, internal::traits<Kernel>::NumDimensions, internal::traits<Kernel>::Layout, TensorIndex> > kern(kernel);
  const TensorRef<const Tensor<typename internal::traits<OutputBackward>::Scalar, internal::traits<OutputBackward>::NumDimensions, internal::traits<OutputBackward>::Layout, TensorIndex> > out(output_backward);

  EIGEN_STATIC_ASSERT(internal::traits<Kernel>::Layout == internal::traits<OutputBackward>::Layout, YOU_MADE_A_PROGRAMMING_MISTAKE);

  static const bool isColMajor = (internal::traits<OutputBackward>::Layout == ColMajor);

  static const int NumDims = internal::traits<OutputBackward>::NumDimensions;

  // Number of filters to apply. This is the same as the output depth of the result
  const TensorIndex kernelFilters = isColMajor ? kern.dimensions()[0] : kern.dimensions()[4];
  // Number of channels. This is the same as the input depth.
  const TensorIndex kernelChannels = isColMajor ? kern.dimensions()[1] : kern.dimensions()[3];
  const TensorIndex kernelPlanes = isColMajor ? kern.dimensions()[2] : kern.dimensions()[2];
  const TensorIndex kernelRows = isColMajor ? kern.dimensions()[3] : kern.dimensions()[1];
  const TensorIndex kernelCols = isColMajor ? kern.dimensions()[4] : kern.dimensions()[0];

  const TensorIndex outputPlanes = isColMajor ? out.dimensions()[1] : out.dimensions()[NumDims - 2];
  const TensorIndex outputRows = isColMajor ? out.dimensions()[2] : out.dimensions()[NumDims - 3];
  const TensorIndex outputCols = isColMajor ? out.dimensions()[3] : out.dimensions()[NumDims - 4];

  TensorIndex forward_pad_z, forward_pad_y, forward_pad_x;
  const TensorIndex size_z = ceil(inputPlanes / static_cast<float>(stridePlanes));
  const TensorIndex size_y = ceil(inputRows / static_cast<float>(strideRows));
  const TensorIndex size_x = ceil(inputCols / static_cast<float>(strideCols));

  // Infer padding type.
  if (size_z == outputPlanes && size_y == outputRows && size_x == outputCols) {
    // SAME padding.
    const TensorIndex dz = size_z * stridePlanes + kernelPlanes - 1 - inputPlanes;
    const TensorIndex dy = size_y * strideRows + kernelRows - 1 - inputRows;
    const TensorIndex dx = size_x * strideCols + kernelCols - 1 - inputCols;

    forward_pad_z = dz - dz / 2;
    forward_pad_y = dy - dy / 2;
    forward_pad_x = dx - dx / 2;
  } else {
    // VALID padding.
    forward_pad_z = 0;
    forward_pad_y = 0;
    forward_pad_x = 0;
  }
  const TensorIndex padding_ztop = kernelPlanes - 1 - forward_pad_z;
  const TensorIndex padding_top = kernelRows - 1 - forward_pad_y;
  const TensorIndex padding_left = kernelCols - 1 - forward_pad_x;

  const TensorIndex padding_zbottom = inputPlanes + kernelPlanes - 1 - (outputPlanes - 1) * stridePlanes - 1 - padding_ztop;
  const TensorIndex padding_bottom = inputRows + kernelRows - 1 - (outputRows - 1) * strideRows - 1 - padding_top;
  const TensorIndex padding_right = inputCols + kernelCols - 1 - (outputCols - 1) * strideCols - 1 - padding_left;

  eigen_assert(padding_ztop >= 0);
  eigen_assert(padding_zbottom >= 0);
  eigen_assert(padding_top >= 0);
  eigen_assert(padding_left >= 0);
  eigen_assert(padding_bottom >= 0);
  eigen_assert(padding_right >= 0);

  // The kernel has dimensions filters X channels X patch_planes X patch_rows X patch_cols.
  // We need to reverse the kernel along the spatial dimensions.
  array<bool, 5> kernel_reverse;
  if (isColMajor) {
    kernel_reverse[0] = false;
    kernel_reverse[1] = false;
    kernel_reverse[2] = true;
    kernel_reverse[3] = true;
    kernel_reverse[4] = true;
  } else {
    kernel_reverse[0] = true;
    kernel_reverse[1] = true;
    kernel_reverse[2] = true;
    kernel_reverse[3] = false;
    kernel_reverse[4] = false;
  }

  DSizes<TensorIndex, 3> kernel_dims;
  if (isColMajor) {
    kernel_dims[0] = kernelFilters;
    kernel_dims[1] = kernelChannels;
    kernel_dims[2] = kernelRows * kernelCols * kernelPlanes;
  } else {
    kernel_dims[0] = kernelRows * kernelCols * kernelPlanes;
    kernel_dims[1] = kernelChannels;
    kernel_dims[2] = kernelFilters;
  }

  // The output_backward has dimensions out_depth X out_planes X out_rows X out_cols X OTHERS
  // When we extract the image patches from output_backward, it will have dimensions:
  //   out_depth X (patch_planes * patch_rows * patch_cols) X (input_planes * input_rows * input_cols * OTHERS)
  DSizes<TensorIndex, 3> pre_contract_dims;
  if (isColMajor) {
    pre_contract_dims[0] = kernelFilters;
    pre_contract_dims[1] = kernelRows * kernelCols * kernelPlanes;
    pre_contract_dims[2] = inputRows * inputCols * inputPlanes;
    for (int i = 4; i < NumDims; ++i) {
      pre_contract_dims[2] *= out.dimension(i);
    }
  } else {
    pre_contract_dims[2] = kernelFilters;
    pre_contract_dims[1] = kernelRows * kernelCols * kernelPlanes;
    pre_contract_dims[0] = inputRows * inputCols * inputPlanes;
    for (int i = 0; i < NumDims - 4; ++i) {
      pre_contract_dims[0] *= out.dimension(i);
    }
  }

  // We will contract along dimensions (0, 2) in kernel and (0, 1) in
  // output_backward, if this is col-major, and
  // dimensions (0, 2) in kernel and (1, 2) in output_backward, if this row-major.
  array<IndexPair<TensorIndex>, 2> contract_dims;
  if (isColMajor) {
    // col-major: kernel.contract(output.patches)
    contract_dims[0] = IndexPair<TensorIndex>(0, 0);
    contract_dims[1] = IndexPair<TensorIndex>(2, 1);
  } else {
    // row-major: output.patches.contract(kernel)
    contract_dims[0] = IndexPair<TensorIndex>(1, 0);
    contract_dims[1] = IndexPair<TensorIndex>(2, 2);
  }

  // Post contraction, the dimensions of the input_backprop is
  //  channels X input_planes X input_rows X input_cols X OTHERS
  DSizes<TensorIndex, NumDims> post_contract_dims;
  if (isColMajor) {
    post_contract_dims[0] = kernelChannels;
    post_contract_dims[1] = inputPlanes;
    post_contract_dims[2] = inputRows;
    post_contract_dims[3] = inputCols;
    for (int i = 4; i < NumDims; ++i) {
      post_contract_dims[i] = out.dimension(i);
    }
  } else {
    post_contract_dims[NumDims - 1] = kernelChannels;
    post_contract_dims[NumDims - 2] = inputPlanes;
    post_contract_dims[NumDims - 3] = inputRows;
    post_contract_dims[NumDims - 4] = inputCols;
    for (int i = 0; i < NumDims - 4; ++i) {
      post_contract_dims[i] = out.dimension(i);
    }
  }

  DSizes<TensorIndex, NumDims> strides;
  for (int i = 0; i < NumDims; i++) {
    strides[i] = 1;
  }
  if (isColMajor) {
    strides[1] = stridePlanes;
    strides[2] = strideRows;
    strides[3] = strideCols;
  } else {
    strides[NumDims - 2] = stridePlanes;
    strides[NumDims - 3] = strideRows;
    strides[NumDims - 4] = strideCols;
  }

  return choose(
      Cond<internal::traits<OutputBackward>::Layout == ColMajor>(),
      kernel.reverse(kernel_reverse)
          .reshape(kernel_dims)
          .contract(
              output_backward.extract_volume_patches(kernelPlanes, kernelRows, kernelCols,
                                                     1, 1, 1, stridePlanes, strideRows, strideCols,
                               padding_ztop, padding_zbottom,
                               padding_top, padding_bottom,
                               padding_left, padding_right)
                  .reshape(pre_contract_dims),
              contract_dims)
          .reshape(post_contract_dims),
      output_backward.extract_volume_patches(kernelPlanes, kernelRows, kernelCols,
                                             1, 1, 1, stridePlanes, strideRows, strideCols,
                       padding_ztop, padding_zbottom,
                       padding_top, padding_bottom,
                       padding_left, padding_right)
          .reshape(pre_contract_dims)
          .contract(kernel.reverse(kernel_reverse).reshape(kernel_dims),
                    contract_dims)
          .reshape(post_contract_dims));
}


/** CuboidConvolutionBackwardKernel
  * \ingroup CXX11_NeuralNetworks_Module
  *
  * \brief Computes the backprop for the filter of a 3D convolution.
  *
  * The output_backward parameter is expected to be a tensor with a rank of 4 or more (channels, depth, height, width, and optionally others)
  * The kernel parameter is expected to be a 4D tensor (filters, channels, kernel_depth, kernel_height, kernel_width)
  * output_backward and kernel have to be in the same layout.
  *
  * The dimensions of the result will be filters, depth, height, width (and others if applicable).
  *
  * It is possible to swap the order of the depth, width and height dimensions provided that the same order is used in the input, the kernel, and the output.
  *
  * All dimension orders above are given for col-major, and should be reversed for row-major.
  */
template <typename OutputBackward, typename Input>
EIGEN_ALWAYS_INLINE static const typename internal::conditional<
    internal::traits<OutputBackward>::Layout == ColMajor,
    const TensorShufflingOp<
        const array<typename internal::traits<OutputBackward>::Index, 5>,
        const TensorReverseOp<
            const array<bool, 5>,
            const TensorReshapingOp<
                const DSizes<typename internal::traits<OutputBackward>::Index, 5>,
                const TensorContractionOp<
                    const array< IndexPair<typename internal::traits<Input>::Index>, 2>,
                    const TensorReshapingOp<
                        const DSizes<typename internal::traits<Input>::Index, 3>,
                        const Input>,
                    const TensorReshapingOp<
                        const DSizes< typename internal::traits<OutputBackward>::Index, 4>,
                        const TensorVolumePatchOp<Dynamic, Dynamic, Dynamic, const OutputBackward>
                    >
                >
            >
        >
    >,
    const TensorShufflingOp<
        const array<typename internal::traits<OutputBackward>::Index, 5>,
        const TensorReverseOp<
            const array<bool, 5>,
            const TensorReshapingOp<
                const DSizes<typename internal::traits<OutputBackward>::Index, 5>,
                const TensorContractionOp<
                    const array< IndexPair<typename internal::traits<Input>::Index>, 2>,
                    const TensorReshapingOp<
                        const DSizes< typename internal::traits<OutputBackward>::Index, 4>,
                        const TensorVolumePatchOp<Dynamic, Dynamic, Dynamic, const OutputBackward>
                    >,
                    const TensorReshapingOp<
                        const DSizes<typename internal::traits<Input>::Index, 3>,
                        const Input
                    >
                >
            >
        >
    >
>::type
CuboidConvolutionBackwardKernel(
    const Input& input, const OutputBackward& output_backward,
    typename internal::traits<Input>::Index kernelPlanes,
    typename internal::traits<Input>::Index kernelRows,
    typename internal::traits<Input>::Index kernelCols,
    const DenseIndex stridePlanes = 1,
    const DenseIndex strideRows = 1,
    const DenseIndex strideCols = 1) {
  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);
  TensorRef<Tensor<typename internal::traits<OutputBackward>::Scalar, internal::traits<OutputBackward>::NumDimensions, internal::traits<OutputBackward>::Layout, TensorIndex> > out(output_backward);

  EIGEN_STATIC_ASSERT(internal::traits<Input>::Layout == internal::traits<OutputBackward>::Layout, YOU_MADE_A_PROGRAMMING_MISTAKE);

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

  static const int NumDims = internal::traits<Input>::NumDimensions;
  EIGEN_STATIC_ASSERT(internal::traits<Input>::NumDimensions == internal::traits<OutputBackward>::NumDimensions, YOU_MADE_A_PROGRAMMING_MISTAKE);

  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);

  const TensorIndex outputPlanes = isColMajor ? out.dimension(1) : out.dimension(NumDims - 2);
  const TensorIndex outputRows = isColMajor ? out.dimension(2) : out.dimension(NumDims - 3);
  const TensorIndex outputCols = isColMajor ? out.dimension(3) : out.dimension(NumDims - 4);

  const TensorIndex kernelFilters = isColMajor ? out.dimension(0) : out.dimension(NumDims - 1);
  const TensorIndex kernelChannels = isColMajor ? in.dimension(0) : in.dimension(NumDims - 1);

  TensorIndex forward_pad_z, forward_pad_y, forward_pad_x;
  const TensorIndex size_z = ceil(inputPlanes / static_cast<float>(stridePlanes));
  const TensorIndex size_y = ceil(inputRows / static_cast<float>(strideRows));
  const TensorIndex size_x = ceil(inputCols / static_cast<float>(strideCols));

  // Infer padding type.
  if (size_z == outputPlanes && size_y == outputRows && size_x == outputCols) {
    // SAME padding.
    const TensorIndex dz = size_z * stridePlanes + kernelPlanes - 1 - inputPlanes;
    const TensorIndex dy = size_y * strideRows + kernelRows - 1 - inputRows;
    const TensorIndex dx = size_x * strideCols + kernelCols - 1 - inputCols;

    forward_pad_z = dz - dz / 2;
    forward_pad_y = dy - dy / 2;
    forward_pad_x = dx - dx / 2;
  } else {
    // VALID padding.
    forward_pad_z = 0;
    forward_pad_y = 0;
    forward_pad_x = 0;
  }

  const TensorIndex padding_ztop = kernelPlanes - 1 - forward_pad_z;
  const TensorIndex padding_top = kernelRows - 1 - forward_pad_y;
  const TensorIndex padding_left = kernelCols - 1 - forward_pad_x;

  const TensorIndex padding_zbottom = inputPlanes + kernelPlanes - 1 - (outputPlanes - 1) * stridePlanes - 1 - padding_ztop;
  const TensorIndex padding_bottom = inputRows + kernelRows - 1 - (outputRows - 1) * strideRows - 1 - padding_top;
  const TensorIndex padding_right = inputCols + kernelCols - 1 - (outputCols - 1) * strideCols - 1 - padding_left;

  eigen_assert(padding_ztop >= 0);
  eigen_assert(padding_zbottom >= 0);
  eigen_assert(padding_top >= 0);
  eigen_assert(padding_left >= 0);
  eigen_assert(padding_bottom >= 0);
  eigen_assert(padding_right >= 0);

  // The output_backward has dimensions out_depth X out_plaens X out_rows X out_cols X OTHERS
  // When we extract the image patches from output_backward (with input as the
  // kernel), it will have dimensions
  //  (out_depth) X (input_planes * input_rows * input_cols) X (kernel_planes * kernel_rows * kernel_cols) X OTHERS
  DSizes<TensorIndex, 4> pre_contract_dims;
  if (isColMajor) {
    pre_contract_dims[0] = kernelFilters;
    pre_contract_dims[1] = inputRows * inputCols * inputPlanes;
    pre_contract_dims[2] = kernelRows * kernelCols * kernelPlanes;
    pre_contract_dims[3] = 1;
    for (int i = 4; i < NumDims; ++i) {
      pre_contract_dims[3] *= out.dimension(i);
    }
  } else {
    pre_contract_dims[3] = kernelFilters;
    pre_contract_dims[2] = inputRows * inputCols * inputPlanes;
    pre_contract_dims[1] = kernelRows * kernelCols * kernelPlanes;
    pre_contract_dims[0] = 1;
    for (int i = 0; i < NumDims - 4; ++i) {
      pre_contract_dims[0] *= out.dimension(i);
    }
  }

  // The input has dimensions in_depth X (input_planes * input_rows * input_cols) X OTHERS
  DSizes<TensorIndex, 3> input_dims;
  if (isColMajor) {
    input_dims[0] = kernelChannels;
    input_dims[1] = inputRows * inputCols * inputPlanes;
    input_dims[2] = 1;
    for (int i = 4; i < NumDims; ++i) {
      input_dims[2] *= in.dimension(i);
    }
    eigen_assert(input_dims[2] == pre_contract_dims[3]);
  } else {
    input_dims[2] = kernelChannels;
    input_dims[1] = inputRows * inputCols * inputPlanes;
    input_dims[0] = 1;
    for (int i = 0; i < NumDims - 4; ++i) {
      input_dims[0] *= in.dimension(i);
    }
    eigen_assert(input_dims[0] == pre_contract_dims[0]);
  }

  // We will contract along dimensions (1, 2) in in and (1, 3) in out, if
  // this is col-major.
  // For row-major, it's dimensions (0, 1) in in and (0, 2) in out.
  array<IndexPair<TensorIndex>, 2> contract_dims;
  if (isColMajor) {
    // col-major: in.contract(output.patches)
    contract_dims[0] = IndexPair<TensorIndex>(1, 1);
    contract_dims[1] = IndexPair<TensorIndex>(2, 3);
  } else {
    // row-major: output.patches.contract(in)
    contract_dims[0] = IndexPair<TensorIndex>(0, 0);
    contract_dims[1] = IndexPair<TensorIndex>(2, 1);
  }

  // After the contraction, the kernel will have dimension
  //   in_depth X out_depth X kernel_patches X kernel_rows X kernel_cols
  // We will need to shuffle the first two dimensions and reverse the spatial dimensions.
  // The end shape is:
  //   out_depth X in_shape X kernel_planes X kernel_rows X kernel_cols

  // This is the shape of the kernel *before* the shuffling.
  DSizes<TensorIndex, 5> kernel_dims;
  if (isColMajor) {
    kernel_dims[0] = kernelChannels;
    kernel_dims[1] = kernelFilters;
    kernel_dims[2] = kernelPlanes;
    kernel_dims[3] = kernelRows;
    kernel_dims[4] = kernelCols;
  } else {
    kernel_dims[0] = kernelCols;
    kernel_dims[1] = kernelRows;
    kernel_dims[2] = kernelPlanes;
    kernel_dims[3] = kernelFilters;
    kernel_dims[4] = kernelChannels;
  }

  // Flip filters and channels.
  array<TensorIndex, 5> kernel_shuffle;
  if (isColMajor) {
    kernel_shuffle[0] = 1;
    kernel_shuffle[1] = 0;
    kernel_shuffle[2] = 2;
    kernel_shuffle[3] = 3;
    kernel_shuffle[4] = 4;
  } else {
    kernel_shuffle[0] = 0;
    kernel_shuffle[1] = 1;
    kernel_shuffle[2] = 2;
    kernel_shuffle[3] = 4;
    kernel_shuffle[4] = 3;
  }

  // Reverse the spatial dimensions.
  array<bool, 5> kernel_reverse;
  if (isColMajor) {
    kernel_reverse[0] = false;
    kernel_reverse[1] = false;
    kernel_reverse[2] = true;
    kernel_reverse[3] = true;
    kernel_reverse[4] = true;
  } else {
    kernel_reverse[0] = true;
    kernel_reverse[1] = true;
    kernel_reverse[2] = true;
    kernel_reverse[3] = false;
    kernel_reverse[4] = false;
  }

  DSizes<TensorIndex, NumDims> strides;
  for (int i = 0; i < NumDims; i++) {
    strides[i] = 1;
  }
  if (isColMajor) {
    strides[1] = stridePlanes;
    strides[2] = strideRows;
    strides[3] = strideCols;
  } else {
    strides[NumDims - 2] = stridePlanes;
    strides[NumDims - 3] = strideRows;
    strides[NumDims - 4] = strideCols;
  }
  return choose(
      Cond<internal::traits<Input>::Layout == ColMajor>(),
      input.reshape(input_dims)
          .contract(
              output_backward.extract_volume_patches(
                                 inputPlanes, inputRows, inputCols, 1,
                                 1, 1, stridePlanes, strideRows, strideCols,

                                 padding_ztop, padding_zbottom, padding_top,
                                 padding_bottom, padding_left, padding_right)
                  .reshape(pre_contract_dims),
              contract_dims)
          .reshape(kernel_dims)
          .reverse(kernel_reverse)
          .shuffle(kernel_shuffle),
      output_backward.extract_volume_patches(
                         inputPlanes, inputRows, inputCols, 1, 1, 1,
                         stridePlanes, strideRows, strideCols, padding_ztop,
                         padding_zbottom, padding_top, padding_bottom,
                         padding_left, padding_right)
          .reshape(pre_contract_dims)
          .contract(input.reshape(input_dims), contract_dims)
          .reshape(kernel_dims)
          .reverse(kernel_reverse)
          .shuffle(kernel_shuffle));
}

} // end namespace Eigen

#endif // EIGEN_CXX11_NEURAL_NETWORKS_BACKWARD_CUBOID_CONVOLUTIONS_H
