# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "opencv-python",
#     "pupil-labs-video",
#     "pupil-labs-neon-recording",
#     "ultralytics",
# ]
# ///

from datetime import datetime, timezone
from pathlib import Path

import cv2
from tqdm import tqdm
from ultralytics import YOLO

import pupil_labs.neon_recording as nr

# RECORDING_DIR = Path("path/to/recording")
RECORDING_DIR = Path("/Users/mgg/Desktop/recs/Skii_2025-02-10-16-35-19-a89e7118")


def main():
    # Get nano version of YOLO 11
    detector = YOLO("yolo11n-seg.pt")

    # Load recording
    recording = nr.load(RECORDING_DIR)
    combined_data = zip(
        recording.scene.ts,
        recording.scene.sample(recording.scene.ts),
        recording.gaze.sample(recording.scene.ts),
        strict=False,
    )
    for ts, scene_frame, gaze_datum in tqdm(
        combined_data, total=len(recording.scene.ts)
    ):
        frame_pixels = scene_frame.bgr
        detection = detector(frame_pixels)
        frame_with_detection = detection[0].plot()

        final_frame = cv2.circle(
            frame_with_detection,
            (int(gaze_datum.x), int(gaze_datum.y)),
            10,
            (0, 0, 255),
            5,
        )
        # Add timestamp text on the bottom left
        font = cv2.FONT_HERSHEY_SIMPLEX
        bottom_left_corner_of_text = (10, final_frame.shape[0] - 10)
        font_scale = 1
        font_color = (255, 255, 255)
        line_type = 2

        time_str = datetime.fromtimestamp(ts // 1e9, timezone.utc).strftime(
            "%Y-%m-%d %H:%M:%S.%f"
        )

        cv2.putText(
            final_frame,
            f"{time_str}",
            bottom_left_corner_of_text,
            font,
            font_scale,
            font_color,
            line_type,
        )
        cv2.imshow("frame", final_frame)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break


if __name__ == "__main__":
    main()