# app.py - Labasni Cloth Detection
# Compatible Gradio 5.9.1 avec API REST fonctionnelle

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import gradio as gr
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense
from tensorflow.keras.applications import MobileNetV2
from ultralytics import YOLO
from PIL import Image
import numpy as np
from sklearn.cluster import KMeans
import json

print("=" * 60)
print("🚀 DÉMARRAGE LABASNI CLOTH DETECTION")
print("=" * 60)

# ============================================
# CHARGEMENT DES MODÈLES
# ============================================

print("\n🔄 Chargement du modèle YOLO...")
try:
    yolo_model = YOLO("best.pt")
    print("✅ YOLO chargé avec succès")
except Exception as e:
    print(f"❌ Erreur YOLO: {e}")
    raise

print("\n🔄 Chargement du modèle Style/Season...")
try:
    base_model = MobileNetV2(
        input_shape=(224, 224, 3),
        include_top=False,
        weights=None
    )
    x = base_model.output
    x = GlobalAveragePooling2D()(x)
    style_output = Dense(4, activation='softmax', name='style_output')(x)
    season_output = Dense(4, activation='softmax', name='season_output')(x)
    
    style_model = Model(
        inputs=base_model.input, 
        outputs=[style_output, season_output]
    )
    
    style_model.load_weights(
        "style_season_model.h5", 
        by_name=True, 
        skip_mismatch=True
    )
    print("✅ Style/Season model chargé avec succès")
except Exception as e:
    print(f"❌ Erreur Style/Season: {e}")
    raise

print("\n" + "=" * 60)
print("✅ TOUS LES MODÈLES CHARGÉS")
print("=" * 60 + "\n")

# ============================================
# FONCTIONS UTILITAIRES
# ============================================

def get_dominant_color(crop):
    """Extrait la couleur dominante via K-Means"""
    try:
        img_array = np.array(crop).reshape(-1, 3)
        img_array = img_array[np.any(img_array != [0, 0, 0], axis=1)]
        
        if len(img_array) == 0:
            return "#808080"
        
        kmeans = KMeans(n_clusters=3, random_state=0, n_init=10)
        kmeans.fit(img_array)
        dominant = kmeans.cluster_centers_[0]
        
        return '#{:02x}{:02x}{:02x}'.format(
            int(dominant[0]), 
            int(dominant[1]), 
            int(dominant[2])
        ).upper()
    except:
        return "#808080"

# ============================================
# FONCTION PRINCIPALE DE DÉTECTION
# ============================================

def detect_cloth(image):
    """Détecte type, couleur, style et saison d'un vêtement"""
    
    print("\n" + "="*60)
    print("🔍 NOUVELLE DÉTECTION")
    print("="*60)
    
    try:
        # Validation
        if image is None:
            error_result = {"error": "Aucune image fournie"}
            print(f"❌ {error_result['error']}")
            return error_result
        
        # Conversion PIL
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)
        
        img = image.convert('RGB')
        w, h = img.size
        print(f"📸 Image: {w}x{h} pixels")
        
        # YOLO Detection
        print("🤖 Détection YOLO...")
        results = yolo_model(img, verbose=False)[0]
        boxes = results.boxes
        
        if not boxes or len(boxes) == 0:
            error_result = {"error": "Aucun vêtement détecté"}
            print(f"❌ {error_result['error']}")
            return error_result
        
        best_idx = boxes.conf.argmax()
        confidence = float(boxes.conf[best_idx])
        print(f"✅ Confiance: {confidence:.2%}")
        
        # Extraction
        x1, y1, x2, y2 = map(int, boxes.xyxy[best_idx])
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(w, x2), min(h, y2)
        
        if x2 <= x1 or y2 <= y1:
            error_result = {"error": "Boîte de détection invalide"}
            print(f"❌ {error_result['error']}")
            return error_result
        
        cropped = img.crop((x1, y1, x2, y2))
        
        # Couleur
        print("🎨 Extraction couleur...")
        hex_color = get_dominant_color(cropped)
        print(f"✅ Couleur: {hex_color}")
        
        # Style & Saison
        print("👔 Classification...")
        resized = cropped.resize((224, 224))
        arr = np.array(resized) / 255.0
        arr = np.expand_dims(arr, axis=0)
        
        style_pred, season_pred = style_model.predict(arr, verbose=0)
        
        styles = ["casual", "formal", "sport"]
        seasons = ["summer", "winter", "fall", "spring"]
        
        style = styles[np.argmax(style_pred)]
        season = seasons[np.argmax(season_pred)]
        
        style_conf = float(np.max(style_pred))
        season_conf = float(np.max(season_pred))
        
        print(f"✅ Style: {style} ({style_conf:.2%})")
        print(f"✅ Saison: {season} ({season_conf:.2%})")
        
        # Type
        type_vetement = results.names[int(boxes.cls[best_idx])]
        print(f"✅ Type: {type_vetement}")
        
        # Résultat
        result = {
            "success": True,
            "detection": {
                "type": type_vetement,
                "color": hex_color,
                "style": style,
                "season": season
            },
            "confidence": {
                "detection": f"{confidence:.2%}",
                "style": f"{style_conf:.2%}",
                "season": f"{season_conf:.2%}"
            }
        }
        
        print("="*60)
        print("✅ DÉTECTION RÉUSSIE")
        print("="*60)
        print(json.dumps(result, indent=2))
        
        return result
        
    except Exception as e:
        error_result = {"error": f"Erreur: {str(e)}"}
        print(f"❌ {error_result['error']}")
        import traceback
        traceback.print_exc()
        return error_result

