from ops import *
import csv, math, statistics as st, sys
label=sys.argv[1]; folder=HERE/label
def stats(v):
    v=sorted(v)
    return dict(n=len(v),mean=st.mean(v),median=st.median(v),p95=v[math.ceil(.95*len(v))-1],maximum=v[-1],minimum=v[0]) if v else None
def diag(path):return path.get('performance',{}).get('pipeline_diagnostic',{}).get('providers',{}).get('pose',{})
def delivery(row):return next(iter(row.get('publisher_metrics',{}).get('delivery',{}).values()),{}).get('pose',{})
def delta(a,b,k):return b.get(k,0)-a.get(k,0)
stages=[json.loads(p.read_text()) for p in sorted(folder.glob('stage-*.json'))]
records={};sample_rows=[]
for name in ('startup.jsonl','samples.jsonl'):
    file=folder/name
    if not file.exists():continue
    for line in file.read_text().splitlines():
        s=json.loads(line)
        if name=='samples.jsonl':sample_rows.append(s)
        for c,r in s['channels'].items():
            for p,v in r['providers'].items():
                for f in diag(v.get('metadata_path') or {}).get('input_timeline',[]):
                    records[(c,r['run_id'],f['pts'])]=dict(camera=c,run=r['run_id'],provider=p,**f)
points=[];anchors={};segments={};last={};resets=[]
for row in sorted(records.values(),key=lambda f:f['input']['ingress_ns']):
    x=row['input'];key=(row['camera'],row['run']);old=last.get(key)
    ticks=(x['rtp']-old['rtp']+2**31)%2**32-2**31 if old else 0
    if old and (x['ssrc']!=old['ssrc'] or x['rate']!=old['rate'] or ticks<0):
        resets.append(dict(camera=row['camera'],run=row['run'],pts=row['pts']))
        segments[key]=segments.get(key,0)+1;old=None
    seg=(key,segments.get(key,0))
    unwrapped=old['unwrapped']+ticks if old else x['rtp']
    last[key]=dict(x,unwrapped=unwrapped)
    if seg not in anchors:anchors[seg]=dict(x,unwrapped=unwrapped)
    origin=anchors[seg];media=(unwrapped-origin['unwrapped'])/x['rate']
    point=dict(camera=row['camera'],run=row['run'],pts=row['pts'],rtp=x['rtp'],segment=seg[1],rate=x['rate'],
        t=x['ingress_ns']/1e9,media_elapsed_s=media,ingress_wall_ns=x['ingress_wall_ns'])
    for name in ('depay_first_ns','source_exit_ns','ingress_ns'):
        point[name+'_relative_s']=(x[name]-origin[name])/1e9-media
    point['depay_source_ms']=(x['source_exit_ns']-x['depay_first_ns'])/1e6
    point['source_ingress_ms']=(x['ingress_ns']-x['source_exit_ns'])/1e6
    for name,a,b in (('preprocess_ms','preprocess_start','preprocess_output'),
                     ('infer_queue_ms','preprocess_output','preprocess'),
                     ('infer_output_ms','preprocess','infer_output'),('callback_copy_ms','callback','copied')):
        point[name]=(row[b]-row[a])/1e6 if a in row and b in row else None
    points.append(point)
lookup={(r['camera'],r['run'],r['pts']):r for r in points}
redis_rows=[];redis_origin={};unmatched=0
for line in (HERE/(label+'-redis.jsonl')).read_text().splitlines():
    row=json.loads(line)
    if row.get('kind')!='frame':continue
    key=(str(row['camera']),row['run'],row['pts']);native=lookup.get(key)
    if not native:unmatched+=1;continue
    wall=int(row['redis_id'].split('-')[0]);mono=row['read_mono_ns']/1e9+(wall-row['read_wall_ms'])/1000
    group=(key[0],key[1],native['segment'])
    if group not in redis_origin:redis_origin[group]=(mono,native['media_elapsed_s'])
    first=redis_origin[group]
    redis_rows.append(dict(camera=key[0],run=key[1],pts=key[2],t=mono,
        redis_relative_s=mono-first[0]-(native['media_elapsed_s']-first[1]),
        ingress_to_redis_ms=wall-row['ingest_ms'],visible=row['visible']))
