import copy
import unittest
from geometry import Spatial,Geometry,bodies,relation
from motion import extract,pair_features
from decision import Candidate,passes
from seed_follow import SeedFollow
import sys
sys.path.insert(0,'/baseline')
from features import Settings
from state import Episode


def person(tid,x,score=.25):
    joints=[[x+50,60,.1] for _ in range(17)]
    for j,xy in {0:(50,15),5:(30,40),6:(70,40),7:(20,80),8:(80,80),9:(10,110),10:(90,110),
                 11:(40,110),12:(60,110),13:(40,150),14:(60,150),15:(40,195),16:(60,195)}.items():
        joints[j]=[x+xy[0],xy[1],.1]
    return dict(track_id=tid,query=tid or 0,bbox=[x,0,x+100,200],score=score,keypoints=joints,valid_joints=0)


def row(t,people):return dict(seconds=t,sample=round(t*10),source_frame=round(t*30)+1,width=1200,height=800,persons=people)


def feature(**kw):
    v=dict(kind='hand',vertical=.2,speed=.8,near=.02,approach=.08,elapsed=.2,cosine=.8,extension=.04,actor=0)
    v.update(kw);return v


CFG=dict(speed_min=.5,near_max=.1,approach_speed_min=.2,cosine_min=.45,
         extension_speed_min=.05,foot_vertical_max=.4,min_samples=2,window=1.)


def decision_row(t,features,reset=False):
    return dict(seconds=t,ambiguous=0,pairs_over_cap=0,pairs=[dict(token='S1',features=features,refs=['T1','T2'],
        orders=[0,1],mode='STABLE_T_AND_GEOMETRY',candidate_state='CURRENT_NEAR',hold_age=0,reset=reset)])


