#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::rapl::Rapl::Create();

#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(devID[dev])});
#elif defined(_ENERGY_AMD_)
	                     pmt::rocm::AMD::Create(devID[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;
    }

  const std::string tag{std::string{label}};

  // check if the label already exists
  if (state_cpu.count(tag))
    {
      state_cpu[tag].start = sensor_cpu->Read();
    }
  else
    {
      // create new EnergyState
      const EnergyState newState{sensor_cpu->Read(),
                                 static_cast<pmt::State>(0),
                                 static_cast<double>(0),
                                 static_cast<double>(0),
                                 static_cast<double>(0),
                                 static_cast<unsigned int>(0)};

      // insert the key and initialize the counters
      state_cpu.insert(std::pair<std::string, EnergyState>(tag, newState));
    }

  return;
}

void Stop_PMT_CPU(const char *label)
{
  if (PMT_ERROR)
    {
      PMT_err();
      return;
    }

  const std::string tag{std::string{label}};
  
  // check if the label already exists
  // if not error
  if (!state_cpu.count(tag))
    {
      PMT_ERROR = true;
      PMT_err();
      return;
    }
  else
    {
      // get the energy state
      EnergyState &State = state_cpu[tag];
      
      // read the counter
      State.stop = sensor_cpu->Read();

      // update quantities
      State.seconds += sensor_cpu->seconds(State.start, State.stop);
      
      State.joules  += sensor_cpu->joules(State.start, State.stop);

      State.watts   += sensor_cpu->watts(State.start, State.stop);

      State.count++;
    }
  
  return;
}

void Show_PMT_CPU(const char *label)
{
  const std::string tag{std::string{label}};
  
  if (PMT_ERROR || !state_cpu.count(tag))
    {
      PMT_err();
      return;
    }
  else
    {
      std::cout << "\n\t CPU Kernel:" << tag << ":" << std::endl;
      std::cout << "\t\t" << state_cpu[tag].seconds << " [S]" << std::endl;
      std::cout << "\t\t" << state_cpu[tag].joules  << " [J]" << std::endl;
      std::cout << "\t\t" << state_cpu[tag].watts / state_cpu[tag].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(std::pair<int, std::map<std::string, EnergyState>>(devID, {}));
    }

  const std::string tag{std::string{label}};
  
  // check if the label already exists
  if (state_gpu[devID].count(tag))
    {
      // read the sensor
      state_gpu[devID][tag].start = sensor_gpu[devID]->Read();
    }
  else
    {
      // insert the label and initialize the counters
      const EnergyState newState{sensor_gpu[devID]->Read(),
                                 static_cast<pmt::State>(0),
                                 static_cast<double>(0),
                                 static_cast<double>(0),
                                 static_cast<double>(0),
                                 static_cast<unsigned int>(0)};

      state_gpu[devID].insert(std::pair<std::string, EnergyState>(tag, newState));      
    }

  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
    {
      const std::string tag{std::string{label}};
      
      // check if the label already exists
      // if not error
      if (!state_gpu[devID].count(tag))
	{
	  PMT_ERROR = true;
	  PMT_err();
	  return;
	}
      else
	{
	  EnergyState &State = state_gpu[devID][tag];
	  
	  // read the counter
	  State.stop = sensor_gpu[devID]->Read();

	  // update quantities
	  State.seconds +=
	    sensor_gpu[devID]->seconds(State.start,
				       State.stop);
      
	  State.joules +=
	    sensor_gpu[devID]->joules(State.start,
				      State.stop);

	  State.watts +=
	    sensor_gpu[devID]->watts(State.start,
				     State.stop);
      
	  State.count++;
	}
    }
  
  return;
}

void Show_PMT_GPU(const char *label)
{
  if (PMT_ERROR)
    {
      PMT_err();
      return;
    }
  else
    {
      const std::string tag{std::string{label}};

      for (const auto &[key, value]: state_gpu)
	{
	  if (value.count(tag))
	    {
	      std::cout << "\n\t GPU [" << key << "] kernel:" << tag << ":" << std::endl;
	      std::cout << "\t\t" << value.at(tag).seconds << " [s]" << std::endl;
	      std::cout << "\t\t" << value.at(tag).joules  << " [J]" << std::endl;
	      std::cout << "\t\t" << value.at(tag).watts / value.at(tag).count  << " [W]" << "\n" << std::endl;
	    }
	}
    }
  
  return;
}

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

#endif // _ENERGY_PMT_
