#ifndef CAFFE2_OPERATORS_ELEMENTWISE_LOGICAL_OPS_H_
#define CAFFE2_OPERATORS_ELEMENTWISE_LOGICAL_OPS_H_

#include "caffe2/core/common_omp.h"
#include "caffe2/core/context.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/operator.h"
#include "caffe2/operators/elementwise_ops.h"

#include <unordered_set>

namespace caffe2 {

template <class Context>
class WhereOp final : public Operator<Context> {
 public:
  USE_OPERATOR_FUNCTIONS(Context);
  USE_DISPATCH_HELPER;

  template <class... Args>
  explicit WhereOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...),
        OP_SINGLE_ARG(bool, "broadcast_on_rows", enable_broadcast_, 0) {}

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

  template <typename T>
  bool DoRunWithType() {
    auto& select = Input(0);
    auto& left = Input(1);
    auto& right = Input(2);

    if (enable_broadcast_) {
      CAFFE_ENFORCE_EQ(select.dim(), 1);
      CAFFE_ENFORCE_EQ(select.size(0), right.size(0));
      CAFFE_ENFORCE_EQ(left.sizes(), right.sizes());
    } else {
      CAFFE_ENFORCE_EQ(select.sizes(), left.sizes());
      CAFFE_ENFORCE_EQ(select.sizes(), right.sizes());
    }
    auto* output = Output(0, left.sizes(), at::dtype<T>());

    const bool* select_data = select.template data<bool>();
    const T* left_data = left.template data<T>();
    const T* right_data = right.template data<T>();
    T* output_data = output->template mutable_data<T>();

    if (enable_broadcast_) {
      size_t block_size = left.size_from_dim(1);
      for (int i = 0; i < select.numel(); i++) {
        size_t offset = i * block_size;
        if (select_data[i]) {
          context_.CopyItemsSameDevice(
              output->dtype(),
              block_size,
              left_data + offset,
              output_data + offset);
        } else {
          context_.CopyItemsSameDevice(
              output->dtype(),
              block_size,
              right_data + offset,
              output_data + offset);
        }
      }
    } else {
      for (int i = 0; i < select.numel(); ++i) {
        output_data[i] = select_data[i] ? left_data[i] : right_data[i];
      }
    }
    return true;
  }

 private:
  bool enable_broadcast_;
};

class IsMemberOfValueHolder {
  std::unordered_set<int32_t> int32_values_;
  std::unordered_set<int64_t> int64_values_;
  std::unordered_set<bool> bool_values_;
  std::unordered_set<std::string> string_values_;
  bool has_values_ = false;

 public:
  template <typename T>
  std::unordered_set<T>& get();

  template <typename T>
  void set(const std::vector<T>& args) {
    has_values_ = true;
    auto& values = get<T>();
    values.insert(args.begin(), args.end());
  }

  bool has_values() {
    return has_values_;
  }
};

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

  static constexpr const char* VALUE_TAG = "value";

 public:
  using TestableTypes = TensorTypes<int32_t, int64_t, bool, std::string>;

  template <class... Args>
  explicit IsMemberOfOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    auto dtype =
        static_cast<TensorProto_DataType>(this->template GetSingleArgument<int>(
            "dtype", TensorProto_DataType_UNDEFINED));
    switch (dtype) {
      case TensorProto_DataType_INT32:
        values_.set(this->template GetRepeatedArgument<int32_t>(VALUE_TAG));
        break;
      case TensorProto_DataType_INT64:
        values_.set(this->template GetRepeatedArgument<int64_t>(VALUE_TAG));
        break;
      case TensorProto_DataType_BOOL:
        values_.set(this->template GetRepeatedArgument<bool>(VALUE_TAG));
        break;
      case TensorProto_DataType_STRING:
        values_.set(this->template GetRepeatedArgument<std::string>(VALUE_TAG));
        break;
      case TensorProto_DataType_UNDEFINED:
        // If dtype is not provided, values_ will be filled the first time that
        // DoRunWithType is called.
        break;
      default:
        CAFFE_THROW("Unexpected 'dtype' argument value: ", dtype);
    }
  }
  virtual ~IsMemberOfOp() noexcept {}

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

  template <typename T>
  bool DoRunWithType() {
    auto& input = Input(0);

    auto* output = Output(0, input.sizes(), at::dtype<bool>());

    if (!values_.has_values()) {
      values_.set(this->template GetRepeatedArgument<T>(VALUE_TAG));
    }
    const auto& values = values_.get<T>();

    const T* input_data = input.template data<T>();
    bool* output_data = output->template mutable_data<bool>();
    for (int i = 0; i < input.numel(); ++i) {
      output_data[i] = values.find(input_data[i]) != values.end();
    }
    return true;
  }

 protected:
  IsMemberOfValueHolder values_;
};

} // namespace caffe2

#endif // CAFFE2_OPERATORS_ELEMENTWISE_LOGICAL_OPS_H_