#ifndef CAFFE2_OPERATORS_PACK_RNN_SEQUENCE_OP_H_
#define CAFFE2_OPERATORS_PACK_RNN_SEQUENCE_OP_H_

#include <algorithm>
#include <vector>
#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"
#include "caffe2/utils/math.h"

namespace caffe2 {

template <class Context, bool Forward>
class PackRNNSequenceOpBase : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit PackRNNSequenceOpBase(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {}

  bool RunOnDevice() override {
    return DispatchHelper<TensorTypes<int32_t, int64_t, float, double>>::call(
        this, Input(0));
  }

  template <typename ValT>
  bool DoRunWithType() {
    // The value is copied from the sequence to the pack
    // if Forward is true, and vice versa
    int dim_offset = Forward ? 1 : 2;
    auto& values = Input(0);
    CAFFE_ENFORCE_GT(values.dim(), dim_offset);

    // block_size is the size for each individual feature
    int64_t block_size = values.size_from_dim(dim_offset);
    auto values_vec = values.template data<ValT>();

    auto& lengths = Input(LENGTHS);
    CAFFE_ENFORCE_EQ(lengths.dim(), 1);
    const auto cols = lengths.numel();
    const int32_t* lengths_vec = lengths.template data<int32_t>();
    // the total number of rows is defined as the max number from lengths
    // if when the lengths is empty, we set rows = 0 to support zero lengths
    const auto rows =
        cols ? *std::max_element(lengths_vec, lengths_vec + cols) : 0;
    CAFFE_ENFORCE_GE(rows, 0);
    int length_sum = 0;
    if (cols > 0) {
      math::Sum<int, Context>(cols, lengths_vec, &length_sum, &context_);
    }

    vector<int64_t> shape;
    // the output shape is rows * cols for the pack,
    // or length_sum for the sequence
    if (Forward) {
      shape.push_back(rows);
      shape.push_back(cols);
    } else {
      shape.push_back(length_sum);
    }
    // insert the dim for the feature
    shape.insert(
        shape.end(), values.sizes().begin() + dim_offset, values.sizes().end());

    auto* output = Output(OUTPUTVALUE, shape, at::dtype<ValT>());

    auto output_data = output->template mutable_data<ValT>();
    // initialize output_data with zero, as it is the default value for padding
    // when certain length is smaller than rows
    math::Set<ValT, Context>(output->numel(), 0, output_data, &context_);

    int32_t offset = 0;
    for (int c = 0; c < cols; c++) {
      for (int r = 0; r < lengths_vec[c]; r++) {
        auto input_offset = Forward ? (offset + r) : (r * cols + c);
        auto output_offset = Forward ? (r * cols + c) : (offset + r);
        context_.CopyItemsSameDevice(
            values.dtype(),
            block_size,
            values_vec + input_offset * block_size,
            output_data + output_offset * block_size);
      }
      offset += lengths_vec[c];
    }
    return true;
  }

 private:
  INPUT_TAGS(INPUTVALUE, LENGTHS);
  OUTPUT_TAGS(OUTPUTVALUE);
};
} // namespace caffe2

#endif // CAFFE2_OPERATORS_PACK_RNN_SEQUENCE_OP_H_