import gradio as gr
from transformers import pipeline
from PIL import Image, ImageDraw, ImageFont
import torch
import spaces
import numpy as np
import cv2

# Modèles optimisés pour le temps réel
REALTIME_MODELS = {
    "YOLOS Tiny (ultra-rapide)": "hustvl/yolos-tiny",
    "DETR ResNet-50": "facebook/detr-resnet-50",
    "YOLOS Small": "hustvl/yolos-small",
    "Conditional DETR": "microsoft/conditional-detr-resnet-50"
}

# Cache global pour le modèle
current_detector = None
current_model_name = None

def load_detector(model_name):
    """Charge le détecteur avec cache"""
    global current_detector, current_model_name
    
    if current_model_name != model_name:
        print(f"🔄 Chargement du modèle: {model_name}")
        model_id = REALTIME_MODELS[model_name]
        current_detector = pipeline(
            "object-detection",
            model=model_id,
            device=0 if torch.cuda.is_available() else -1
        )
        current_model_name = model_name
        print(f"✅ Modèle chargé: {model_name}")
    
    return current_detector

@spaces.GPU
def process_webcam_frame(frame, model_choice, confidence_threshold):
    """
    Traite chaque frame de la webcam en temps réel
    Cette fonction est appelée automatiquement pour chaque frame
    """
    if frame is None:
        return frame
    
    try:
        # Charger le détecteur
        detector = load_detector(model_choice)
        
        # Convertir numpy array en PIL Image si nécessaire
        if isinstance(frame, np.ndarray):
            # Gradio webcam donne du RGB
            pil_image = Image.fromarray(frame)
        else:
            pil_image = frame
        
        # Redimensionner pour accélérer le traitement
        original_size = pil_image.size
        max_size = 640  # Réduire la taille pour plus de vitesse
        
        if max(original_size) > max_size:
            ratio = max_size / max(original_size)
            new_size = (int(original_size[0] * ratio), int(original_size[1] * ratio))
            resized_image = pil_image.resize(new_size)
        else:
            resized_image = pil_image
            ratio = 1.0
        
        # Détection sur l'image redimensionnée
        detections = detector(resized_image)
        
        # Filtrer par confiance
        filtered_detections = [
            det for det in detections 
            if det['score'] >= confidence_threshold
        ]
        
        # Redimensionner les coordonnées vers la taille originale
        for det in filtered_detections:
            if ratio != 1.0:
                det['box']['xmin'] = int(det['box']['xmin'] / ratio)
                det['box']['ymin'] = int(det['box']['ymin'] / ratio)
                det['box']['xmax'] = int(det['box']['xmax'] / ratio)
                det['box']['ymax'] = int(det['box']['ymax'] / ratio)
        
        # Dessiner les détections sur l'image originale
        annotated_image = draw_detections_fast(pil_image, filtered_detections)
        
        # Convertir back en numpy pour Gradio
        return np.array(annotated_image)
        
    except Exception as e:
        print(f"❌ Erreur de traitement: {e}")
        return frame

def draw_detections_fast(image, detections):
    """Version optimisée pour dessiner les détections"""
    if not detections:
        return image
    
    draw = ImageDraw.Draw(image)
    
    # Police par défaut pour la vitesse
    try:
        font = ImageFont.load_default()
    except:
        font = None
    
    colors = ["#FF6B6B", "#4ECDC4", "#45B7D1", "#96CEB4", "#FECA57"]
    
    for i, detection in enumerate(detections):
        box = detection['box']
        label = detection['label']
        score = detection['score']
        
        # Coordonnées
        x1, y1 = box['xmin'], box['ymin']
        x2, y2 = box['xmax'], box['ymax']
        
        # Couleur
        color = colors[i % len(colors)]
        
        # Boîte
        draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
        
        # Label avec score
        text = f"{label} {score:.2f}"
        
        # Fond du texte (simplifié)
        if font:
            bbox = draw.textbbox((x1, y1-20), text, font=font)
            draw.rectangle(bbox, fill=color)
            draw.text((x1, y1-20), text, fill="white", font=font)
        else:
            draw.text((x1, y1-15), text, fill=color)
    
    return image

# Interface Gradio avec streaming
with gr.Blocks(title="🎥 Détection Live", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🎥 Détection d'Objets en Temps Réel
    
    **Activez votre webcam** et voyez la détection se faire en direct !
    
    ⚡ **Optimisé pour la vitesse** avec des modèles légers
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            # Composant webcam avec streaming
            webcam = gr.Interface(
                fn=process_webcam_frame,
                inputs=[
                    gr.Image(sources=["webcam"], streaming=True, type="numpy"),
                    gr.Dropdown(
                        choices=list(REALTIME_MODELS.keys()),
                        value="YOLOS Tiny (ultra-rapide)",
                        label="🤖 Modèle (changement en direct)"
                    ),
                    gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.5,
                        step=0.1,
                        label="🎯 Seuil de confiance"
                    )
                ],
                outputs=gr.Image(type="numpy", streaming=True),
                live=True,  # ⭐ CRUCIAL: Active le mode live
                title=None
            )
        
        with gr.Column(scale=1):
            gr.Markdown("""
            ## 📊 Informations Live
            
            ### 🎛️ Contrôles en temps réel:
            - **Modèle**: Change instantanément
            - **Confiance**: Ajuste le filtrage
            - **Streaming**: Traitement frame par frame
            
            ### ⚡ Optimisations:
            - Images redimensionnées à 640px
            - Modèles légers prioritaires
            - Cache intelligent des modèles
            - Dessin optimisé
            
            ### 🎯 Modèles recommandés:
            - **YOLOS Tiny**: Maximum de vitesse
            - **DETR ResNet-50**: Bon équilibre
            """)
    
    # Version alternative avec Interface simple
    gr.Markdown("---")
    gr.Markdown("## 🎥 Version Alternative (Interface Simple)")
    
    alternative_interface = gr.Interface(
        fn=process_webcam_frame,
        inputs=[
            gr.Image(sources=["webcam"], streaming=True),
            gr.Dropdown(
                choices=list(REALTIME_MODELS.keys()),
                value="YOLOS Tiny (ultra-rapide)"
            ),
            gr.Slider(0.1, 1.0, 0.5, step=0.1)
        ],
        outputs=gr.Image(streaming=True),
        live=True,  # ⭐ Mode live activé
        title="Détection Webcam Live",
        description="Cliquez sur la webcam pour démarrer le streaming live!"
    )

if __name__ == "__main__":
    demo.launch()