from pynvml import *
import time


class GPUEnergyTracker:
    """
    GPU energy monitor using NVIDIA NVML.

    Measures instantaneous power and integrates
    it over time to estimate total GPU energy
    consumption.
    Units
    Power  : Watts
    Energy : Joules
    """

    def __init__(self, gpu_index=0):
        nvmlInit()

        self.handle = nvmlDeviceGetHandleByIndex(gpu_index)

        self.energy_joules = 0.0

        self.last_time = time.time()

        self.last_power = self._read_power()

    def _read_power(self):
        power_mw = nvmlDeviceGetPowerUsage(self.handle)

        return power_mw / 1000.0

    def step(self):
        current_time = time.time()

        power = self._read_power()

        dt = current_time - self.last_time

        # trapezoidal integration
        self.energy_joules += 0.5 * (self.last_power + power) * dt

        self.last_power = power
        self.last_time = current_time

    def total_energy(self):
        return self.energy_joules

    def current_power(self):
        return self.last_power