from edge.models import Person, BBox, Keypoint
from edge.tracking import IouTracker


def _person(x1, y1, x2, y2):
    kps = [Keypoint((x1 + x2) / 2, (y1 + y2) / 2, 0.9) for _ in range(17)]
    return Person(BBox(x1, y1, x2, y2), kps)


def test_new_person_gets_track_id():
    tr = IouTracker()
    tracks = tr.update([_person(0, 0, 10, 20)], t=0.0)
    assert len(tracks) == 1
    assert tracks[0].track_id == 1


def test_overlapping_person_keeps_same_id():
    tr = IouTracker(iou_thr=0.3)
    tr.update([_person(0, 0, 10, 20)], t=0.0)
    tracks = tr.update([_person(1, 1, 11, 21)], t=0.1)  # high overlap
    assert tracks[0].track_id == 1


def test_disappeared_track_is_reported_lost():
    tr = IouTracker(max_age=0.2)
    tr.update([_person(0, 0, 10, 20)], t=0.0)
    tr.update([], t=0.5)  # nothing detected, past max_age
    lost = tr.lost()
    assert len(lost) == 1
    assert lost[0].track_id == 1


def test_history_accumulates_frames():
    tr = IouTracker()
    tr.update([_person(0, 0, 10, 20)], t=0.0)
    tracks = tr.update([_person(0, 0, 10, 20)], t=0.1)
    assert len(tracks[0].frames) == 2
    assert tracks[0].latest.t == 0.1