#if defined(_ENERGY_PMT_)

#include <pmt.h>
#include <assert.h>
#include <iostream>
#include <vector>
#include <string>
#include <map>

struct EnergyState
{
  pmt::State   start;
  pmt::State   stop;
  double       joules;
  double       watts;
  double       seconds;
  unsigned int count;
};

static bool PMT_ERROR{false};

void PMT_err(void)
{
  std::cout << "\n\t PMT Error \n" << std::endl;
  
  return;
}

#if defined(_ENERGY_RAPL_)
   static std::unique_ptr<pmt::PMT> sensor_cpu;
   static std::map<std::string, EnergyState> state_cpu;
#endif // _ENERGY_RAPL_

#if defined(_ENERGY_NVIDIA_) || defined(_ENERGY_AMD_)
   static std::vector<std::unique_ptr<pmt::PMT>> sensor_gpu;
   static std::vector<std::map<std::string, EnergyState>> state_gpu;
#endif // defined(_ENERGY_NVIDIA_) || defined(_ENERGY_AMD_)

extern "C"
{
#include "energy_pmt_methods.h"
}

#if defined(_ENERGY_RAPL_) || defined(_ENERGY_NVIDIA_) || defined(_ENERGY_AMD_)
void Create_PMT(const unsigned int numGPUs)
{
#if defined(_ENERGY_RAPL_)
  
  sensor_cpu = pmt::Create("rapl");

#endif // _ENERGY_RAPL_
  
  if (numGPUs > 0)
    {
      for (unsigned int dev=0 ; dev<numGPUs ; dev++)
	{
#if defined(_ENERGY_NVIDIA_)
	  sensor_gpu.push_back(pmt::nvml::NVML::Create(dev));
#endif // _ENERGY_NVIDIA_

#if defined(_ENERGY_AMD_)
	  sensor_gpu.push_back(pmt::rocm::AMD::Create(dev));
#endif // _ENERGY_AMD_
	} // numGPUs
    }

  return;
}
#endif // defined(_ENERGY_RAPL_) || defined(_ENERGY_NVIDIA_) || defined(_ENERGY_AMD_)

#if defined(_ENERGY_RAPL_)
void Start_PMT_CPU(const char *label)
{
  if (PMT_ERROR)
    {
      PMT_err();
      return;
    }

  // check if the label already exists
  if (state_cpu.count(std::string{label}))
    {
      state_cpu[std::string{label}].start = sensor_cpu->Read();
    }
  else
    {
      // insert the key and initialize the counters
      state_cpu.insert({std::string{label},
			{sensor_cpu->Read(),
			 0,
			 0.0, 0.0, 0.0, 0
			}});
    }

  return;
}

void Stop_PMT_CPU(const char *label)
{
  if (PMT_ERROR)
    {
      PMT_err();
      return;
    }
  
  // check if the label already exists
  // if not error
  if (!state_cpu.count(std::string{label}))
    {
      PMT_ERROR = true;
      PMT_err();
      return;
    }
  else
    {
      // read the counter
      state_cpu[std::string{label}].stop = sensor_cpu->Read();

      // update quantities
      state_cpu[std::string{label}].seconds +=
	sensor_cpu->seconds(state_cpu[std::string{label}].start,
			    state_cpu[std::string{label}].stop);
      
      state_cpu[std::string{label}].joules +=
	sensor_cpu->joules(state_cpu[std::string{label}].start,
			   state_cpu[std::string{label}].stop);

      state_cpu[std::string{label}].watts +=
	sensor_cpu->watts(state_cpu[std::string{label}].start,
			  state_cpu[std::string{label}].stop);

      state_cpu[std::string{label}].count++;
    }
  
  return;
}

void Show_PMT_CPU(const char *label)
{
  if (PMT_ERROR || !state_cpu.count(std::string{label}))
    {
      PMT_err();
      return;
    }
  else
    {
      std::cout << "\n\t Kernel:" << std::string{label} << ":" << std::endl;
      std::cout << "\t\t" << state_cpu[std::string{label}].seconds << " [S]" << std::endl;
      std::cout << "\t\t" << state_cpu[std::string{label}].joules  << " [J]" << std::endl;
      std::cout << "\t\t" << state_cpu[std::string{label}].watts / state_cpu[std::string{label}].count  << " [W]" << "\n" << std::endl;
    }
  
  return;
}
#endif // _ENERGY_RAPL_

#endif // _ENERGY_PMT_
