/*!
 * Copyright (c) 2015 by Contributors
 * \file caffe_data_iter.cc
 * \brief register mnist iterator
*/
#include <sys/time.h>
#include <caffe/proto/caffe.pb.h>
#include <dmlc/parameter.h>
#include <atomic>

#include "caffe_common.h"
#include "caffe_stream.h"
#include "caffe_fieldentry.h"
#include "caffe_blob.h"
#include "../../src/io/inst_vector.h"
#include "../../src/io/iter_prefetcher.h"

#define CHECK_NEXT_TIMING

#ifdef CHECK_NEXT_TIMING
#define IF_CHECK_TIMING(__t$) __t$
#else
#define IF_CHECK_TIMING(__t$)
#endif

namespace mxnet {
namespace io {

struct CaffeDataParam : public dmlc::Parameter<CaffeDataParam> {
  /*! \brief protobuf text */
  ::caffe::LayerParameter prototxt;
  /*! \brief number of iterations per epoch */
  int num_examples;
  /*! \brief data mode */
  bool flat;

  DMLC_DECLARE_PARAMETER(CaffeDataParam) {
    DMLC_DECLARE_FIELD(prototxt).set_default("layer{}")
      .describe("Caffe's layer parameter");
    DMLC_DECLARE_FIELD(flat).set_default(false)
      .describe("Augmentation Param: Whether to flat the data into 1D.");
    DMLC_DECLARE_FIELD(num_examples).set_lower_bound(1).set_default(10000)
      .describe("Number of examples in the epoch.");
  }
};

template<typename Dtype>
class CaffeDataIter : public IIterator<TBlobBatch> {
 public:
  explicit CaffeDataIter(int type_flag) : batch_size_(0), channels_(1), width_(1), height_(1)
                               , type_flag_(type_flag), loc_(0)
  {}
  virtual ~CaffeDataIter(void) {}

  // intialize iterator loads data in
  virtual void Init(const std::vector<std::pair<std::string, std::string> >& kwargs) {
    std::map<std::string, std::string> kmap(kwargs.begin(), kwargs.end());
    param_.InitAllowUnknown(kmap);

    // Caffe seems to understand phase inside an "include {}" block
    if (!param_.prototxt.has_phase()) {
      if (param_.prototxt.include().size()) {
        if (param_.prototxt.include(0).has_phase()) {
          param_.prototxt.set_phase(param_.prototxt.include(0).phase());
        }
      }
    }

    std::string type = param_.prototxt.type();
    caffe_data_layer_ = caffe::LayerRegistry<Dtype>::CreateLayer(param_.prototxt);
    CHECK(caffe_data_layer_ != nullptr) << "Failed creating caffe data layer";
    const size_t top_size = param_.prototxt.top_size();
    if (top_size > 0) {
      if (top_size > NR_SUPPORTED_TOP_ITEMS) {
        LOG(WARNING)
          << "Too may \"top\" items, only two (one data, one label) are currently supported";
      }
      top_.reserve(top_size);
      for (size_t x = 0; x < top_size; ++x) {
        ::caffe::Blob<Dtype> *blob = new ::caffe::Blob<Dtype>();
        cleanup_blobs_.push_back(std::unique_ptr<::caffe::Blob<Dtype>>(blob));
        top_.push_back(blob);
      }
      caffe_data_layer_->SetUp(bottom_, top_);
      const std::vector<int> &shape = top_[DATA]->shape();
      const size_t shapeDimCount = shape.size();
      if (shapeDimCount > 0) {
        batch_size_ = shape[0];
        if (shapeDimCount > 1) {
          channels_ = shape[1];
          if (shapeDimCount > 2) {
            width_ = shape[2];
            if (shapeDimCount > 3) {
              height_ = shape[3];
            }
          }
        }
      }

      if (top_size > DATA) {
        if (param_.flat) {
          batch_data_ = TBlob(nullptr, mshadow::Shape2(batch_size_,
                                                       channels_ * width_ * height_),
                              cpu::kDevCPU, type_flag_);
        } else {
          batch_data_ = TBlob(nullptr, mxnet::TShape(top_[DATA]->shape().begin(),
                                                     top_[DATA]->shape().end()),
                              cpu::kDevCPU, type_flag_);
        }
      }
      out_.data.clear();
      if (top_size > LABEL) {
          batch_label_ = TBlob(nullptr, mxnet::TShape(top_[LABEL]->shape().begin(),
                                                      top_[LABEL]->shape().end()),
                               cpu::kDevCPU, type_flag_);
      }
      out_.batch_size = batch_size_;
    }
  }

  virtual void BeforeFirst(void) {
    loc_ = 0;
  }

