import torch
from ultralytics import YOLO
from PIL import Image, ImageDraw
import numpy as np
import os

# Model yolları
YOLO_MODEL_PATH = os.path.join(os.path.dirname(__file__), "yolov8n.pt")

# Tehlikeli obje sınıfları (COCO veri setindeki indeksler)
DANGEROUS_CLASSES = {
    "bookcase": 0.9,       # Kitaplık
    "cabinet": 0.8,        # Dolap
    "refrigerator": 0.7,   # Buzdolabı
    "oven": 0.6,           # Fırın
    "tv": 0.5,             # Televizyon
    "microwave": 0.5,      # Mikrodalga
    "chair": 0.3,          # Sandalye
    "couch": 0.4,          # Kanepe
    "bed": 0.4,            # Yatak
    "dining table": 0.6,   # Yemek masası
    "toilet": 0.3,         # Tuvalet
    "sink": 0.3,           # Lavabo
    "clock": 0.2,          # Saat
    "vase": 0.2,           # Vazo
    "scissors": 0.1,       # Makas
    "hair drier": 0.1,     # Saç kurutma makinesi
    "toothbrush": 0.1,     # Diş fırçası
}

class DangerDetector:
    def __init__(self):
        # GPU bellek optimizasyonu
        torch.cuda.empty_cache()
        
        # Model yükleme
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Using device: {self.device}")
        
        # YOLOv8 modelini yükle
        if not os.path.exists(YOLO_MODEL_PATH):
            print("Downloading YOLOv8 model...")
            self.model = YOLO("yolov8n.pt")  # Otomatik indirecek
            self.model.save(YOLO_MODEL_PATH)
        else:
            self.model = YOLO(YOLO_MODEL_PATH)
        
        # Modeli GPU'ya taşı ve optimize et
        self.model.to(self.device)
        
        print("Danger detector initialized")
    
    def detect_dangers(self, image):
        """Görüntüdeki tehlikeli objeleri tespit et"""
        # Görüntüyü modele ver
        results = self.model(image, verbose=False)
        
        # Sonuçları işle
        detected_objects = []
        for result in results:
            boxes = result.boxes
            for box in boxes:
                class_id = int(box.cls.item())
                class_name = result.names[class_id]
                
                # Sadece tehlikeli sınıfları filtrele
                if class_name in DANGEROUS_CLASSES:
                    confidence = box.conf.item()
                    danger_score = DANGEROUS_CLASSES[class_name] * confidence
                    x1, y1, x2, y2 = box.xyxy[0].tolist()
                    
                    detected_objects.append({
                        "class": class_name,
                        "box": [x1, y1, x2, y2],
                        "confidence": confidence,
                        "danger_score": danger_score
                    })
        
        return detected_objects
    
    def mark_dangers(self, image):
        """Görüntüdeki tehlikeli objeleri işaretle ve en güvenli noktayı göster"""
        # NumPy dizisine dönüştür
        if isinstance(image, str):
            image = Image.open(image)
        elif isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # Tehlikeli objeleri tespit et
        dangers = self.detect_dangers(np.array(image))
        
        # Görüntü üzerine çizim yapmak için kopya oluştur
        result_image = image.copy()
        draw = ImageDraw.Draw(result_image)
        
        # Tehlikeli objeleri işaretle
        for obj in dangers:
            x1, y1, x2, y2 = obj["box"]
            class_name = obj["class"]
            danger_score = obj["danger_score"]
            
            # Tehlike skoruna göre renk belirle (kırmızı -> sarı)
            color = (255, int(255 * (1 - danger_score)), 0)
            
            # Kutuyu çiz
            draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
            
            # Etiketi çiz
            draw.text((x1, y1-15), f"{class_name}: {danger_score:.2f}", fill=color)
        
        # Basit güvenli nokta hesaplama
        safe_point = self._find_safe_point(np.array(image), dangers)
        
        # Güvenli noktayı işaretle
        if safe_point:
            x, y = safe_point
            radius = 20
            draw.ellipse((x-radius, y-radius, x+radius, y+radius), fill=(0, 255, 0, 128))
            draw.text((x+radius+5, y), "EN GÜVENLİ NOKTA", fill=(0, 255, 0))
        
        return np.array(result_image)
    
    def _find_safe_point(self, image, dangers):
        """En güvenli noktayı bul (basit versiyon)"""
        h, w = image.shape[:2]
        
        # Tehlike haritası oluştur
        danger_map = np.zeros((h, w), dtype=np.float32)
        
        # Her tehlikeli obje için tehlike alanı oluştur
        for obj in dangers:
            x1, y1, x2, y2 = obj["box"]
            danger_score = obj["danger_score"]
            
            # Objenin merkezi
            center_x = (x1 + x2) / 2
            center_y = (y1 + y2) / 2
            
            # Tüm pikseller için tehlike skorunu hesapla
            for y in range(0, h, 10):  # Performans için her 10 pikselde bir hesapla
                for x in range(0, w, 10):
                    # Mesafe hesapla
                    distance = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                    
                    # Mesafeye bağlı tehlike skoru (uzaklık arttıkça azalır)
                    pixel_danger = danger_score * (1000 / (distance + 10))
                    
                    # Tehlike haritasını güncelle
                    danger_map[y:min(y+10, h), x:min(x+10, w)] += pixel_danger
        
        # En güvenli nokta (en düşük tehlike skoru)
        min_danger = np.min(danger_map)
        safe_indices = np.where(danger_map == min_danger)
        
        if len(safe_indices[0]) > 0:
            # Birden fazla güvenli nokta varsa, rastgele birini seç
            idx = np.random.randint(0, len(safe_indices[0]))
            safe_y = safe_indices[0][idx]
            safe_x = safe_indices[1][idx]
            return (safe_x, safe_y)
        
        return None

# Singleton örneği
detector = None

def get_detector():
    """Detector singleton örneğini döndür"""
    global detector
    if detector is None:
        detector = DangerDetector()
    return detector

def mark_dangers(image):
    """Gradio arayüzü için wrapper fonksiyon"""
    detector = get_detector()
    return detector.mark_dangers(image)

# Test
if __name__ == "__main__":
    import matplotlib.pyplot as plt
    
    # Test görüntüsü
    test_image = "test_room.jpg"
    
    # Tehlikeleri işaretle
    detector = DangerDetector()
    result = detector.mark_dangers(test_image)
    
    # Sonucu göster
    plt.figure(figsize=(12, 8))
    plt.imshow(result)
    plt.axis('off')
    plt.title("Tehlike Analizi")
    plt.show()