# ============================================
# WRAPPER POUR FORMATER LE JSON EN STRING (UI uniquement)
# ============================================

def detect_cloth_wrapper(image):
    """Wrapper qui convertit le résultat en JSON string
    Évite le bug gr.JSON dans Gradio 5.x (TypeError avec additionalProperties)
    L'API REST retournera une string JSON que les clients devront parser
    """
    result = detect_cloth(image)
    # Convertir le dictionnaire en JSON string formaté
    return json.dumps(result, indent=2, ensure_ascii=False)

# ============================================
# INTERFACE GRADIO 5.x - AVEC API
# ============================================

# CSS personnalisé
custom_css = """
.gradio-container {
    font-family: 'IBM Plex Sans', sans-serif;
}
.output-json {
    font-family: 'Courier New', monospace;
    font-size: 14px;
}
"""

# Création de l'interface avec Gradio Blocks (meilleur contrôle UI et API)
with gr.Blocks(
    title="🧥 Labasni - Détection de Vêtements",
    theme=gr.themes.Soft(),
    css=custom_css
) as demo:
    
    gr.Markdown("""
    # 👔 Labasni - Détection de Vêtements
    
    ## 👔 Labasni - Détection Automatique de Vêtements
    
    Cette API utilise l'intelligence artificielle pour analyser vos photos de vêtements et détecter :
    
    - 👕 **Type** : Haut, pantalon, robe, chaussures, etc.
    - 🎨 **Couleur dominante** : Code hexadécimal de la couleur principale
    - 👔 **Style** : Casual, Formel, Sport
    - 🌤️ **Saison** : Été, Hiver, Automne, Printemps
    
    ## 🚀 Comment utiliser
    
    1. **Via l'interface** : Uploadez une image ci-dessous
    2. **Via l'API** : Utilisez l'endpoint `/api/predict`
    
    ## 📡 Exemple d'appel API
    
    ```bash
    curl -X POST "https://syleto-labasni-detection.hf.space/api/predict" \\
      -H "Content-Type: application/json" \\
      -d '{"data": ["<BASE64_IMAGE>"]}'
    ```
    
    ## ⚙️ Modèles utilisés
    
    - **YOLO** (best.pt) : Détection des vêtements
    - **MobileNetV2** : Classification style/saison
    - **K-Means** : Extraction couleur dominante
    """)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(
                label="📸 Photo du vêtement",
                type="pil",
                height=400
            )
            
            with gr.Row():
                clear_btn = gr.Button("Clear", variant="secondary")
                submit_btn = gr.Button("Submit", variant="primary")
        
        with gr.Column():
            # ✅ SOLUTION : Utiliser gr.Textbox au lieu de gr.JSON pour éviter le bug Gradio 5.x
            # Le bug se produit dans la génération du schéma API pour gr.JSON
            result_output = gr.Textbox(
                label="📊 Résultat de la détection",
                elem_classes=["output-json"],
                lines=15,
                max_lines=20,
                interactive=False
            )
    
    # ✅ IMPORTANT : Configurer l'API explicitement pour Gradio 5.x
    # SOLUTION : Utiliser detect_cloth_wrapper qui retourne une string JSON
    # Cela évite le bug gr.JSON (TypeError avec additionalProperties)
    # L'API REST retournera une string JSON que les clients devront parser
    # C'est acceptable car cela résout le bug "No API found"
    submit_btn.click(
        fn=detect_cloth_wrapper,  # ✅ Wrapper qui retourne JSON string formaté
        inputs=[image_input],
        outputs=[result_output],
        api_name="predict"  # Crée l'endpoint /api/predict
    )
    
    clear_btn.click(
        fn=lambda: (None, ""),  # ✅ Retourner "" au lieu de None pour Textbox
        inputs=[],
        outputs=[image_input, result_output]
    )
    
    gr.Markdown("""
    ---
    **Développé par** : Aziz Ben Ammar & Équipe Labasni  
    **Version** : 1.0.0  
    **Contact** : [Hugging Face Space](https://huggingface.co/spaces/Syleto/labasni-detection)
    """)

# ============================================
# LANCEMENT
# ============================================

if __name__ == "__main__":
    print("\n" + "="*60)
    print("🌐 LANCEMENT DE L'INTERFACE GRADIO")
    print("="*60 + "\n")
    
    # ✅ IMPORTANT : Dans Gradio 5.x, l'API est automatiquement exposée avec api_name
    # L'interface API séparée garantit que l'API REST est accessible
    # Pas besoin de show_api (paramètre n'existe pas dans Gradio 5.x)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        share=False
    )