import numpy as np

from behavision.detection import Detection
from behavision.tracking import IouTracker

KPS = np.zeros((5, 2), dtype=np.float32)


def det(box, score=0.9, quality=0.5):
    return Detection(box=box, kps=KPS, score=score, quality=quality)


def test_same_face_keeps_one_track():
    tracker = IouTracker(iou_threshold=0.3, max_misses=3)
    for i in range(5):
        active, ended = tracker.update([det((100 + i, 100, 200 + i, 200))])
    assert len(active) == 1
    assert active[0].hits == 5
    assert not ended


def test_two_faces_two_tracks():
    tracker = IouTracker()
    active, _ = tracker.update([det((0, 0, 50, 50)), det((300, 300, 350, 350))])
    assert len(active) == 2
    assert active[0].id != active[1].id


def test_track_ends_after_max_misses():
    tracker = IouTracker(max_misses=2)
    tracker.update([det((0, 0, 50, 50))])
    ended_all = []
    for _ in range(4):
        _, ended = tracker.update([])
        ended_all += ended
    assert len(ended_all) == 1
    assert not tracker.tracks


def test_best_quality_is_retained():
    tracker = IouTracker()
    tracker.update([det((0, 0, 50, 50), quality=0.8)])
    active, _ = tracker.update([det((1, 1, 51, 51), quality=0.3)])
    assert active[0].quality == 0.3
    assert active[0].best_quality == 0.8