# robot_utils/camera.py

from ultralytics import YOLO
import cv2
import pybullet as p
import numpy as np

class Camera:
    def __init__(self, robot_id, physics_client_id, yolo_model_path, classes_path, img_width=640, img_height=480):
        self.robot_id = robot_id
        self.physics_client_id = physics_client_id
        self.active = False
        self.model = YOLO(yolo_model_path)  # Load YOLOv8 model
        self.img_width = img_width
        self.img_height = img_height
        
        with open(classes_path, 'r') as f:
            self.classes = [line.strip() for line in f.readlines()]

    def activate(self):
        self.active = True
        print("Camera activated.")

    def deactivate(self):
        self.active = False
        print("Camera deactivated.")

    def capture_image(self, robot_id, physics_client_id):
        if not self.active:
            print("Camera is off.")
            return None

        # Get robot's position and orientation
        robot_pos, robot_orn = p.getBasePositionAndOrientation(robot_id, physicsClientId=physics_client_id)
        yaw = p.getEulerFromQuaternion(robot_orn)[-1]

        # Set camera parameters
        distance = 2.5
        camera_eye_position = [robot_pos[0], robot_pos[1], robot_pos[2] + 1]
        camera_target_position = [
            robot_pos[0] + distance * np.cos(yaw),
            robot_pos[1] + distance * np.sin(yaw),
            robot_pos[2]
        ]

        # Compute view and projection matrices
        view_matrix = p.computeViewMatrix(
            cameraEyePosition=camera_eye_position,
            cameraTargetPosition=camera_target_position,
            cameraUpVector=[0, 1, 0],
            physicsClientId=physics_client_id
        )
        projection_matrix = p.computeProjectionMatrixFOV(
            fov=90,
            aspect=self.img_width / self.img_height,
            nearVal=0.1,
            farVal=100,
            physicsClientId=physics_client_id
        )

        # Capture the image
        _, _, img_arr, _, _ = p.getCameraImage(
            width=self.img_width,
            height=self.img_height,
            viewMatrix=view_matrix,
            projectionMatrix=projection_matrix,
            physicsClientId=physics_client_id
        )

        # Convert the image array to a proper format
        img = np.array(img_arr, dtype=np.uint8).reshape((self.img_height, self.img_width, 4))
        img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
        return img

    def detect_objects(self):
        img = self.capture_image()
        if img is None:
            return []
        
        results = self.model(img)  # Perform object detection
        
        # Extract detected classes
        detected_classes = [self.model.names[int(cls)] for result in results for cls in result.boxes.cls]
        unique_classes = list(set(detected_classes))
        
        print(f"Detected objects: {unique_classes}")
        return unique_classes

    def get_annotated_frame(self):
        img = self.capture_image(self.robot_id, self.physics_client_id)
        if img is None:
            return None
        
        results = self.model(img)  # Perform object detection
        annotated_frame = img.copy()

        for result in results:
            for box in result.boxes:
                cls_id = int(box.cls)
                label = self.model.names[cls_id]
                conf = box.conf
                x1, y1, x2, y2 = map(int, box.xyxy)
                cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(annotated_frame, f"{label} {conf:.2f}", (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        return annotated_frame


Traceback (most recent call last):
  File "/Users/apple/Music/major project/multimodal_robot_simulation/main.py", line 137, in <module>
    main()
  File "/Users/apple/Music/major project/multimodal_robot_simulation/main.py", line 120, in main
    execute_instructions(robot.functions, instructions)
  File "/Users/apple/Music/major project/multimodal_robot_simulation/main.py", line 20, in execute_instructions
    robot_functions.move_forward(params['distance'])
  File "/Users/apple/Music/major project/multimodal_robot_simulation/functions.py", line 49, in move_forward
    self._move(distance, forward=True)
  File "/Users/apple/Music/major project/multimodal_robot_simulation/functions.py", line 63, in _move
    self.display_annotated_frame()
  File "/Users/apple/Music/major project/multimodal_robot_simulation/functions.py", line 39, in display_annotated_frame
    annotated_frame = self.camera.get_annotated_frame()
  File "/Users/apple/Music/major project/multimodal_robot_simulation/robot_utils/camera.py", line 102, in get_annotated_frame
    x1, y1, x2, y2 = map(int, box.xyxy)
ValueError: only one element tensors can be converted to Python scalars