#!/usr/bin/env python3
"""Minimal NVIDIA GPU energy wrapper for Critical Load UAW measurements.
Runs a command and records cumulative NVML energy-counter deltas plus sampled power.
Requires: pip install nvidia-ml-py
"""
import argparse, csv, os, subprocess, threading, time
from statistics import mean
import pynvml

p=argparse.ArgumentParser()
p.add_argument('--out', required=True)
p.add_argument('--sample-ms', type=int, default=100)
p.add_argument('command', nargs=argparse.REMAINDER)
a=p.parse_args()
if not a.command: raise SystemExit('Provide command after --')
if a.command and a.command[0]=='--': a.command=a.command[1:]

pynvml.nvmlInit(); n=pynvml.nvmlDeviceGetCount(); hs=[pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(n)]

def energy_mj(h): return pynvml.nvmlDeviceGetTotalEnergyConsumption(h)
start_e=[energy_mj(h) for h in hs]
samples=[]; stop=False

def sampler():
    while not stop:
        ts=time.time()
        for i,h in enumerate(hs):
            try: pw=pynvml.nvmlDeviceGetPowerUsage(h)/1000.0
            except: pw=float('nan')
            samples.append((ts,i,pw))
        time.sleep(a.sample_ms/1000)

t=threading.Thread(target=sampler,daemon=True); t.start(); t0=time.time()
proc=subprocess.run(a.command); t1=time.time(); stop=True; t.join(timeout=1)
end_e=[energy_mj(h) for h in hs]

with open(a.out,'w',newline='') as f:
    w=csv.writer(f); w.writerow(['gpu_index','energy_delta_j','avg_sampled_power_w','duration_s','return_code'])
    for i,(s,e) in enumerate(zip(start_e,end_e)):
        pw=[x[2] for x in samples if x[1]==i and x[2]==x[2]]
        w.writerow([i,(e-s)/1000.0,mean(pw) if pw else '',t1-t0,proc.returncode])
pynvml.nvmlShutdown(); raise SystemExit(proc.returncode)