#pragma once

#include <torch/csrc/jit/ir.h>
#include <ATen/ATen.h>
#include <ATen/core/ivalue.h>
#include <ATen/core/jit_type.h>
#include <ATen/core/stack.h>

#include <list>
#include <vector>

namespace torch {
namespace jit {

using ::c10::ProfiledTensorTypePtr;

struct ProfilingRecord {
  // N.B. ProfilingRecord's copy and move c-tor are disabled, so we won't
  // end up accidentally copying or moving ProfilingRecords whose addresses
  // are captured in callbacks_
  ProfilingRecord(const ProfilingRecord&) = delete;
  ProfilingRecord(ProfilingRecord&&) noexcept = delete;
  static ProfiledTensorTypePtr toProfiledTensorTypePtr(const IValue& ival);
  static std::unique_ptr<ProfilingRecord> instrumentGraph(
      const std::shared_ptr<Graph>& graph);

  std::shared_ptr<Graph> profiled_graph_;
  std::mutex mutex_;
  size_t profiling_count_;

 private:
  Node* createProfileNode(
      const std::function<void(Stack&)>& fp,
      at::ArrayRef<Value*> inputs);
  void instrumentBlock(Block* block);
  ProfilingRecord(std::shared_ptr<Graph> g);
  // N.B. list is used to make sure that std::function objs
  // aren't moved anywhere since we are stashing pointers
  // to those in nodes' attributes
  std::list<std::function<void(Stack&)>> callbacks_;
};

} // jit
} // torch
