#!/usr/bin/env python3
import csv, math, sys
from collections import defaultdict

if len(sys.argv) < 2:
    print('Usage: compute_metrics.py run_log.csv [output.csv]')
    sys.exit(1)
input_path = sys.argv[1]
out_path = sys.argv[2] if len(sys.argv) > 2 else 'uaw_summary.csv'

def f(x):
    try: return float(x)
    except: return 0.0

def b(x):
    return str(x).strip().lower() in {'1','true','yes','y'}

rows=list(csv.DictReader(open(input_path, newline='', encoding='utf-8')))
groups=defaultdict(list)
for r in rows:
    key=(r['task_id'], r['model'], r['model_version'], r['hardware'], r['engine'], r['precision'], r['measurement_boundary'])
    groups[key].append(r)

out=[]
for key, rs in groups.items():
    accepted=[r for r in rs if b(r['accepted'])]
    attempts=len(rs); nacc=len(accepted)
    energy_field='node_energy_j' if any(f(r.get('node_energy_j'))>0 for r in rs) else 'gpu_energy_j'
    total_energy=sum(f(r.get(energy_field)) for r in rs)
    success=nacc/attempts if attempts else math.nan
    e_acc=total_energy/nacc if nacc else math.nan
    first_energy=0.0; retry_energy=0.0
    by_run=defaultdict(list)
    for r in rs: by_run[r['run_id']].append(r)
    for _, ars in by_run.items():
        ars=sorted(ars,key=lambda x:int(x.get('attempt_index') or 0))
        if ars:
            first_energy += f(ars[0].get(energy_field))
            retry_energy += sum(f(x.get(energy_field)) for x in ars[1:])
    retry_share=retry_energy/total_energy if total_energy else 0.0
    avg_quality=sum(f(r.get('quality_score')) for r in accepted)/nacc if nacc else math.nan
    avg_latency=sum(f(r.get('latency_s')) for r in rs)/attempts if attempts else math.nan
    out.append({
        'task_id':key[0],'model':key[1],'model_version':key[2],'hardware':key[3],
        'engine':key[4],'precision':key[5],'measurement_boundary':key[6],
        'attempts':attempts,'accepted_tasks':nacc,'success_rate':success,
        'energy_field_used':energy_field,'total_energy_j':total_energy,
        'energy_per_accepted_task_j':e_acc,'retry_energy_share':retry_share,
        'mean_accepted_quality':avg_quality,'mean_attempt_latency_s':avg_latency
    })

fields=list(out[0].keys()) if out else []
with open(out_path,'w',newline='',encoding='utf-8') as fh:
    w=csv.DictWriter(fh,fieldnames=fields); w.writeheader(); w.writerows(out)
print(f'Wrote {len(out)} summary rows to {out_path}')