text: #pragma once
  
  #include <set>
  #include <any>
  #include <typeindex>
  #include <sstream>
  
  #include "nlohmann/json.hpp"
  
  #include "megbrain/tensor.h"
  
  #include "./events.h"
  
  namespace mgb::imperative::profiler {
  
  using StackManager = interpreter::intl::StackManager;
  
  struct ProfileTensorState {
      uint64_t id = 0;
      std::optional<uint64_t> source;
      TensorLayout layout;
      CompNode device;
      std::string name;
      profiler::HostTime produced = profiler::HostTime::min();
      profiler::Duration living_time = profiler::Duration::zero();
  
      size_t size_in_bytes() const {
          if (!layout.dtype.valid()) {
              return 0;
          }
          return layout.dtype.size(layout.total_nr_elems());
      }
  
      std::string info(HostTime current_time) {
          std::string shape = layout.TensorShape::to_string();
          std::string dtype = layout.dtype.name();
          return ssprintf("%s(%s:%s:%s)", name.c_str(), shape.c_str(), dtype.c_str(), device.to_string().c_str());
      }
  
      nlohmann::json detail(HostTime current_time) {
          nlohmann::json args;
          args["id"] = id;
          args["name"] = name;
          args["shape"] = layout.TensorShape::to_string();
          args["dtype"] = layout.dtype.name();
          args["nr_elements"] = layout.total_nr_elems();
          args["device"] = device.to_string();
          if (produced != produced.min()) {
              double ms_count = std::chrono::duration_cast<std::chrono::duration<double, std::micro>>(current_time - produced + living_time).count();
              args["living_time"] = ssprintf("%lf ms", ms_count);
          }
          return args;
      }
  };
  
  struct ProfileOperatorState {
      uint64_t id = 0;
      std::string name;
      OpParams params;
      SmallVector<uint64_t> inputs;
      SmallVector<uint64_t> outputs;
      CompNode device;
      Trace trace;
  
      struct Execution {
          std::string reason;
          profiler::HostTime begin;
          profiler::HostTime end;
      };
  
      SmallVector<Execution> executions;
  
      nlohmann::json detail() {
          nlohmann::json args;
          for (auto&& [name, value]: params) {
              args[name] = value;
          }
          args["__id__"] = id;
          args["__name__"] = name;
          args["__device__"] = device.to_string();
          return args;
      }
  };
  
  template <typename TProp>
  struct ProfileTensorPropPair {
      uint64_t id;
      TProp value;
  
      bool operator<(const ProfileTensorPropPair& lhs) const {
          return value == lhs.value ? id < lhs.id : value < lhs.value;
      }
  
      bool operator==(const ProfileTensorPropPair& lhs) const {
          return id == lhs.id && value == lhs.value;
      }
  
      bool operator>(const ProfileTensorPropPair& lhs) const {
          return value == lhs.value ? id > lhs.id : value > lhs.value;
      }
  };
  
  using ProfileTensorSizePair = ProfileTensorPropPair<size_t>;
  using ProfileTensorProducedPair = ProfileTensorPropPair<uint64_t>;
  
  struct ProfileState {
      std::unordered_map<uint64_t, ProfileTensorState> tensors;
      std::unordered_map<uint64_t, ProfileOperatorState> operators;
      std::unordered_map<std::string, uint64_t> tensor_name_counter;
      std::set<ProfileTensorSizePair> tensors_by_size;
      std::set<ProfileTensorSizePair> tensors_by_produced;
  
      std::vector<uint64_t> top_k_tensor_in_device(CompNode device, size_t k) {
          std::vector<uint64_t> results;
          for (auto iter = tensors_by_size.rbegin(); iter != tensors_by_size.rend(); ++iter) {
              if (!k) {
                  break;
              }
              if (tensors[iter->id].device == device) {
                  results.push_back(iter->id);
                  --k;
              }
          }
          return results;
      }
  };
  
  template<typename T, typename = void>
  struct is_op_event : std::false_type { };
  
  template<typename T>
  struct is_op_event<T, decltype(std::declval<T>().op_id, void())> : std::true_type { };
  
  template<typename T, typename = void>
  struct is_tensor_event : std::false_type { };
  
  template<typename T>
  struct is_tensor_event<T, decltype(std::declval<T>().tensor_id, void())> : std::true_type { };
  template<typename T, typename = void>
  struct is_trace_event : std::false_type { };
  template<typename T>
  struct is_trace_event<T, decltype(std::declval<T>().trace, void())> : std::true_type { };
  
  template <typename... TItems>
  class AnyToVariantConverter {
  public:
      using any_t = AnyPtr;
      using variant_t = std::variant<TItems...>;
  private:
      std::unordered_map<std::type_index, std::function<variant_t(const any_t&)>> m_table;
  
      template <typename TItem>
      void register_converter() {
          m_table[typeid(TItem)] = [](const any_t& input) {
              return variant_t(*input.as<TItem>());
          };
      }
  public:
      AnyToVariantConverter() {
          (register_converter<TItems>(), ...);
      }
      variant_t operator()(const any_t& input) {
          return m_table[input.type()](std::move(input));
      }
  };
  
  template <typename TSelf>
  class EventVisitor {
  private:
      std::unordered_map<size_t, ProfileOperatorState> m_operators;
      std::unordered_map<size_t, ProfileTensorState> m_tensors;
      std::unordered_map<size_t, std::vector<Profiler::Record>> m_duration_stack;
      HostTime m_start_time;
      CompNode::UnorderedMap<size_t> m_device_tid_table;
      std::unordered_map<std::thread::id, size_t> m_host_tid_table;
      CompNode::UnorderedMap<std::map<profiler::HostTime, profiler::RealDuration>> m_device_timeline;
      std::unordered_map<std::thread::id, std::vector<Trace>> m_trace_stack;
      std::unordered_map<std::string, int64_t> m_counter_table;
  protected:
      Profiler::Record* current;
      ProfileOperatorState* current_op;
      ProfileTensorState* current_tensor;
  protected:
      profiler::Duration since_start(profiler::HostTime time) {
          return time - m_start_time;
      }
  
      profiler::HostTime to_device_time(profiler::HostTime time, CompNode device) {
          auto& device_timeline = m_device_timeline[device];
          auto upper = device_timeline.lower_bound(time);
          if (upper == device_timeline.end()) {
              if (upper == device_timeline.begin()) {
                  return time;
              } else {
                  --upper;
                  return time + std::chrono::duration_cast<profiler::Duration>(upper->second);
              }
          } else if (upper->first == time) {
              return time + std::chrono::duration_cast<profiler::Duration>(upper->second);
          } else if (upper == device_timeline.begin()) {
              return time + std::chrono::duration_cast<profiler::Duration>(upper->second);
          }
          auto lower = upper;
          -- lower;
          double ratio = ((double)(time - lower->first).count() / (double)(upper->first - lower->first).count());
          mgb_assert(ratio > 0 && ratio < 1, "invalid ratio");
          mgb_assert(lower->first + lower->second <= upper->first + upper->second, "device time corr");
          auto shift = lower->second + ratio * (upper->second - lower->second);
          auto result = time + std::chrono::duration_cast<profiler::Duration>(shift);
          return result;
      }
  
      size_t to_tid(std::thread::id host_tid) {
          return m_host_tid_table.at(host_tid);
      }
  
      size_t to_tid(CompNode device) {
          return m_device_tid_table.at(device);
      }
  
      SmallVector<std::thread::id> host_threads() {
          SmallVector<std::thread::id> host_threads;
          for (auto&& [host, _]: m_host_tid_table) {
              host_threads.push_back(host);
          }
          return host_threads;
      }
  
      SmallVector<CompNode> devices() {
          SmallVector<CompNode> devices;
          for (auto&& [device, _]: m_device_tid_table) {
              devices.push_back(device);
          }
          return devices;
      }
  
      void inc_counter(const char* key, int64_t delta) {
          if (!m_counter_table.count(key)) {
              m_counter_table[key] = 0;
          }
          auto& value = m_counter_table[key];
          static_cast<TSelf&>(*this).notify_counter(key, value, value + delta);
          value += delta;
      }
  public:
      void process_events(Profiler::bundle_t& bundle) {
          m_start_time = bundle.start_at;
  
          auto& self = static_cast<TSelf&>(*this);
          AnyToVariantConverter<OpDispatchEvent, OpExecuteEvent, OpExecuteFinishEvent,
                  KernelLaunchEvent, KernelLaunchFinishEvent,
                  OpInputEvent, OpInputFinishEvent, OpOutputEvent, OpOutputFinishEvent,
                  TensorDeclareEvent, TensorProduceEvent, TensorUsageEvent, TensorReleaseEvent, TensorEraseEvent,
                  TensorGetPropEvent, TensorNotifyPropEvent, TensorWaitPropEvent, TensorWaitPropFinishEvent,
                  SampleDeviceEvent, SampleDeviceFinishEvent, WorkerExceptionEvent, ShapeInferEvent, SyncEvent, SyncFinishEvent,
                  StartProfileEvent, StartProfileFinishEvent, StopProfileEvent, StopProfileFinishEvent,
                  TensorCommandEvent, TensorCommandFinishEvent, AutoEvictEvent, AutoEvictFinishEvent,
                  CustomEvent, CustomFinishEvent, RecordDeviceEvent, ScopeEvent, ScopeFinishEvent,
                  HostToDeviceEvent, HostToDeviceFinishEvent> converter;
  
          auto for_each_entry = [&](auto&& handler) {
              for (auto& entry: bundle.entries) {
                  current = &entry;
                  std::visit(handler, converter(entry.data));
              }
              current = nullptr;
          };
  
          // build device timeline
          struct DeviceStartPair {
              profiler::HostTime host;
              std::shared_ptr<CompNode::Event> device;
          };
          CompNode::UnorderedMap<DeviceStartPair> device_start_table;
  
          for_each_entry([&](auto&& event){
              using T = std::decay_t<decltype(event)>;
              if constexpr (std::is_same_v<T, RecordDeviceEvent>) {
                  using namespace std::chrono_literals;
                  DeviceStartPair& device_start = device_start_table[event.event->comp_node()];
                  if (!device_start.device) {
                      device_start = { current->time, event.event };
                  }
                  event.event->host_wait();
                  auto device_time = (device_start.host - current->time) + std::chrono::duration_cast<profiler::RealDuration>(device_start.device->elapsed_time_until(*event.event) * 1s);
                  m_device_timeline[event.event->comp_node()][current->time] = device_time;
              }
          });
  
          // register host threads
          for_each_entry([&](auto&& event){
              if (!m_host_tid_table.count(current->tid)) {
                  m_host_tid_table[current->tid] = {m_device_tid_table.size() + m_host_tid_table.size()};
              }
          });
  
          for_each_entry([&](auto&& event){
              using T = std::decay_t<decltype(event)>;
              if constexpr (std::is_same_v<T, OpDispatchEvent>) {
                  auto& op = m_operators[event.op_id];
                  mgb_assert(op.id == 0, "duplicate operator id");
                  op.id = event.op_id;
                  op.name = event.op_name;
                  op.params = event.op_params();
                  op.inputs = event.inputs;
                  op.outputs = event.outputs;
                  op.trace = event.trace;
                  for (auto&& output: event.outputs) {
                      m_tensors[output].source = op.id;
                  }
              } else if constexpr (std::is_same_v<T, TensorDeclareEvent>) {
                  auto& tensor = m_tensors[event.tensor_id];
                  mgb_assert(tensor.id == 0, "duplicated tensor id");
                  tensor.id = event.tensor_id;
                  tensor.name = event.name;
              } else if constexpr (std::is_same_v<T, TensorProduceEvent>) {
                  auto& tensor = m_tensors[event.tensor_id];
                  if (!m_device_tid_table.count(event.device)) {
                      m_device_tid_table[event.device] = {m_device_tid_table.size() + m_host_tid_table.size()};
                  }
                  tensor.device = event.device;
                  tensor.layout = event.layout;
              }
          });
  
          // replay execution
          using namespace std::placeholders;
          for_each_entry([&](auto&& event){
              using T = std::decay_t<decltype(event)>;
              // update current_op/tensor
              if constexpr (is_op_event<T>::value) {
                  current_op = &m_operators[event.op_id];
                  if (current_op->id == 0) {
                      current_op->id = event.op_id;
                      current_op->name = "UnknownOperator";
                  }
              } else if constexpr (is_tensor_event<T>::value) {
                  current_tensor = &m_tensors[event.tensor_id];
                  if (current_tensor->id == 0) {
                      current_tensor->id = event.tensor_id;
                      current_tensor->name = "UnknownTensor";
                  }
              }
              if constexpr (std::is_same_v<T, OpExecuteEvent>) {
                  current_op->executions.emplace_back();
                  current_op->executions.back().reason = event.reason;
                  current_op->executions.back().begin = current->time;
              } else if constexpr (std::is_same_v<T, OpExecuteFinishEvent>) {
                  current_op->executions.back().end = current->time;
              }
              // update counters
              if constexpr (std::is_same_v<T, OpDispatchEvent>) {
                  inc_counter("nr_op_pending", 1);
              } else if constexpr (std::is_same_v<T, OpExecuteEvent>) {
                  inc_counter("nr_op_pending", -1);
              } else if constexpr (std::is_same_v<T, TensorProduceEvent>) {
                  inc_counter("nr_alive_tensor", 1);
              } else if constexpr (std::is_same_v<T, TensorReleaseEvent>) {
                  inc_counter("nr_alive_tensor", -1);
              } else if constexpr (std::is_same_v<T, TensorEraseEvent>) {
                  if (event.use_count == 0) {
                      inc_counter("nr_redunant_tensor", 1);
                  }
              } else if constexpr (std::is_same_v<T, ShapeInferEvent>) {
                  if (!event.success) {
                      inc_counter("nr_shape_infer_failure", 1);
                  }
              } else if constexpr (std::is_same_v<T, WorkerExceptionEvent>) {
                  inc_counter("nr_exception", 1);
              } else if constexpr (std::is_same_v<T, KernelLaunchFinishEvent>) {
                  auto& execution = current_op->executions.back();
                  if (execution.reason == "dtr") {
                      auto overhead = to_device_time(current->time, event.device) - to_device_time(execution.begin, event.device);
                      inc_counter("dtr_overhead_us", std::chrono::duration_cast<std::chrono::microseconds>(overhead).count());
                  }
              }
              // visit_event_impl
              self.visit_event(event);
              // reset current_op/tensor
              if constexpr (is_op_event<T>::value) {
                  current_op = nullptr;
              } else if constexpr (is_tensor_event<T>::value) {
                  current_tensor = nullptr;
              }
          });
      }
  };
  
  }
  