#ifndef CAFFE2_OPERATORS_FEATURE_MAPS_OPS_H_
#define CAFFE2_OPERATORS_FEATURE_MAPS_OPS_H_

#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"

namespace caffe2 {

template <class Context>
class MergeSingleScalarFeatureTensorsOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit MergeSingleScalarFeatureTensorsOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    numInputs_ = InputSize() / kNumTensorsPerInput;
    featureIDs_ = this->template GetRepeatedArgument<int64_t>("feature_ids");
  }
  virtual ~MergeSingleScalarFeatureTensorsOp() noexcept {}

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

  template <typename T>
  bool DoRunWithType() {
    int numExamples = Input(0).numel();
    int totalNumFeatures = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      const bool* inPresenceData =
          Input(kNumTensorsPerInput * inputIndex + 1).template data<bool>();
      for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
        if (inPresenceData[exampleIndex]) {
          ++totalNumFeatures;
        }
      }
    }

    auto* outLengths = Output(0, {numExamples}, at::dtype<int32_t>());
    auto* outKeys = Output(1, {totalNumFeatures}, at::dtype<int64_t>());
    auto* outValues = Output(2, {totalNumFeatures}, at::dtype<T>());

    int32_t* outLengthsData = outLengths->template mutable_data<int32_t>();
    int64_t* outKeysData = outKeys->template mutable_data<int64_t>();
    T* outValuesData = outValues->template mutable_data<T>();

    int keysOffset = 0;
    for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
      outLengthsData[exampleIndex] = 0;
      for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
        const T* inData =
            Input(kNumTensorsPerInput * inputIndex).template data<T>();
        const bool* inPresenceData =
            Input(kNumTensorsPerInput * inputIndex + 1).template data<bool>();
        if (inPresenceData[exampleIndex]) {
          ++outLengthsData[exampleIndex];
          outKeysData[keysOffset] = featureIDs_[inputIndex];
          outValuesData[keysOffset] = inData[exampleIndex];
          ++keysOffset;
        }
      }
    }
    return true;
  }

 private:
  const int kNumTensorsPerInput = 2;
  int numInputs_;
  std::vector<int64_t> featureIDs_;
};

template <class Context>
class MergeSingleScalarFeatureTensorsGradientOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit MergeSingleScalarFeatureTensorsGradientOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    numFeatureInputs_ = InputSize() - 1; // Everything other than values_grad
  }
  virtual ~MergeSingleScalarFeatureTensorsGradientOp() noexcept {}

  bool RunOnDevice() override {
    return DispatchHelper<
        TensorTypes<bool, int32_t, int64_t, float, double, std::string>>::
        call(this, Input(InputSize() - 1));
  }

  template <typename T>
  bool DoRunWithType() {
    int numExamples = Input(0).numel();
    for (int inputIndex = 0; inputIndex < numFeatureInputs_; ++inputIndex) {
      Output(inputIndex)->ResizeLike(Input(inputIndex));
    }

    const T* inValuesGradData = Input(InputSize() - 1).template data<T>();

    T default_value = T();
    int valuesOffset = 0;
    for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
      for (int inputIndex = 0; inputIndex < numFeatureInputs_; ++inputIndex) {
        const bool* inPresenceData = Input(inputIndex).template data<bool>();
        T* outFeatureData = Output(inputIndex)->template mutable_data<T>();
        if (inPresenceData[exampleIndex]) {
          outFeatureData[exampleIndex] = inValuesGradData[valuesOffset];
          ++valuesOffset;
        } else {
          outFeatureData[exampleIndex] = default_value;
        }
      }
    }
    return true;
  }

 private:
  int numFeatureInputs_;
};

