import cv2
import argparse
import socket

from ultralytics import YOLO

def parse_argument() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="YOLOv8 live")
    parser.add_argument(
        "--webcam-resolution",
        default=[640, 480],
        nargs=2, 
        type=int
    )
    args = parser.parse_args()
    return args

def send_data_to_server(host, port, data):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((host, port))
        s.sendall(data.encode())
        print(f"Data sent: {data}")

def main():
    args = parse_argument()
    frame_width, frame_height = args.webcam_resolution

    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, frame_width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, frame_height)

    model_path = r"C:\Users\guill\OneDrive\Documenten\Project_one\miniproject\2023-2024-projectone-ctai-GuillemynRune\Dataset\runs\detect\train\weights\best.pt"
    model = YOLO(model_path)

    raspberry_pi_ip = '192.168.168.167'  # Replace with your Raspberry Pi's IP address
    port = 65432

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

        result = model(frame, agnostic_nms=True)[0]
        detections = result.boxes  # Get the boxes directly from the result
        labels = [
            f"{model.names[int(box.cls.item())]} {box.conf.item():.2f}"
            for box in detections
        ]

        for box in detections:
            # Draw the box on the frame
            x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
            conf = box.conf.item()
            cls = int(box.cls.item())
            label = f"{model.names[cls]} {conf:.2f}"
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)

        # If there are detections, send them to the server
        if labels:
            data_to_send = "; ".join(labels)
            send_data_to_server(raspberry_pi_ip, port, data_to_send)

        cv2.imshow("YOLOv8", frame)
        
        if (cv2.waitKey(30) == 27):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()