from ops import *
import statistics,sys
phase=sys.argv[1];d=json.loads((HERE/(phase+'.json')).read_text());dt=d['seconds']
channels=[];groups={}
for a,b in zip(d['samples'][0]['channels'],d['samples'][-1]['channels']):
    am,bm=a['metadata_path'],b['metadata_path']
    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={},channels=0,diagnostic_missed=0))
        g['channels']+=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)
        durations={}
        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)
            durations[k]=round(ns/dn/1e6,2) if dn else None
            total=g['durations'].setdefault(k,dict(count=0,ns=0));total['count']+=dn;total['ns']+=ns
        reasons={k:v-am.get(k,0) for k,v in bm.items() if k.startswith('rejected_'+provider+'_')}
        for k,v in reasons.items():g['reasons'][k]=g['reasons'].get(k,0)+v
        channels.append(dict(camera=b['camera'],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,mean_ms=durations,reasons=reasons))
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'];a,b=ta['stats'],tb['stats'];stats={}
    for k,v in b['inference_stats'].items():
        n=int(v['count'])-int(a['inference_stats'][k]['count']);ns=int(v['ns'])-int(a['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 a['batch_stats']};batch=[]
    for v in b['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(b['execution_count'])-int(a['execution_count']);inf=int(b['inference_count'])-int(a['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 d['cpu_after'].items():
    if k not in d['cpu_before']:continue
    old=d['cpu_before'][k];r=cpu.setdefault(v['channel'],dict(cpu_s=0.,runnable_wait_s=0.))
    r['cpu_s']+=(v['values'][0]-old['values'][0])/1e9;r['runnable_wait_s']+=(v['values'][1]-old['values'][1])/1e9
for r in cpu.values():
    for k in r:r[k]=round(r[k],3)
a=[int(x) for x in d['samples'][0]['cpu'].split()[1:9]];b=[int(x) for x in d['samples'][-1]['cpu'].split()[1:9]];dc=[y-x for x,y in zip(a,b)]
summary=dict(phase=phase,seconds=round(dt,3),groups=groups,channels=channels,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))
save(phase+'-summary.json',summary)
print(json.dumps({k:v for k,v in summary.items() if k not in ('channels','cpu')}))