"""A: ID-keyed pair episodes. B: bounded spatial correspondence + ROI episode.

This is a review prototype, not a drop-in production engine. No I/O or public IDs.
"""
from itertools import combinations
from math import dist,isfinite
from features import Settings,bodies,match_cost
from state import PairState,Episode


class Engine:
    def __init__(self,mode,cfg=None):
        assert mode in ('A','B')
        self.mode=mode;self.cfg=cfg or Settings();self.states={};self.serial=0
        self.roi=Episode(self.cfg);self.last_time=None;self.last_sample=None;self.geometry=None
        self.last_event=None;self.event_count=0

    def _new(self):
        self.serial+=1;token='S'+str(self.serial)
        state=PairState(token,self.cfg);self.states[token]=state;return state

    def step(self,row,prepared=None):
        t=row['seconds'];sample=row['sample'];cfg=self.cfg
        if not isfinite(t) or t<0 or (self.last_time is not None and (t<=self.last_time or sample<=self.last_sample)):
            raise ValueError('NON_MONOTONIC_OBSERVATION')
        geometry=(row['width'],row['height'])
        if self.geometry is not None and geometry!=self.geometry:raise ValueError('GEOMETRY_CHANGED_CREATE_NEW_ENGINE')
        self.geometry=geometry;self.last_time=t;self.last_sample=sample
        retention=cfg.pair_retention if self.mode=='A' else cfg.pair_match_grace+1e-8
        expired=[key for key,s in self.states.items() if s.last_seen is not None and t-s.last_seen>retention]
        for key in expired:del self.states[key]
        bs,rejected=bodies(row,cfg) if prepared is None else prepared
        rejected=list(rejected)
        if self.mode=='A':
            rejected += [dict(order=b['order'],reason='A_REQUIRES_TRACK_ID') for b in bs if b['tid'] is None]
            bs=sorted((b for b in bs if b['tid'] is not None),key=lambda b:b['tid'])
            assert len({b['tid'] for b in bs})==len(bs)
        pairs=[];overcap=0
        for pair in combinations(bs,2):
            if dist(pair[0]['center'],pair[1]['center'])/((pair[0]['scale']+pair[1]['scale'])/2)>cfg.pair_distance_ratio:continue
            if len(pairs)>=cfg.max_pairs:overcap+=1;continue
            pairs.append(pair)
        assignments={};ambiguous=0
        if self.mode=='A':
            for i,pair in enumerate(pairs):
                key=tuple(b['tid'] for b in pair)
                if key not in self.states:
                    if len(self.states)>=cfg.max_states:continue
                    self.states[key]=PairState('T'+str(key[0])+'+T'+str(key[1]),cfg)
                assignments[i]=(self.states[key],pair)
        else:
            options={};reverse={}
            for i,pair in enumerate(pairs):
                choices=[]
                for key,s in self.states.items():
                    if s.last_seen is None or t-s.last_seen>cfg.pair_match_grace+1e-8:continue
                    orient=[]
                    for current in (pair,pair[::-1]):
                        cost=match_cost(s.last_bodies,current,cfg)
                        if cost is not None:orient.append((cost,current))
                    orient.sort(key=lambda x:x[0])
                    if not orient or (len(orient)>1 and orient[1][0]-orient[0][0]<cfg.spatial_ambiguity_margin):continue
                    cost,current=orient[0];choices.append((cost,key,current));reverse.setdefault(key,[]).append((cost,i))
                choices.sort(key=lambda x:(x[0],x[1]));options[i]=choices
            for i,pair in enumerate(pairs):
                choices=options[i];chosen=None
                if choices:
                    cost,key,current=choices[0];contenders=sorted(reverse[key])
                    certain=(len(choices)==1 or choices[1][0]-cost>=cfg.spatial_ambiguity_margin)
                    mutual=contenders[0][1]==i and (len(contenders)==1 or contenders[1][0]-contenders[0][0]>=cfg.spatial_ambiguity_margin)
                    if certain and mutual:chosen=(self.states[key],current)
                    else:ambiguous+=1
                if chosen is None and len(self.states)<cfg.max_states:chosen=(self._new(),pair)
                if chosen:assignments[i]=chosen
        decisions=[];events=[];used=set()
        for i in range(len(pairs)):
            if i not in assignments:continue
            state,pair=assignments[i];used.add(id(state));d=state.step(t,pair)
            if self.mode=='A':
                allowed=self.last_event is None or t-self.last_event>=cfg.event_cooldown-1e-8
                ep=state.episode.step(t,d['hit'],d['quiet_valid'],allowed);d['episode']=ep
                if ep['event']:
                    self.last_event=t;self.event_count+=1
                    events.append(dict(test_event=self.event_count,pair=state.token,rules=d['rules'],seconds=t))
            decisions.append(d)
        if self.mode=='A':
            for s in self.states.values():
                if id(s) not in used:s.episode.step(t,False,False)
            hit=any(d['hit'] for d in decisions);display=any(d['episode']['display'] for d in decisions)
            latched=any(s.episode.latched for s in self.states.values())
            valid=bool(decisions) and all(d['quiet_valid'] for d in decisions)
            status='ALERT' if events else 'RULE_HIT' if hit else 'DISPLAY_HOLD' if display else 'LATCHED_WAIT' if latched else 'MONITORING' if valid else 'WARMUP_OR_INSUFFICIENT'
            global_ep=dict(status=status,event=bool(events),display=display,latched=latched)
        else:
            hit=any(d['hit'] for d in decisions)
            # Unresolved spatial candidates cannot certify observed quiet.
            valid=bool(decisions) and not ambiguous and not overcap and len(decisions)==len(pairs) and all(d['quiet_valid'] for d in decisions)
            global_ep=self.roi.step(t,hit,valid)
            if global_ep['event']:
                self.event_count+=1;events=[dict(test_event=self.event_count,pair=None,rules=sorted({r for d in decisions for r in d['rules']}),seconds=t)]
        return dict(sample=sample,source_frame=row['source_frame'],seconds=t,mode=self.mode,episode=global_ep,
            current_rules=sorted({r for d in decisions for r in d['rules']}),events=events,event_count=self.event_count,
            eligible_bodies=len(bs),candidate_pairs=len(pairs),pairs_over_cap=overcap,state_count=len(self.states),
            expired_states=len(expired),ambiguous_matches=ambiguous,unassigned_pairs=len(pairs)-len(decisions),
            rejected=rejected,pairs=decisions,causal=True)