template <class Context>
class MergeSingleListFeatureTensorsOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit MergeSingleListFeatureTensorsOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    numInputs_ = InputSize() / kNumTensorsPerInput;
    inValuesOffset_.resize(numInputs_);
    featureIDs_ = this->template GetRepeatedArgument<int64_t>("feature_ids");
  }
  virtual ~MergeSingleListFeatureTensorsOp() noexcept {}

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

  template <typename T>
  bool DoRunWithType() {
    int numExamples = Input(0).numel();
    int totalNumFeatures = 0;
    int totalNumValues = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      const int32_t* inLengthsData =
          Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
      const bool* inPresenceData =
          Input(kNumTensorsPerInput * inputIndex + 2).template data<bool>();
      for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
        if (inPresenceData[exampleIndex]) {
          ++totalNumFeatures;
          totalNumValues += inLengthsData[exampleIndex];
        }
      }
    }

    auto* outLengths = Output(0, {numExamples}, at::dtype<int32_t>());
    auto* outKeys = Output(1, {totalNumFeatures}, at::dtype<int64_t>());
    auto* outValuesLengths =
        Output(2, {totalNumFeatures}, at::dtype<int32_t>());
    auto* outValuesValues = Output(3, {totalNumValues}, at::dtype<T>());

    int32_t* outLengthsData = outLengths->template mutable_data<int32_t>();
    int64_t* outKeysData = outKeys->template mutable_data<int64_t>();
    int32_t* outValuesLengthsData =
        outValuesLengths->template mutable_data<int32_t>();
    T* outValuesValuesData = outValuesValues->template mutable_data<T>();

    int keysOffset = 0;
    int valuesOffset = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      inValuesOffset_[inputIndex] = 0;
    }
    for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
      outLengthsData[exampleIndex] = 0;
      for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
        const int32_t* inLengthsData =
            Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
        const auto& inValues = Input(kNumTensorsPerInput * inputIndex + 1);
        const bool* inPresenceData =
            Input(kNumTensorsPerInput * inputIndex + 2).template data<bool>();
        if (inPresenceData[exampleIndex]) {
          ++outLengthsData[exampleIndex];
          outKeysData[keysOffset] = featureIDs_[inputIndex];
          outValuesLengthsData[keysOffset] = inLengthsData[exampleIndex];
          context_.CopyItemsSameDevice(
              inValues.dtype(),
              inLengthsData[exampleIndex],
              &inValues.template data<T>()[inValuesOffset_[inputIndex]],
              &outValuesValuesData[valuesOffset]);
          valuesOffset += inLengthsData[exampleIndex];
          inValuesOffset_[inputIndex] += inLengthsData[exampleIndex];
          ++keysOffset;
        }
      }
    }
    return true;
  }

 private:
  const int kNumTensorsPerInput = 3;
  int numInputs_;
  std::vector<int> inValuesOffset_;
  std::vector<int64_t> featureIDs_;
};

template <class Context>
class MergeSingleListOrMapFeatureTensorsGradientOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit MergeSingleListOrMapFeatureTensorsGradientOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    numFeatureInputs_ = (InputSize() - 1) / kNumTensorsPerInput;
  }
  virtual ~MergeSingleListOrMapFeatureTensorsGradientOp() noexcept {}

  bool RunOnDevice() override {
    return DispatchHelper<
        TensorTypes<bool, int32_t, int64_t, float, double, std::string>>::
        call(this, Input(InputSize() - 1));
  }

  template <typename T>
  bool DoRunWithType() {
    int numExamples = Input(0).numel();
    std::vector<int> outValuesOffset(numFeatureInputs_);
    for (int inputIndex = 0; inputIndex < numFeatureInputs_; ++inputIndex) {
      int inputNumValues = 0;
      const int32_t* inLengthsData =
          Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
      const bool* inPresenceData =
          Input(kNumTensorsPerInput * inputIndex + 1).template data<bool>();
      for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
        if (inPresenceData[exampleIndex]) {
          inputNumValues += inLengthsData[exampleIndex];
        }
      }
      Output(inputIndex)->Resize(inputNumValues);
    }

    const auto& inValuesValuesGrad = Input(InputSize() - 1);
    const T* inValuesValuesGradData = inValuesValuesGrad.template data<T>();

    int inValuesValuesOffset = 0;
    for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
      for (int inputIndex = 0; inputIndex < numFeatureInputs_; ++inputIndex) {
        const int32_t* inLengthsData =
            Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
        const bool* inPresenceData =
            Input(kNumTensorsPerInput * inputIndex + 1).template data<bool>();
        if (inPresenceData[exampleIndex]) {
          T* outFeatureValues = Output(inputIndex)->template mutable_data<T>();
          context_.CopyItemsSameDevice(
              inValuesValuesGrad.dtype(),
              inLengthsData[exampleIndex],
              &inValuesValuesGradData[inValuesValuesOffset],
              &outFeatureValues[outValuesOffset[inputIndex]]);
          outValuesOffset[inputIndex] += inLengthsData[exampleIndex];
          inValuesValuesOffset += inLengthsData[exampleIndex];
        }
      }
    }
    return true;
  }

 private:
  const int kNumTensorsPerInput = 2;
  int numFeatureInputs_;
};

