from ops import *
import csv,math,statistics
ROOT=HERE/'attempt2'
ORDER=[3,5,6,8,9,10,11,12,13,14,15,16,17,18,19,20,21]
def delta(a,b,key):return b.get(key,0)-a.get(key,0)
def p95(hist,bounds):
    total=sum(hist)
    if not total:return None
    target=math.ceil(total*.95);seen=0
    for i,n in enumerate(hist):
        seen+=n
        if seen>=target:return bounds[i] if i<len(bounds) else '>'+str(bounds[-1])
def delivery(row,p):
    return next(iter(row.get('publisher_metrics',{}).get('delivery',{}).values()),{}).get(p,{})
def diag(path,p):return path.get('performance',{}).get('pipeline_diagnostic',{}).get('providers',{}).get(p,{})
result=json.loads((ROOT/'trial-result.json').read_text());summary=[];lanes=[];windows=[];previous=0
for stage in result['stages']:
    if len(stage['samples'])<2:continue
    a,b=stage['samples'][0],stage['samples'][-1];dt=b['t']-a['t'];level=stage['level'];groups={}
    for camera,bc in b['channels'].items():
        ac=a['channels'][camera]
        assert ac['run_id']==bc['run_id']
        for p,bprovider in bc['providers'].items():
            ap=ac['providers'][p]['metadata_path'];bp=bprovider['metadata_path']
            an=ap['providers'][p];bn=bp['providers'][p]
            completed=delta(an,bn,'responses');late=delta(an,bn,'late');accepted=delta(ap,bp,'accepted_result')
            assert completed-late==accepted
            ad,bd=diag(ap,p),diag(bp,p)
            same_window=bool(ad.get('enabled') and ad.get('started_ns')==bd.get('started_ns'))
            if not same_window:ad={}
            metric_valid=bool(bd.get('enabled') and bd.get('collecting'))
            diag_dt=((bp['performance']['sampled_ns']-ap['performance']['sampled_ns'])/1e9 if same_window else
                     (bp['performance']['sampled_ns']-bd.get('started_ns',bp['performance']['sampled_ns']))/1e9)
            durations={}
            for name,value in bd.get('durations',{}).items():
                old=ad.get('durations',{}).get(name,{});n=delta(old,value,'count');ns=delta(old,value,'sum_ns')
                histogram=[x-y for x,y in zip(value['histogram'],old.get('histogram',[0]*len(value['histogram'])))]
                durations[name]=dict(count=n,sum_ns=ns,mean_ms=ns/n/1e6 if n else None,
                    p95_bucket_upper_ms=p95(histogram,bd['histogram_bounds_ms']),histogram=histogram,
                    diagnostic_window_max_ms=value['max_ns']/1e6)
            pub0,pub1=delivery(ac,p),delivery(bc,p);sent=delta(pub0,pub1,'count')
            lane=dict(level=level,phase_status=stage['status'],camera=int(camera),provider=p,
                cohort='existing' if int(camera) in ORDER[:previous] else 'new',seconds=dt,
                completed=completed,accepted=accepted,sent=sent,late=late,
                completed_fps=completed/dt,accepted_fps=accepted/dt,sent_fps=sent/dt,
                rejected_pct=late*100/completed if completed else None,diagnostic_complete=metric_valid,
                diagnostic_scope='sample_difference' if same_window else 'configure_to_last_child_snapshot',
                diagnostic_seconds=diag_dt,
                selected=delta(ap['performance']['counts'],bp['performance']['counts'],'selected'),
                infer_entries=delta(ad.get('counts',{}),bd.get('counts',{}),'preprocess') if metric_valid else None,
                infer_entry_fps=(delta(ad.get('counts',{}),bd.get('counts',{}),'preprocess')/diag_dt) if metric_valid and diag_dt>0 else None,
                pre_infer_expired=delta(ad.get('counts',{}),bd.get('counts',{}),'pre_infer_expired') if metric_valid else None,
                diagnostic_missed=delta(ad,bd,'missed'),durations=durations,
                rejection_reasons={k:delta(ap,bp,k) for k in bp if k.startswith('rejected_')},
                first_response_s=bn['first_response_at']-result['starts'][camera]['sent'] if bn.get('first_response_at') else None,
                first_redis_commit_s=pub1.get('first_at_ns',0)/1e9-result['starts'][camera]['sent'] if pub1 else None,
                maximum_sampled_silence_ms=max(diag(s['channels'][camera]['providers'][p]['metadata_path'],p).get('acceptance_silence_ns',0)/1e6 for s in stage['samples']),
                redis_commit_p95_bucket_upper_ms=p95([v-x for v,x in zip(pub1.get('histogram',[]),pub0.get('histogram',[0]*11))],(20,50,75,100,150,200,300,500,1000,2000)))
            lanes.append(lane);groups.setdefault(p,[]).append(lane)
    models={}
    for m,tb in b['resources']['triton'].items():
        ta=a['resources']['triton'][m];stats={}
        for name,stat in tb['inference_stats'].items():
            old=ta['inference_stats'][name];n=int(stat['count'])-int(old['count']);ns=int(stat['ns'])-int(old['ns'])
            stats[name]=dict(count=n,mean_ms=ns/n/1e6 if n else None)
        old={r['batch_size']:r for r in ta['batch_stats']};batch={}
        for r in tb['batch_stats']:
            batch[str(r['batch_size'])]=int(r['compute_infer']['count'])-int(old.get(r['batch_size'],{}).get('compute_infer',{}).get('count',0))
        executions=int(tb['execution_count'])-int(ta['execution_count']);inferences=int(tb['inference_count'])-int(ta['inference_count'])
        models[m]=dict(stats=stats,batch_executions=batch,mean_batch=inferences/executions if executions else None,
            executions=executions,inferences=inferences)
    cpu={};ra,rb=a['resources'],b['resources'];rdt=rb['started_monotonic']-ra['started_monotonic']
    for label,tasks in rb['cpu'].items():
        old=ra['cpu'].get(label,{})
        cohort=label.rsplit(':',1)[0];v=cpu.setdefault(cohort,dict(cpu_s=0.,runnable_wait_thread_s=0.,matched_threads=0,unmatched_threads=0))
        for tid,row in tasks.items():
            if tid not in old:v['unmatched_threads']+=1;continue
            v['matched_threads']+=1;v['cpu_s']+=(row[0]-old[tid][0])/1e9;v['runnable_wait_thread_s']+=(row[1]-old[tid][1])/1e9
    for v in cpu.values():v['mean_cores']=v['cpu_s']/rdt
    acg=dict(l.split() for l in ra['cgroup']['cpu.stat'].splitlines());bcg=dict(l.split() for l in rb['cgroup']['cpu.stat'].splitlines())
    hx=[int(v) for v in ra['host_cpu'].split()[1:9]];hy=[int(v) for v in rb['host_cpu'].split()[1:9]];hd=[v-x for x,v in zip(hx,hy)]
    gpu=[[float(x.strip()) for x in s['gpu'].split(',')] for s in stage['samples']]
    aggregated={}
    for p,entries in groups.items():
        duration_totals={}
        for lane in entries:
            for name,d in lane['durations'].items():
                v=duration_totals.setdefault(name,dict(count=0,sum_ns=0,histogram=[0]*16))
                v['count']+=d['count'];v['sum_ns']+=d['sum_ns'];v['histogram']=[x+y for x,y in zip(v['histogram'],d['histogram'])]
        for v in duration_totals.values():
            v['mean_ms']=v['sum_ns']/v['count']/1e6 if v['count'] else None
            v['p95_bucket_upper_ms']=p95(v['histogram'],[1,2,5,10,20,50,75,100,150,200,300,500,1000,2000,5000])
        n=sum(v['completed'] for v in entries);lost=sum(v['late'] for v in entries)
        aggregated[p]=dict(completed_fps=sum(v['completed_fps'] for v in entries),accepted_fps=sum(v['accepted_fps'] for v in entries),
            per_channel_completed_mean=statistics.mean(v['completed_fps'] for v in entries),
            per_channel_accepted_mean=statistics.mean(v['accepted_fps'] for v in entries),
            per_channel_accepted_min=min(v['accepted_fps'] for v in entries),per_channel_accepted_max=max(v['accepted_fps'] for v in entries),
            per_channel_accepted_std=statistics.pstdev(v['accepted_fps'] for v in entries),rejected_pct=lost*100/n if n else None,
            sent_fps=sum(v['sent_fps'] for v in entries),diagnostic_complete_channels=sum(v['diagnostic_complete'] for v in entries),
            selected_fps=sum(v['selected'] or 0 for v in entries)/dt,infer_entry_fps=sum(v['infer_entry_fps'] or 0 for v in entries),
            pre_infer_expired=sum(v['pre_infer_expired'] or 0 for v in entries),durations=duration_totals,
            existing_accepted_mean=statistics.mean(v['accepted_fps'] for v in entries if v['cohort']=='existing') if previous else None,
            new_accepted_mean=statistics.mean(v['accepted_fps'] for v in entries if v['cohort']=='new'))
    value=dict(level=level,status=stage['status'],seconds=dt,providers=aggregated,models=models,cpu=cpu,
        host_cpu_busy_pct=100*(sum(hd)-hd[3]-hd[4])/sum(hd),container_mean_cores=(int(bcg['usage_usec'])-int(acg['usage_usec']))/1e6/rdt,
        throttled_usec=int(bcg.get('throttled_usec',0))-int(acg.get('throttled_usec',0)),cpu_quota=rb['cgroup']['cpu.max'],
        gpu_util_mean=statistics.mean(v[0] for v in gpu),gpu_memory_util_mean=statistics.mean(v[1] for v in gpu),
        gpu_memory_peak_mib=max(v[2] for v in gpu),gpu_power_mean_w=statistics.mean(v[3] for v in gpu),
        ram_peak_mib=max(int(s['resources']['cgroup']['memory.current']) for s in stage['samples'])/1024**2,
        host_available_min_mib=min(s['resources']['mem']['MemAvailable'] for s in stage['samples'])/1024**2,
        expected_children=level*3,observed_children=len([v for r in b['channels'].values() for v in r['providers'].values() if v.get('child_pid')]),
        renderer_peak_sampled=max(s['artifacts']['rendering'] for s in stage['samples']),
        artifact_count_delta={k:delta(a['artifacts']['counts'],b['artifacts']['counts'],k) for k in b['artifacts']['counts']},
        max_probe_seconds=max(s['finished']-s['t'] for s in stage['samples']),metric_errors=stage['metric_errors'])
    summary.append(value);previous=level
    for first,last in zip(stage['samples'],stage['samples'][1:]):
        windows.append(dict(level=level,begin_wall_ms=first['wall_ms'],end_wall_ms=last['wall_ms'],seconds=last['t']-first['t'],
            renderer_first=first['artifacts']['rendering'],renderer_last=last['artifacts']['rendering'],gpu_last=last['gpu'],
            channel_provider=[dict(camera=c,provider=p,completed=delta(first['channels'][c]['providers'][p]['metadata_path']['providers'][p],v['metadata_path']['providers'][p],'responses'),
                                  late=delta(first['channels'][c]['providers'][p]['metadata_path']['providers'][p],v['metadata_path']['providers'][p],'late')) for c,r in last['channels'].items() for p,v in r['providers'].items()]))
save('summary-v2.json',summary);save('lanes-v2.json',lanes);save('windows-v2.json',windows)
flat=[{k:v for k,v in r.items() if k not in ('durations','rejection_reasons')} for r in lanes]
with (HERE/'lanes-v2.csv').open('x',newline='') as f:
    w=csv.DictWriter(f,fieldnames=list(flat[0]));w.writeheader();w.writerows(flat)
print(json.dumps([dict(level=s['level'],state=s['status'],gpu=s['gpu_util_mean'],cpu=s['container_mean_cores'],
    ram=s['ram_peak_mib'],providers={p:{k:v for k,v in r.items() if k in ('completed_fps','accepted_fps','per_channel_accepted_mean','rejected_pct','sent_fps')} for p,r in s['providers'].items()}) for s in summary]))