import cv2
import numpy as np
from ultralytics import YOLO
from typing import List, Tuple, Dict, Any

# Load YOLOv8 model for general object detection
model = YOLO("models/yolov8n.pt")

def detect_objects(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
    """
    Detect cars, bikes, humans, dogs, and other objects in the frame using YOLOv8.
    Args:
        frame: Input frame as a numpy array.
    Returns:
        Tuple of (list of detections, annotated frame).
    """
    # Perform inference
    results = model(frame, conf=0.5)  # Detect all classes with confidence > 0.5
    
    detections = []
    for i, r in enumerate(results[0].boxes):
        x_min, y_min, x_max, y_max = map(int, r.xyxy[0])
        conf = float(r.conf)
        cls = int(r.cls)
        # Map YOLOv8 class IDs to labels (based on COCO dataset)
        label_map = {
            0: "person",  # Human
            1: "bicycle", # Bike (approximation)
            2: "car",
            3: "motorcycle", # Bike
            5: "bus",     # Treat as car
            7: "truck",   # Treat as car
            16: "dog",
        }
        dtype = label_map.get(cls, "object")  # Default to "object" for unmapped classes
        if dtype in ["bicycle", "motorcycle"]:
            dtype = "bike"
        elif dtype in ["bus", "truck"]:
            dtype = "car"
        label = f"{dtype.capitalize()} {i+1}"
        
        # Determine severity (not used for objects, but included for consistency)
        area = (x_max - x_min) * (y_max - y_min)
        severity = "Moderate"  # Default for objects
        
        detections.append({
            "box": [x_min, y_min, x_max, y_max],
            "label": label,
            "type": dtype,
            "confidence": conf,
            "severity": severity
        })
    
    return detections, frame