template <class Context>
class MergeSingleMapFeatureTensorsOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit MergeSingleMapFeatureTensorsOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    numInputs_ = InputSize() / kNumTensorsPerInput;
    inValuesOffset_.resize(numInputs_);
    featureIDs_ = this->template GetRepeatedArgument<int64_t>("feature_ids");
  }
  virtual ~MergeSingleMapFeatureTensorsOp() noexcept {}

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

  template <typename K>
  bool DoRunWithType() {
    return DispatchHelper<
        TensorTypes2<bool, int32_t, int64_t, float, double, std::string>,
        K>::call(this, Input(2));
  }

  template <typename K, typename V>
  bool DoRunWithType2() {
    int numExamples = Input(0).numel();
    int totalNumFeatures = 0;
    int totalNumValues = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      const int32_t* inLengthsData =
          Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
      const bool* inPresenceData =
          Input(kNumTensorsPerInput * inputIndex + 3).template data<bool>();
      for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
        if (inPresenceData[exampleIndex]) {
          ++totalNumFeatures;
          totalNumValues += inLengthsData[exampleIndex];
        }
      }
    }

    auto* outLengths = Output(0, {numExamples}, at::dtype<int32_t>());
    auto* outKeys = Output(1, {totalNumFeatures}, at::dtype<int64_t>());
    auto* outValuesLengths =
        Output(2, {totalNumFeatures}, at::dtype<int32_t>());
    auto* outValuesKeys = Output(3, {totalNumValues}, at::dtype<K>());
    auto* outValuesValues = Output(4, {totalNumValues}, at::dtype<V>());

    int32_t* outLengthsData = outLengths->template mutable_data<int32_t>();
    int64_t* outKeysData = outKeys->template mutable_data<int64_t>();
    int32_t* outValuesLengthsData =
        outValuesLengths->template mutable_data<int32_t>();
    K* outValuesKeysData = outValuesKeys->template mutable_data<K>();
    V* outValuesValuesData = outValuesValues->template mutable_data<V>();

    int keysOffset = 0;
    int valuesOffset = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      inValuesOffset_[inputIndex] = 0;
    }
    for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
      outLengthsData[exampleIndex] = 0;
      for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
        const int32_t* inLengthsData =
            Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
        const auto& inKeys = Input(kNumTensorsPerInput * inputIndex + 1);
        const auto& inValues = Input(kNumTensorsPerInput * inputIndex + 2);
        const bool* inPresenceData =
            Input(kNumTensorsPerInput * inputIndex + 3).template data<bool>();
        if (inPresenceData[exampleIndex]) {
          ++outLengthsData[exampleIndex];
          outKeysData[keysOffset] = featureIDs_[inputIndex];
          outValuesLengthsData[keysOffset] = inLengthsData[exampleIndex];
          context_.CopyItemsSameDevice(
              inKeys.dtype(),
              inLengthsData[exampleIndex],
              &inKeys.template data<K>()[inValuesOffset_[inputIndex]],
              &outValuesKeysData[valuesOffset]);
          context_.CopyItemsSameDevice(
              inValues.dtype(),
              inLengthsData[exampleIndex],
              &inValues.template data<V>()[inValuesOffset_[inputIndex]],
              &outValuesValuesData[valuesOffset]);
          valuesOffset += inLengthsData[exampleIndex];
          inValuesOffset_[inputIndex] += inLengthsData[exampleIndex];
          ++keysOffset;
        }
      }
    }
    return true;
  }

 private:
  const int kNumTensorsPerInput = 4;
  int numInputs_;
  std::vector<int> inValuesOffset_;
  std::vector<int64_t> featureIDs_;
};

