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

import os
import gradio as gr
from PIL import Image
from improved_mae_classifier import ImprovedMAEWasteClassifier

print("🚀 Initializing Improved MAE waste classifier...")
try:
    # Load the improved classifier with optimized settings
    classifier = ImprovedMAEWasteClassifier(
        hf_model_id="ysfad/mae-waste-classifier",
        temperature=2.5,  # Reduced overconfidence
        cardboard_penalty=0.8  # Reduced cardboard bias
    )
    print("✅ Improved MAE Classifier ready!")
except Exception as e:
    print(f"❌ Error loading improved classifier: {e}")
    raise

def classify_waste(image):
    """Classify waste item and provide disposal instructions with improved handling."""
    if image is None:
        return "Please upload an image.", "", "", ""
    
    try:
        # Classify the image using ensemble prediction for better accuracy
        result = classifier.classify_image(image, top_k=5, use_ensemble=True)
        
        if not result['success']:
            return f"Error: {result['error']}", "", "", ""
        
        predicted_class = result['predicted_class']
        confidence = result['confidence']
        top_predictions = result['top_predictions']
        
        # Format prediction result with confidence handling
        if predicted_class == "Uncertain":
            prediction_text = f"🤔 **Uncertain Classification**\n\nConfidence too low for reliable prediction ({confidence:.1%})\n\n💡 **Suggestions:**\n- Try a clearer photo\n- Better lighting\n- Different angle\n- Remove background clutter"
            confidence_text = f"Highest confidence: {confidence:.1%} (below threshold)"
        else:
            prediction_text = f"🎯 **{predicted_class}**\n\nConfidence: {confidence:.1%}"
            confidence_text = f"Confidence: {confidence:.1%}"
        
        # Get disposal instructions
        instructions = classifier.get_disposal_instructions(predicted_class)
        
        # Create detailed predictions table
        predictions_table = "| Rank | Class | Confidence |\n|------|-------|------------|\n"
        for i, pred in enumerate(top_predictions, 1):
            conf_percent = pred['confidence'] * 100
            predictions_table += f"| {i} | {pred['class']} | {conf_percent:.1f}% |\n"
        
        # Model information
        model_info = classifier.get_model_info()
        info_text = f"""**Model:** {model_info['model_name']}
**Architecture:** {model_info['architecture']}
**Classes:** {model_info['num_classes']}
**Device:** {model_info['device']}
**Improvements:** Temperature scaling, bias correction, ensemble prediction"""
        
        return prediction_text, confidence_text, instructions, predictions_table, info_text
        
    except Exception as e:
        return f"Error processing image: {str(e)}", "", "", "", ""

