#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::map<int, std::unique_ptr<pmt::PMT>> sensor_gpu;
   static std::map<int, std::map<std::string, EnergyState>> state_gpu;
#endif // defined(_ENERGY_NVIDIA_) || defined(_ENERGY_AMD_)

#if !defined(__NVCC__) && !defined(__NVCOMPILER)
   extern "C"
      {
         #include "energy_pmt_methods.h"
      }
#endif // !defined(__NVCC__) && !defined(__NVCOMPILER)

#if defined(_ENERGY_RAPL_) || defined(_ENERGY_NVIDIA_) || defined(_ENERGY_AMD_)
void Create_PMT(int *devID,
		const unsigned int numGPUs)
{
#if defined(_ENERGY_RAPL_)

  sensor_cpu = pmt::Create("rapl");

#endif // _ENERGY_RAPL_
  
  if ((numGPUs > 0) && (devID != nullptr))
    {
      for (unsigned int dev=0 ; dev<numGPUs ; dev++)
	{
#if defined(_ENERGY_NVIDIA_) || defined(_ENERGY_AMD_)
	  sensor_gpu.insert({devID[dev],

#if defined(_ENERGY_NVIDIA_)
	                     pmt::nvml::NVML::Create(dev)});
#elif defined(_ENERGY_AMD_)
	                     pmt::rocm::AMD::Create(dev)});
#endif

#endif // defined(_ENERGY_NVIDIA_) || defined(_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 CPU 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_

#if defined(_ENERGY_NVIDIA_) || defined(_ENERGY_AMD_)
void Start_PMT_GPU(const char *label,
		   const int  devID)
{
  if (PMT_ERROR || !sensor_gpu.count(devID))
    {
      PMT_err();
      return;
    }

  // check if the devID already exists
  if (!state_gpu.count(devID))
    {
      // insert devID
      state_gpu.insert({devID, {}});
    }
  // check if the label already exists
  if (state_gpu[devID].count(std::string{label}))
    {
      // read the sensor
      state_gpu[devID][std::string{label}].start = sensor_gpu[devID]->Read();
    }
  else
    {
      // insert the label and initialize the counters
      state_gpu[devID].insert({std::string{label},
			       {
				 sensor_gpu[devID]->Read(),
				 0,
				 0.0, 0.0, 0.0, 0
			       }});
    }

  return;
}

void Stop_PMT_GPU(const char *label,
		  const int   devID)
{
  // check if the devID already exists
  // if not error
  if (!state_gpu.count(devID) || PMT_ERROR || !sensor_gpu.count(devID))
    {
      PMT_ERROR = true;
      PMT_err();
      return;
    }
  else
    {
      // check if the label already exists
      // if not error
      if (!state_gpu[devID].count(std::string{label}))
	{
	  PMT_ERROR = true;
	  PMT_err();
	  return;
	}
      else
	{
	  // read the counter
	  state_gpu[devID][std::string{label}].stop = sensor_gpu[devID]->Read();

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

	  state_gpu[devID][std::string{label}].watts +=
	    sensor_gpu[devID]->watts(state_gpu[devID][std::string{label}].start,
				     state_gpu[devID][std::string{label}].stop);
      
	  state_gpu[devID][std::string{label}].count++;
	}
    }
  
  return;
}

void Show_PMT_GPU(const char *label)
{
  if (PMT_ERROR)
    {
      PMT_err();
      return;
    }
  else
    {
      // show quantities for all devices
      for (const auto& [key, value]: state_gpu)
	{
	  std::cout << "\n\t GPU [" << key << "] kernel:" << std::string{label} << ":" << std::endl;
	  std::cout << "\t\t" << value.at(std::string{label}).seconds << " [s]" << std::endl;
	  std::cout << "\t\t" << value.at(std::string{label}).joules  << " [J]" << std::endl;
	  std::cout << "\t\t" << value.at(std::string{label}).watts / value.at(std::string{label}).count  << " [W]" << "\n" << std::endl;
	}
    }
  
  return;
}

#endif // defined(_ENERGY_NVIDIA_) || defined(_ENERGY_AMD_)

#endif // _ENERGY_PMT_