class Tests(unittest.TestCase):
    def test_score_inclusive(self):
        bs,_=bodies(row(0,[person(1,100)]),Geometry());self.assertEqual(len(bs),1);self.assertEqual(len(bs[0]['joints']),17)
        self.assertFalse(bodies(row(0,[person(1,100,.249999)]),Geometry())[0])

    def test_joint_below_invalid(self):
        p=person(1,100);p['keypoints'][9][2]=.099999;p['keypoints'][10][0]=float('nan')
        bs,_=bodies(row(0,[p]),Geometry());self.assertNotIn(9,bs[0]['joints']);self.assertNotIn(10,bs[0]['joints'])

    def test_single_person(self):self.assertFalse(Spatial().step(row(0,[person(1,100)]))['pairs'])

    def test_duplicate_id(self):self.assertFalse(Spatial().step(row(0,[person(1,100),person(1,130)]))['pairs'])

    def test_same_t_reordered(self):
        e=Spatial();a=e.step(row(0,[person(1,100),person(2,230)]))['pairs'][0]
        b=e.step(row(.1,[person(2,230),person(1,100)]))['pairs'][0]
        self.assertEqual(a['token'],b['token']);self.assertEqual([p['tid'] for p in b['bodies']],[1,2])

    def test_t_changed_spatial(self):
        e=Spatial();a=e.step(row(0,[person(1,100),person(2,230)]))['pairs'][0]
        b=e.step(row(.1,[person(3,100),person(4,230)]))['pairs'][0]
        self.assertEqual(a['token'],b['token']);self.assertEqual(b['mode'],'UNAMBIGUOUS_SPATIAL')

    def test_hold_expires(self):
        e=Spatial();e.step(row(0,[person(1,100),person(2,230)]))
        p=e.step(row(.1,[person(1,100),person(2,290)]))['pairs'][0]
        self.assertEqual(p['candidate_state'],'GEOMETRY_HOLD')
        self.assertTrue(e.step(row(.9,[person(1,100),person(2,290)]))['pairs'])
        self.assertFalse(e.step(row(1.1,[person(1,100),person(2,290)]))['pairs'])

    def test_far_release(self):
        e=Spatial();e.step(row(0,[person(1,100),person(2,230)]))
        v=e.step(row(.1,[person(1,100),person(2,500)]));self.assertFalse(v['pairs']);self.assertEqual(v['expired'][0]['reason'],'BEYOND_EXIT')

    def test_gap_no_motion(self):
        e=Spatial();e.step(row(0,[person(1,100),person(2,230)]))
        p=e.step(row(.3,[person(1,100),person(2,230)]))['pairs'][0]
        self.assertTrue(p['reset']);self.assertFalse(pair_features(p,.3)[0])

    def test_no_future(self):
        e=Spatial()
        for t in (0,.1,.2,.3):
            p=e.step(row(t,[person(1,100),person(2,230)]))['pairs'][0]
            self.assertTrue(all(f['input_start']<f['input_end']==t for f in pair_features(p,t)[0]))

    def test_translation_zero(self):
        a=bodies(row(0,[person(1,100),person(2,230)]),Geometry())[0]
        b=bodies(row(.1,[person(1,110),person(2,240)]),Geometry())[0]
        fs,_=extract(a,b,.1);self.assertTrue(fs);self.assertTrue(all(f['motion']<1e-9 for f in fs))

    def test_scale_zero(self):
        a=bodies(row(0,[person(1,100),person(2,230)]),Geometry())[0];b=copy.deepcopy(a)
        for p in b:
            p['height']*=1.1;p['bbox']=[x*1.1 for x in p['bbox']]
            p['joints']={j:[k[0]*1.1,k[1]*1.1,k[2]] for j,k in p['joints'].items()}
        fs,_=extract(a,b,.1);self.assertTrue(all(f['motion']<1e-8 for f in fs))

    def test_missing_not_zero(self):
        a=bodies(row(0,[person(1,100),person(2,230)]),Geometry())[0];b=copy.deepcopy(a)
        for p in b:p['joints']={}
        fs,why=extract(a,b,.1);self.assertFalse(fs);self.assertTrue(why)

    def test_duplicate_lags_one_sample(self):
        e=Candidate(CFG);v=e.step(decision_row(0,[feature(),feature()]));self.assertFalse(v['hit']);self.assertEqual(v['pairs'][0]['counts'],[1,0])
        self.assertTrue(e.step(decision_row(.1,[feature()]))['hit'])

    def test_other_actor_not_combined(self):
        e=Candidate(CFG);e.step(decision_row(0,[feature(actor=0)]))
        self.assertFalse(e.step(decision_row(.1,[feature(actor=1)]))['hit'])

    def test_current_evidence_required(self):
        e=Candidate(CFG);e.step(decision_row(0,[feature()]));e.step(decision_row(.1,[feature()]))
        self.assertFalse(e.step(decision_row(.2,[]))['hit'])

    def test_reset_clears(self):
        e=Candidate(CFG);e.step(decision_row(0,[feature()]));self.assertFalse(e.step(decision_row(.3,[feature()],True))['hit'])

    def test_expiry(self):
        e=Candidate(CFG);e.step(decision_row(0,[feature()]));self.assertFalse(e.step(decision_row(1.1,[feature()]))['hit'])

    def test_threshold_boundaries(self):
        self.assertTrue(passes(feature(speed=.5,near=.1,approach=.04,extension=.01),CFG))
        self.assertFalse(passes(feature(speed=.499999),CFG));self.assertFalse(passes(feature(near=.100001),CFG))

    def test_incidental_nearness_not_enough(self):
        self.assertFalse(passes(feature(speed=0,near=0,approach=0,extension=0),CFG))
        self.assertFalse(passes(feature(approach=-.1),CFG))

    def test_foot_geometry(self):
        self.assertFalse(passes(feature(kind='foot',vertical=.5),CFG))
        self.assertTrue(passes(feature(kind='foot',vertical=.3),CFG))

    def test_seed_follow_requires_real_current_motion(self):
        e=SeedFollow(dict(CFG,seed_seconds=.6))
        seed=feature(limb='LW',delta=[1.,0.])
        e.step(decision_row(0,[seed]))
        self.assertFalse(e.step(decision_row(.1,[]))['hit'])
        back=feature(limb='LW',delta=[-1.,0.],approach=-.05,cosine=-.8)
        v=e.step(decision_row(.2,[back]));self.assertTrue(v['hit'])
        self.assertEqual(v['pairs'][0]['passing'][0]['evidence_role'],'OBSERVED_RETURN_OR_CONTINUATION')

    def test_seed_follow_other_limb_not_seeded(self):
        e=SeedFollow(dict(CFG,seed_seconds=.6))
        e.step(decision_row(0,[feature(limb='LW',delta=[1.,0.])]))
        for actor,limb in ((0,'RW'),(1,'LW')):
            v=e.step(decision_row(.1 if actor==0 else .2,[feature(actor=actor,limb=limb,delta=[-1.,0.],approach=-.1,cosine=-.8)]))
            self.assertFalse(v['raw'])

    def test_seed_follow_expiry_not_renewed_by_returns(self):
        e=SeedFollow(dict(CFG,seed_seconds=.6))
        e.step(decision_row(0,[feature(limb='LW',delta=[1.,0.])]))
        back=feature(limb='LW',delta=[-1.,0.],approach=-.1,cosine=-.8)
        self.assertTrue(e.step(decision_row(.5,[back]))['raw'])
        self.assertFalse(e.step(decision_row(.7,[back]))['raw'])

    def test_seed_follow_reset(self):
        e=SeedFollow(dict(CFG,seed_seconds=.6))
        e.step(decision_row(0,[feature(limb='LW',delta=[1.,0.])]))
        self.assertFalse(e.step(decision_row(.1,[feature(limb='LW',delta=[-1.,0.],approach=-.1,cosine=-.8)],True))['raw'])

    def test_event_duplicate_display_and_unknown(self):
        e=Episode(Settings())
        self.assertTrue(e.step(0,True,False)['event'])
        self.assertFalse(e.step(.1,True,False)['event'])
        hold=e.step(.2,False,False)
        self.assertTrue(hold['display']);self.assertEqual(hold['quiet_seconds'],0)
        self.assertFalse(e.step(.6,False,False)['display'])
        expired=e.step(5.2,False,False)
        self.assertEqual(expired['transition'],'EXPIRED_NOT_NORMAL')
        self.assertFalse(expired['latched'])

    def test_decision_prefix_invariant(self):
        rows=[decision_row(t/10,[feature(limb='LW',delta=[1.,0.])]) for t in range(10)]
        for cls in (Candidate,SeedFollow):
            a=cls(dict(CFG,seed_seconds=.6));b=cls(dict(CFG,seed_seconds=.6))
            full=[a.step(r) for r in rows]
            prefix=[b.step(r) for r in rows[:5]]
            self.assertEqual(full[:5],prefix)


if __name__=='__main__':unittest.main(verbosity=2)