#!/usr/bin/env python3
"""Gradio app for waste classification using finetuned MAE ViT-Base model."""

import os
import gradio as gr
from PIL import Image
from mae_waste_classifier import MAEWasteClassifier

print("🚀 Initializing MAE waste classifier...")
try:
    # Load the finetuned MAE model from Hugging Face Hub
    classifier = MAEWasteClassifier(hf_model_id="ysfad/mae-waste-classifier")
    print("✅ MAE Classifier ready!")
except Exception as e:
    print(f"❌ Error loading MAE classifier: {e}")
    raise

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 not result['success']:
            return f"Error: {result['error']}", "", "", ""
        
        # Get model info
        model_info = classifier.get_model_info()
        
        # Format main prediction
        main_prediction = f"""
        **🎯 Predicted Class:** {result['predicted_class']}
        **🎲 Confidence:** {result['confidence']:.3f}
        **🤖 Model:** {model_info['model_name']}
        **🏆 Validation Accuracy:** 93.27%
        """
        
        # Get disposal instructions
        disposal_text = classifier.get_disposal_instructions(result['predicted_class'])
        
        # Format detailed results table
        if result['top_predictions']:
            table_rows = []
            for i, pred in enumerate(result['top_predictions'], 1):
                table_rows.append([
                    str(i),
                    pred['class'],
                    f"{pred['confidence']:.3f}"
                ])
            
            # Create HTML table
            table_html = f"""
            <div style="margin-top: 15px;">
                <h4>🔍 Top {len(result['top_predictions'])} 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;">Class</th>
                            <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">Confidence</th>
                        </tr>
                    </thead>
                    <tbody>
            """
            
            for row in table_rows:
                # Color coding based on confidence
                confidence_val = float(row[2])
                if confidence_val > 0.7:
                    row_color = "#e8f5e8"  # Light green
                elif confidence_val > 0.4:
                    row_color = "#fff3cd"  # Light yellow
                else:
                    row_color = "#f8d7da"  # Light red
                
                table_html += f"""
                        <tr style="background-color: {row_color};">
                            <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>
                        </tr>
                """
            
            table_html += """
                    </tbody>
                </table>
            </div>
            """
        else:
            table_html = "<p>No predictions available.</p>"
        
        # Format model info
        model_info_text = f"""
        **Architecture:** {model_info['architecture']} 
        **Pretrained:** {model_info['pretrained']}
        **Classes:** {model_info['num_classes']} waste categories
        **Device:** {model_info['device'].upper()}
        **Training:** Finetuned on RealWaste dataset (4,752 images)
        **Performance:** 93.27% validation accuracy
        **Model Hub:** [ysfad/mae-waste-classifier](https://huggingface.co/ysfad/mae-waste-classifier)
        """
        
        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="🗂️ MAE Waste Classifier", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🗂️ MAE Waste Classification System
    
    Upload an image of waste item to get **classification** and **disposal instructions**.
    
    Uses a **finetuned MAE ViT-Base model** achieving **93.27% validation accuracy** on 9 waste categories!
    
    **Model:** [ysfad/mae-waste-classifier](https://huggingface.co/ysfad/mae-waste-classifier)
    """)
    
    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 (if available)
    if os.path.exists("examples"):
        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"]
            ],
            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 **MAE (Masked Autoencoder) ViT-Base** model finetuned on the RealWaste dataset.
    The model was pretrained with MAE self-supervised learning and then finetuned for waste classification.
    
    **⚡ Performance:** Achieved **93.27% validation accuracy** on 9 waste categories with 4,752 training images.
    
    **📊 Categories:** Cardboard, Food Organics, Glass, Metal, Miscellaneous Trash, Paper, Plastic, Textile Trash, Vegetation
    
    **🤗 Model:** [ysfad/mae-waste-classifier](https://huggingface.co/ysfad/mae-waste-classifier)
    """)

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