# ==========================================
# EMOTION DETECTION WEB APP
# Model: koyelog/face
# Compatible with Gradio 6.1.0
# Created by: Koyeliya Ghosh
# ==========================================

import gradio as gr
import torch
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import numpy as np
import warnings
warnings.filterwarnings('ignore')

print("="*70)
print("🎭 AI EMOTION DETECTOR - INITIALIZING")
print("="*70)

# ===== CONFIGURATION =====
MODEL_ID = "koyelog/face"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

print(f"\n📦 Model ID: {MODEL_ID}")
print(f"🖥️  Device: {DEVICE}")
print(f"💾 PyTorch Version: {torch.__version__}")

# ===== LOAD MODEL & PROCESSOR =====
print("\n⏳ Loading model from HuggingFace...")

try:
    model = ViTForImageClassification.from_pretrained(MODEL_ID)
    processor = ViTImageProcessor.from_pretrained(MODEL_ID)
    model.to(DEVICE)
    model.eval()
    print("✅ Model loaded successfully!")
    print(f"📊 Model Parameters: {sum(p.numel() for p in model.parameters()):,}")
    
except Exception as e:
    print(f"❌ ERROR loading model: {e}")
    raise

# ===== EMOTION CONFIGURATION =====
EMOTIONS = {
    0: {'name': 'Angry', 'emoji': '😠', 'color': '#ff4444', 'description': 'Showing anger or frustration'},
    1: {'name': 'Disgust', 'emoji': '🤢', 'color': '#44ff44', 'description': 'Expressing disgust or dislike'},
    2: {'name': 'Fear', 'emoji': '😨', 'color': '#9944ff', 'description': 'Showing fear or anxiety'},
    3: {'name': 'Happy', 'emoji': '😊', 'color': '#ffdd44', 'description': 'Expressing happiness or joy'},
    4: {'name': 'Sad', 'emoji': '😢', 'color': '#4444ff', 'description': 'Showing sadness or sorrow'},
    5: {'name': 'Surprise', 'emoji': '😲', 'color': '#ff44ff', 'description': 'Expressing surprise or shock'},
    6: {'name': 'Neutral', 'emoji': '😐', 'color': '#888888', 'description': 'No strong emotion detected'}
}

print(f"\n🎭 Loaded {len(EMOTIONS)} emotion classes:")
for idx, emo in EMOTIONS.items():
    print(f"   {idx}: {emo['emoji']} {emo['name']}")

# ===== PREDICTION FUNCTION =====
@torch.no_grad()
def predict_emotion(image):
    """
    Predict emotion from image
    Args:
        image: PIL Image or numpy array
    Returns:
        results_dict: Dictionary for Gradio Label
        html_output: Formatted HTML result
    """
    
    if image is None:
        return None, """
        <div style='text-align: center; padding: 40px; color: #ff4444; background: #ffe6e6; border-radius: 15px;'>
            <h2>⚠️ No Image Provided</h2>
            <p style='color: #666;'>Please upload an image or use webcam to capture!</p>
        </div>
        """
    
    try:
        # Convert numpy to PIL if needed
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # Convert to RGB
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        print(f"\n📸 Processing image: {image.size[0]}x{image.size[1]}")
        
        # Preprocess
        inputs = processor(images=image, return_tensors="pt")
        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
        
        # Inference
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.nn.functional.softmax(logits, dim=-1)[0].cpu()
        
        # Get predictions
        predicted_id = torch.argmax(probs).item()
        confidence = probs[predicted_id].item()
        
        # Get emotion details
        emotion = EMOTIONS[predicted_id]
        
        print(f"🎯 Prediction: {emotion['emoji']} {emotion['name']}")
        print(f"📊 Confidence: {confidence*100:.2f}%")
        
        # Format results for Gradio Label component
        results = {
            f"{EMOTIONS[i]['emoji']} {EMOTIONS[i]['name']}": float(probs[i])
            for i in range(len(EMOTIONS))
        }
        
        # Generate HTML output
        html = generate_result_html(
            emotion['name'],
            emotion['emoji'],
            emotion['color'],
            emotion['description'],
            confidence,
            probs
        )
        
        return results, html
        
    except Exception as e:
        print(f"❌ ERROR during prediction: {e}")
        
        error_html = f"""
        <div style='text-align: center; padding: 40px; background: #ffe6e6; border-radius: 15px;'>
            <h2 style='color: #ff4444;'>❌ Prediction Error</h2>
            <p style='color: #666;'>{str(e)}</p>
            <p style='color: #999; font-size: 0.9em;'>Please try a different image with a clear face</p>
        </div>
        """
        return None, error_html

