#ifndef CAFFE2_OPERATORS_FILLER_OP_H_
#define CAFFE2_OPERATORS_FILLER_OP_H_

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

namespace caffe2 {

// FillerOp takes in either zero or one input.
//
// If the number of input is 1, the shape will be identical to that of the input
// at run time with optional additional dimensions appended at the end as
// specified by "extra_shape" argument. In that case the "shape" parameter
// should not be set.
//
// If the number of inputs is 0, the full shape must be provided via "shape"
// argument
template <class Context>
class FillerOp : public Operator<Context> {
 public:
  template <class... Args>
  explicit FillerOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...),
        shape_(this->template GetRepeatedArgument<int64_t>("shape")),
        extra_shape_(ToVectorint64_t(
            this->template GetRepeatedArgument<int>("extra_shape"))),
        input_as_shape_(
            this->template GetSingleArgument<bool>("input_as_shape", false)) {
    if (InputSize()) {
      if (shape_.size() != 0) {
        CAFFE_THROW(
            "Cannot set the shape argument and pass in an input at "
            "the same time");
      }
    } else {
      if (!extra_shape_.empty()) {
        CAFFE_THROW("Cannot set extra_shape when there is no input");
      }
      if (input_as_shape_) {
        CAFFE_THROW("An input must be given if input_as_shape is true");
      }
      if (shape_.size() == 0 &&
          this->template HasSingleArgumentOfType<int>("shape")) {
        CAFFE_THROW("Fill 'shape' argument was a scalar, list expected");
      }
    }
  }

  virtual ~FillerOp() {}
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  bool RunOnDevice() override {
    auto* output = Operator<Context>::Output(0);
    if (InputSize()) {
      auto shape = vector<int64_t>{};
      if (input_as_shape_) {
        if (this->InputIsTensorType(0, CPU)) {
          // originally, shape input must be in CPU context
          auto& input = this->template Input<Tensor>(0, CPU);
          CAFFE_ENFORCE_EQ(
              input.dim(),
              1,
              "When input_as_shape is true, the input must be a 1D tensor of "
              "data type int64_t");
          CAFFE_ENFORCE(input.numel() > 0);
          auto* shape_data = input.template data<int64_t>();
          shape.insert(shape.end(), shape_data, shape_data + input.dim32(0));
        } else {
          // in ONNX case, we allow shape to be in CUDA context
          auto& input = Input(0);
          CAFFE_ENFORCE_EQ(
              input.dim(),
              1,
              "When input_as_shape is true, the input must be a 1D tensor of "
              "data type int64_t");
          CAFFE_ENFORCE(input.numel() > 0);
          auto* shape_data = input.template data<int64_t>();
          std::unique_ptr<int64_t[]> shape_data_copy = caffe2::make_unique<int64_t[]>(input.dim32(0));
          context_.template CopyToCPU<int64_t>(input.dim32(0), shape_data, shape_data_copy.get());
          shape.insert(shape.end(), shape_data_copy.get(), shape_data_copy.get() + input.dim32(0));
        }
      } else {
        auto& input = Input(0);
        shape.insert(shape.end(), input.sizes().begin(), input.sizes().end());
      }
      shape.insert(shape.end(), extra_shape_.begin(), extra_shape_.end());
      output->Resize(shape);
    } else {
      output->Resize(shape_);
    }
    return Fill(output);
  }

  virtual bool Fill(Tensor* output) = 0;

 protected:
  vector<int64_t> shape_;
  vector<int64_t> extra_shape_;
  bool input_as_shape_;
};

template <typename T, class Context>
class UniformFillOp final : public FillerOp<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit UniformFillOp(Args&&... args)
      : FillerOp<Context>(std::forward<Args>(args)...),
        min_(this->template GetSingleArgument<T>("min", 0)),
        max_(this->template GetSingleArgument<T>("max", 1)) {
    if (InputSize() == 3) {
      CAFFE_ENFORCE(
          !this->template HasSingleArgumentOfType<T>("min"),
          "Cannot set both min arg and min input blob");
      CAFFE_ENFORCE(
          !this->template HasSingleArgumentOfType<T>("max"),
          "Cannot set both max arg and max input blob");
    } else {
      CAFFE_ENFORCE_LT(
          min_, max_, "Max value should be bigger than min value.");
    }
  }

  bool Fill(Tensor* output) override {
    T min = min_;
    T max = max_;
    if (InputSize() == 3) {
      CAFFE_ENFORCE_EQ(1, Input(1).numel(), "min blob must be scalar");
      CAFFE_ENFORCE_EQ(1, Input(2).numel(), "max blob must be scalar");
      min = *Input(1).template data<T>();
      max = *Input(2).template data<T>();
      if (min > max) {
        auto shape = output->sizes().vec();
        shape[0] = 0;
        output->Resize(shape);
        output->template mutable_data<T>();
        return true;
      }
    }
    math::RandUniform<T, Context>(
        output->numel(),
        min,
        max,
        output->template mutable_data<T>(),
        &context_);
    return true;
  }

 private:
  T min_;
  T max_;
};

