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


def test_crossing_keeps_ids_via_global_best_match():
    # Track-iteration-order greedy let track 1 steal the box that overlaps track 2 far more
    # (0.48 vs 0.74), leaving track 2 unmatched and its person re-created as a NEW id —
    # FALLEN state and frozen baseline would ride along to the wrong person. Global greedy
    # must assign the highest-IoU pair first.
    tr = IouTracker(iou_thr=0.3)
    tr.update([_person(0, 0, 10, 20), _person(5, 0, 15, 20)], t=0.0)
    tracks = tr.update([_person(3.5, 0, 13.5, 20), _person(-4, 0, 6, 20)], t=0.1)
    by_id = {t.track_id: t for t in tracks}
    assert set(by_id) == {1, 2}          # no phantom new track
    assert by_id[2].latest.bbox.x1 == 3.5
    assert by_id[1].latest.bbox.x1 == -4


def test_contended_match_is_flagged_ambiguous():
    # One person overlapping BOTH tracks (a crossing): whichever track wins, the match may
    # be a swap, so every contended track must be reported for a safety-state reset.
    tr = IouTracker(iou_thr=0.3)
    tr.update([_person(0, 0, 10, 20), _person(5, 0, 15, 20)], t=0.0)
    tr.update([_person(2, 0, 12, 20)], t=0.1)
    assert tr.ambiguous() == {1, 2}


def test_clean_match_is_not_ambiguous():
    tr = IouTracker(iou_thr=0.3)
    tr.update([_person(0, 0, 10, 20), _person(100, 0, 110, 20)], t=0.0)
    tr.update([_person(1, 1, 11, 21), _person(101, 1, 111, 21)], t=0.1)
    assert tr.ambiguous() == set()