# app.py - Optimized for Hugging Face Spaces
import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image, ImageDraw, ImageFilter, ImageEnhance
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
import io
import time
from datetime import datetime
import base64

# Load the trained model
print("🔄 Loading the trained model...")
try:
    model = tf.keras.models.load_model('model.h5')
    print("✅ Model loaded successfully!")
except Exception as e:
    print(f"❌ Error loading model: {e}")
    raise

class DigitRecognizer:
    def __init__(self, model):
        self.model = model
        self.preprocessing_time = 0
        self.prediction_time = 0
    
    def preprocess_image(self, image):
        """Optimized preprocessing without OpenCV"""
        start_time = time.time()
        
        try:
            # Convert to grayscale if needed
            if image.mode != 'L':
                image = image.convert('L')
            
            # Enhance contrast
            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(2.0)
            
            # Apply simple blur for noise reduction (using PIL)
            image = image.filter(ImageFilter.SMOOTH_MORE)
            
            # Resize to 28x28 pixels
            image = image.resize((28, 28), Image.Resampling.LANCZOS)
            
            # Convert to numpy array
            img_array = np.array(image)
            
            # Normalize to 0-1
            img_array = img_array.astype('float32') / 255.0
            
            # Simple background detection and inversion
            if np.mean(img_array) > 0.5:
                img_array = 1.0 - img_array
            
            # Reshape for model
            img_array = img_array.reshape(1, 28, 28, 1)
            
            self.preprocessing_time = time.time() - start_time
            return img_array
            
        except Exception as e:
            print(f"Preprocessing error: {e}")
            return None
    
    def predict(self, image):
        """Prediction with confidence analysis"""
        start_time = time.time()
        
        try:
            if image is None:
                return None, None, None
            
            # Preprocess image
            processed_image = self.preprocess_image(image)
            if processed_image is None:
                return None, None, None
            
            # Make prediction
            predictions = self.model.predict(processed_image, verbose=0)[0]
            
            # Get top prediction
            predicted_digit = np.argmax(predictions)
            confidence = np.max(predictions)
            
            # Create confidence data
            confidence_data = []
            for i, prob in enumerate(predictions):
                confidence_data.append({
                    'digit': i,
                    'confidence': float(prob),
                    'percentage': f"{prob:.2%}",
                    'is_top': i == predicted_digit
                })
            
            # Sort by confidence
            confidence_data.sort(key=lambda x: x['confidence'], reverse=True)
            
            self.prediction_time = time.time() - start_time
            
            return predicted_digit, confidence, confidence_data
            
        except Exception as e:
            print(f"Prediction error: {e}")
            return None, None, None

# Initialize recognizer
recognizer = DigitRecognizer(model)

def create_confidence_chart(confidence_data):
    """Create confidence visualization"""
    if not confidence_data:
        return None
    
    digits = [str(item['digit']) for item in confidence_data]
    confidences = [item['confidence'] for item in confidence_data]
    colors = ['#FF4B4B' if item['is_top'] else '#4A90E2' for item in confidence_data]
    
    plt.figure(figsize=(10, 5))
    bars = plt.bar(digits, confidences, color=colors, alpha=0.8, edgecolor='white', linewidth=2)
    
    # Add value labels
    for bar, conf in zip(bars, confidences):
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                f'{conf:.1%}', ha='center', va='bottom', fontweight='bold', fontsize=10)
    
    plt.ylabel('Confidence Score', fontweight='bold', fontsize=12)
    plt.xlabel('Digits', fontweight='bold', fontsize=12)
    plt.title('Digit Recognition Confidence', fontsize=14, fontweight='bold', pad=20)
    plt.ylim(0, 1)
    plt.grid(axis='y', alpha=0.3)
    plt.tight_layout()
    
    # Convert to image
    buf = io.BytesIO()
    plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
    buf.seek(0)
    plt.close()
    
    return Image.open(buf)

def create_confidence_table(confidence_data):
    """Create HTML confidence table"""
    if not confidence_data:
        return ""
    
    html = """
    <div style="background: linear-gradient(135deg, #667eea20, #764ba220); padding: 20px; border-radius: 15px; border: 1px solid #e0e0e0;">
        <h3 style="margin-top: 0; margin-bottom: 20px; color: #333;">🎯 Top Predictions</h3>
    """
    
    for item in confidence_data[:3]:  # Show top 3
        bar_width = int(item['confidence'] * 100)
        bar_color = "#FF4B4B" if item['is_top'] else "#4A90E2"
        star = "⭐" if item['is_top'] else ""
        
        html += f"""
        <div style="margin-bottom: 15px;">
            <div style="display: flex; justify-content: space-between; margin-bottom: 5px;">
                <span style="font-weight: bold; color: #333;">{star} Digit {item['digit']}</span>
                <span style="color: {bar_color}; font-weight: bold;">{item['percentage']}</span>
            </div>
            <div style="background: #e0e0e0; border-radius: 10px; height: 25px; overflow: hidden;">
                <div style="background: {bar_color}; height: 100%; width: {bar_width}%; border-radius: 10px; transition: width 0.3s ease;"></div>
            </div>
        </div>
        """
    
    html += "</div>"
    return html

