from ops import *
import sys,statistics
mode=sys.argv[1];d=json.loads((HERE/mode/'result.json').read_text())
assert d['measurement']=='PASS'
a,b=d['samples'][0],d['samples'][-1];dt=b['t']-a['t'];groups={};lanes=[]
for label,bs in b['groups'].items():
    am=a['groups'][label]['metadata_path'];bm=bs['metadata_path']
    assert a['groups'][label]['run_id']==bs['run_id']
    for provider,bp in bm['providers'].items():
        ap=am['providers'][provider];n=bp['responses']-ap['responses'];lost=bp['late']-ap['late']
        ad=am['performance']['pipeline_diagnostic']['providers'][provider];bd=bm['performance']['pipeline_diagnostic']['providers'][provider]
        g=groups.setdefault(provider,dict(responses=0,late=0,selected=0,pre_infer_expired=0,durations={},reasons={},lanes=0,diagnostic_missed=0))
        g['lanes']+=1;g['responses']+=n;g['late']+=lost;g['diagnostic_missed']+=bd['missed']-ad['missed']
        for k in ('selected','pre_infer_expired'):g[k]+=bd['counts'].get(k,0)-ad['counts'].get(k,0)
        for k,v in bd['durations'].items():
            old=ad['durations'].get(k,{});dn=v['count']-old.get('count',0);ns=v['sum_ns']-old.get('sum_ns',0)
            total=g['durations'].setdefault(k,dict(count=0,ns=0));total['count']+=dn;total['ns']+=ns
        for k,v in bm.items():
            if k.startswith('rejected_'+provider+'_'):g['reasons'][k]=g['reasons'].get(k,0)+v-am.get(k,0)
        lanes.append(dict(label=label,provider=provider,completed_fps=round(n/dt,2),accepted_fps=round((n-lost)/dt,2),rejected_pct=round(100*lost/n,2) if n else None))
for g in groups.values():
    g['completed_fps']=round(g['responses']/dt,2);g['accepted_fps']=round((g['responses']-g['late'])/dt,2)
    g['rejected_pct']=round(g['late']*100/g['responses'],2) if g['responses'] else None
    g['mean_ms']={k:round(v['ns']/v['count']/1e6,2) if v['count'] else None for k,v in g.pop('durations').items()}
models={}
for m,ta in d['triton_before'].items():
    tb=d['triton_after'][m];assert ta['config']==tb['config'];x,y=ta['stats'],tb['stats'];stats={}
    for k,v in y['inference_stats'].items():
        n=int(v['count'])-int(x['inference_stats'][k]['count']);ns=int(v['ns'])-int(x['inference_stats'][k]['ns'])
        stats[k]=dict(count=n,ms=round(ns/n/1e6,2) if n else None)
    old={v['batch_size']:v for v in x['batch_stats']};batch=[]
    for v in y['batch_stats']:
        n=v['compute_infer']['count']-old.get(v['batch_size'],{}).get('compute_infer',{}).get('count',0);batch.append([v['batch_size'],n])
    ex=int(y['execution_count'])-int(x['execution_count']);inf=int(y['inference_count'])-int(x['inference_count'])
    models[m]=dict(stats=stats,batches=batch,mean_batch=round(inf/ex,3) if ex else None,executions=ex,inferences=inf)
cpu={}
for k,v in b['thread_cpu'].items():
    if k not in a['thread_cpu']:continue
    old=a['thread_cpu'][k];r=cpu.setdefault(k.split(':')[0],dict(cpu_s=0.,runnable_wait_s=0.))
    r['cpu_s']+=(v[0]-old[0])/1e9;r['runnable_wait_s']+=(v[1]-old[1])/1e9
for r in cpu.values():
    for k in r:r[k]=round(r[k],3)
x=[int(i) for i in a['host_cpu'].split()[1:9]];y=[int(i) for i in b['host_cpu'].split()[1:9]];dc=[j-i for i,j in zip(x,y)]
cg0=dict(line.split() for line in a['cgroup']['cpu.stat'].splitlines());cg1=dict(line.split() for line in b['cgroup']['cpu.stat'].splitlines())
summary=dict(mode=mode,seconds=round(dt,3),groups=groups,lanes=lanes,models=models,cpu=cpu,
    host_cpu_busy_pct=round(100*(sum(dc)-dc[3]-dc[4])/sum(dc),1),
    gpu_mean_pct=round(statistics.mean(float(s['gpu'].split(',')[0]) for s in d['samples']),1),
    gpu_memory_peak_mib=max(float(s['gpu'].split(',')[1]) for s in d['samples']),
    cgroup_memory_peak_mib=round(max(int(s['cgroup']['memory.current']) for s in d['samples'])/1024**2),
    cpu_quota=b['cgroup']['cpu.max'].strip(),throttled_usec=int(cg1.get('throttled_usec',0))-int(cg0.get('throttled_usec',0)),
    cleanup={label:s['last_cleanup'] for label,s in d['after_stop']['groups'].items()})
save(mode+'-summary.json',summary)
print(json.dumps({k:v for k,v in summary.items() if k not in ('lanes','cpu','cleanup')}))