"""Loose spatial candidates and short pair correspondence, not a replacement Tracker."""
from collections import deque
from dataclasses import dataclass,asdict
from itertools import combinations
from math import dist,hypot,isfinite


@dataclass(frozen=True)
class Geometry:
    person_score:float=.25
    joint_score:float=.10
    enter_gap:float=.35
    exit_gap:float=.60
    enter_center:float=1.25
    exit_center:float=1.60
    candidate_hold:float=1.0
    spatial_grace:float=.30
    step_gap:float=.15
    correspondence_distance:float=.45
    correspondence_size:float=1.8
    ambiguity:float=.06
    window:float=1.0
    max_people:int=32
    max_pairs:int=128
    max_states:int=256

    def payload(self):return asdict(self)


def point(k,w,h,threshold=.10):
    return len(k)==3 and all(isfinite(x) for x in k) and 0<=k[0]<=w and 0<=k[1]<=h and threshold<=k[2]<=1


def bodies(row,cfg):
    result=[];reject=[]
    for order,p in enumerate(row['persons']):
        b=p['bbox'];s=p['score']
        if not isfinite(s) or not 0<=s<=1:reason='INVALID_SCORE'
        elif s<cfg.person_score:reason='PERSON_SCORE_BELOW'
        elif len(b)!=4 or not all(isfinite(x) for x in b) or b[0]>=b[2] or b[1]>=b[3] or b[2]<=0 or b[3]<=0 or b[0]>=row['width'] or b[1]>=row['height']:reason='INVALID_BBOX'
        elif len(result)>=cfg.max_people:reason='PEOPLE_CAP'
        else:reason=None
        if reason:reject.append(dict(order=order,reason=reason));continue
        joints={i:list(k) for i,k in enumerate(p['keypoints']) if point(k,row['width'],row['height'],cfg.joint_score)}
        result.append(dict(order=order,tid=p['track_id'],query=p['query'],bbox=list(b),score=s,
            center=[(b[0]+b[2])/2,(b[1]+b[3])/2],height=b[3]-b[1],joints=joints,
            ref='T'+str(p['track_id']) if p['track_id'] is not None else 'q'+str(p['query'])))
    return result,reject


def relation(a,b):
    x,y=a['bbox'],b['bbox'];unit=max(1,min(a['height'],b['height']))
    dx=max(0,x[0]-y[2],y[0]-x[2]);dy=max(0,x[1]-y[3],y[1]-x[3])
    return dict(gap=hypot(dx,dy)/unit,center=dist(a['center'],b['center'])/max(1,(a['height']+b['height'])/2),
                height_ratio=max(a['height'],b['height'])/unit)


def near(rel,cfg):return rel['gap']<=cfg.enter_gap and rel['center']<=cfg.enter_center
def inside_exit(rel,cfg):return rel['gap']<=cfg.exit_gap and rel['center']<=cfg.exit_center
def valid_ids(pair):return all(p['tid'] is not None for p in pair) and pair[0]['tid']!=pair[1]['tid']


def cost(old,current,cfg):
    distances=[]
    for a,b in zip(old,current):
        if min(a['height'],b['height'])<=0 or max(a['height']/b['height'],b['height']/a['height'])>cfg.correspondence_size:return None
        d=dist(a['center'],b['center'])/max(1,(a['height']+b['height'])/2)
        if d>cfg.correspondence_distance:return None
        distances.append(d)
    return max(distances)


class Spatial:
    def __init__(self,cfg=None):
        self.cfg=cfg or Geometry();self.states={};self.serial=0;self.last_time=None

    def step(self,row):
        t=row['seconds'];cfg=self.cfg
        assert isfinite(t) and (self.last_time is None or t>self.last_time)
        self.last_time=t;bs,reject=bodies(row,cfg)
        ids={p['tid']:p for p in bs if p['tid'] is not None}
        expired=[]
        for token,s in list(self.states.items()):
            old=s['bodies'];far=False
            if valid_ids(old) and all(p['tid'] in ids for p in old):
                far=not inside_exit(relation(*(ids[p['tid']] for p in old)),cfg)
            if far or t-s['last_near']>cfg.candidate_hold+1e-8:
                expired.append(dict(token=token,reason='BEYOND_EXIT' if far else 'CANDIDATE_EXPIRED'))
                del self.states[token]
        offered=[];overcap=0
        for pair in combinations(bs,2):
            if pair[0]['tid'] is not None and pair[0]['tid']==pair[1]['tid']:
                reject.append(dict(orders=[p['order'] for p in pair],reason='SAME_T_NOT_TWO_PEOPLE'));continue
            rel=relation(*pair)
            if not inside_exit(rel,cfg):continue
            if len(offered)>=cfg.max_pairs:overcap+=1;continue
            offered.append((pair,rel))
        choices={}
        for i,(pair,rel) in enumerate(offered):
            opts=[]
            for token,s in self.states.items():
                old=s['bodies'];stable=valid_ids(old) and valid_ids(pair) and {x['tid'] for x in old}=={x['tid'] for x in pair}
                if not stable and t-s['last_seen']>cfg.spatial_grace+1e-8:continue
                orientations=[]
                for current in (pair,pair[::-1]):
                    if stable and [p['tid'] for p in current]!=[p['tid'] for p in old]:continue
                    c=cost(old,current,cfg)
                    if c is not None:orientations.append((c,current))
                orientations.sort(key=lambda x:x[0])
                if not orientations or (len(orientations)>1 and orientations[1][0]-orientations[0][0]<cfg.ambiguity):continue
                c,current=orientations[0];opts.append((0 if stable else 1,c,token,current))
            opts.sort(key=lambda x:(x[0],x[1],x[2]));choices[i]=opts
        reverse={}
        for i,opts in choices.items():
            for rank,c,token,_ in opts:reverse.setdefault(token,[]).append((rank,c,i))
        output=[];used=set();ambiguous=0
        for i,(pair,rel) in enumerate(offered):
            opts=choices[i];match=None
            if opts:
                rank,c,token,current=opts[0];contenders=sorted(reverse[token])
                own=len(opts)==1 or opts[1][0]>rank or opts[1][1]-c>=cfg.ambiguity
                mutual=contenders[0][2]==i and (len(contenders)==1 or contenders[1][0]>rank or contenders[1][1]-c>=cfg.ambiguity)
                if own and mutual and token not in used:match=(token,current,'STABLE_T_AND_GEOMETRY' if rank==0 else 'UNAMBIGUOUS_SPATIAL')
                else:ambiguous+=1
            entering=near(rel,cfg)
            if match is None:
                if not entering or len(self.states)>=cfg.max_states:continue
                self.serial+=1;token='S'+str(self.serial);current=pair;mode='NEW_CANDIDATE'
                self.states[token]=dict(bodies=current,last_seen=t,last_near=t,history=deque(maxlen=40))
                reset=True
            else:
                token,current,mode=match;s=self.states[token]
                reset=t-s['last_seen']>cfg.step_gap+1e-8
                if reset:s['history'].clear()
            s=self.states[token]
            if entering:s['last_near']=t
            used.add(token)
            history=list(s['history'])
            output.append(dict(token=token,bodies=current,relation=rel,mode=mode,reset=reset,
                candidate_state='CURRENT_NEAR' if entering else 'GEOMETRY_HOLD',
                hold_age=t-s['last_near'],history=history))
            s['history'].append((t,current));s.update(bodies=current,last_seen=t)
        return dict(bodies=bs,rejected=reject,pairs=output,expired=expired,ambiguous=ambiguous,
                    pairs_over_cap=overcap,state_count=len(self.states),retained_absent=len(set(self.states)-used))