/*
 * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 */

#pragma once

#include <cuda_runtime.h>

#include <cstdio>
#include <string>
#include <unordered_map>

/**
 * Contains utility methods to measure elapsed time on GPU-related functions
 *
 * To just the timer just for once:
 * @code{.cpp}
 * using namespace cugraph::ops::utils::gpu;
 * tic("some-gpu-work");
 * some_gpu_work();
 * toc("some-gpu-work");
 * auto& timer = get_timer("some-gpu-work");
 * std::cout << "Time (in s): " << timer.elapsed() << std::endl;
 * @endcode
 *
 * In case you want to count the same logic multiple times (eg: in a loop):
 * @code{.cpp}
 * using namespace cugraph::ops::utils::gpu;
 * while (!done) {
 *   tic("some-gpu-work");
 *   some_gpu_work();
 *   toc("some-gpu-work");
 * }
 * auto& timer = get_timer("some-gpu-work");
 * std::cout << "Average time (in s): " << timer.avg_elapsed() << std::endl;
 * std::cout << "Total time (in s): " << timer.elapsed() << std::endl;
 * std::cout << "Number of tic/toc's: " << timer.count() << std::endl;
 * @endcode
 */

namespace cugraph::ops::utils::gpu {

/**
 * @brief Timer class to be used while measuring GPU elapsed times
 */
class timer_t {
 public:
  explicit timer_t(cudaStream_t s = nullptr);
  ~timer_t();
  /** default copy/copy assignment */
  timer_t(const timer_t& other)            = default;
  timer_t& operator=(const timer_t& other) = default;
  /** (Re)Starts the timer */
  void tic();
  /** stops the timer */
  void toc();
  /** elapsed time (in s) */
  [[nodiscard]] double elapsed() const { return elapsed_; }
  /** number of times the timer was stopped */
  [[nodiscard]] int count() const { return count_; }
  /** average elapsed time per stop (in s) */
  [[nodiscard]] double avg_elapsed() const { return elapsed_ / count_; }

 private:
  double elapsed_{0.0};
  int count_{0};
  bool running_{false};
  cudaStream_t stream_{nullptr};
  cudaEvent_t start_{nullptr};
  cudaEvent_t stop_{nullptr};
};  // class timer_t

/**
 * @brief get the timer object of the given name
 *
 * @param[in] name name of the timer to be queried for and returned
 *
 * @return a reference to the timer object, if it exists, else an exception will
 *         be raised
 */
timer_t& get_timer(const std::string& name);

/** start a named timer */
void tic(const std::string& name);

/** stop the previously started named timer */
void toc(const std::string& name);

/** remove all the named timers created so far */
void reset_timers();

/**
 * @brief Start/Stop the timer and print timing info after stopping
 *
 * @param[in] lev  log level
 * @param[in] name name of the timer to be stopped
 *
 * @note As as side-effect, the timer is not run if the input log level is not
 *       verbose enough to be printed out.
 *
 * @{
 */
void tic_for_log(int lev, const std::string& name);
void toc_and_log(int lev, const std::string& name);
/** @} */

}  // namespace cugraph::ops::utils::gpu
