#ifndef CAFFE2_FILLER_H_
#define CAFFE2_FILLER_H_

#include <sstream>

#include "caffe2/core/logging.h"
#include "caffe2/core/tensor.h"
#include "caffe2/utils/math.h"

namespace caffe2 {

// TODO: replace filler distribution enum with a better abstraction
enum FillerDistribution { FD_UNIFORM, FD_FIXEDSUM, FD_SYNTHETIC };

class TensorFiller {
 public:
  template <class Type, class Context>
  void Fill(Tensor* tensor, Context* context) const {
    CAFFE_ENFORCE(context, "context is null");
    CAFFE_ENFORCE(tensor, "tensor is null");
    auto min = (min_ < std::numeric_limits<Type>::min())
        ? std::numeric_limits<Type>::min()
        : static_cast<Type>(min_);
    auto max = (max_ > std::numeric_limits<Type>::max())
        ? std::numeric_limits<Type>::max()
        : static_cast<Type>(max_);
    CAFFE_ENFORCE_LE(min, max);

    Tensor temp_tensor(shape_, Context::GetDeviceType());
    std::swap(*tensor, temp_tensor);
    Type* data = tensor->template mutable_data<Type>();

    // select distribution
    switch (dist_) {
      case FD_UNIFORM: {
        math::RandUniform<Type, Context>(
            tensor->numel(), min, max, data, context);
        break;
      }
      case FD_FIXEDSUM: {
        auto fixed_sum = static_cast<Type>(fixed_sum_);
        CAFFE_ENFORCE_LE(min * tensor->numel(), fixed_sum);
        CAFFE_ENFORCE_GE(max * tensor->numel(), fixed_sum);
        math::RandFixedSum<Type, Context>(
            tensor->numel(), min, max, fixed_sum_, data, context);
        break;
      }
      case FD_SYNTHETIC: {
        math::RandSyntheticData<Type, Context>(
            tensor->numel(), min, max, data, context);
        break;
      }
    }
  }

  TensorFiller& Dist(FillerDistribution dist) {
    dist_ = dist;
    return *this;
  }

  template <class Type>
  TensorFiller& Min(Type min) {
    min_ = (double)min;
    return *this;
  }

  template <class Type>
  TensorFiller& Max(Type max) {
    max_ = (double)max;
    return *this;
  }

  template <class Type>
  TensorFiller& FixedSum(Type fixed_sum) {
    dist_ = FD_FIXEDSUM;
    fixed_sum_ = (double)fixed_sum;
    return *this;
  }

  // A helper function to construct the lengths vector for sparse features
  // We try to pad least one index per batch unless the total_length is 0
  template <class Type>
  TensorFiller& SparseLengths(Type total_length) {
    return FixedSum(total_length)
        .Min(std::min(static_cast<Type>(1), total_length))
        .Max(total_length);
  }

  // a helper function to construct the segments vector for sparse features
  template <class Type>
  TensorFiller& SparseSegments(Type max_segment) {
    CAFFE_ENFORCE(dist_ != FD_FIXEDSUM);
    return Min(0).Max(max_segment).Dist(FD_SYNTHETIC);
  }

  TensorFiller& Shape(const std::vector<int64_t>& shape) {
    shape_ = shape;
    return *this;
  }

  template <class Type>
  TensorFiller(const std::vector<int64_t>& shape, Type fixed_sum)
      : shape_(shape), dist_(FD_FIXEDSUM), fixed_sum_((double)fixed_sum) {}

  TensorFiller(const std::vector<int64_t>& shape)
      : shape_(shape), dist_(FD_UNIFORM), fixed_sum_(0) {}

  TensorFiller() : TensorFiller(std::vector<int64_t>()) {}

  std::string DebugString() const {
    std::stringstream stream;
    stream << "shape = [" << shape_ << "]; min = " << min_
           << "; max = " << max_;
    switch (dist_) {
      case FD_FIXEDSUM:
        stream << "; dist = FD_FIXEDSUM";
        break;
      case FD_SYNTHETIC:
        stream << "; dist = FD_SYNTHETIC";
        break;
      default:
        stream << "; dist = FD_UNIFORM";
        break;
    }
    return stream.str();
  }

 private:
  std::vector<int64_t> shape_;
  // TODO: type is unknown until a user starts to fill data;
  // cast everything to double for now.
  double min_ = 0.0;
  double max_ = 1.0;
  FillerDistribution dist_;
  double fixed_sum_;
};

} // namespace caffe2

#endif // CAFFE2_FILLER_H_