#ifndef CAFFE2_OPERATORS_LENGTH_SPLIT_OP_H_
#define CAFFE2_OPERATORS_LENGTH_SPLIT_OP_H_

#include "caffe2/core/common_omp.h"
#include "caffe2/core/context.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/operator.h"
#include "caffe2/utils/math.h"

namespace caffe2 {

template <class Context>
class LengthsSplitOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit LengthsSplitOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...),
        n_split_(OperatorBase::GetSingleArgument<int32_t>("n_split", 0)) {
    if (InputSize() == 1) {
      // If not specified, then must have this argument
      CAFFE_ENFORCE(
          OperatorBase::HasArgument("n_split"),
          "Argument `n_split` is missing and was not specified as input.");
      CAFFE_ENFORCE(
          n_split_ > 0,
          "`n_split` must contain a positive value for defined behavior.");
    }
  }
  ~LengthsSplitOp() {}

  bool RunOnDevice() override {
    const auto& L = Input(0);
    CAFFE_ENFORCE_EQ(L.dim(), 1, "Input `LENGTHS` should be a 1D vector.");

    if (InputSize() > 1) {
      // We potentially have n_split specified as inputs as well
      CAFFE_ENFORCE(
          Input(1).dim() == 1 && Input(1).numel() == 1,
          "Input `n_split` should be a vector of size 1.");

      const auto& input1 = Input(1);
      context_.template CopyItems<Context, CPUContext>(
          input1.dtype(), 1, input1.raw_data(), &n_split_);
    }

    CAFFE_ENFORCE(
        n_split_ > 0,
        "`n_split` must contain a positive value for defined behavior.");
    const auto M = L.numel();

    auto* Y = Output(0, {M * n_split_}, at::dtype<int32_t>());

    const int32_t* Ldata = L.template data<int32_t>();
    int32_t* Ydata = Y->template mutable_data<int32_t>();

    for (int i = 0; i < M; i++) {
      int32_t mod = Ldata[i] % n_split_;
      int32_t res =
          mod != 0 ? math::DivUp(Ldata[i], n_split_) : Ldata[i] / n_split_ + 1;
      for (int j = 0; j < n_split_; j++) {
        Ydata[(i * n_split_) + j] = mod-- > 0 ? res : res - 1;
      }
    }
    return true;
  }

 private:
  int32_t n_split_;
};

} // namespace caffe2

#endif // CAFFE2_OPERATORS_LENGTH_SPLIT_OP_H_