#ifndef CAFFE2_OPERATORS_MERGE_ID_LISTS_OP_H_
#define CAFFE2_OPERATORS_MERGE_ID_LISTS_OP_H_

#include <set>
#include <vector>
#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"

namespace caffe2 {

template <class Context>
class MergeIdListsOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  USE_SIMPLE_CTOR_DTOR(MergeIdListsOp);

  template <typename T>
  bool DoRunWithType() {
    auto& first_lengths = Input(0);
    CAFFE_ENFORCE_EQ(first_lengths.dim(), 1, "LENGTHS should be 1-D");
    const auto batch_size = first_lengths.numel();

    auto* out_lengths = Output(0, first_lengths.sizes(), at::dtype<int32_t>());

    auto* out_lengths_data = out_lengths->template mutable_data<int32_t>();

    /**
     * Loop to figure out how much space to reserve for output
     * and perform checks.
     */
    auto M = 0;
    for (size_t i = 0; i < InputSize(); i += 2) {
      auto& lengths = Input(i);
      CAFFE_ENFORCE_EQ(lengths.dim(), 1, "LENGTHS should be 1-D");
      CAFFE_ENFORCE_EQ(lengths.numel(), batch_size, "LENGTHS should be equal");
      auto& values = Input(i + 1);
      CAFFE_ENFORCE_EQ(values.dim(), 1, "VALUES should be 1-D");
      M += values.numel();
    }

    auto* out_values = Output(1, {M}, at::dtype<T>());

    T* out_values_data = out_values->template mutable_data<T>();
    auto pos = 0;

    // TODO(badri): Use unordered_set if performance is an issue
    std::set<T> deduped;
    std::vector<int> offsets(InputSize(), 0);
    for (auto sample = 0; sample < batch_size; sample++) {
      for (size_t i = 0; i < InputSize(); i += 2) {
        auto& lengths = Input(i);
        const auto* lengths_data = lengths.template data<int32_t>();

        auto& values = Input(i + 1);
        const T* values_data = values.template data<T>();
        const auto length = lengths_data[sample];

        for (auto j = offsets[i]; j < offsets[i] + length; j++) {
          deduped.insert(values_data[j]);
        }
        offsets[i] += length;
      }
      for (auto val : deduped) {
        out_values_data[pos++] = val;
      }
      out_lengths_data[sample] = deduped.size();
      deduped.clear();
    }
    out_values->Resize(pos);
    return true;
  }

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

} // namespace caffe2

#endif // CAFFE2_OPERATORS_MERGE_ID_LISTS_OP_H_