#include <utils/ExecutionStatManager.hpp>

#include <utils/Exceptions.hpp>
#include <utils/Messenger.hpp>

#include <cmath>
#include <rang.hpp>
#include <sys/resource.h>

ExecutionStatManager* ExecutionStatManager::m_instance = nullptr;

void
ExecutionStatManager::_printMaxResidentMemory() const
{
  class Memory
  {
   private:
    double m_value;

   public:
    PUGS_INLINE const double&
    value() const
    {
      return m_value;
    }

    std::string
    prettyPrint() const
    {
      const std::vector<std::string> units = {"B", "KB", "MB", "GB", "TB", "PB", "EB"};

      double local_memory = m_value;
      size_t i_unit       = 0;
      while ((local_memory >= 1024) and (i_unit < units.size())) {
        ++i_unit;
        local_memory /= 1024;
      }
      std::ostringstream os;
      os << local_memory << units[i_unit];
      return os.str();
    }

    Memory()
    {
      rusage u;
      getrusage(RUSAGE_SELF, &u);
      m_value = u.ru_maxrss * 1024;
    }

    Memory(double value) : m_value{value} {}
  };

  Memory memory;
  std::cout << "Memory: " << rang::style::bold << Memory{parallel::allReduceSum(memory.value())}.prettyPrint()
            << rang::style::reset;
  if (parallel::size() > 1) {
    std::cout << " (over " << parallel::size() << " processes)";
    std::cout << " Avg: " << rang::style::bold
              << Memory{parallel::allReduceSum(memory.value()) / parallel::size()}.prettyPrint() << rang::style::reset;
    std::cout << " Min: " << rang::style::bold << Memory{parallel::allReduceMin(memory.value())}.prettyPrint()
              << rang::style::reset;
    std::cout << " Max: " << rang::style::bold << Memory{parallel::allReduceMax(memory.value())}.prettyPrint()
              << rang::style::reset;
  }
  std::cout << '\n';
}

void
ExecutionStatManager::_printElapseTime() const
{
  std::cout << "Execution: " << rang::style::bold << m_instance->m_elapse_time.seconds() << 's' << rang::style::reset
            << '\n';
}

void
ExecutionStatManager::_printTotalCPUTime() const
{
  rusage u;
  getrusage(RUSAGE_SELF, &u);

  const double total_cpu_time =
    u.ru_utime.tv_sec + u.ru_stime.tv_sec + (u.ru_utime.tv_usec + u.ru_stime.tv_usec) * 1E-6;

  std::cout << "Total CPU: " << rang::style::bold << parallel::allReduceSum(total_cpu_time) << 's'
            << rang::style::reset;
  std::cout << " (" << parallel::allReduceSum(Kokkos::DefaultHostExecutionSpace::concurrency()) << " threads over "
            << parallel::size() << " processes)";
  if (total_cpu_time > 60) {
    size_t seconds    = std::floor(total_cpu_time);
    const size_t days = seconds / (24 * 3600);
    seconds -= days * (24 * 3600);
    const size_t hours = seconds / 3600;
    seconds -= hours * 3600;
    const size_t minutes = seconds / 60;
    seconds -= minutes * 60;
    std::cout << " " << rang::style::bold;
    bool print = false;
    if (days > 0) {
      print = true;
      std::cout << days << "d" << ' ';
    }
    if (print or (hours > 0)) {
      print = true;
      std::cout << std::setw(2) << std::setfill('0') << hours << "h";
    }
    if (print or (minutes > 0)) {
      print = true;
      std::cout << std::setw(2) << std::setfill('0') << minutes << "mn";
    }
    if (print) {
      std::cout << rang::style::bold << std::setw(2) << std::setfill('0') << seconds << "s";
    }
    std::cout << rang::style::reset;
  }
  std::cout << '\n';
}

void
ExecutionStatManager::printInfo()
{
  if (ExecutionStatManager::getInstance().doPrint()) {
    std::cout << "----------------- " << rang::fg::green << "pugs exec stats" << rang::fg::reset
              << " ---------------------\n";

    ExecutionStatManager::getInstance()._printElapseTime();
    ExecutionStatManager::getInstance()._printTotalCPUTime();
    ExecutionStatManager::getInstance()._printMaxResidentMemory();
  }
}

void
ExecutionStatManager::create()
{
  if (ExecutionStatManager::m_instance == nullptr) {
    ExecutionStatManager::m_instance = new ExecutionStatManager;
  } else {
    throw UnexpectedError("ExecutionStatManager already created");
  }
}

void
ExecutionStatManager::destroy()
{
  // One allows multiple destruction to handle unexpected code exit
  if (ExecutionStatManager::m_instance != nullptr) {
    delete ExecutionStatManager::m_instance;
    ExecutionStatManager::m_instance = nullptr;
  }
}