"""Three bounded candidate variants, using frozen features and stored baseline."""
import bootstrap
from bootstrap import SOURCE
from collections import Counter
from dataclasses import replace
import itertools
import json
import os
from data_robust import OUT,read,paths,lines,reference
from refined import Refined,augment
from features import Settings
from state import Episode

options=[dict(id='R1',compact_hand=.12),dict(id='R2',compact_hand=.14,self_head=.08),
         dict(id='R3',compact_hand=.14,self_head=.08,mutual_hand=.18)]
suffix=''
if os.environ.get('CLIP')=='weak':
    options=[dict(id='R4',compact_hand=.12,compact_max_speed=1.0,compact_max_motion=.10),
             dict(id='R5',compact_hand=.12,compact_max_speed=1.0)]
    suffix='_weak'
cfg=read(SOURCE/'candidates.json')['candidates'][0]['config']
feedback=read(OUT/'feedback.json');results={v['id']:dict(config=v,clips={}) for v in options}
for clip in read(OUT/'specs.json'):
    fp,bp,_=paths(clip)
    engines=[Refined(cfg,v) for v in options]
    episodes=[Episode(replace(Settings(),display_hold=1.0)) for _ in options]
    stats=[dict(labels={},feedback={},new_hit_frames=[],new_negative_events=[],lost_positive_frames=[],remaining_negative_frames=[]) for _ in options]
    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)
        tag,ids=reference(clip,old,feedback)
        for i,engine in enumerate(engines):
            r=engine.step(row);ep=episodes[i].step(row['seconds'],r['hit'],r['quiet'])
            if not row['displayed']:continue
            d=stats[i];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']
            if r['hit'] and not old['hit']:d['new_hit_frames'].append(row['source_frame'])
            if tag=='positive_reference' and old['hit'] and not r['hit']:d['lost_positive_frames'].append(row['source_frame'])
            if tag=='negative_reference' and r['hit']:d['remaining_negative_frames'].append(row['source_frame'])
            if tag=='negative_reference' and ep['event'] and not old['episode']['event']:d['new_negative_events'].append(row['source_frame'])
            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']
    for i,opt in enumerate(options):results[opt['id']]['clips'][clip]=stats[i]
    print('PROBE',clip,{v['id']:{k:(x['base_hit'],x['new_hit']) for k,x in stats[i]['labels'].items()} for i,v in enumerate(options)},flush=True)
(OUT/('reports/probe_robust'+suffix+'.json')).write_text(json.dumps(results,indent=2))
print('PROBE_COMPLETE',len(results),'variants',flush=True)