#ifndef CAFFE2_OPERATORS_SUMMARIZE_OP_H_
#define CAFFE2_OPERATORS_SUMMARIZE_OP_H_

#include <fstream>

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

namespace caffe2 {

constexpr char kSummaryzeOpExtension[] = ".summary";

template <typename T, class Context>
class SummarizeOp final : public Operator<Context> {
 public:
  explicit SummarizeOp(const OperatorDef& def, Workspace* ws)
      : Operator<Context>(def, ws),
        to_file_(this->template GetSingleArgument<int>("to_file", 0)) {
    if (to_file_) {
      // We will output to file instead of printing on screen.
      const string& target_folder = ws->RootFolder();
      // We will write each individual tensor to its individual file.
      // Also, since the namescope is currently represented by "/", we will
      // need to replace it with a symbol that does not conflict with the
      // folder separator in Linux.
      string proper_name = def.input(0);
      std::replace(proper_name.begin(), proper_name.end(), '/', '#');
      log_file_.reset(new std::ofstream(
          target_folder + "/" + proper_name + kSummaryzeOpExtension,
          std::ofstream::out | std::ofstream::trunc));
      CAFFE_ENFORCE(
          log_file_->good(),
          "Failed to open summarize file for tensor ",
          def.input(0),
          ". rdstate() = ",
          log_file_->rdstate());
    }
  }
  ~SummarizeOp() {
    if (to_file_)
      log_file_->close();
  }
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  bool RunOnDevice() override;

  static constexpr int MIN_IDX = 0;
  static constexpr int MAX_IDX = 1;
  static constexpr int MEAN_IDX = 2;
  static constexpr int STD_IDX = 3;

  static constexpr int NUM_STATS = 4;

 protected:
  bool to_file_;
  std::unique_ptr<std::ofstream> log_file_;
};

} // namespace caffe2

#endif // CAFFE2_OPERATORS_SUMMARIZE_OP_H_