summary=[];lanes=[];windows=[]
for stage in stages:
    a,b=stage['samples'][0],stage['samples'][-1];seconds=b['t']-a['t'];level=stage['level']
    channels=[];totals={}
    for c,br in b['channels'].items():
        ar=a['channels'][c];av=ar['providers']['pose'];bv=br['providers']['pose'];ap=av['metadata_path'];bp=bv['metadata_path']
        ad,bd=diag(ap),diag(bp);pub0,pub1=delivery(av),delivery(bv)
        durs={}
        for k,v in bd.get('durations',{}).items():
            old=ad.get('durations',{}).get(k,{});n=delta(old,v,'count');ns=delta(old,v,'sum_ns')
            durs[k]=dict(n=n,mean_ms=ns/n/1e6 if n else None)
            agg=totals.setdefault(k,[0,0]);agg[0]+=n;agg[1]+=ns
        p=[x for x in points if x['camera']==c and a['t']<=x['t']<=b['t']]
        rr=[x for x in redis_rows if x['camera']==c and a['t']<=x['t']<=b['t']]
        lane=dict(stage=stage['index'],level=level,camera=c,seconds=seconds,
            complete_fps=delta(ap['providers']['pose'],bp['providers']['pose'],'responses')/seconds,
            accepted_fps=delta(ap,bp,'accepted_result')/seconds,
            finalized_fps=delta(ap,bp,'finalized')/seconds,redis_fps=delta(pub0,pub1,'count')/seconds,
            queue_expired=delta(ap,bp,'completed_queue_expired'),capacity_drop=delta(ap,bp,'completed_queue_drop'),
            input_timing_frames=len(p),redis_joined_frames=len(rr),durations=durs,
            source_drift=stats([x['source_exit_ns_relative_s'] for x in p]),
            ingress_drift=stats([x['ingress_ns_relative_s'] for x in p]),
            depay_drift=stats([x['depay_first_ns_relative_s'] for x in p]),
            redis_drift=stats([x['redis_relative_s'] for x in rr]),processing_ms=stats([x['ingress_to_redis_ms'] for x in rr]),
            queue_mean_at_probes={k:delta(ad.get('queue_levels',{}).get(k,{}),v,'total')/max(1,delta(ad.get('queue_levels',{}).get(k,{}),v,'count')) for k,v in bd.get('queue_levels',{}).items()},
            queue_max_since_start={k:v['maximum'] for k,v in bd.get('queue_levels',{}).items()},
            frame_counts={k:delta(ap['performance'].get('counts',{}),bp['performance'].get('counts',{}),k)
                          for k in ('decoded','selected','sampling_drop','source_queue_gap','source_analysis_gap')},
            sender_counters={k:delta(av['publisher_metrics'],bv['publisher_metrics'],k)
                             for k in ('sent_frames','replaced','expired','send_error')})
        channels.append(lane);lanes.append(lane)
        for i in range(math.ceil(seconds/5)):
            lo=a['t']+i*5;hi=min(b['t'],lo+5)
            pp=[x for x in p if lo<=x['t']<hi];rs=[x for x in rr if lo<=x['t']<hi]
            windows.append(dict(stage=stage['index'],level=level,camera=c,window=i,seconds=hi-lo,
                ingress_relative_mean_s=st.mean(x['ingress_ns_relative_s'] for x in pp) if pp else None,
                source_relative_mean_s=st.mean(x['source_exit_ns_relative_s'] for x in pp) if pp else None,
                depay_relative_mean_s=st.mean(x['depay_first_ns_relative_s'] for x in pp) if pp else None,
                redis_relative_mean_s=st.mean(x['redis_relative_s'] for x in rs) if rs else None))
    ra,rb=a['resources'],b['resources'];cpu=[]
    for group,tasks in rb['cpu'].items():
        old=ra['cpu'].get(group,{});pairs=[(old[t],v) for t,v in tasks.items() if t in old]
        cpu.append(dict(group=group,cpu_s=sum(v[0]-u[0] for u,v in pairs)/1e9,
            runnable_wait_thread_s=sum(v[1]-u[1] for u,v in pairs)/1e9,matched_threads=len(pairs)))
    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-u for u,v in zip(hx,hy)]
    ac=dict(x.split() for x in ra['cgroup']['cpu.stat'].splitlines());bc=dict(x.split() for x in rb['cgroup']['cpu.stat'].splitlines())
    ta=ra['triton']['RFDETRKeypointPreview'];tb=rb['triton']['RFDETRKeypointPreview']
    old={r['batch_size']:r for r in ta['batch_stats']}
    batch={str(r['batch_size']):int(r['compute_infer']['count'])-int(old.get(r['batch_size'],{}).get('compute_infer',{}).get('count',0)) for r in tb['batch_stats']}
    triton={}
    for k,v in tb['inference_stats'].items():
        old=ta['inference_stats'][k];n=int(v['count'])-int(old['count']);ns=int(v['ns'])-int(old['ns'])
        triton[k]=dict(count=n,mean_ms=ns/n/1e6 if n else None)
    summary.append(dict(stage=stage['index'],level=level,seconds=seconds,
        complete_mean=st.mean(x['complete_fps'] for x in channels),redis_mean=st.mean(x['redis_fps'] for x in channels),
        redis_min=min(x['redis_fps'] for x in channels),redis_std=st.pstdev(x['redis_fps'] for x in channels),
        expired=sum(x['queue_expired'] for x in channels),capacity_drop=sum(x['capacity_drop'] for x in channels),
        durations={k:dict(count=v[0],mean_ms=v[1]/v[0]/1e6 if v[0] else None) for k,v in totals.items()},
        cpu_busy_pct=100*(sum(hd)-hd[3]-hd[4])/sum(hd),cpu_steal_pct=100*hd[7]/sum(hd),
        quota=rb['cgroup']['cpu.max'],throttled_us=int(bc.get('throttled_usec',0))-int(ac.get('throttled_usec',0)),
        gpu_util_mean=st.mean(float(s['gpu'].split(',')[0]) for s in stage['samples']),cpu=cpu,
        batch_executions=batch,triton=triton,
        artifact_counts={k:delta(a['artifacts']['counts'],b['artifacts']['counts'],k) for k in b['artifacts']['counts']},
        renderer_peak_sampled=max(s['artifacts']['rendering'] for s in stage['samples'])))
