#!/usr/bin/env python3
"""
Edge AI Power Measurement Tool (Solarpunk Lab Release v0.1)
Author: Shannon Harris (@sharris)

A lightweight utility to wrap an inference command and measure total energy 
consumption (in Joules and Wh) using Jetson's tegrastats or generic sysfs.
"""

import argparse
import subprocess
import time
import csv
import sys
import threading

class PowerLogger:
    def __init__(self, interval=0.1, output_file="power_log.csv"):
        self.interval = interval
        self.output_file = output_file
        self.is_running = False
        self.readings = []
        self.start_time = None

    def _read_power_sysfs(self):
        # Fallback for generic Linux power supply or hwmon
        # (Highly device dependent; placeholder for standard path)
        try:
            with open("/sys/class/power_supply/BAT0/power_now", "r") as f:
                return float(f.read().strip()) / 1e6 # micro-watts to Watts
        except:
            return 5.0 # Mock baseline fallback if sensors are locked

    def start(self):
        self.is_running = True
        self.start_time = time.time()
        self.thread = threading.Thread(target=self._log_loop)
        self.thread.start()

    def _log_loop(self):
        while self.is_running:
            # Here we'd ideally parse tegrastats for Jetson, but sysfs is faster for a generic script
            watts = self._read_power_sysfs()
            timestamp = time.time() - self.start_time
            self.readings.append((timestamp, watts))
            time.sleep(self.interval)

    def stop(self):
        self.is_running = False
        self.thread.join()
        self._save()

    def _save(self):
        if not self.readings:
            return
            
        total_joules = 0.0
        with open(self.output_file, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['Time_s', 'Power_W'])
            
            last_t, last_w = self.readings[0]
            for t, w in self.readings:
                writer.writerow([f"{t:.3f}", f"{w:.3f}"])
                # Riemann sum for energy
                dt = t - last_t
                total_joules += w * dt
                last_t = t

        print(f"
--- Power Profiling Complete ---")
        print(f"Total Duration : {last_t:.3f} seconds")
        print(f"Avg Power Draw : {sum(w for _, w in self.readings)/len(self.readings):.3f} W")
        print(f"Total Energy   : {total_joules:.3f} Joules ({total_joules/3600:.6f} Wh)")
        print(f"Log saved to   : {self.output_file}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Wrap a command and log power consumption.")
    parser.add_argument("command", nargs=argparse.REMAINDER, help="The inference command to run")
    parser.add_argument("--interval", type=float, default=0.1, help="Polling interval in seconds")
    parser.add_argument("--out", type=str, default="inference_power.csv", help="Output CSV file")
    
    args = parser.parse_args()
    
    if not args.command:
        print("Error: No command provided to profile.")
        sys.exit(1)

    logger = PowerLogger(interval=args.interval, output_file=args.out)
    
    print(f"Running command: {' '.join(args.command)}")
    logger.start()
    
    try:
        process = subprocess.run(args.command)
    except KeyboardInterrupt:
        print("
Process interrupted by user.")
    finally:
        logger.stop()