from ultralytics import YOLO
import cv2
from pathlib import Path
import logging

logger = logging.getLogger(__name__)


class MemoryDetector:
    def __init__(self, model_path):
        try:
            self.model = YOLO(model_path)
            logger.info(f"Loaded model from {model_path}")
        except Exception as e:
            logger.error(f"Model loading failed: {str(e)}")
            raise

    def detect(self, image_path):
        try:
            # Run inference
            results = self.model.predict(image_path, imgsz=416, conf=0.5)
            
            # Extract results
            boxes = results[0].boxes.xyxy.cpu().numpy()
            confidences = results[0].boxes.conf.cpu().numpy()
            
            # Convert to list of [x1, y1, x2, y2, confidence]
            detections = []
            for box, conf in zip(boxes, confidences):
                detections.append({
                    'box': [int(x) for x in box],
                    'confidence': float(conf)
                })
            
            # Annotate image
            annotated_img = self._draw_boxes(image_path, detections)
            
            return {
                'detections': detections,
                'annotated_image': annotated_img
            }
            
        except Exception as e:
            logger.error(f"Detection failed: {str(e)}")
            raise

    def _draw_boxes(self, image_path, detections):
        img = cv2.imread(str(image_path))
        for det in detections:
            x1, y1, x2, y2 = det['box']
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(img, f"{det['confidence']:.2f}", 
                       (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1)
        return img