from ultralytics import YOLO
from PIL import Image
import numpy as np
from typing import Tuple, List, Dict

class MemoryDetector:
    def __init__(self, 
                 model_path='model/weights/best.pt',
                 conf_threshold=0.25,
                 iou_threshold=0.45):
        """
        Initialize the detector with the trained model.
        
        Args:
            model_path (str): Path to the trained model weights
            conf_threshold (float): Confidence threshold for detections
            iou_threshold (float): IoU threshold for NMS
        """
        self.model = YOLO(model_path)
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold

    def detect(self, 
              image: Image.Image,
              conf_threshold: float = None,
              iou_threshold: float = None) -> Tuple[Image.Image, List[Dict]]:
        """
        Detect memory modules in the given image.
        
        Args:
            image (PIL.Image): Input image to process
            conf_threshold (float, optional): Override default confidence threshold
            iou_threshold (float, optional): Override default IoU threshold
            
        Returns:
            Tuple[PIL.Image, List[Dict]]: Annotated image and list of detections
        """
        # Use provided thresholds or defaults
        conf = conf_threshold if conf_threshold is not None else self.conf_threshold
        iou = iou_threshold if iou_threshold is not None else self.iou_threshold
        
        # Run inference
        results = self.model.predict(
            source=image,
            conf=conf,
            iou=iou,
            max_det=10,
            verbose=False
        )
        
        # Get the annotated image
        annotated_img = results[0].plot()
        
        # Extract detection information
        detections = []
        for box in results[0].boxes:
            detection = {
                'xyxy': box.xyxy[0].tolist(),  # Bounding box coordinates
                'confidence': float(box.conf[0]),  # Detection confidence
                'class': int(box.cls[0])  # Class ID
            }
            detections.append(detection)
        
        return Image.fromarray(annotated_img), detections

    def optimize_thresholds(self, validation_images: List[Image.Image]) -> Tuple[float, float]:
        """
        Find optimal confidence and IoU thresholds using validation images.
        
        Args:
            validation_images (List[Image.Image]): List of validation images
            
        Returns:
            Tuple[float, float]: Optimal confidence and IoU thresholds
        """
        best_conf = 0.25
        best_iou = 0.45
        
        # Grid search for best parameters
        conf_range = [0.15, 0.2, 0.25, 0.3, 0.35]
        iou_range = [0.35, 0.4, 0.45, 0.5, 0.55]
        
        best_score = 0
        
        for conf in conf_range:
            for iou in iou_range:
                total_score = 0
                for img in validation_images:
                    _, detections = self.detect(img, conf, iou)
                    # Score based on number of detections and confidence
                    score = sum([d['confidence'] for d in detections])
                    total_score += score
                
                if total_score > best_score:
                    best_score = total_score
                    best_conf = conf
                    best_iou = iou
        
        return best_conf, best_iou