template <class Context>
class MergeMultiScalarFeatureTensorsOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit MergeMultiScalarFeatureTensorsOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    numInputs_ = InputSize() / kNumTensorsPerInput;
    inKeysOffset_.resize(numInputs_);
  }
  virtual ~MergeMultiScalarFeatureTensorsOp() noexcept {}

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

  template <typename T>
  bool DoRunWithType() {
    int numExamples = Input(0).numel();
    int totalNumFeatures = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      totalNumFeatures += Input(kNumTensorsPerInput * inputIndex + 1).numel();
    }

    auto* outLengths = Output(0, {numExamples}, at::dtype<int32_t>());
    auto* outKeys = Output(1, {totalNumFeatures}, at::dtype<int64_t>());
    auto* outValues = Output(2, {totalNumFeatures}, at::dtype<T>());

    int32_t* outLengthsData = outLengths->template mutable_data<int32_t>();
    int64_t* outKeysData = outKeys->template mutable_data<int64_t>();
    T* outValuesData = outValues->template mutable_data<T>();

    int outKeysOffset = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      inKeysOffset_[inputIndex] = 0;
    }
    for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
      outLengthsData[exampleIndex] = 0;
      for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
        const int32_t* inLengthsData =
            Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
        const int64_t* inKeysData = Input(kNumTensorsPerInput * inputIndex + 1)
                                        .template data<int64_t>();
        const T* inValuesData =
            Input(kNumTensorsPerInput * inputIndex + 2).template data<T>();
        outLengthsData[exampleIndex] += inLengthsData[exampleIndex];
        for (int featureIndex = 0; featureIndex < inLengthsData[exampleIndex];
             ++featureIndex) {
          outKeysData[outKeysOffset] = inKeysData[inKeysOffset_[inputIndex]];
          outValuesData[outKeysOffset] =
              inValuesData[inKeysOffset_[inputIndex]];
          ++outKeysOffset;
          ++inKeysOffset_[inputIndex];
        }
      }
    }

    return true;
  }

 private:
  const int kNumTensorsPerInput = 3;
  int numInputs_;
  std::vector<int> inKeysOffset_;
};

template <class Context>
class MergeMultiScalarFeatureTensorsGradientOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit MergeMultiScalarFeatureTensorsGradientOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    numFeatureInputs_ = (InputSize() - 1) / kNumTensorsPerInput;
  }
  virtual ~MergeMultiScalarFeatureTensorsGradientOp() noexcept {}

  bool RunOnDevice() override {
    return DispatchHelper<
        TensorTypes<bool, int32_t, int64_t, float, double, std::string>>::
        call(this, Input(InputSize() - 1));
  }

  template <typename T>
  bool DoRunWithType() {
    int numExamples = Input(0).numel();
    std::vector<int> outValuesOffset(numFeatureInputs_);
    for (int inputIndex = 0; inputIndex < numFeatureInputs_; ++inputIndex) {
      int inputNumValues = 0;
      const int32_t* inLengthsData =
          Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
      for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
        inputNumValues += inLengthsData[exampleIndex];
      }
      Output(inputIndex)->Resize(inputNumValues);
    }

    const auto& inValuesGrad = Input(InputSize() - 1);
    const T* inValuesGradData = inValuesGrad.template data<T>();

    int inValuesOffset = 0;
    for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
      for (int inputIndex = 0; inputIndex < numFeatureInputs_; ++inputIndex) {
        const int32_t* inLengthsData =
            Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
        if (inLengthsData[exampleIndex] > 0) {
          T* outFeatureValues = Output(inputIndex)->template mutable_data<T>();
          context_.CopyItemsSameDevice(
              inValuesGrad.dtype(),
              inLengthsData[exampleIndex],
              &inValuesGradData[inValuesOffset],
              &outFeatureValues[outValuesOffset[inputIndex]]);
          outValuesOffset[inputIndex] += inLengthsData[exampleIndex];
          inValuesOffset += inLengthsData[exampleIndex];
        }
      }
    }
    return true;
  }

 private:
  int kNumTensorsPerInput = 1;
  int numFeatureInputs_;
};

