"""
Object detection capabilities for FMUS-VID.

This module provides AI-powered object detection using popular models.
"""

import torch
import numpy as np
from typing import List, Dict, Tuple, Optional, Union
from pathlib import Path
import logging
from dataclasses import dataclass

logger = logging.getLogger(__name__)

@dataclass
class Detection:
    """Represents a detected object."""
    label: str
    confidence: float
    bbox: Tuple[float, float, float, float]  # x1, y1, x2, y2 (normalized)

class ObjectDetector:
    """AI-powered object detection."""

    def __init__(self, model_name: str = "yolov5s", device: str = "auto"):
        """
        Initialize object detector.

        Args:
            model_name: Name of detection model to use
            device: Device to run inference on ("cpu", "cuda", or "auto")
        """
        self.model_name = model_name

        # Determine device
        if device == "auto":
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        # Load model
        try:
            if model_name.startswith("yolov5"):
                self.model = torch.hub.load("ultralytics/yolov5", model_name)
                self.model.to(self.device)
            else:
                raise ValueError(f"Unsupported model: {model_name}")
        except Exception as e:
            logger.error(f"Failed to load model {model_name}: {e}")
            raise

    def detect(self, frame: np.ndarray, conf_threshold: float = 0.5) -> List[Detection]:
        """
        Detect objects in a frame.

        Args:
            frame: Input frame (RGB format)
            conf_threshold: Confidence threshold for detections

        Returns:
            List of Detection objects
        """
        # Convert frame to tensor
        img = torch.from_numpy(frame).to(self.device)
        img = img.permute(2, 0, 1)  # HWC to CHW

        # Run inference
        results = self.model([img])

        # Process detections
        detections = []
        for pred in results.pred[0]:
            if pred[4] >= conf_threshold:  # Check confidence
                # Get coordinates
                x1, y1, x2, y2 = pred[:4].cpu().numpy()

                # Normalize coordinates
                height, width = frame.shape[:2]
                x1, x2 = x1/width, x2/width
                y1, y2 = y1/height, y2/height

                # Get class info
                cls_id = int(pred[5])
                label = self.model.names[cls_id]
                confidence = float(pred[4])

                detections.append(Detection(
                    label=label,
                    confidence=confidence,
                    bbox=(x1, y1, x2, y2)
                ))

        return detections

    def annotate_frame(self, frame: np.ndarray, detections: List[Detection]) -> np.ndarray:
        """
        Draw detection boxes on frame.

        Args:
            frame: Input frame
            detections: List of detections

        Returns:
            Annotated frame
        """
        import cv2

        result = frame.copy()
        height, width = frame.shape[:2]

        for det in detections:
            # Convert normalized coordinates to pixel coordinates
            x1 = int(det.bbox[0] * width)
            y1 = int(det.bbox[1] * height)
            x2 = int(det.bbox[2] * width)
            y2 = int(det.bbox[3] * height)

            # Draw box
            cv2.rectangle(result, (x1, y1), (x2, y2), (0, 255, 0), 2)

            # Draw label
            label = f"{det.label} {det.confidence:.2f}"
            cv2.putText(result, label, (x1, y1-10),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        return result