# ===== HTML GENERATOR =====
def generate_result_html(name, emoji, color, description, confidence, probs):
    """Generate beautiful HTML result display"""
    
    # Calculate probability bars HTML
    bars_html = ""
    for idx in sorted(range(len(EMOTIONS)), key=lambda i: probs[i], reverse=True):
        emo = EMOTIONS[idx]
        prob = probs[idx].item()
        percentage = prob * 100
        bar_width = min(percentage, 100)
        
        bars_html += f"""
        <div style='margin: 12px 0;'>
            <div style='display: flex; justify-content: space-between; align-items: center; margin-bottom: 6px;'>
                <div style='display: flex; align-items: center; gap: 10px;'>
                    <span style='font-size: 1.8em;'>{emo['emoji']}</span>
                    <span style='font-weight: 600; color: #333;'>{emo['name']}</span>
                </div>
                <span style='font-weight: 700; color: {emo['color']}; font-size: 1.1em;'>{percentage:.1f}%</span>
            </div>
            <div style='width: 100%; background: #e9ecef; border-radius: 10px; height: 12px; overflow: hidden; box-shadow: inset 0 2px 4px rgba(0,0,0,0.06);'>
                <div style='width: {bar_width}%; background: linear-gradient(90deg, {emo['color']}, {emo['color']}dd); height: 100%; transition: width 0.8s ease; border-radius: 10px;'></div>
            </div>
        </div>
        """
    
    # Main HTML
    html = f"""
    <div style='font-family: "Segoe UI", -apple-system, BlinkMacSystemFont, sans-serif; max-width: 100%;'>
        
        <!-- Main Result Card -->
        <div style='
            text-align: center;
            padding: 50px 30px;
            background: linear-gradient(135deg, {color}18 0%, {color}30 100%);
            border-radius: 25px;
            box-shadow: 0 10px 40px rgba(0,0,0,0.12);
            margin-bottom: 30px;
            border: 2px solid {color}40;
        '>
            <div style='
                font-size: 120px;
                margin: 0 0 20px 0;
                display: inline-block;
                animation: bounce 1s ease infinite;
            '>
                {emoji}
            </div>
            
            <h1 style='
                color: {color};
                font-size: 3.5em;
                margin: 20px 0 10px 0;
                font-weight: 800;
                text-shadow: 2px 2px 8px rgba(0,0,0,0.1);
            '>
                {name}
            </h1>
            
            <p style='
                font-size: 1.3em;
                color: #555;
                margin: 15px 0;
                font-weight: 500;
            '>
                {description}
            </p>
            
            <div style='
                display: inline-flex;
                align-items: center;
                gap: 15px;
                margin: 25px 0;
                padding: 15px 35px;
                background: white;
                border-radius: 50px;
                box-shadow: 0 4px 20px rgba(0,0,0,0.1);
            '>
                <span style='font-size: 1.2em; color: #666;'>Confidence:</span>
                <span style='font-size: 2em; font-weight: 800; color: {color};'>{confidence*100:.1f}%</span>
            </div>
            
            <!-- Confidence Bar -->
            <div style='
                width: 100%;
                max-width: 500px;
                height: 50px;
                background: #e9ecef;
                border-radius: 25px;
                overflow: hidden;
                margin: 30px auto 0;
                box-shadow: inset 0 4px 8px rgba(0,0,0,0.1);
            '>
                <div style='
                    width: {confidence*100}%;
                    height: 100%;
                    background: linear-gradient(90deg, {color}, {color}cc);
                    border-radius: 25px;
                    transition: width 1.5s ease;
                    display: flex;
                    align-items: center;
                    justify-content: center;
                '>
                    <span style='
                        color: white;
                        font-weight: 800;
                        font-size: 1.3em;
                        text-shadow: 0 2px 4px rgba(0,0,0,0.3);
                    '>
                        {confidence*100:.1f}%
                    </span>
                </div>
            </div>
        </div>
        
        <!-- Detailed Breakdown -->
        <div style='
            background: white;
            padding: 35px;
            border-radius: 20px;
            box-shadow: 0 8px 32px rgba(0,0,0,0.08);
            border: 1px solid #e9ecef;
        '>
            <h2 style='
                margin: 0 0 25px 0;
                color: #333;
                font-size: 1.8em;
                font-weight: 700;
            '>
                📊 Detailed Emotion Analysis
            </h2>
            
            {bars_html}
        </div>
        
        <!-- Model Info -->
        <div style='
            margin-top: 25px;
            padding: 20px;
            background: linear-gradient(135deg, #f8f9fa, #e9ecef);
            border-radius: 15px;
            text-align: center;
            font-size: 0.9em;
            color: #666;
        '>
            <p style='margin: 5px 0;'>
                <strong>Model:</strong> koyelog/face | 
                <strong>Accuracy:</strong> 98.80% | 
                <strong>Parameters:</strong> 85.8M
            </p>
        </div>
    </div>
    
    <style>
        @keyframes bounce {{
            0%, 100% {{ transform: translateY(0); }}
            50% {{ transform: translateY(-20px); }}
        }}
    </style>
    """
    
    return html