template <class Context>
class UniqueUniformFillOp final : public FillerOp<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit UniqueUniformFillOp(Args&&... args)
      : FillerOp<Context>(std::forward<Args>(args)...) {
    TensorProto_DataType dtype =
        static_cast<TensorProto_DataType>(this->template GetSingleArgument<int>(
            "dtype", TensorProto_DataType_INT32));

    switch (dtype) {
      case TensorProto_DataType_INT32:
        CheckRange<int>();
        body_ = &UniqueUniformFillOp::FillWithType<int>;
        break;
      case TensorProto_DataType_INT64:
        CheckRange<int64_t>();
        body_ = &UniqueUniformFillOp::FillWithType<int64_t>;
        break;
      case TensorProto_DataType_UNDEFINED:
        CAFFE_THROW(
            "UniqueUniformFill op cannot have undefined 'dtype' argument");
      // break;
      default:
        CAFFE_THROW("Unexpected 'dtype' argument value: ", dtype);
    }
  }

  bool Fill(Tensor* output) override {
    return (this->*body_)(output);
  }

 private:
  template <typename T>
  void CheckRange() {
    CAFFE_ENFORCE(this->template HasSingleArgumentOfType<T>("min"));
    CAFFE_ENFORCE(this->template HasSingleArgumentOfType<T>("max"));
    CAFFE_ENFORCE_LT(
        this->template GetSingleArgument<T>("min", 0),
        this->template GetSingleArgument<T>("max", 0),
        "Max value should be bigger than min value.");
  }

  template <typename T>
  bool FillWithType(Tensor* output) {
    T min = this->template GetSingleArgument<T>("min", 0);
    T max = this->template GetSingleArgument<T>("max", 0);

    const T* avoid_data = nullptr;
    size_t avoid_size = 0;
    if (InputSize() >= 2) {
      auto& avoid = Input(1);
      avoid_data = avoid.template data<T>();
      avoid_size = avoid.numel();
    }
    math::RandUniformUnique<T, Context>(
        output->numel(),
        min,
        max,
        output->template mutable_data<T>(),
        avoid_size,
        avoid_data,
        &context_);
    return true;
  }

  bool (UniqueUniformFillOp::*body_)(Tensor* output);
};

template <class Context>
class ConstantFillOp final : public FillerOp<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit ConstantFillOp(Args&&... args)
      : FillerOp<Context>(std::forward<Args>(args)...) {
    TensorProto_DataType dtype =
        static_cast<TensorProto_DataType>(this->template GetSingleArgument<int>(
            "dtype", TensorProto_DataType_FLOAT));

    if (!OperatorBase::HasArgument("dtype") &&
        OperatorBase::HasArgument("value")) {
      // If 'dtype' is not provided, infer type based on the type of 'value'
      // Currently, single argument contains either float, int64 or bytes
      if (this->template HasSingleArgumentOfType<float>("value")) {
        dtype = TensorProto_DataType_FLOAT;
      } else if (this->template HasSingleArgumentOfType<int64_t>("value")) {
        dtype = TensorProto_DataType_INT64;
      } else {
        CAFFE_THROW("Argument 'value' is of unexpected type");
      }
      VLOG(1) << "Argument 'dtype' is not provided. Assume the data type is "
              << "the same as that of argument 'value': " << dtype;
    }

    switch (dtype) {
      case TensorProto_DataType_FLOAT:
        body_ = &ConstantFillOp::FillWithType<float>;
        break;
      case TensorProto_DataType_DOUBLE:
        body_ = &ConstantFillOp::FillWithType<double>;
        break;
      case TensorProto_DataType_BOOL:
        body_ = &ConstantFillOp::FillWithType<bool>;
        break;
      case TensorProto_DataType_INT8:
        body_ = &ConstantFillOp::FillWithType<int8_t>;
        break;
      case TensorProto_DataType_INT16:
        body_ = &ConstantFillOp::FillWithType<int16_t>;
        break;
      case TensorProto_DataType_INT32:
        body_ = &ConstantFillOp::FillWithType<int>;
        break;
      case TensorProto_DataType_INT64:
        body_ = &ConstantFillOp::FillWithType<int64_t>;
        break;
      case TensorProto_DataType_UINT8:
        body_ = &ConstantFillOp::FillWithType<uint8_t>;
        break;
      case TensorProto_DataType_UINT16:
        body_ = &ConstantFillOp::FillWithType<uint16_t>;
        break;
      case TensorProto_DataType_STRING:
        body_ = &ConstantFillOp::FillWithString;
        break;
      case TensorProto_DataType_UNDEFINED:
        CAFFE_THROW("ConstantFill op cannot have undefined 'dtype' argument");
      // break;
      default:
        CAFFE_THROW("Unexpected 'dtype' argument value: ", dtype);
    }
  }

  bool Fill(Tensor* output) override {
    return (this->*body_)(output);
  }

  template <typename T>
  bool FillWithType(Tensor* output) {
    T value = this->template GetSingleArgument<T>("value", 0);
    auto* data = output->template mutable_data<T>();
    if (output->numel()) {
      math::Set<T, Context>(output->numel(), value, data, &context_);
    }
    return true;
  }

  bool FillWithString(Tensor* output) {
    auto value = this->template GetSingleArgument<std::string>("value", "");
    auto* data = output->template mutable_data<std::string>();
    for (int i = 0; i < output->numel(); ++i) {
      data[i] = value;
    }
    return true;
  }

 private:
  bool (ConstantFillOp::*body_)(Tensor* output);
};

