import supervision as sv
from pathlib import Path
from ultralytics import YOLO

VIDEO_FILE = Path('data', 'test-traffic.mp4')
VIDEO_OUT = Path('data', 'test-traffic-annotated.mp4')


# Only a regression test. PIL will never be used to open video files.
trace_annotator = sv.LabelAnnotator()

model_det = YOLO("yolov8n")
video_info = sv.VideoInfo.from_video_path(video_path=str(VIDEO_FILE))
frames_generator = sv.get_video_frames_generator(source_path=str(VIDEO_FILE))
tracker = sv.ByteTrack()
smoother = sv.DetectionsSmoother()

with sv.VideoSink(target_path=str(VIDEO_OUT), video_info=video_info) as sink:
    for frame in frames_generator:
        result = model_det(frame)[0]
        detections = sv.Detections.from_ultralytics(result)
        detections = tracker.update_with_detections(detections)
        smoother.update_with_detections(detections)
        detections = smoother.get_smoothed_detections()

        # print(detections.tracker_id, detections.class_id)
        labels = [
            f"#{tracker_id} {class_id}"
            for tracker_id, class_id
            in zip(detections.tracker_id, detections.class_id)
        ]

        annotated_frame = trace_annotator.annotate(
            scene=frame.copy(),
            detections=detections,
            labels=labels
        )
        sink.write_frame(frame=annotated_frame)