# ===== GRADIO INTERFACE =====
print("\n🎨 Building Gradio interface...")

# Custom CSS
custom_css = """
.gradio-container {
    font-family: 'Segoe UI', -apple-system, BlinkMacSystemFont, sans-serif !important;
    max-width: 1400px !important;
}

.gr-button-primary {
    background: linear-gradient(135deg, #667eea, #764ba2) !important;
    border: none !important;
    font-size: 18px !important;
    font-weight: 600 !important;
    padding: 16px 40px !important;
    border-radius: 12px !important;
}

footer {
    visibility: hidden !important;
}
"""

# Create Gradio Interface
with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="purple",
        secondary_hue="pink"
    ),
    css=custom_css,
    title="🎭 AI Emotion Detector | koyelog"
) as demo:
    
    # Header
    gr.HTML("""
        <div style='
            text-align: center;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 60px 30px;
            border-radius: 25px;
            margin-bottom: 40px;
            box-shadow: 0 15px 50px rgba(102, 126, 234, 0.3);
        '>
            <h1 style='font-size: 4em; margin: 0; font-weight: 900;'>
                🎭 AI Emotion Detector
            </h1>
            <p style='font-size: 1.5em; margin: 20px 0 10px; opacity: 0.95;'>
                Powered by Vision Transformer | 98.80% Accuracy
            </p>
            <p style='font-size: 1.1em; opacity: 0.85;'>
                Model: <strong>koyelog/face</strong> | 85.8M Parameters
            </p>
            <div style='margin-top: 20px; display: flex; gap: 15px; justify-content: center; flex-wrap: wrap;'>
                <span style='background: rgba(255,255,255,0.25); padding: 10px 25px; border-radius: 25px;'>😠 Angry</span>
                <span style='background: rgba(255,255,255,0.25); padding: 10px 25px; border-radius: 25px;'>🤢 Disgust</span>
                <span style='background: rgba(255,255,255,0.25); padding: 10px 25px; border-radius: 25px;'>😨 Fear</span>
                <span style='background: rgba(255,255,255,0.25); padding: 10px 25px; border-radius: 25px;'>😊 Happy</span>
                <span style='background: rgba(255,255,255,0.25); padding: 10px 25px; border-radius: 25px;'>😢 Sad</span>
                <span style='background: rgba(255,255,255,0.25); padding: 10px 25px; border-radius: 25px;'>😲 Surprise</span>
                <span style='background: rgba(255,255,255,0.25); padding: 10px 25px; border-radius: 25px;'>😐 Neutral</span>
            </div>
        </div>
    """)
    
    with gr.Tabs():
        
        # TAB 1: WEBCAM
        with gr.Tab("📹 Live Webcam"):
            gr.Markdown("""
            ### 🎥 Capture Your Face
            Use your webcam to capture your face and detect your emotion in real-time!
            """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    webcam_input = gr.Image(
                        sources=["webcam"],
                        type="pil",
                        label="📸 Your Face"
                    )
                    webcam_button = gr.Button(
                        "🔍 Detect My Emotion",
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column(scale=1):
                    webcam_html = gr.HTML(label="🎯 Result")
                    webcam_label = gr.Label(
                        label="📊 Probabilities",
                        num_top_classes=7
                    )
            
            webcam_button.click(
                fn=predict_emotion,
                inputs=webcam_input,
                outputs=[webcam_label, webcam_html]
            )
        
        # TAB 2: UPLOAD
        with gr.Tab("🖼️ Upload Image"):
            gr.Markdown("""
            ### 📤 Upload Face Image
            Upload or drag & drop an image to detect emotions!
            """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    image_input = gr.Image(
                        type="pil",
                        label="🖼️ Upload Image",
                        sources=["upload", "clipboard"]
                    )
                    image_button = gr.Button(
                        "🔍 Detect Emotion",
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column(scale=1):
                    image_html = gr.HTML(label="🎯 Result")
                    image_label = gr.Label(
                        label="📊 Probabilities",
                        num_top_classes=7
                    )
            
            image_button.click(
                fn=predict_emotion,
                inputs=image_input,
                outputs=[image_label, image_html]
            )
    
    # Footer
    gr.HTML("""
        <div style='
            text-align: center;
            margin-top: 60px;
            padding: 50px 30px;
            background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
            border-radius: 25px;
        '>
            <h2 style='color: #333; margin-bottom: 30px;'>📊 Model Information</h2>
            
            <div style='
                display: grid;
                grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
                gap: 25px;
                margin: 30px 0;
            '>
                <div style='background: white; padding: 25px; border-radius: 15px;'>
                    <p style='font-weight: 700; color: #667eea; margin-bottom: 10px;'>Model ID</p>
                    <p style='font-size: 1.2em; color: #333;'>koyelog/face</p>
                </div>
                <div style='background: white; padding: 25px; border-radius: 15px;'>
                    <p style='font-weight: 700; color: #667eea; margin-bottom: 10px;'>Architecture</p>
                    <p style='font-size: 1.2em; color: #333;'>Vision Transformer</p>
                </div>
                <div style='background: white; padding: 25px; border-radius: 15px;'>
                    <p style='font-weight: 700; color: #667eea; margin-bottom: 10px;'>Parameters</p>
                    <p style='font-size: 1.2em; color: #333;'>85.8 Million</p>
                </div>
                <div style='background: white; padding: 25px; border-radius: 15px;'>
                    <p style='font-weight: 700; color: #667eea; margin-bottom: 10px;'>Accuracy</p>
                    <p style='font-size: 1.2em; color: #333;'>98.80%</p>
                </div>
            </div>
            
            <p style='color: #666; margin-top: 30px;'>
                Created by <strong style='color: #667eea;'>Koyeliya Ghosh</strong><br>
                <a href='https://huggingface.co/koyelog/face' target='_blank' style='color: #667eea;'>
                    View Model →
                </a>
            </p>
        </div>
    """)

# ===== LAUNCH =====
if __name__ == "__main__":
    print("\n" + "="*70)
    print("🚀 LAUNCHING APP")
    print("="*70)
    print("✅ Model ready")
    print("✅ Starting server...\n")
    
    demo.launch()