"""
detection.py - Dog Detection using YOLOv8
"""

import cv2
import numpy as np
import torch
from ultralytics import YOLO
from typing import List, Optional
from dataclasses import dataclass


@dataclass
class Detection:
    bbox: List[float]
    confidence: float
    image_crop: Optional[np.ndarray] = None


class DogDetector:
    def __init__(self, confidence_threshold: float = 0.50, device: str = 'cuda'):
        self.confidence_threshold = confidence_threshold
        self.device = device if torch.cuda.is_available() else 'cpu'
        self.model = YOLO('yolov8m.pt')
        self.model.to(self.device)
        self.dog_class_id = 16
        self.min_detection_area = 900
        
        print(f"✅ Dog Detector initialized on {self.device}")
        print(f"   Confidence threshold: {self.confidence_threshold:.2f}")
    
    def detect(self, frame: np.ndarray) -> List[Detection]:
        if frame is None or frame.size == 0:
            return []
        
        results = self.model(
            frame,
            conf=self.confidence_threshold,
            classes=[self.dog_class_id],
            verbose=False
        )
        
        detections = []
        
        if results and len(results) > 0:
            result = results[0]
            
            if result.boxes is not None:
                boxes = result.boxes
                
                for i in range(len(boxes)):
                    x1, y1, x2, y2 = boxes.xyxy[i].cpu().numpy()
                    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                    
                    h, w = frame.shape[:2]
                    x1 = max(0, min(w-1, x1))
                    y1 = max(0, min(h-1, y1))
                    x2 = max(0, min(w, x2))
                    y2 = max(0, min(h, y2))
                    
                    if x2 <= x1 or y2 <= y1:
                        continue
                    
                    area = (x2 - x1) * (y2 - y1)
                    if area < self.min_detection_area:
                        continue
                    
                    confidence = float(boxes.conf[i])
                    dog_crop = frame[y1:y2, x1:x2].copy()
                    
                    detection = Detection(
                        bbox=[x1, y1, x2, y2],
                        confidence=confidence,
                        image_crop=dog_crop
                    )
                    
                    detections.append(detection)
        
        return detections
    
    def set_confidence(self, threshold: float):
        self.confidence_threshold = max(0.1, min(1.0, threshold))
        print(f"Detection confidence updated: {self.confidence_threshold:.2f}")


EnhancedDogDetector = DogDetector