import cv2
import dlib
import time  # To measure time for FPS calculation

# GStreamer pipeline to acquire video from the CSI camera (Jetson Nano)
def gstreamer_pipeline(
    sensor_id=0,
    capture_width=1280,
    capture_height=720,
    display_width=640,
    display_height=360,
    framerate=15,
    flip_method=0,
):
    return (
        f"nvarguscamerasrc sensor-id={sensor_id} ! "
        f"video/x-raw(memory:NVMM), width=(int){capture_width}, height=(int){capture_height}, "
        f"format=(string)NV12, framerate=(fraction){framerate}/1 ! "
        f"nvvidconv flip-method={flip_method} ! "
        f"video/x-raw, width=(int){display_width}, height=(int){display_height}, format=(string)BGRx ! "
        f"videoconvert ! video/x-raw, format=(string)BGR ! appsink"
    )

# Capture from the camera using OpenCV and start tracking
cap = cv2.VideoCapture(gstreamer_pipeline(), cv2.CAP_GSTREAMER)

if not cap.isOpened():
    print("Unable to open camera")
else:
    # Read the first frame to allow user to draw a bounding box
    ret, frame = cap.read()
    if not ret:
        print("Failed to grab the initial frame")
        cap.release()

    # Let the user draw the initial bounding box
    bbox = cv2.selectROI("Select ROI", frame, fromCenter=False, showCrosshair=True)

    # Initialize Dlib correlation tracker
    tracker = dlib.correlation_tracker()
    tracker.start_track(frame, dlib.rectangle(bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]))

    img_counter = 0  # To count frames and name images
    frame_counter = 0  # To track the number of frames
    start_time = time.time()  # Start time for FPS calculation

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                print("Failed to grab frame")
                break

            # Start measuring time per frame
            frame_start_time = time.time()

            # Update the tracker with the current frame
            tracker.update(frame)
            pos = tracker.get_position()

            # Extract the bounding box from the tracker
            x1, y1, x2, y2 = int(pos.left()), int(pos.top()), int(pos.right()), int(pos.bottom())

            # Draw the updated tracking bounding box on the frame
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

            # Show the frame with the tracking rectangle
            cv2.imshow("Tracking", frame)

            # Save every 10th frame
            if frame_counter % 10 == 0:
                img_name = f"tracked_frame_{img_counter}.jpg"
                cv2.imwrite(img_name, frame)
                print(f"Saved {img_name}")
                img_counter += 1

            # Calculate and print FPS
            frame_time = time.time() - frame_start_time
            fps = 1.0 / frame_time
            print(f"FPS: {fps:.2f}")

            frame_counter += 1

            # Exit if 'q' is pressed
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    except KeyboardInterrupt:
        print("Tracking interrupted by user")

    cap.release()
    cv2.destroyAllWindows()

# Calculate and print average FPS over all frames
end_time = time.time()
total_time = end_time - start_time
average_fps = frame_counter / total_time
print(f"Average FPS: {average_fps:.2f}")

print("Finished tracking")