template <class Context>
class DiagonalFillOp final : public FillerOp<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit DiagonalFillOp(Args&&... args)
      : FillerOp<Context>(std::forward<Args>(args)...) {
    TensorProto_DataType dtype =
        static_cast<TensorProto_DataType>(this->template GetSingleArgument<int>(
            "dtype", TensorProto_DataType_FLOAT));

    if (!OperatorBase::HasArgument("dtype") &&
        OperatorBase::HasArgument("value")) {
      // If 'dtype' is not provided, infer type based on the type of 'value'
      // Currently, single argument contains either float, int64 or bytes
      if (this->template HasSingleArgumentOfType<float>("value")) {
        dtype = TensorProto_DataType_FLOAT;
      } else if (this->template HasSingleArgumentOfType<int64_t>("value")) {
        dtype = TensorProto_DataType_INT64;
      } else {
        CAFFE_THROW("Argument 'value' is of unexpected type");
      }
      VLOG(1) << "Argument 'dtype' is not provided. Assume the data type is "
              << "the same as that of argument 'value': " << dtype;
    }

    switch (dtype) {
      case TensorProto_DataType_FLOAT:
        body_ = &DiagonalFillOp::FillWithType<float>;
        break;
      case TensorProto_DataType_DOUBLE:
        body_ = &DiagonalFillOp::FillWithType<double>;
        break;
      case TensorProto_DataType_BOOL:
        body_ = &DiagonalFillOp::FillWithType<bool>;
        break;
      case TensorProto_DataType_INT8:
        body_ = &DiagonalFillOp::FillWithType<int8_t>;
        break;
      case TensorProto_DataType_INT16:
        body_ = &DiagonalFillOp::FillWithType<int16_t>;
        break;
      case TensorProto_DataType_INT32:
        body_ = &DiagonalFillOp::FillWithType<int>;
        break;
      case TensorProto_DataType_INT64:
        body_ = &DiagonalFillOp::FillWithType<int64_t>;
        break;
      case TensorProto_DataType_UINT8:
        body_ = &DiagonalFillOp::FillWithType<uint8_t>;
        break;
      case TensorProto_DataType_UINT16:
        body_ = &DiagonalFillOp::FillWithType<uint16_t>;
        break;
      case TensorProto_DataType_UNDEFINED:
        CAFFE_THROW("Cannot have undefined 'dtype' argument");
      default:
        CAFFE_THROW("Unexpected 'dtype' argument value: ", dtype);
    }
  }

  bool Fill(Tensor* output) override {
    return (this->*body_)(output);
  }

  template <typename T>
  bool FillWithType(Tensor* output);

 private:
  void VerifyOutputShape(Tensor* output) {
    CAFFE_ENFORCE(output->dim() >= 2, "Input shape must be >= 2D");
  }

  int64_t GetStepSize(Tensor* output) {
    int64_t step;
    if (output->dim() == 2) {
      step = output->size(1) + 1;
    } else {
      int64_t prev_i = output->size(0);
      for (auto i : output->sizes()) {
        if (i != prev_i) {
          CAFFE_THROW("All dimensions of input must be of equal length");
        }
      }
      vector<int64_t> cumprod(output->dim());
      auto dims = output->sizes();
      std::partial_sum(
          dims.begin(),
          dims.end() - 1,
          cumprod.begin(),
          std::multiplies<int64_t>());
      step = 1 +
          std::accumulate(
                 cumprod.begin(), cumprod.end(), static_cast<int64_t>(0));
      VLOG(0) << step;
    }
    return step;
  }

  bool (DiagonalFillOp::*body_)(Tensor* output);
};

