# Attempting live visual processing and detection
# Importing computervision library
import cv2
import numpy as np
from time import time

conf_threshold = 0.5 # Threshold for confidence in prediction to show box
prev_frame_time = 0 # Previous frame time to track frames per second
new_frame_time = 0 # New frame time
colors = np.random.uniform(0, 255, size=(len(model.names), 3)) # Random colors to more easily differentiate between cards

# Function to run YOLOv8 prediction and draw bounding boxes on the frame
def draw_boxes(frame):
    global conf_threshold, prev_frame_time, new_frame_time, colors # Import global variables / constants
    
    # Time tracking for fps calculation
    new_frame_time = time()
    fps = 1/(new_frame_time-prev_frame_time)
    prev_frame_time = new_frame_time

    # Predicting on frame
    results = model.predict(frame)  # Run prediction
    results = results[0] if isinstance(results, list) else results
    boxes = results.boxes.xyxy.cpu().numpy()
    confArr = results.boxes.conf.cpu().numpy()
    clsArr = results.boxes.cls.cpu().numpy()

    # Define a list of distinct colors
    for i in range(len(boxes)):
        x1, y1, x2, y2 = boxes[i]
        conf = confArr[i]
        cls = clsArr[i]
        color = colors[int(cls)]
        if conf > conf_threshold:  # Only draw boxes with confidence greater than the set threshold
            # Draw bounding box
            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
            # Put label and confidence
            label = f'{model.names[int(cls)]}: {conf:.2f}'
            cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
    cv2.putText(frame, str(round(fps)), (7, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (100, 100, 100), 2, 2)
    return frame

# Opening live feed from usb camera
video = cv2.VideoCapture(0)
# Check if video feed is accessible
if not video.isOpened():
    print("Error: Could not open video stream.")
    exit()

# Loop through each frame from the video until q is pressed
while True:
    ret, frame = video.read()
    # Check if frame is read correctly
    if not ret:
        print("Error: Could not read frame.")
        break
    # Show prediction
    # cv2.resize(frame, (1080,1920))
    cv2.imshow('Live Video', draw_boxes(frame))
    # key q stops live video
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
# Release the video capture object and close all OpenCV windows
video.release()
cv2.destroyAllWindows()