#!/usr/bin/env python3
"""Gradio app for waste classification using finetuned CLIP model."""

import os
import gradio as gr
from PIL import Image
from clip_waste_classifier.finetuned_classifier import FinetunedCLIPWasteClassifier

# Initialize classifier with Hugging Face model
# Replace with your actual HF model ID after uploading
HF_MODEL_ID = "ysfad/openclip-finetune-waste"  # Your HF repository

print("🚀 Initializing CLIP waste classifier...")
try:
    # Try to load finetuned model from HF Hub, fallback to pretrained
    classifier = FinetunedCLIPWasteClassifier(hf_model_id=HF_MODEL_ID)
    print("✅ Classifier ready!")
except Exception as e:
    print(f"⚠️ Error loading classifier: {e}")
    print("🔄 Loading fallback classifier...")
    classifier = FinetunedCLIPWasteClassifier()

def classify_waste(image):
    """Classify waste item and provide disposal instructions."""
    if image is None:
        return "Please upload an image.", "", "", ""
    
    try:
        # Classify the image
        result = classifier.classify_image(image, top_k=5)
        
        if "error" in result:
            return f"Error: {result['error']}", "", "", ""
        
        # Get model info
        model_info = classifier.get_model_info()
        model_type = result.get('model_type', 'unknown')
        
        # Format main prediction
        main_prediction = f"""
        **🎯 Predicted Item:** {result['predicted_item']}
        **📂 Category:** {result['predicted_category']}
        **🎲 Confidence:** {result['best_confidence']:.3f}
        **🤖 Model:** {model_type.title()} CLIP ({model_info['model_name']})
        """
        
        # Format disposal instructions
        best_match = result['top_items'][0] if result['top_items'] else None
        disposal_text = best_match['disposal_method'] if best_match else "No instructions available"
        
        # Format detailed results table
        if result['top_items']:
            table_rows = []
            for i, item in enumerate(result['top_items'][:5], 1):
                table_rows.append([
                    str(i),
                    item['item'],
                    item['category'],
                    f"{item['confidence']:.3f}"
                ])
            
            # Create HTML table
            table_html = f"""
            <div style="margin-top: 15px;">
                <h4>🔍 Top 5 Predictions</h4>
                <table style="width: 100%; border-collapse: collapse;">
                    <thead>
                        <tr style="background-color: #f0f0f0;">
                            <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">#</th>
                            <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">Item</th>
                            <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">Category</th>
                            <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">Confidence</th>
                        </tr>
                    </thead>
                    <tbody>
            """
            
            for row in table_rows:
                table_html += f"""
                        <tr>
                            <td style="border: 1px solid #ddd; padding: 8px;">{row[0]}</td>
                            <td style="border: 1px solid #ddd; padding: 8px;"><strong>{row[1]}</strong></td>
                            <td style="border: 1px solid #ddd; padding: 8px;">{row[2]}</td>
                            <td style="border: 1px solid #ddd; padding: 8px;">{row[3]}</td>
                        </tr>
                """
            
            table_html += """
                    </tbody>
                </table>
            </div>
            """
        else:
            table_html = "<p>No predictions available.</p>"
        
        # Format model info
        model_info_text = f"""
        **Architecture:** {model_info['model_name']} 
        **Pretrained:** {model_info['pretrained']}
        **Classes:** {model_info['num_classes']} waste categories
        **Device:** {model_info['device'].upper()}
        **Type:** {model_type.title()} Model
        """
        
        return main_prediction, disposal_text, table_html, model_info_text
        
    except Exception as e:
        return f"Error during classification: {str(e)}", "", "", ""

# Create Gradio interface
with gr.Blocks(title="🗂️ AI Waste Classifier", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🗂️ AI Waste Classification System
    
    Upload an image of waste item to get **classification** and **disposal instructions**.
    
    Uses a **finetuned CLIP model** with **fluid OpenCLIP matching** against 2,205+ waste items!
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Input section
            gr.Markdown("### 📸 Upload Image")
            image_input = gr.Image(
                type="pil",
                label="Upload waste item image",
                height=300
            )
            
            classify_btn = gr.Button(
                "🔍 Classify Waste",
                variant="primary",
                size="lg"
            )
            
            # Model info section
            gr.Markdown("### 🤖 Model Information")
            model_info_output = gr.Markdown("")
        
        with gr.Column(scale=1):
            # Results section
            gr.Markdown("### 🎯 Classification Results")
            prediction_output = gr.Markdown("")
            
            gr.Markdown("### ♻️ Disposal Instructions")
            disposal_output = gr.Textbox(
                label="How to dispose of this item",
                lines=4,
                interactive=False
            )
            
            # Detailed results
            gr.Markdown("### 📊 Detailed Results")
            detailed_output = gr.HTML("")
    
    # Example images section
    gr.Markdown("### 💡 Try these examples:")
    gr.Examples(
        examples=[
            ["examples/plastic_bottle.jpg"],
            ["examples/cardboard_box.jpg"],
            ["examples/aluminum_can.jpg"],
            ["examples/glass_bottle.jpg"],
            ["examples/battery.jpg"]
        ] if os.path.exists("examples") else [],
        inputs=image_input,
        outputs=[prediction_output, disposal_output, detailed_output, model_info_output],
        fn=classify_waste,
        cache_examples=False
    )
    
    # Event handlers
    classify_btn.click(
        fn=classify_waste,
        inputs=image_input,
        outputs=[prediction_output, disposal_output, detailed_output, model_info_output]
    )
    
    image_input.change(
        fn=classify_waste,
        inputs=image_input,
        outputs=[prediction_output, disposal_output, detailed_output, model_info_output]
    )
    
    # Footer
    gr.Markdown("""
    ---
    **🔬 About:** This system uses a **finetuned CLIP** (ViT-B-16) model with **fluid OpenCLIP matching** 
    against a comprehensive database of 2,205+ waste items from Toronto municipal waste guidelines.
    
    **⚡ Performance:** The model encoder was finetuned on 30 waste categories achieving 91.33% accuracy, 
    then applied with fluid matching for maximum flexibility.
    """)

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