template <typename T, class Context>
class GaussianFillOp final : public FillerOp<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit GaussianFillOp(Args&&... args)
      : FillerOp<Context>(std::forward<Args>(args)...),
        mean_(this->template GetSingleArgument<float>("mean", 0)),
        std_(this->template GetSingleArgument<float>("std", 1)) {
    DCHECK_GT(std_, 0) << "Standard deviation should be nonnegative.";
  }

  bool Fill(Tensor* output) override {
    math::RandGaussian<T, Context>(
        output->numel(),
        mean_,
        std_,
        output->template mutable_data<T>(),
        &context_);
    return true;
  }

 private:
  T mean_;
  T std_;
};

template <typename T, class Context>
class XavierFillOp final : public FillerOp<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit XavierFillOp(Args&&... args)
      : FillerOp<Context>(std::forward<Args>(args)...) {}

  bool Fill(Tensor* output) override {
    const int fan_in = output->numel() / output->dim32(0);
    T scale = std::sqrt(T(3) / fan_in);
    math::RandUniform<T, Context>(
        output->numel(),
        -scale,
        scale,
        output->template mutable_data<T>(),
        &context_);
    return true;
  }
};

template <typename T, class Context>
class MSRAFillOp final : public FillerOp<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit MSRAFillOp(Args&&... args)
      : FillerOp<Context>(std::forward<Args>(args)...) {}

  bool Fill(Tensor* output) override {
    const int fan_out = output->numel() / output->dim32(1);
    T scale = std::sqrt(T(2) / fan_out);
    math::RandGaussian<T, Context>(
        output->numel(),
        0.0,
        scale,
        output->template mutable_data<T>(),
        &context_);
    return true;
  }
};

// This is mostly used just as a debugging purpose stuff: it fills a tensor
// sequentially with values 0, 1, 2..., which can then be used to check e.g.
// reshape operations by allowing one to read the indices more easily.
template <typename T, class Context>
class RangeFillOp final : public FillerOp<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit RangeFillOp(Args&&... args)
      : FillerOp<Context>(std::forward<Args>(args)...) {}

  bool Fill(Tensor* output) override;
};

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

  bool RunOnDevice() override {
    auto& input = Input(0);

    auto* input_data = input.template data<int32_t>();

    CAFFE_ENFORCE_EQ(input.dim(), 1, "Input must be a vector.");

    auto len_sum = std::accumulate(input_data, input_data + input.numel(), 0);

    auto* output = Output(0, {len_sum}, at::dtype<int32_t>());
    auto* output_data = output->template mutable_data<int32_t>();

    int32_t offset = 0;
    for (int i = 0; i < input.numel(); ++i) {
      auto len = input_data[i];
      auto start = output_data + offset;
      std::iota(
          start,
          start + len,
          0); // make the third argument the arg of this operator
      offset += len;
    }
    return true;
  }
};

template <int VALUE_TYPE = TensorProto_DataType_FLOAT>
inline std::vector<TensorShape> FillerTensorInference(
    const OperatorDef& def,
    const vector<TensorShape>& in) {
  vector<TensorShape> out(1);
  ArgumentHelper helper(def);
  out[0].set_data_type(static_cast<TensorProto_DataType>(
      helper.GetSingleArgument<int>("dtype", VALUE_TYPE)));

  if (in.size()) {
    // TODO
    bool input_as_shape =
        helper.GetSingleArgument<bool>("input_as_shape", false);
    if (input_as_shape) {
      out[0].set_unknown_shape(true);
      return out;
    }
    for (auto d : in[0].dims()) {
      out[0].add_dims(d);
    }
  } else {
    auto shape = helper.GetRepeatedArgument<int64_t>("shape");
    for (auto d : shape) {
      out[0].add_dims(d);
    }
  }
  return out;
}

} // namespace caffe2

#endif // CAFFE2_OPERATORS_FILLER_OP_H_