#ifndef CAFFE2_OPERATORS_PIECEWISE_LINEAR_TRANSFORM_OP_H_
#define CAFFE2_OPERATORS_PIECEWISE_LINEAR_TRANSFORM_OP_H_

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

C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10(PiecewiseLinearTransform);

namespace caffe2 {

template <typename T, class Context>
class PiecewiseLinearTransformOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit PiecewiseLinearTransformOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    binary_ = this->template GetSingleArgument<bool>("binary", false);

    // Retrieve transform params (i.e., the linear functions).
    bounds_from_arg_ = this->template GetRepeatedArgument<T>("bounds");
    slopes_from_arg_ = this->template GetRepeatedArgument<T>("slopes");
    intercepts_from_arg_ = this->template GetRepeatedArgument<T>("intercepts");
    transform_param_from_arg_ = CheckTransParamFromArg();
  }

  bool RunOnDevice() override {
    return binary_ ? TransformBinary() : TransformGeneral();
  }

 private:
  // num_func_per_group is the number of pieces of linear functions of
  // each group.
  // num_group: The number of groups of linear functions. Each group is for
  // transforming one column of predictions.
  void InferNumFunctionsPerGroup(
      const int64_t num_bounds,
      const int64_t num_slopes,
      const int64_t num_intercepts,
      int64_t* num_func_per_group,
      int64_t* num_group) {
    CAFFE_ENFORCE_EQ(num_slopes, num_intercepts);

    // This is based on the facts:
    // 1. in each group, the num of bounds minus the num of slopes is 1;
    // 2. each group has the same number of pieces.
    *num_group = num_bounds - num_slopes;
    CAFFE_ENFORCE_GT(*num_group, 0);
    if (binary_) {
      CAFFE_ENFORCE_EQ(*num_group, 1);
    }
    *num_func_per_group = num_slopes / *num_group;
    CAFFE_ENFORCE_GT(*num_func_per_group, 0);
    CAFFE_ENFORCE_EQ(num_slopes % *num_group, 0);
  }

  bool CheckBoundsSorted(
      const T* bounds,
      const int64_t num_bounds_per_group,
      const int64_t num_group) {
    const T* start = bounds;
    for (int64_t i = 0; i < num_group; i++) {
      if (!std::is_sorted(start, start + num_bounds_per_group)) {
        return false;
      }
      start += num_bounds_per_group;
    }
    return true;
  }

  // Returns true if the transform params from arg are valid.
  // Otherwise, we will assume the transform params will pass from Input blobs.
  bool CheckTransParamFromArg() {
    int good_param = 0;
    good_param += bounds_from_arg_.size() > 0;
    good_param += slopes_from_arg_.size() > 0;
    good_param += intercepts_from_arg_.size() > 0;
    CAFFE_ENFORCE(
        good_param == 0 || good_param == 3,
        "bounds, slopes, intercepts must be all set or all not set");
    if (good_param == 3) {
      int64_t num_func_per_group;
      int64_t num_group;
      InferNumFunctionsPerGroup(
          bounds_from_arg_.size(),
          slopes_from_arg_.size(),
          intercepts_from_arg_.size(),
          &num_func_per_group,
          &num_group);
      CAFFE_ENFORCE(
          CheckBoundsSorted(
              bounds_from_arg_.data(), num_func_per_group + 1, num_group),
          "bounds must be sorted for each group");
    }

    return good_param == 3;
  }

  void setUpTensors(int64_t& num_func_per_group, int64_t& num_group, int64_t M);

  void GetTransParamData(
      const T** bounds,
      const T** slopes,
      const T** intercepts,
      int64_t* num_func_per_group,
      int64_t* num_group) {
    int64_t num_bounds;
    int64_t num_slopes;
    int64_t num_intercepts;

    if (transform_param_from_arg_) {
      CAFFE_ENFORCE_EQ(InputSize(), 1);
      *bounds = bounds_from_arg_.data();
      *slopes = slopes_from_arg_.data();
      *intercepts = intercepts_from_arg_.data();
      num_bounds = bounds_from_arg_.size();
      num_slopes = slopes_from_arg_.size();
      num_intercepts = intercepts_from_arg_.size();
    } else {
      CAFFE_ENFORCE_EQ(InputSize(), 4);
      auto& bounds_input = Input(BOUNDS);
      auto& slopes_input = Input(SLOPES);
      auto& intercepts_input = Input(INTERCEPTS);
      *bounds = bounds_input.template data<T>();
      *slopes = slopes_input.template data<T>();
      *intercepts = intercepts_input.template data<T>();
      num_bounds = bounds_input.numel();
      num_slopes = slopes_input.numel();
      num_intercepts = intercepts_input.numel();
    }
    InferNumFunctionsPerGroup(
        num_bounds, num_slopes, num_intercepts, num_func_per_group, num_group);
  }

  bool TransformGeneral() {
    auto& X = Input(0);

    CAFFE_ENFORCE_EQ(X.dim(), 2);
    int64_t N = X.dim32(0);
    int64_t M = X.dim32(1);
    auto* Y = Output(0, X.sizes(), at::dtype<T>());
    const auto* Xdata = X.template data<T>();
    T* Ydata = Y->template mutable_data<T>();

    const T* bounds;
    const T* slopes;
    const T* intercepts;
    int64_t num_func_per_group;
    int64_t num_group;
    GetTransParamData(
        &bounds, &slopes, &intercepts, &num_func_per_group, &num_group);
    CAFFE_ENFORCE_EQ(num_group, M);

    for (int64_t j = 0; j < M; ++j) {
      const T* bounds_group = bounds + j * (num_func_per_group + 1);
      const T* slopes_group = slopes + j * num_func_per_group;
      const T* intercepts_group = intercepts + j * num_func_per_group;
      for (int64_t i = 0; i < N; ++i) {
        Ydata[i * M + j] = PiecewiseLinearTransform(
            Xdata[i * M + j],
            bounds_group,
            slopes_group,
            intercepts_group,
            num_func_per_group);
      }
    }
    return true;
  }

  bool TransformBinary() {
    auto& X = Input(PREDICTIONS);

    CAFFE_ENFORCE(X.dim() == 1 || X.dim() == 2);
    int64_t N = X.dim32(0);
    int64_t M = X.dim() == 2 ? X.dim32(1) : 1;
    CAFFE_ENFORCE(
        M == 1 || M == 2,
        "If binary is set to true, the input must be Nx2 or Nx1 tensor");
    auto* Y = Output(0, X.sizes(), at::dtype<T>());
    const auto* Xdata = X.template data<T>();
    T* Ydata = Y->template mutable_data<T>();

    const T* bounds;
    const T* slopes;
    const T* intercepts;
    int64_t num_func_per_group;
    int64_t num_group;
    GetTransParamData(
        &bounds, &slopes, &intercepts, &num_func_per_group, &num_group);
    CAFFE_ENFORCE_EQ(num_group, 1);

    if (M == 1) {
      for (int64_t i = 0; i < N; ++i) {
        Ydata[i] = PiecewiseLinearTransform(
            Xdata[i], bounds, slopes, intercepts, num_func_per_group);
      }
    } else {
      for (int64_t i = 0; i < N; ++i) {
        Ydata[i * M + 1] = PiecewiseLinearTransform(
            Xdata[i * M + 1], bounds, slopes, intercepts, num_func_per_group);
        Ydata[i * M] = 1.0f - Ydata[i * M + 1];
      }
    }

    return true;
  }

  T PiecewiseLinearTransform(
      const T x,
      const T* bounds,
      const T* slopes,
      const T* intercepts,
      const int64_t num_func_per_group) {
    T y = 0;
    // deal with samples out of bounds
    // make it the same as the upper/lower bound value
    if (x <= bounds[0]) {
      y = slopes[0] * bounds[0] + intercepts[0];
    } else if (x >= bounds[num_func_per_group]) {
      y = slopes[num_func_per_group - 1] * bounds[num_func_per_group] +
          intercepts[num_func_per_group - 1];
    } else {
      auto low_bound =
          std::lower_bound(bounds, bounds + num_func_per_group + 1, x);
      int bounds_idx = low_bound - bounds - 1;
      // compute the piecewise linear transformation as Y
      y = slopes[bounds_idx] * x + intercepts[bounds_idx];
    }
    return y;
  }

 private:
  bool binary_;
  vector<T> bounds_from_arg_;
  vector<T> slopes_from_arg_;
  vector<T> intercepts_from_arg_;

  Tensor bounds_device_{Context::GetDeviceType()};
  Tensor intercepts_device_{Context::GetDeviceType()};
  Tensor slopes_device_{Context::GetDeviceType()};
  bool gpu_copied_ = false;

  // If true, the piecewise linear functions are passed through args,
  // otherwise, they are passed through Input blobs.
  bool transform_param_from_arg_;

  INPUT_TAGS(PREDICTIONS, BOUNDS, SLOPES, INTERCEPTS);
};

} // namespace caffe2

#endif // CAFFE2_OPERATORS_PIECEWISE_LINEAR_TRANSFORM_OP_H_