import bootstrap
from bootstrap import PRIOR,SOURCE
from dataclasses import replace
import json,os,itertools
from data_robust import OUT,read,paths,lines,reference,sha
from refined import Refined,augment
from features import Settings
from state import Episode
from evaluate_integrated import group_states

clip=os.environ['CLIP']; fp,bp,source=paths(clip)
expected=read(PRIOR/'reports/probe_robust_weak.json')['R5']
cfg=read(SOURCE/'candidates.json')['candidates'][0]['config']
engine=Refined(cfg,expected['config']); episode=Episode(replace(Settings(),display_hold=1.0)); holds={}
feedback=read(OUT/'feedback.json')
d=dict(labels={},feedback={},new_hit_frames=[],new_negative_events=[],lost_positive_frames=[],remaining_negative_frames=[])
dest=OUT/'download'/clip/'data';dest.mkdir(parents=True,exist_ok=True)
events=[];reds=[];changes=[]
with (dest/'R5_diagnostics.jsonl').open('w') as out:
 for row,old in itertools.zip_longest(lines(fp),lines(bp)):
  assert row is not None and old is not None
  assert all(row[k]==old[k] for k in ('sample','source_frame','seconds','displayed'))
  row=augment(row);r=engine.step(row);ep=episode.step(row['seconds'],r['hit'],r['quiet'])
  tag,ids=reference(clip,old,feedback)
  result=dict(sample=row['sample'],source_frame=row['source_frame'],seconds=row['seconds'],displayed=row['displayed'],review_frame=row['review_frame'],reference=tag,episode=ep,groups=group_states(row,r,row['seconds'],holds),**r)
  out.write(json.dumps(result,separators=(',',':'))+'\n')
  if not row['displayed']:continue
  s=d['labels'].setdefault(tag,dict(N=0,observable=0,base_hit=0,new_hit=0,base_events=0,new_events=0))
  s['N']+=1;s['observable']+=r['observable'];s['base_hit']+=old['hit'];s['new_hit']+=r['hit'];s['base_events']+=old['episode']['event'];s['new_events']+=ep['event']
  frame=row['source_frame']
  if r['hit']:reds.append(frame)
  if ep['event']:events.append(frame)
  if r['hit'] and not old['hit']:d['new_hit_frames'].append(frame)
  if tag=='positive_reference' and old['hit'] and not r['hit']:d['lost_positive_frames'].append(frame)
  if tag=='negative_reference' and r['hit']:d['remaining_negative_frames'].append(frame)
  if tag=='negative_reference' and ep['event'] and not old['episode']['event']:d['new_negative_events'].append(frame)
  if r['hit']!=old['hit'] or ep['event']!=old['episode']['event']:changes.append(dict(frame=frame,seconds=row['seconds'],baseline_hit=old['hit'],r5_hit=r['hit'],baseline_event=old['episode']['event'],r5_event=ep['event']))
  for key in ids:
   f=d['feedback'].setdefault(key,dict(N=0,base_hit=0,new_hit=0,base_events=0,new_events=0))
   f['N']+=1;f['base_hit']+=old['hit'];f['new_hit']+=r['hit'];f['base_events']+=old['episode']['event'];f['new_events']+=ep['event']
assert d==expected['clips'][clip],(clip,'R5 reproduction mismatch')
(dest/'replay_validation.json').write_text(json.dumps(dict(status='PASS',comparison=d,red_frames=reds,event_frames=events,changes=changes,trace_sha256=sha(dest/'R5_diagnostics.jsonl')),indent=2))
print('REPLAY_PASS',clip,flush=True)