from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image

class RoadsideDetector:
    def __init__(self, model_path='yol_kenari_tespit_best.pt'):
        """
        Initialize the detector with the trained YOLO model.
        """
        print("🤖 Loading AI model...")
        try:
            self.model = YOLO(model_path)  # Load YOLO model
            print("✅ Model loaded successfully!")
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            self.model = None
        
        # Class names based on your model in Turkish
        self.classes = ['çöp', 'çöp kutusu', 'taşan çöp']

    def detect_items(self, image, confidence=0.25):
        """
        Detects items in an image based on the trained YOLO model.
        """
        if self.model is None:
            print("❌ Model not loaded!")
            return []

        print("🔍 Detecting illegal dumping...")
        results = self.model(image, conf=confidence)
        
        detections = []
        for result in results:
            boxes = result.boxes
            if boxes is not None:
                for i in range(len(boxes)):
                    x1, y1, x2, y2 = boxes.xyxy[i].cpu().numpy()
                    confidence = float(boxes.conf[i])
                    class_id = int(boxes.cls[i])
                    class_name = self.classes[class_id]  # Use Turkish classes
                    
                    detection = {
                        'type': class_name,
                        'confidence': confidence,
                        'bbox': [float(x1), float(y1), float(x2), float(y2)],
                        'area': (x2 - x1) * (y2 - y1)
                    }
                    detections.append(detection)
                    
                    print(f"  Found: {class_name} ({confidence:.1%} sure)")

        print(f"✅ Found {len(detections)} items in total")
        return detections

    def draw_boxes(self, image, detections):
        """
        Draw bounding boxes around detected objects.
        """
        img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        for det in detections:
            x1, y1, x2, y2 = [int(x) for x in det['bbox']]
            color = (255, 0, 0)  # Red color for bounding box
            cv2.rectangle(img_cv, (x1, y1), (x2, y2), color, 2)
            
            label = f"{det['type']} {det['confidence']:.1%}"
            label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
            cv2.rectangle(img_cv, 
                         (x1, y1 - label_size[1] - 10),
                         (x1 + label_size[0], y1),
                         color, -1)  # -1 means fill
            
            cv2.putText(img_cv, label,
                       (x1, y1 - 5),
                       cv2.FONT_HERSHEY_SIMPLEX,
                       0.5, (255, 255, 255), 2)
        
        img_pil = Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB))
        return img_pil

    def calculate_severity(self, detections):
        """
        Calculate the severity of illegal dumping based on detected items.
        """
        count = len(detections)
        severity = "NONE"
        color = "🟢"
        message = "Area is clean"
        if count >= 5:
            severity = "HIGH"
            color = "🔴"
            message = "Urgent cleanup needed!"
        elif count >= 2:
            severity = "MEDIUM"
            color = "🟠"
            message = "Schedule cleanup soon"
        elif count == 1:
            severity = "LOW"
            color = "🟡"
            message = "Minor issue"
        
        return {
            'level': severity,
            'color': color,
            'count': count,
            'types': {det['type']: count for det in detections},
            'message': message
        }

    def process_image(self, image_path):
        """
        Complete pipeline: detect, draw bounding boxes, calculate severity.
        """
        image = Image.open(image_path) if isinstance(image_path, str) else image_path
        detections = self.detect_items(image)
        annotated_image = self.draw_boxes(image, detections)
        severity = self.calculate_severity(detections)
        summary = f"{severity['color']} {severity['level']} - "
        summary += f"Found {severity['count']} items\n"
        
        if severity['types']:
            summary += "Details:\n"
            for item_type, count in severity['types'].items():
                summary += f"  • {count}x {item_type}\n"
        
        summary += f"\n{severity['message']}"
        
        return {
            'image': annotated_image,
            'detections': detections,
            'severity': severity,
            'summary': summary
        }

# Example usage (no need to import `RoadsideDetector` from itself)
if __name__ == "__main__":
    detector = RoadsideDetector('yol_kenari_tespit_best.pt')
    image_path = 'path_to_image.jpg'  # Replace with an actual image path
    result = detector.process_image(image_path)
    print(result['summary'])
    result['image'].show()