template <class Context>
class MergeMultiListFeatureTensorsOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit MergeMultiListFeatureTensorsOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    numInputs_ = InputSize() / kNumTensorsPerInput;
    inKeysOffset_.resize(numInputs_);
    inValuesValuesOffset_.resize(numInputs_);
  }
  virtual ~MergeMultiListFeatureTensorsOp() noexcept {}

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

  template <typename T>
  bool DoRunWithType() {
    int numExamples = Input(0).numel();
    int totalNumFeatures = 0;
    int totalNumValues = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      totalNumFeatures += Input(kNumTensorsPerInput * inputIndex + 1).numel();
      totalNumValues += Input(kNumTensorsPerInput * inputIndex + 3).numel();
    }

    auto* outLengths = Output(0, {numExamples}, at::dtype<int32_t>());
    auto* outKeys = Output(1, {totalNumFeatures}, at::dtype<int64_t>());
    auto* outValuesLengths =
        Output(2, {totalNumFeatures}, at::dtype<int32_t>());
    auto* outValuesValues = Output(3, {totalNumValues}, at::dtype<T>());

    int32_t* outLengthsData = outLengths->template mutable_data<int32_t>();
    int64_t* outKeysData = outKeys->template mutable_data<int64_t>();
    int32_t* outValuesLengthsData =
        outValuesLengths->template mutable_data<int32_t>();
    T* outValuesValuesData = outValuesValues->template mutable_data<T>();

    int outKeysOffset = 0;
    int outValuesValuesOffset = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      inKeysOffset_[inputIndex] = 0;
      inValuesValuesOffset_[inputIndex] = 0;
    }
    for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
      outLengthsData[exampleIndex] = 0;
      for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
        const int32_t* inLengthsData =
            Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
        const int64_t* inKeysData = Input(kNumTensorsPerInput * inputIndex + 1)
                                        .template data<int64_t>();
        const int32_t* inValuesLengthsData =
            Input(kNumTensorsPerInput * inputIndex + 2)
                .template data<int32_t>();
        const auto& inValuesValues =
            Input(kNumTensorsPerInput * inputIndex + 3);
        outLengthsData[exampleIndex] += inLengthsData[exampleIndex];
        for (int featureIndex = 0; featureIndex < inLengthsData[exampleIndex];
             ++featureIndex) {
          outKeysData[outKeysOffset] = inKeysData[inKeysOffset_[inputIndex]];
          outValuesLengthsData[outKeysOffset] =
              inValuesLengthsData[inKeysOffset_[inputIndex]];
          context_.CopyItemsSameDevice(
              inValuesValues.dtype(),
              inValuesLengthsData[inKeysOffset_[inputIndex]],
              &inValuesValues
                   .template data<T>()[inValuesValuesOffset_[inputIndex]],
              &outValuesValuesData[outValuesValuesOffset]);
          outValuesValuesOffset +=
              inValuesLengthsData[inKeysOffset_[inputIndex]];
          inValuesValuesOffset_[inputIndex] +=
              inValuesLengthsData[inKeysOffset_[inputIndex]];
          ++outKeysOffset;
          ++inKeysOffset_[inputIndex];
        }
      }
    }

    return true;
  }

 private:
  const int kNumTensorsPerInput = 4;
  int numInputs_;
  std::vector<int> inKeysOffset_;
  std::vector<int> inValuesValuesOffset_;
};