value=dict(label=label,summary=summary,lanes=lanes,resets=resets,redis_unmatched=unmatched,
    notes=['Source probe is depay sink AFTER jitterbuffer, not NIC receive. Source exit is decoded source-bin ghost pad BEFORE mux.',
           'Drift anchors are first observed valid samples per run/segment, not absolute source capture time.',
           'Queue levels are probe-weighted. Maxima are since process start, not stable-window maxima.',
           ('Private timeline is bounded64 in baseline, bounded8/8KiB after IPC fix; source/Redis joined distributions exclude unmatched frames.')],
    points=len(points),redis_joined=len(redis_rows))
for name,data in [('analysis.json',value),('input-points.json',points),('redis-points.json',redis_rows),('windows.json',windows)]:
    with (folder/name).open('w') as f:json.dump(data,f,indent=2)
with (folder/'windows.csv').open('w',newline='') as f:
    writer=csv.DictWriter(f,fieldnames=list(windows[0]) if windows else ['stage']);writer.writeheader();writer.writerows(windows)
print(json.dumps([dict(level=s['level'],stage=s['stage'],complete=s['complete_mean'],redis=s['redis_mean'],expired=s['expired'],cpu=s['cpu_busy_pct'],gpu=s['gpu_util_mean'],
    durations={k:v['mean_ms'] for k,v in s['durations'].items() if k in ('depay_to_source_exit','source_exit_to_ingress','preprocess_start_to_preprocess_output','preprocess_output_to_infer_entry','preprocess_to_infer_output','result_total_wall')}) for s in summary]))