#include <ATen/ATen.h>
#include <ATen/PadNd.h>

#include <c10/util/irange.h>

namespace at { namespace native {

Tensor constant_pad_nd(const Tensor& self, IntArrayRef pad, const Scalar& value) {
    TORCH_CHECK(pad.size() % 2 == 0, "Length of pad must be even but instead it equals ",
             pad.size());

    auto input_sizes = self.sizes();
    auto l_inp = self.dim();

    auto l_pad = pad.size() / 2;
    auto l_diff = l_inp - l_pad;
    TORCH_CHECK(l_inp >= (int64_t)l_pad, "Length of pad should be no more than twice the number of "
             "dimensions of the input. Pad length is ", pad.size(), "while the input has ",
             l_inp, "dimensions.");

    std::vector<int64_t> new_shape;

    bool all_pads_non_positive = true;

    auto c_input = self;
    for (const auto i : c10::irange(l_diff, l_inp)) {
        auto pad_idx = 2 * (l_inp - i - 1);
        if (pad[pad_idx] < 0) {
            c_input = c_input.narrow(i, -pad[pad_idx], c_input.size(i) + pad[pad_idx]);
        } else if (pad[pad_idx] != 0) {
            all_pads_non_positive = false;
        }
        if (pad[pad_idx + 1] < 0) {
            c_input = c_input.narrow(i, 0, c_input.size(i) + pad[pad_idx + 1]);
        } else if (pad[pad_idx + 1] != 0) {
            all_pads_non_positive = false;
        }
    }

    // if none of the pads are positive we can optimize and just return the result
    // of calling .narrow() on the input
    if (all_pads_non_positive) {
        return c_input.clone();
    }


    for (size_t i = 0; i < (size_t)l_diff; i ++) {
        new_shape.emplace_back(input_sizes[i]);
    }

    for (const auto i : c10::irange((size_t)l_pad)) {
        auto pad_idx = pad.size() - ((i + 1) * 2);
        auto new_dim = input_sizes[l_diff + i] + pad[pad_idx] + pad[pad_idx + 1];
        TORCH_CHECK(new_dim > 0, "The input size ", input_sizes[l_diff + i], ", plus negative padding ",
                 pad[pad_idx], " and ", pad[pad_idx + 1], " resulted in a negative output size, "
                 "which is invalid. Check dimension ", l_diff + i, " of your input.");
        new_shape.emplace_back(new_dim);
    }

    at::Tensor output;
    const auto memory_format = self.suggest_memory_format();
    if (self.is_quantized()) {
        const auto qscheme = self.qscheme();
        TORCH_CHECK(qscheme == kPerTensorAffine || qscheme == kPerTensorSymmetric,
                    "Only per-tensor padding is supported.");
        output = at::_empty_affine_quantized(
            new_shape, self.options().memory_format(memory_format),
            self.q_scale(), self.q_zero_point(), c10::nullopt);
    } else {
        output = at::empty(new_shape, self.options().memory_format(memory_format));
    }
    output.fill_(value);

    auto c_output = output;
    for (const auto i : c10::irange(l_diff, l_inp)) {
        auto pad_idx = 2 * (l_inp - i - 1);
        if (pad[pad_idx] > 0) {
            c_output = c_output.narrow(i, pad[pad_idx], c_output.size(i) - pad[pad_idx]);
        }
        if (pad[pad_idx + 1] > 0) {
            c_output = c_output.narrow(i, 0, c_output.size(i) - pad[pad_idx + 1]);
        }
    }
    c_output.copy_(c_input);
    return output;
}

Tensor _pad_circular(const Tensor &self, IntArrayRef padding) {
  const auto in_shape = self.sizes();
  const auto ndim = static_cast<int64_t>(in_shape.size()) - 2;
  TORCH_CHECK(padding.size() + 4 == in_shape.size() * 2,
              "Invalid padding size, expected ", ndim * 2, " but got ", padding.size());

  DimVector out_shape(in_shape.size());
  out_shape[0] = in_shape[0];
  out_shape[1] = in_shape[1];

  // Get shape of padded tensor
  for (const auto i : c10::irange(ndim)) {
    const auto pad_l = padding[2 * (ndim - i - 1) + 0];
    const auto pad_r = padding[2 * (ndim - i - 1) + 1];
    const auto size = in_shape[2 + i];
    out_shape[2 + i] = size + pad_l + pad_r;

    TORCH_CHECK(
        pad_l <= size && pad_r <= size,
        "Padding value causes wrapping around more than once.");
    TORCH_CHECK(
        out_shape[2 + i] >= 0,
        "Negative padding value is resulting in an empty dimension");
  }

  auto out = self.new_empty(out_shape, self.options());

  // Put original array into the padded array
  Tensor out_slice = out;
  Tensor in_slice = self;
  constexpr int64_t zero = 0;
  for (const auto i : c10::irange(ndim)) {
    const auto dim = ndim - i + 1;
    const auto pad_l = padding[2*i + 0];
    const auto pad_r = padding[2*i + 1];
    out_slice = out_slice.slice(dim, std::max(pad_l, zero), out_shape[dim] - std::max(pad_r, zero));
    in_slice = in_slice.slice(dim, std::max(-pad_l, zero), in_shape[dim] - std::max(-pad_r, zero));
  }
  out_slice.copy_(in_slice);

  // The following steps first pad the beginning of the tensor (left side),
  // and then pad the end of the tensor (right side).
  // Note: Corners will be written more than once when ndim > 1.
  //
  // Only in cases where padding values are > 0 are when additional copying
  // is required.
  for (const auto i : c10::irange(ndim)) {
    const auto dim = ndim - i + 1;
    const auto pad_l = padding[2*i + 0];
    const auto pad_r = padding[2*i + 1];

    if (pad_l > 0) {
      out_slice = out.slice(dim, 0, pad_l);
      in_slice = out.slice(dim,
                           out_shape[dim] - pad_l - std::max(pad_r, zero),
                           out_shape[dim] - std::max(pad_r, zero));
      out_slice.copy_(in_slice);
    }

    if (pad_r > 0) {
      out_slice = out.slice(dim, out_shape[dim] - pad_r, out_shape[dim]);
      in_slice = out.slice(dim, std::max(pad_l, zero), std::max(pad_l, zero) + pad_r);
      out_slice.copy_(in_slice);
    }
  }

  return out;
}

Tensor _pad_enum(const Tensor &self, IntArrayRef pad, int64_t mode_int, c10::optional<double> value) {
  const auto input_dim = self.dim();
  TORCH_CHECK(pad.size() % 2 == 0, "Padding length must be divisible by 2");
  TORCH_CHECK(static_cast<int64_t>(pad.size()) <= input_dim * 2, "Padding length too large");
  auto mode = static_cast<at::padding_mode>(mode_int);

  if (mode == at::padding_mode::constant) {
    return at::constant_pad_nd(self, pad, value.value_or(0.0));
  }
  TORCH_CHECK(!value.has_value() || *value == 0,
              "Padding mode \"", padding_mode_string(mode),
              "\" doesn't take in value argument");

  if (pad.size() == 2 && (input_dim == 2 || input_dim == 3)) {
    switch (mode) {
      case at::padding_mode::reflect: return at::reflection_pad1d(self, pad);
      case at::padding_mode::replicate: return at::replication_pad1d(self, pad);
      case at::padding_mode::circular: return at::_pad_circular(self, pad);
      default: {}
    }
  } else if(pad.size() == 4 && (input_dim == 3 || input_dim == 4)) {
    switch (mode) {
      case at::padding_mode::reflect: return at::reflection_pad2d(self, pad);
      case at::padding_mode::replicate: return at::replication_pad2d(self, pad);
      case at::padding_mode::circular: return at::_pad_circular(self, pad);
      default: {}
    }
  } else if (pad.size() == 6 && (input_dim == 4 || input_dim == 5)) {
    switch (mode) {
      case at::padding_mode::reflect: return at::reflection_pad3d(self, pad);
      case at::padding_mode::replicate: return at::replication_pad3d(self, pad);
      case at::padding_mode::circular: return at::_pad_circular(self, pad);
      default: {}
    }
  }
  C10_THROW_ERROR(NotImplementedError,
      "Only 2D, 3D, 4D, 5D padding with non-constant padding are supported for now");
}

Tensor pad(const Tensor &self, IntArrayRef pad, c10::string_view mode, c10::optional<double> value) {
  const auto mode_enum = [&] {
    if (mode == "reflect") {
      return at::padding_mode::reflect;
    } else if (mode == "constant") {
      return at::padding_mode::constant;
    } else if (mode == "replicate") {
      return at::padding_mode::replicate;
    } else if (mode == "circular") {
      return at::padding_mode::circular;
    }
    C10_THROW_ERROR(NotImplementedError,
                    c10::str("Unrecognised padding mode ", mode));
  }();
  return at::native::_pad_enum(self, pad, static_cast<int64_t>(mode_enum), value);
}

}}  // namespace at::native