template <class Context>
class MergeMultiMapFeatureTensorsOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit MergeMultiMapFeatureTensorsOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    numInputs_ = InputSize() / kNumTensorsPerInput;
    inKeysOffset_.resize(numInputs_);
    inValuesValuesOffset_.resize(numInputs_);
  }
  virtual ~MergeMultiMapFeatureTensorsOp() noexcept {}

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

  template <typename K>
  bool DoRunWithType() {
    return DispatchHelper<
        TensorTypes2<bool, int32_t, int64_t, float, double, std::string>,
        K>::call(this, Input(4));
  }

  template <typename K, typename V>
  bool DoRunWithType2() {
    int numExamples = Input(0).numel();
    int totalNumFeatures = 0;
    int totalNumValues = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      totalNumFeatures += Input(kNumTensorsPerInput * inputIndex + 1).numel();
      totalNumValues += Input(kNumTensorsPerInput * inputIndex + 4).numel();
    }

    auto* outLengths = Output(0, {numExamples}, at::dtype<int32_t>());
    auto* outKeys = Output(1, {totalNumFeatures}, at::dtype<int64_t>());
    auto* outValuesLengths =
        Output(2, {totalNumFeatures}, at::dtype<int32_t>());
    auto* outValuesKeys = Output(3, {totalNumValues}, at::dtype<K>());
    auto* outValuesValues = Output(4, {totalNumValues}, at::dtype<V>());

    int32_t* outLengthsData = outLengths->template mutable_data<int32_t>();
    int64_t* outKeysData = outKeys->template mutable_data<int64_t>();
    int32_t* outValuesLengthsData =
        outValuesLengths->template mutable_data<int32_t>();
    K* outValuesKeysData = outValuesKeys->template mutable_data<K>();
    V* outValuesValuesData = outValuesValues->template mutable_data<V>();

    int outKeysOffset = 0;
    int outValuesValuesOffset = 0;
    for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
      inKeysOffset_[inputIndex] = 0;
      inValuesValuesOffset_[inputIndex] = 0;
    }
    for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
      outLengthsData[exampleIndex] = 0;
      for (int inputIndex = 0; inputIndex < numInputs_; ++inputIndex) {
        const int32_t* inLengthsData =
            Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
        const int64_t* inKeysData = Input(kNumTensorsPerInput * inputIndex + 1)
                                        .template data<int64_t>();
        const int32_t* inValuesLengthsData =
            Input(kNumTensorsPerInput * inputIndex + 2)
                .template data<int32_t>();
        const auto& inValuesKeys = Input(kNumTensorsPerInput * inputIndex + 3);
        const auto& inValuesValues =
            Input(kNumTensorsPerInput * inputIndex + 4);
        outLengthsData[exampleIndex] += inLengthsData[exampleIndex];
        for (int featureIndex = 0; featureIndex < inLengthsData[exampleIndex];
             ++featureIndex) {
          outKeysData[outKeysOffset] = inKeysData[inKeysOffset_[inputIndex]];
          outValuesLengthsData[outKeysOffset] =
              inValuesLengthsData[inKeysOffset_[inputIndex]];
          context_.CopyItemsSameDevice(
              inValuesKeys.dtype(),
              inValuesLengthsData[inKeysOffset_[inputIndex]],
              &inValuesKeys
                   .template data<K>()[inValuesValuesOffset_[inputIndex]],
              &outValuesKeysData[outValuesValuesOffset]);
          context_.CopyItemsSameDevice(
              inValuesValues.dtype(),
              inValuesLengthsData[inKeysOffset_[inputIndex]],
              &inValuesValues
                   .template data<V>()[inValuesValuesOffset_[inputIndex]],
              &outValuesValuesData[outValuesValuesOffset]);
          outValuesValuesOffset +=
              inValuesLengthsData[inKeysOffset_[inputIndex]];
          inValuesValuesOffset_[inputIndex] +=
              inValuesLengthsData[inKeysOffset_[inputIndex]];
          ++outKeysOffset;
          ++inKeysOffset_[inputIndex];
        }
      }
    }

    return true;
  }

 private:
  const int kNumTensorsPerInput = 5;
  int numInputs_;
  std::vector<int> inKeysOffset_;
  std::vector<int> inValuesValuesOffset_;
};