# Create Gradio interface with improved design
with gr.Blocks(
    title="🗂️ Improved MAE Waste Classifier",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        max-width: 1200px !important;
    }
    .header {
        text-align: center;
        padding: 20px;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        border-radius: 10px;
        margin-bottom: 20px;
    }
    .improvement-box {
        background: #e8f5e8;
        border: 2px solid #4caf50;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    .warning-box {
        background: #fff3cd;
        border: 2px solid #ffc107;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    """
) as demo:
    
    # Header
    gr.HTML("""
    <div class="header">
        <h1>🗂️ Improved MAE Waste Classifier</h1>
        <p>Enhanced AI-powered waste classification with bias correction and uncertainty handling</p>
        <p><strong>✨ New Features:</strong> Temperature scaling • Cardboard bias reduction • Uncertainty detection • Ensemble predictions</p>
    </div>
    """)
    
    # Improvements notice
    gr.HTML("""
    <div class="improvement-box">
        <h3>🎉 Recent Improvements</h3>
        <ul>
            <li><strong>✅ Reduced Cardboard Bias:</strong> From 83% to 17% false cardboard predictions</li>
            <li><strong>✅ Better Confidence:</strong> 39% reduction in overconfident predictions</li>
            <li><strong>✅ Uncertainty Handling:</strong> Shows "Uncertain" for low-confidence predictions</li>
            <li><strong>✅ Ensemble Predictions:</strong> Uses multiple augmentations for stability</li>
        </ul>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Image input
            image_input = gr.Image(
                label="📸 Upload Waste Image",
                type="pil",
                height=400
            )
            
            # Classification button
            classify_btn = gr.Button(
                "🔍 Classify Waste",
                variant="primary",
                size="lg"
            )
            
            # Quick tips
            gr.HTML("""
            <div class="warning-box">
                <h4>📋 Tips for Better Results:</h4>
                <ul>
                    <li>Use clear, well-lit photos</li>
                    <li>Center the item in frame</li>
                    <li>Avoid cluttered backgrounds</li>
                    <li>Try different angles if uncertain</li>
                </ul>
            </div>
            """)
        
        with gr.Column(scale=2):
            # Results section
            with gr.Group():
                gr.HTML("<h3>🎯 Classification Results</h3>")
                
                prediction_output = gr.Markdown(
                    label="Prediction",
                    value="Upload an image to get started!"
                )
                
                confidence_output = gr.Textbox(
                    label="📊 Confidence Score",
                    interactive=False
                )
                
                instructions_output = gr.Textbox(
                    label="♻️ Disposal Instructions",
                    lines=3,
                    interactive=False
                )
    
    # Detailed results section
    with gr.Row():
        with gr.Column():
            gr.HTML("<h3>📊 Detailed Predictions</h3>")
            predictions_table = gr.Markdown(
                label="Top 5 Predictions",
                value="| Rank | Class | Confidence |\n|------|-------|------------|\n| - | Upload image first | - |"
            )
        
        with gr.Column():
            gr.HTML("<h3>🤖 Model Information</h3>")
            model_info_output = gr.Markdown(
                label="Model Details",
                value="Model information will appear here after classification."
            )
    
    # About section
    with gr.Accordion("ℹ️ About This Improved Model", open=False):
        gr.HTML("""
        <div style="padding: 20px;">
            <h4>🧠 Model Architecture</h4>
            <p>This classifier uses a <strong>Vision Transformer (ViT-Base)</strong> pre-trained with <strong>Masked Autoencoder (MAE)</strong> and fine-tuned on the RealWaste dataset.</p>
            
            <h4>✨ Key Improvements</h4>
            <ul>
                <li><strong>Temperature Scaling (T=2.5):</strong> Reduces overconfident predictions</li>
                <li><strong>Cardboard Bias Correction:</strong> Applies 0.8x penalty to cardboard predictions</li>
                <li><strong>Class-specific Thresholds:</strong> Higher threshold (0.8) for cardboard, lower (0.4) for textile</li>
                <li><strong>Ensemble Prediction:</strong> Averages 5 augmented predictions for stability</li>
                <li><strong>Uncertainty Detection:</strong> Shows "Uncertain" when confidence is too low</li>
            </ul>
            
            <h4>📊 Performance Metrics</h4>
            <ul>
                <li><strong>Original Validation Accuracy:</strong> 93.27%</li>
                <li><strong>Cardboard Bias Reduction:</strong> 66.6% improvement</li>
                <li><strong>Confidence Calibration:</strong> 38.7% reduction in overconfidence</li>
                <li><strong>Classes:</strong> 9 waste categories</li>
            </ul>
            
            <h4>🗂️ Waste Categories</h4>
            <p><strong>Cardboard, Food Organics, Glass, Metal, Miscellaneous Trash, Paper, Plastic, Textile Trash, Vegetation</strong></p>
        </div>
        """)
    
    # Event handlers
    classify_btn.click(
        fn=classify_waste,
        inputs=[image_input],
        outputs=[
            prediction_output,
            confidence_output, 
            instructions_output,
            predictions_table,
            model_info_output
        ]
    )
    
    # Auto-classify on image upload
    image_input.change(
        fn=classify_waste,
        inputs=[image_input],
        outputs=[
            prediction_output,
            confidence_output,
            instructions_output,
            predictions_table,
            model_info_output
        ]
    )

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