#ifndef CAFFE2_OPERATORS_CREATE_SCOPE_OP_H_
#define CAFFE2_OPERATORS_CREATE_SCOPE_OP_H_

#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "caffe2/core/context.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/operator.h"
#include "caffe2/proto/caffe2_pb.h"

C10_DECLARE_bool(caffe2_workspace_stack_debug);

namespace caffe2 {
namespace detail {

/*
 * Keeps track of forward and backward gradient workspaces in stack,
 * reuses previously created workspaces, non-thread safe
 */
class CAFFE2_API WorkspaceStack {
 public:
  explicit WorkspaceStack() : parent_ws_(nullptr), top_(-1) {}

  std::shared_ptr<Workspace> pushForwardWorkspace(Workspace* parent_ws) {
    return pushForwardWorkspace(
        parent_ws, std::unordered_map<std::string, std::string>());
  }

  std::shared_ptr<Workspace> pushForwardWorkspace(
      Workspace* parent_ws,
      const std::unordered_map<std::string, std::string>& blob_bindings) {
    checkStack();
    if (FLAGS_caffe2_workspace_stack_debug) {
      if (parent_ws_) {
        CAFFE_ENFORCE_EQ(parent_ws_, parent_ws, "Parent workspace mismatch");
      } else {
        parent_ws_ = parent_ws;
      }
      if (!blob_bindings_.empty()) {
        checkBindingsMatch(blob_bindings_, blob_bindings);
      } else {
        blob_bindings_ = blob_bindings;
      }
    }

    if (top_ == workspaces_.size() - 1) {
      workspaces_.push_back(
          std::make_shared<Workspace>(parent_ws, blob_bindings));
    } else {
      // when reusing workspace, make sure copies of external blobs are
      // removed and blob bindings are set
      auto& workspace = workspaces_[top_ + 1];
      const auto& local_blobs = workspace->LocalBlobs();
      std::unordered_set<std::string> local_blobs_set;
      local_blobs_set.insert(local_blobs.begin(), local_blobs.end());
      bool found_local_copy = false;
      for (const auto& blob_pair : blob_bindings) {
        if (local_blobs_set.count(blob_pair.first)) {
          workspace->RemoveBlob(blob_pair.first);
          found_local_copy = true;
        }
      }
      if (found_local_copy) {
        workspace->AddBlobMapping(parent_ws, blob_bindings);
      }
    }

    return workspaces_[++top_];
  }

  std::shared_ptr<Workspace> popGradientWorkspace(
      Workspace* parent_ws,
      const std::unordered_map<std::string, std::string>& grad_blob_bindings) {
    checkStack();
    if (FLAGS_caffe2_workspace_stack_debug) {
      if (parent_ws_) {
        CAFFE_ENFORCE_EQ(parent_ws_, parent_ws, "Parent workspace mismatch");
      } else {
        parent_ws_ = parent_ws;
      }
      if (!grad_blob_bindings_.empty()) {
        checkBindingsMatch(grad_blob_bindings_, grad_blob_bindings);
      } else {
        grad_blob_bindings_ = grad_blob_bindings;
      }
    }

    if (top_ < 0) {
      return nullptr;
    }
    auto& grad_workspace = workspaces_[top_];
    grad_workspace->AddBlobMapping(parent_ws, grad_blob_bindings, true);
    --top_;
    return grad_workspace;
  }

  std::shared_ptr<Workspace> reuseLastForwardWorkspace(Workspace* parent_ws) {
    return reuseLastForwardWorkspace(
        parent_ws, std::unordered_map<std::string, std::string>());
  }

  std::shared_ptr<Workspace> reuseLastForwardWorkspace(
      Workspace* parent_ws,
      const std::unordered_map<std::string, std::string>& blob_bindings) {
    checkStack();
    if (top_ < 0) {
      return nullptr;
    }
    workspaces_[top_]->AddBlobMapping(parent_ws, blob_bindings);
    return workspaces_[top_];
  }

  void clear() {
    checkStack();
    top_ = -1;
  }

  bool empty() const {
    return top_ < 0;
  }

 private:
  void checkStack() const {
    CAFFE_ENFORCE_GT(
        (int)workspaces_.size(), top_, "Corrupted workspaces stack");
  }

  void checkBindingsMatch(
      const std::unordered_map<std::string, std::string>& bindings,
      const std::unordered_map<std::string, std::string>& test_bindings) const {
    CAFFE_ENFORCE_EQ(
        bindings.size(), test_bindings.size(), "Blob bindings mismatch");
    for (const auto& blob_binding : bindings) {
      CAFFE_ENFORCE(
          test_bindings.count(blob_binding.first), "Blob bindings mismatch");
      CAFFE_ENFORCE_EQ(
          test_bindings.at(blob_binding.first),
          blob_binding.second,
          "Blob bindings mismatch");
    }
  }

  std::unordered_map<std::string, std::string> blob_bindings_;
  std::unordered_map<std::string, std::string> grad_blob_bindings_;
  Workspace* parent_ws_;
  int top_;
  std::vector<std::shared_ptr<Workspace>> workspaces_;
};
}

template <class Context>
class CreateScopeOp final : public Operator<Context> {
 public:
  template <class... Args>
  explicit CreateScopeOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {}

  USE_OPERATOR_CONTEXT_FUNCTIONS;
  bool RunOnDevice() override;
};

template <class Context>
class HasScopeOp final : public Operator<Context> {
 public:
  template <class... Args>
  explicit HasScopeOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {}

  USE_OPERATOR_CONTEXT_FUNCTIONS;
  bool RunOnDevice() override;
};

} // namespace caffe2

#endif // CAFFE2_OPERATORS_CREATE_SCOPE_OP_H_