import cv2
from websocket import create_connection
import argparse
import math
from ultralytics import YOLO

# object classes
classNames = ["person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat",
              "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
              "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
              "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
              "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
              "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
              "carrot", "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", "bed",
              "diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
              "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
              "teddy bear", "hair drier", "toothbrush"
              ]


def main(websocket_addr):
    # Connect to websocket server
    # ws = create_connection(websocket_addr)

    # Load YOLO model
    model = YOLO('yolov8s.pt')  # Load an official Detect model

    # Open webcam
    cap = cv2.VideoCapture(0)
    cap.set(3, 640)
    cap.set(4, 480)

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # Perform object detection
        results = model.predict(frame)

        max_x1, max_y1, max_x2, max_y2 = 0

        for r in results:
            boxes = r.boxes
            for box in boxes:
                if box.cls[0].item() == 0:
                    # bounding box
                    x1, y1, x2, y2 = box.xyxy[0]
                    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) # convert to int values

                    # put box in cam
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 255), 3)

                    # class name
                    cls = int(box.cls[0])

                    # object details
                    org = [x1, y1]
                    font = cv2.FONT_HERSHEY_SIMPLEX
                    fontScale = 1
                    color = (255, 0, 0)
                    thickness = 2

                    cv2.putText(frame, classNames[cls], org, font, fontScale, color, thickness)

                    if (x2-x1) * (y2 - y1):
                        max_x1 = x1
                        max_x2 = x2 
                        max_y1 = y1 
                        max_y2 = y2

        print(f"{max_x1}:{max_y1}, {max_x2}:{max_y2}")

        cv2.imshow('Webcam', frame)
        if cv2.waitKey(1) == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Detect people from webcam feed and send data to websocket server')
    parser.add_argument('websocket_addr', type=str, help='Websocket server address')
    args = parser.parse_args()

    main(args.websocket_addr)