  virtual bool Next(void) {
    // MxNet iterator is expected to return CPU-accessible memory
    if (::caffe::Caffe::mode() != ::caffe::Caffe::CPU) {
      ::caffe::Caffe::set_mode(::caffe::Caffe::CPU);
      CHECK_EQ(::caffe::Caffe::mode(), ::caffe::Caffe::CPU);
    }
    caffe_data_layer_->Forward(bottom_, top_);
    CHECK_GT(batch_size_, 0) << "batch size must be greater than zero";
    CHECK_EQ(out_.batch_size, batch_size_) << "Internal Error: batch size mismatch";

    if (loc_ + batch_size_ <= param_.num_examples) {
      batch_data_.dptr_ = top_[DATA]->mutable_cpu_data();
      batch_label_.dptr_ = top_[LABEL]->mutable_cpu_data();

      out_.data.clear();
      out_.data.push_back(batch_data_);
      out_.data.push_back(batch_label_);
      loc_ += batch_size_;
      return true;
    }

    return false;
  }

  virtual const TBlobBatch &Value(void) const {
    return out_;
  }

 private:
  /*! \brief indexes into top_ */
  enum { DATA = 0, LABEL, NR_SUPPORTED_TOP_ITEMS };

  /*! \brief MNISTCass iter params */
  CaffeDataParam param_;
  /*! \brief Shape scalar values */
  index_t batch_size_, channels_, width_, height_;
  /*! \brief Caffe data layer */
  boost::shared_ptr<caffe::Layer<Dtype> >  caffe_data_layer_;
  /*! \brief batch data blob */
  mxnet::TBlob batch_data_;
  /*! \brief batch label blob */
  mxnet::TBlob batch_label_;
  /*! \brief Output blob data for this iteration */
  TBlobBatch out_;
  /*! \brief Bottom and top connection-point blob data */
  std::vector<::caffe::Blob<Dtype>*> bottom_, top_;
  /*! \brief Cleanup these blobs on exit */
  std::list<std::unique_ptr<::caffe::Blob<Dtype>>> cleanup_blobs_;
  /*! \brief type flag of the tensor blob */
  const int type_flag_;
  /*! \brief Blobs done so far */
  std::atomic<size_t>  loc_;
};  // class CaffeDataIter

class CaffeDataIterWrapper : public PrefetcherIter {
 public:
  CaffeDataIterWrapper() : PrefetcherIter(NULL), next_time_(0) {}
  virtual ~CaffeDataIterWrapper() {
    IF_CHECK_TIMING(
      if (next_time_.load() > 0) {
        LOG(WARNING) << "Caffe data loader was blocked for "
                     << next_time_.load()
                     << " ms waiting for incoming data";
      }
    )
  }
  virtual void Init(const std::vector<std::pair<std::string, std::string> >& kwargs) {
    // We need to init prefetcher args in order to get dtype
    this->param_.InitAllowUnknown(kwargs);
    if (!this->param_.dtype) this->param_.dtype = mshadow::kFloat32;
    switch (this->param_.dtype.value()) {
      case mshadow::kFloat32:
        this->loader_.reset(new CaffeDataIter<float>(this->param_.dtype.value()));
        break;
      case mshadow::kFloat64:
        this->loader_.reset(new CaffeDataIter<double>(this->param_.dtype.value()));
        break;
      case mshadow::kFloat16:
        LOG(FATAL) << "float16 layer is not supported by caffe";
        return;
      default:
        LOG(FATAL) << "Unsupported type " << this->param_.dtype.value();
        return;
    }
    PrefetcherIter::Init(kwargs);
    this->param_.prefetch_buffer = 1;
  }
  virtual void BeforeFirst(void) {
    return PrefetcherIter::BeforeFirst();
  }
  virtual bool Next(void) {
    IF_CHECK_TIMING(
      const uint64_t start_time = GetTickCountMS();
    )
    const bool rc = PrefetcherIter::Next();
    IF_CHECK_TIMING(
      const uint64_t diff_time  = GetTickCountMS() - start_time;
      next_time_.fetch_add(diff_time);
    )
    return rc;
  }

 protected:
  IF_CHECK_TIMING(
    static uint64_t GetTickCountMS() {
      struct timeval tv;
      gettimeofday(&tv, 0);
      return uint64_t( tv.tv_sec ) * 1000 + tv.tv_usec / 1000;
    }
  )

  /*! \brief milliseconds spent in Next() */
  std::atomic<uint64_t> next_time_;
};  // class CaffeDataIterWrapper

DMLC_REGISTER_PARAMETER(CaffeDataParam);

MXNET_REGISTER_IO_ITER(CaffeDataIter)
.describe("Create MxNet iterator for a Caffe data layer.")
.add_arguments(CaffeDataParam::__FIELDS__())
.add_arguments(PrefetcherParam::__FIELDS__())
.set_body([]() {
    return new CaffeDataIterWrapper();
});

}  // namespace io
}  // namespace mxnet