from ops import *
import csv,statistics
table=[];overview={};batch=[]
for mode in ('pose','pose-a'):
    root=HERE/mode
    result=json.loads((root/'trial-result.json').read_text())
    summary=json.loads((root/'summary-v2.json').read_text())
    details=json.loads((root/'detail-sot.json').read_text())
    lanes=json.loads((root/'lanes-v2.json').read_text())
    age=json.loads((root/'age-windows.json').read_text());tails=json.loads((root/'age-tail.json').read_text())
    clean=[p['last_cleanup'] for r in result['cleanup']['states'].values() for p in r['providers'].values()]
    overview[mode]=dict(stages=[s['level'] for s in summary if s['status']=='PASS'],errors=result['errors'],
        reaped=result['cleanup']['reaped'],children=len(clean),forced=sum(bool(c['forced']) for c in clean),
        exits=sorted(set(c['diagnostic']['exit_code'] for c in clean)),
        terminal_failed=sum(bool(c['diagnostic'].get('terminal_failed')) for c in clean))
    for s in summary:
        if s['status']!='PASS':continue
        n=s['level']
        for p,v in s['providers'].items():
            rows=[r for r in details if r['level']==n and r['provider']==p]
            ls=[r for r in lanes if r['level']==n and r['provider']==p]
            out=dict(mode=mode,queue_seconds=1 if mode=='pose-a' else 2,channels=n,provider=p,
                seconds=s['seconds'],completed_mean=v['per_channel_completed_mean'],completed_min=min(r['completed_fps'] for r in ls),
                completed_total=v['completed_fps'],redis_mean=v['sent_fps']/n,redis_total=v['sent_fps'],
                redis_min=min(r['sent_fps'] for r in ls),redis_std=statistics.pstdev(r['sent_fps'] for r in ls),
                finalized_mean=statistics.mean(r['finalized_fps'] for r in rows),
                expired=sum(r['queue_expired'] for r in rows),capacity_drop=sum(r['queue_drop'] for r in rows),
                send_replaced=sum(r['counts']['replaced'] for r in rows),send_errors=sum(r['counts']['send_error'] for r in rows),
                max_sampled_depth=max(r['max_queue_depth'] for r in rows),max_sampled_queue_age=max(r['max_queue_age'] for r in rows),
                host_cpu_busy=s['host_cpu_busy_pct'],container_cores=s['container_mean_cores'],gpu_util=s['gpu_util_mean'],
                gpu_memory_mib=s['gpu_memory_peak_mib'],ram_mib=s['ram_peak_mib'],throttled_us=s['throttled_usec'],
                events_created=sum(r['events_created'] for r in rows),renderer_peak_sampled=s['renderer_peak_sampled'],
                phase_wall_ms={k:d['mean_ms'] for k,d in v['durations'].items() if k.startswith('result_') and k.endswith('_wall')},
                phase_cpu_ms={k:d['mean_ms'] for k,d in v['durations'].items() if k.startswith('result_') and k.endswith('_cpu')})
            for name in ('selected_to_redis','completed_to_redis'):
                stats=[r['timings'].get(name,{}) for r in rows];count=sum(x.get('count',0) for x in stats)
                out[name+'_mean_ms']=sum(x['count']*x['mean_ms'] for x in stats if x.get('count'))/count if count else None
                ts=[r for r in tails if r['level']==n and r['provider']==p and r['metric']==name]
                out[name+'_max_ms']=max((x['configured_window_max_ms'] for x in ts),default=None)
                out[name+'_lane_p50_range_ms']=[min(x['recent_p50_ms'] for x in ts),max(x['recent_p50_ms'] for x in ts)] if ts else []
                out[name+'_lane_p95_max_ms']=max((x['recent_p95_ms'] for x in ts),default=None)
                windows=[]
                for i in range(1,7):
                    ws=[r for r in age if r['level']==n and r['provider']==p and r['window']==i and r[name+'_count']>0]
                    count=sum(x[name+'_count'] for x in ws)
                    windows.append(sum(x[name+'_mean_ms']*x[name+'_count'] for x in ws)/count if count else None)
                out[name+'_five_second_means_ms']=windows
            table.append(out)
        for model,m in s['models'].items():
            counts=m['batch_executions'];executions=sum(counts.values())
            batch.append(dict(mode=mode,channels=n,model=model,executions=executions,frames=m['inferences'],
                requests=m['stats']['success']['count'],batch_executions=counts,
                batch8_pct=100*counts.get('8',0)/executions if executions else None,
                batch_gt4_pct=100*sum(v for k,v in counts.items() if int(k)>4)/executions if executions else None,
                queue_ms=m['stats']['queue']['mean_ms'],compute_ms=m['stats']['compute_infer']['mean_ms']))
    start=result['stages'][0]['started_wall_ms']/1000;end=result['finished_wall_ms']/1000
    sql=f"SELECT coalesce(json_agg(t),'[]'::json) FROM (SELECT payload->>'scenario' AS scenario,artifact_state,delivery_state,CASE WHEN error ~ '^[A-Z0-9_:-]+$' THEN left(error,100) ELSE NULL END AS error_code,count(*) AS count FROM contextra_event_artifacts WHERE created_at >= {start} AND created_at <= {end} GROUP BY 1,2,3,4) t;"
    value=json.loads(cmd(['docker','exec','-i','contextra-postgresql-1','sh','-c','exec psql -X -qAt -U "$POSTGRES_USER" -d "$POSTGRES_DB"'],input=sql))
    save(mode+'-artifacts.json',dict(start=start,end=end,groups=value))
for name,data in [('comparison-sot',table),('batch-sot',batch)]:
    save(name+'.json',data)
    if data:
        with (HERE/(name+'.csv')).open('x',newline='') as f:
            w=csv.DictWriter(f,fieldnames=list(data[0]));w.writeheader();w.writerows(data)
save('overview.json',overview)
print(json.dumps(overview))
for r in table:
    print(r['mode'],r['channels'],r['provider'],round(r['completed_mean'],2),round(r['redis_mean'],2),r['expired'],round(r['selected_to_redis_mean_ms'] or 0),round(r['host_cpu_busy'],1))