def recognize_digit(image):
    """Main recognition function"""
    if image is None:
        return "Please draw a digit first", None, "0 ms", "0 ms", ""
    
    digit, confidence, confidence_data = recognizer.predict(image)
    
    if digit is None:
        return "Could not process image", None, "0 ms", "0 ms", ""
    
    # Determine confidence level
    if confidence > 0.95:
        emoji = "🎯"
        message = "Excellent match!"
        color = "#10B981"
    elif confidence > 0.8:
        emoji = "✅"
        message = "Good recognition!"
        color = "#3B82F6"
    elif confidence > 0.6:
        emoji = "⚠️"
        message = "Moderate confidence"
        color = "#F59E0B"
    else:
        emoji = "❓"
        message = "Low confidence"
        color = "#EF4444"
    
    # Format timing
    preprocess_time = f"{recognizer.preprocessing_time*1000:.1f} ms"
    predict_time = f"{recognizer.prediction_time*1000:.1f} ms"
    
    # Create HTML output
    output = f"""
    <div style="background: {color}10; padding: 20px; border-radius: 15px; border-left: 5px solid {color};">
        <div style="display: flex; align-items: center; gap: 10px; margin-bottom: 15px;">
            <span style="font-size: 2em;">{emoji}</span>
            <h2 style="margin: 0; color: {color};">Recognition Result</h2>
        </div>
        
        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 15px; margin-bottom: 15px;">
            <div style="background: white; padding: 20px; border-radius: 12px; text-align: center;">
                <div style="color: #666; margin-bottom: 5px;">Predicted Digit</div>
                <span style="font-size: 3em; font-weight: bold; color: {color};">{digit}</span>
            </div>
            <div style="background: white; padding: 20px; border-radius: 12px; text-align: center;">
                <div style="color: #666; margin-bottom: 5px;">Confidence</div>
                <span style="font-size: 2em; font-weight: bold; color: {color};">{confidence:.1%}</span>
            </div>
        </div>
        
        <div style="background: white; padding: 15px; border-radius: 12px;">
            <div style="display: flex; justify-content: space-between; align-items: center;">
                <span><strong>Status:</strong> {message}</span>
                <span><strong>Time:</strong> {datetime.now().strftime('%H:%M:%S')}</span>
            </div>
        </div>
    </div>
    """
    
    # Create visualizations
    chart_image = create_confidence_chart(confidence_data)
    table_html = create_confidence_table(confidence_data)
    
    return output, chart_image, preprocess_time, predict_time, table_html

def clear_all():
    """Clear all inputs and outputs"""
    return None, None, "0 ms", "0 ms", ""

# Custom CSS
custom_css = """
.gradio-container {
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
}
.main-header {
    text-align: center;
    padding: 2rem;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 20px;
    margin-bottom: 2rem;
    color: white;
}
.drawing-area {
    border: 2px dashed #cbd5e1 !important;
    border-radius: 15px !important;
    background: #f8fafc !important;
    min-height: 350px !important;
}
.predict-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    border: none !important;
    border-radius: 25px !important;
    font-weight: 600 !important;
    padding: 12px 24px !important;
    transition: transform 0.2s !important;
}
.predict-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 10px 20px rgba(102, 126, 234, 0.3) !important;
}
.clear-btn {
    background: white !important;
    border: 2px solid #667eea !important;
    color: #667eea !important;
    border-radius: 25px !important;
    font-weight: 600 !important;
    padding: 12px 24px !important;
}
.stats-card {
    background: linear-gradient(135deg, #f8f9fa 0%, #e9ecef 100%) !important;
    padding: 15px !important;
    border-radius: 12px !important;
    border: none !important;
    box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important;
}
.footer {
    text-align: center;
    padding: 20px;
    color: #666;
    font-size: 0.9em;
    margin-top: 40px;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    
    # Header
    with gr.Row():
        with gr.Column():
            gr.HTML("""
                <div class="main-header">
                    <h1>🔢 Handwritten Digit Recognition</h1>
                    <p style="font-size: 1.2em; opacity: 0.9;">Powered by Deep Learning CNN</p>
                </div>
            """)
    
    with gr.Row():
        # Left Column
        with gr.Column(scale=1):
            gr.Markdown("### 🎨 Draw Your Digit")
            input_image = gr.Image(
                label="Draw here (0-9)",
                image_mode="L",
                type="pil",
                height=350,
                elem_classes=["drawing-area"]
            )
            
            with gr.Row():
                clear_btn = gr.Button(
                    "🗑️ Clear", 
                    elem_classes=["clear-btn"],
                    scale=1
                )
                predict_btn = gr.Button(
                    "🚀 Predict Digit", 
                    elem_classes=["predict-btn"],
                    scale=1
                )
            
            # Performance metrics
            gr.Markdown("### ⚡ Performance")
            with gr.Row():
                preprocess_time = gr.Textbox(
                    label="Preprocessing",
                    value="0 ms",
                    interactive=False,
                    elem_classes=["stats-card"]
                )
                inference_time = gr.Textbox(
                    label="Inference", 
                    value="0 ms",
                    interactive=False,
                    elem_classes=["stats-card"]
                )
        
        # Right Column
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Results")
            output_text = gr.HTML(
                label="Prediction",
                value="<div style='text-align: center; padding: 50px; color: #666;'>🎨 Draw a digit to see results</div>"
            )
            
            gr.Markdown("### 📈 Confidence Distribution")
            confidence_chart = gr.Image(
                label="Confidence Chart",
                interactive=False,
                height=250
            )
            
            confidence_table = gr.HTML(
                label="Confidence Details",
                value=""
            )
    
    # Footer
    gr.HTML("""
        <div class="footer">
            <p>Built with TensorFlow + Gradio | Deployed on Hugging Face Spaces</p>
            <p style="font-size: 0.8em;">Draw any digit 0-9 and see AI recognition in real-time</p>
        </div>
    """)
    
    # Event handlers
    predict_btn.click(
        fn=recognize_digit,
        inputs=[input_image],
        outputs=[output_text, confidence_chart, preprocess_time, inference_time, confidence_table]
    )
    
    clear_btn.click(
        fn=clear_all,
        inputs=[],
        outputs=[input_image, output_text, preprocess_time, inference_time, confidence_table]
    )

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