import gradio as gr
from transformers import pipeline
from PIL import Image
import numpy as np

# Catégories de vêtements bien définies
FASHION_CATEGORIES = [
    "t-shirt", "button-down shirt", "polo shirt", 
    "sweatshirt", "hoodie", "sweater",
    "jacket", "coat", "blazer",
    "dress", "long dress", "short dress",
    "skirt", "long skirt", "short skirt",
    "jeans", "pants", "shorts", "leggings",
    "sneakers", "boots", "heels", "sandals"
]

print("🔧 Loading classification model...")
# Modèle principal pour l'analyse globale
class_pipe = pipeline(
    "zero-shot-image-classification", 
    model="openai/clip-vit-base-patch32"
)

# Modèle de secours pour confirmation
backup_pipe = pipeline(
    "image-classification",
    model="google/vit-base-patch16-224"
)

print("✅ Models loaded successfully!")

def analyze_complete_image(image):
    """Analyse l'image ENTIÈRE sans segmentation"""
    try:
        if image is None:
            return "❌ Please upload an image first", None
            
        # Conversion en PIL Image si nécessaire
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # Réduction de taille pour de meilleures performances
        image = image.resize((512, 512))
        
        # 🔥 ANALYSE PRINCIPALE - Image entière
        predictions = class_pipe(
            image, 
            candidate_labels=FASHION_CATEGORIES,
            hypothesis_template="a complete photo of {}",
            multi_label=False
        )
        
        # 🔥 ANALYSE DE CONFIRMATION avec modèle secondaire
        backup_preds = backup_pipe(image)
        
        # Filtrage des résultats peu confidentiels
        confident_predictions = [
            p for p in predictions 
            if p['score'] > 0.15  # Seuil de confiance augmenté
        ]
        
        if not confident_predictions:
            return "❌ No confident prediction. Try a clearer image.", image
        
        # Formatage des résultats
        result_text = "🎯 **Fashion Analysis Results:**\n\n"
        result_text += "**Main predictions:**\n"
        
        for i, pred in enumerate(confident_predictions[:3]):
            result_text += f"{i+1}. **{pred['label']}**: {pred['score']*100:.1f}%\n"
        
        # Ajouter la prédiction du modèle de secours
        if backup_preds:
            result_text += f"\n**Secondary model suggests**: {backup_preds[0]['label']}\n"
            result_text += f"**Confidence**: {backup_preds[0]['score']*100:.1f}%"
        
        # Conseils basés sur la prédiction
        top_pred = confident_predictions[0]['label']
        result_text += f"\n\n💡 **Tip**: For better accuracy, make sure the {top_pred} is clearly visible and centered."
        
        return result_text, image
        
    except Exception as e:
        return f"❌ Error: {str(e)}", None

# Interface optimisée
with gr.Blocks(
    title="Fashion AI - Complete Image Analysis",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container { max-width: 900px; margin: auto; }
    .header { text-align: center; margin-bottom: 20px; }
    """
) as demo:
    
    gr.Markdown("""
    <div class='header'>
    <h1>👗 Fashion AI - Complete Image Analysis</h1>
    <p>Analyzes the ENTIRE image without cropping or segmentation</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload Image")
            image_input = gr.Image(
                type="pil", 
                label="Upload Complete Image",
                height=300
            )
            analyze_btn = gr.Button(
                "🔍 Analyze Complete Image", 
                variant="primary",
                size="lg"
            )
            
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Analysis Results")
            output_text = gr.Markdown(
                label="Results",
                show_label=False
            )
            output_image = gr.Image(
                label="Original Image (for reference)",
                interactive=False,
                height=300
            )
    
    # Section d'instructions
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            ### 💡 Best Practices:
            - ✅ **Full garment visible** - don't crop
            - ✅ **Good lighting** - no shadows
            - ✅ **Neutral background** - less distraction
            - ✅ **Single item** - one piece per photo
            - ✅ **Clear view** - front angle preferred
            """)
    
    # Section d'exemples
    with gr.Row():
        gr.Markdown("""
        ### 🎯 Examples of good images:
        - Full t-shirt visible on plain background
        - Complete dress without cropping
        - Entire pair of jeans clearly visible
        """)
    
    # Événements
    analyze_btn.click(
        fn=analyze_complete_image,
        inputs=image_input,
        outputs=[output_text, output_image]
    )
    
    image_input.upload(
        fn=analyze_complete_image,
        inputs=image_input,
        outputs=[output_text, output_image]
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )