// Copyright (c) 2022 PaddlePaddle 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.

#pragma once

#include "paddle/common/ddim.h"
#include "paddle/phi/core/meta_tensor.h"

namespace phi {

template <typename T = int>
inline void UpdatePaddingAndDilation(std::vector<T>* paddings,
                                     std::vector<T>* dilation,
                                     const std::string padding_algorithm,
                                     const DDim data_dims,
                                     const std::vector<T>& strides,
                                     const std::vector<T>& ksize) {
  // set padding size == data_dims.size() * 2
  auto data_shape = vectorize<T>(data_dims);
  if (static_cast<int>(paddings->size()) == data_dims.size()) {
    for (int i = 0; i < data_dims.size(); ++i) {
      T copy_pad = *(paddings->begin() + 2 * i);
      paddings->insert(paddings->begin() + 2 * i + 1, copy_pad);
    }
  } else {
    PADDLE_ENFORCE_EQ(
        data_dims.size() * 2,
        paddings->size(),
        common::errors::InvalidArgument(
            "Attribute padding's size should be the same or twice as the "
            "input's dimension. "
            "But received: padding's size is %d, padding is [%s]; input's "
            "dimension is %d, input's shape is [%s].",
            paddings->size(),
            make_ddim(*paddings),
            data_dims.size(),
            data_dims));
  }

  // when padding_algorithm is "VALID" or "SAME"
  if (padding_algorithm == "SAME") {
    for (int i = 0; i < data_dims.size(); ++i) {
      T out_size = (data_dims[i] + strides[i] - 1) / strides[i];
      T pad_sum =
          std::max((out_size - 1) * strides[i] + ksize[i] - data_shape[i],
                   static_cast<T>(0));
      T pad_0 = pad_sum / 2;
      T pad_1 = pad_sum - pad_0;
      *(paddings->begin() + i * 2) = pad_0;
      *(paddings->begin() + i * 2 + 1) = pad_1;

      // dilation
      *(dilation->begin() + i) = 1;
    }

  } else if (padding_algorithm == "VALID") {
    for (auto it = paddings->begin(); it != paddings->end(); it++) {
      *it = 0;
    }
  }
}

inline int ConvOutSize(int input_size,
                       int filter_size,
                       int dilation,
                       int pad_left,
                       int pad_right,
                       int stride) {
  const int64_t dkernel =
      static_cast<int64_t>(dilation) * (filter_size - 1) + 1;
  const int64_t output_size =
      (static_cast<int64_t>(input_size) + pad_left + pad_right - dkernel) /
          stride +
      1;

  PADDLE_ENFORCE_GT(
      output_size,
      0,
      common::errors::InvalidArgument(
          "The output's size is expected to be greater than 0. "
          "But received: output's size is %ld. The output's size is "
          "computed by ((input_size + pad_left + pad_right - "
          "(dilation * (filter_size - 1) + 1)) / stride + 1), where "
          "input_size is %d, padding is (%d, %d), filter_size is %d, "
          "dilation is %d, stride is %d.",
          output_size,
          input_size,
          pad_left,
          pad_right,
          filter_size,
          dilation,
          stride));

  PADDLE_ENFORCE_LE_INT_MAX(output_size, "conv output size");
  return static_cast<int>(output_size);
}

inline std::vector<int64_t> ComputeOutputShape(
    const MetaTensor& input,
    const MetaTensor& filter,
    const MetaTensor& bias,
    const std::vector<int>& strides,
    const std::vector<int>& paddings,
    const std::string& padding_algorithm,
    const std::vector<int>& dilations,
    int groups,
    const std::string& data_format,
    bool channel_last,
    MetaConfig config) {
  auto in_dims = input.dims();
  auto filter_dims = filter.dims();

  int dilation_size = dilations.size();
  for (int i = 0; i < dilation_size; ++i) {
    PADDLE_ENFORCE_GT(
        dilations[i],
        0,
        common::errors::InvalidArgument(
            "The dilation of Op(Conv) should be larger than 0, but received "
            "dilation is %d.",
            dilations[i]));
  }

  PADDLE_ENFORCE_EQ(
      in_dims.size() == 4 || in_dims.size() == 5,
      true,
      common::errors::InvalidArgument(
          "The input of Op(Conv) should be a 4-D or 5-D Tensor. But "
          "received: input's dimension is %u, input's shape is [%s].",
          in_dims.size(),
          in_dims));

  PADDLE_ENFORCE_EQ(
      in_dims.size(),
      filter_dims.size(),
      common::errors::InvalidArgument(
          "The input's dimension and filter's dimension of "
          "Op(Conv) should be equal. But received: the input's shape is "
          "[%s], "
          "the input's dimension is %d; the filter's shape is [%s],  "
          "the filter's dimension is %d.",
          in_dims,
          in_dims.size(),
          filter_dims,
          filter_dims.size()));

  int stride_size = strides.size();
  for (int i = 0; i < stride_size; ++i) {
    PADDLE_ENFORCE_GT(
        strides[i],
        0,
        common::errors::InvalidArgument(
            "The stride of Op(Conv) should be larger than 0, but received "
            "stride is %d.",
            strides[i]));
  }

  PADDLE_ENFORCE_EQ(
      in_dims.size(),
      strides.size() + 2U,
      common::errors::InvalidArgument(
          "The difference of input's dimension and Attr(strides)'s "
          "length must be equal to 2 for Op(Conv). "
          "But received: input's dimension is %d, input's shape is [%s]; "
          "Attr(stride)'s length is %d, Attr(stride) is [%s]; "
          "difference of input's dimension and Attr(strides)'s length = %u.",
          in_dims.size(),
          in_dims,
          strides.size(),
          make_ddim(strides),
          in_dims.size() - stride_size));

  const auto input_channels =
      channel_last ? in_dims[in_dims.size() - 1] : in_dims[1];
  if (config.is_runtime) {
    PADDLE_ENFORCE_EQ(
        input_channels,
        (channel_last ? filter_dims[filter_dims.size() - 1] : filter_dims[1]) *
            groups,
        common::errors::InvalidArgument(
            "The number of input's channels should be equal to filter's "
            "channels "
            "* groups for Op(Conv). But received: the input's channels is %d, "
            "the input's shape is [%s]; the filter's channels is %d, the "
            "filter's shape is [%s]; the groups is %d, the data_format is %s. "
            "The error may come from wrong data_format setting.",
            input_channels,
            in_dims,
            channel_last ? filter_dims[filter_dims.size() - 1] : filter_dims[1],
            filter_dims,
            groups,
            data_format));

    PADDLE_ENFORCE_EQ(
        filter_dims[0] % groups,
        0,
        common::errors::InvalidArgument(
            "The number of output's channels (filter's first dimension) of "
            "Op(Conv) should be divided by groups. But received: "
            "the output channels is %d, the filter's shape is [%s], "
            "the groups is %d.",
            filter_dims[0],
            filter_dims,
            groups));

    PADDLE_ENFORCE_GT(
        filter_dims[0],
        0,
        common::errors::InvalidArgument(
            "the size of filter at axis 0 should be greater than 0"));
  }

  DDim in_data_dims;
  if (channel_last) {
    in_data_dims = slice_ddim(in_dims, 1, in_dims.size() - 1);
  } else {
    in_data_dims = slice_ddim(in_dims, 2, in_dims.size());
  }

  DDim filter_data_dims;
  if (channel_last) {
    filter_data_dims = slice_ddim(filter_dims, 1, filter_dims.size() - 1);
  } else {
    filter_data_dims = slice_ddim(filter_dims, 2, filter_dims.size());
  }

  std::vector<int64_t> ksize = vectorize<int64_t>(filter_data_dims);
  std::vector<int64_t> paddings_vec(paddings.begin(), paddings.end());
  std::vector<int64_t> dilations_vec(dilations.begin(), dilations.end());
  std::vector<int64_t> strides_vec(strides.begin(), strides.end());
  phi::UpdatePaddingAndDilation(&paddings_vec,
                                &dilations_vec,
                                padding_algorithm,
                                in_data_dims,
                                strides_vec,
                                ksize);

  std::vector<int64_t> output_shape({in_dims[0]});
  if (!channel_last) {
    output_shape.push_back(filter_dims[0]);
  }
  for (int i = 0; i < in_data_dims.size(); ++i) {
    if (!config.is_runtime &&
        (in_data_dims[i] <= 0 || filter_dims[i + 2] <= 0)) {
      output_shape.push_back(-1);
    } else {
      PADDLE_ENFORCE_LE_INT_MAX(in_data_dims[i], "conv input size");
      PADDLE_ENFORCE_LE_INT_MAX(filter_data_dims[i], "conv filter size");
      PADDLE_ENFORCE_LE_INT_MAX(paddings_vec[2 * i], "conv padding left");
      PADDLE_ENFORCE_LE_INT_MAX(paddings_vec[2 * i + 1], "conv padding right");
      output_shape.push_back(
          ConvOutSize(static_cast<int>(in_data_dims[i]),
                      static_cast<int>(filter_data_dims[i]),
                      static_cast<int>(dilations_vec[i]),
                      static_cast<int>(paddings_vec[2 * i]),
                      static_cast<int>(paddings_vec[2 * i + 1]),
                      static_cast<int>(strides_vec[i])));
    }
  }
  if (channel_last) {
    output_shape.push_back(filter_dims[0]);
  }

  return output_shape;
}

inline bool IsExpand(const std::vector<int64_t>& filter_dim,
                     const std::vector<int>& strides,
                     const std::vector<int>& paddings,
                     const std::vector<int>& dilations) {
  bool filter_1 = true, strides_1 = true, padding_0 = true, dilation_1 = true;
  for (size_t j = 0; j < strides.size(); ++j) {
    filter_1 = filter_1 && (filter_dim[j + 2] == 1);
    strides_1 = strides_1 && (strides[j] == 1);
    padding_0 = padding_0 && (paddings[j] == 0);
    dilation_1 = dilation_1 && (dilations[j] == 1);
  }
  if (paddings.size() != strides.size()) {
    for (size_t j = 0; j < paddings.size(); ++j) {
      padding_0 = padding_0 && (paddings[j] == 0);
    }
  }
  return !(filter_1 && strides_1 && padding_0 && dilation_1);
}

}  // namespace phi