"""B0-B3 ablation on the synthetic scenarios of ref_impl.py (multi-seed). Synthetic logic verification only."""
import numpy as np, json
import ref_impl as R
FPS=R.FPS
NAMES=['normal','talking','brief_head','single_yawn','repeated_yawns','occlusion','narrow_eyed','sustained_closure','gradual_fatigue']
FAT={'sustained_closure','gradual_fatigue'}
def sec_flags(frame_flag):
    n=len(frame_flag)//FPS
    return frame_flag[:n*FPS].reshape(n,FPS).any(axis=1)
def engine_alarm(e,ear,mar,pit,valid):
    n=len(ear); out=np.zeros(n//FPS,bool)
    for i in range(n):
        st,_=e.step(ear[i],mar[i],pit[i],valid[i])
        if (i+1)%FPS==0: out[(i+1)//FPS-1]=(st=='DROWSY')
    return out
def conditions(name,seed):
    t,ear,mar,pit,valid,onset=R.scenario(name,seed)
    E=np.median(ear[:60*FPS])
    res={}
    res['B0']=sec_flags((ear<0.2)&valid)
    res['B1']=sec_flags((ear<R.P['alpha']*E)&valid)
    e2=R.Engine(); e2.base=dict(E=1.0,pit=0.0,bpm=15.0); e2.state='Alert'   # population constants: EAR<0.2, 15 blinks/min, pitch 0
    res['B2']=engine_alarm(e2,ear,mar,pit,valid)
    res['B3']=engine_alarm(R.Engine(),ear,mar,pit,valid)
    nsec=len(res['B0']); sec=np.arange(nsec)
    lab=np.zeros(nsec,bool)
    if onset is not None: lab[sec>=onset]=True
    return res,lab,onset
def edges(a): return int(np.sum(a[1:]&~a[:-1])+ (1 if a[0] else 0))
def main(nseeds=30):
    conds=['B0','B1','B2','B3']
    per_seed=[]; scen={c:{n:[] for n in NAMES} for c in conds}
    for seed in range(nseeds):
        tp={c:0 for c in conds}; fp=tp.copy(); fn=tp.copy(); fpev={c:0 for c in conds}; negsec=0; lead={c:[] for c in conds}
        for n in NAMES:
            res,lab,onset=conditions(n,seed)
            ev=np.arange(len(lab))>=60            # evaluation period after calibration
            negsec+=int(np.sum(ev&~lab))
            for c in conds:
                a=res[c]
                tp[c]+=int(np.sum(a&lab&ev)); fp[c]+=int(np.sum(a&~lab&ev)); fn[c]+=int(np.sum(~a&lab&ev))
                neg=a&~lab; neg=neg&ev
                fpev[c]+=edges(neg)
                if onset is not None:
                    idx=np.where(a&lab)[0]; lead[c].append(float(idx[0]-onset) if len(idx) else np.nan)
                scen[c][n].append(float(np.mean(a[ev])))
        row={}
        for c in conds:
            p=tp[c]/(tp[c]+fp[c]) if tp[c]+fp[c]>0 else np.nan
            r=tp[c]/(tp[c]+fn[c]); f=2*p*r/(p+r) if (p==p and p+r>0) else np.nan
            row[c]=dict(P=p,R=r,F1=f,FAh=fpev[c]/(negsec/3600),lead=float(np.nanmean(lead[c])) if not np.all(np.isnan(lead[c])) else np.nan,lead_det=int(np.sum(~np.isnan(lead[c]))))
        per_seed.append(row)
    summ={}
    for c in conds:
        summ[c]={}
        for k in ['P','R','F1','FAh','lead']:
            v=np.array([r[c][k] for r in per_seed],float)
            v=v[~np.isnan(v)]
            summ[c][k]=dict(mean=float(v.mean()),lo=float(np.percentile(v,2.5)),hi=float(np.percentile(v,97.5)),n=int(len(v)))
        summ[c]['detected_runs']=int(sum(r[c]['lead_det'] for r in per_seed))
    json.dump(dict(summary=summ,nseeds=nseeds,frac={c:{n:float(np.mean(scen[c][n])) for n in NAMES} for c in conds}),open('ablation_results.json','w'),indent=1)
    for c in conds: print(c,{k:(round(v['mean'],3),round(v['lo'],3),round(v['hi'],3),v['n']) for k,v in summ[c].items() if k!='detected_runs'},summ[c]['detected_runs'])
    for c in conds: print(c,{n:round(float(np.mean(scen[c][n])),3) for n in NAMES})
main()