template <class Context>
class MergeMultiListOrMapFeatureTensorsGradientOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit MergeMultiListOrMapFeatureTensorsGradientOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    numFeatureInputs_ = (InputSize() - 1) / kNumTensorsPerInput;
  }
  virtual ~MergeMultiListOrMapFeatureTensorsGradientOp() noexcept {}

  bool RunOnDevice() override {
    return DispatchHelper<
        TensorTypes<bool, int32_t, int64_t, float, double, std::string>>::
        call(this, Input(InputSize() - 1));
  }

  template <typename T>
  bool DoRunWithType() {
    int numExamples = Input(0).numel();
    std::vector<int> outValuesLengthOffset(numFeatureInputs_);
    std::vector<int> outValuesValuesOffset(numFeatureInputs_);
    for (int inputIndex = 0; inputIndex < numFeatureInputs_; ++inputIndex) {
      int inputNumValues = 0;
      auto& inValuesLength = Input(kNumTensorsPerInput * inputIndex + 1);
      const int32_t* inValuesLengthsData =
          inValuesLength.template data<int32_t>();
      for (int valuesIndex = 0; valuesIndex < inValuesLength.numel();
           ++valuesIndex) {
        inputNumValues += inValuesLengthsData[valuesIndex];
      }
      Output(inputIndex)->Resize(inputNumValues);
    }

    const auto& inValuesValuesGrad = Input(InputSize() - 1);
    const T* inValuesValuesGradData = inValuesValuesGrad.template data<T>();

    int inValuesValuesOffset = 0;
    for (int exampleIndex = 0; exampleIndex < numExamples; ++exampleIndex) {
      for (int inputIndex = 0; inputIndex < numFeatureInputs_; ++inputIndex) {
        const int32_t* inLengthsData =
            Input(kNumTensorsPerInput * inputIndex).template data<int32_t>();
        const int32_t* inValuesLengthsData =
            Input(kNumTensorsPerInput * inputIndex + 1)
                .template data<int32_t>();
        int valuesLengthCopy = 0;
        for (int valuesLengthIndex = 0;
             valuesLengthIndex < inLengthsData[exampleIndex];
             ++valuesLengthIndex) {
          valuesLengthCopy += inValuesLengthsData
              [outValuesLengthOffset[inputIndex] + valuesLengthIndex];
        }
        if (valuesLengthCopy > 0) {
          T* outFeatureValues = Output(inputIndex)->template mutable_data<T>();
          context_.CopyItemsSameDevice(
              inValuesValuesGrad.dtype(),
              valuesLengthCopy,
              &inValuesValuesGradData[inValuesValuesOffset],
              &outFeatureValues[outValuesValuesOffset[inputIndex]]);
        }
        outValuesLengthOffset[inputIndex] += inLengthsData[exampleIndex];
        outValuesValuesOffset[inputIndex] += valuesLengthCopy;
        inValuesValuesOffset += valuesLengthCopy;
      }
    }
    return true;
  }

 private:
  int kNumTensorsPerInput = 2;
  int numFeatureInputs_;
};

} // namespace caffe2

#endif // CAFFE2_OPERATORS_FEATURE_MAPS_OPS_H_