import gradio as gr
import cv2
import numpy as np
from huggingface_hub import hf_hub_download

# Load model
print("Loading model...")
try:
    import tensorflow as tf
    model_path = hf_hub_download(
        repo_id="maftuh-main/meme-emotion-detector",
        filename="model.tflite"
    )
    
    # Load TFLite interpreter
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()
    
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    
    print("✓ Model loaded successfully!")
    MODEL_LOADED = True
except Exception as e:
    print(f"Error loading model: {e}")
    MODEL_LOADED = False

# Emotion labels
EMOTIONS = ['angry', 'disgust', 'fear', 'happy', 'sad', 'surprise', 'neutral']

# Emotion to emoji mapping
EMOJI_MAP = {
    'angry': '😠',
    'disgust': '🤢',
    'fear': '😨',
    'happy': '😊',
    'sad': '😢',
    'surprise': '😲',
    'neutral': '😐'
}

def predict_emotion(image):
    """Predict emotion from image."""
    if not MODEL_LOADED:
        return {"error": 1.0}
    
    try:
        # Convert to grayscale
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image
        
        # Resize to 48x48
        resized = cv2.resize(gray, (48, 48))
        
        # Normalize
        normalized = resized.astype(np.float32) / 255.0
        
        # Add batch and channel dimensions
        input_data = np.expand_dims(normalized, axis=0)
        input_data = np.expand_dims(input_data, axis=-1)
        
        # Run inference
        interpreter.set_tensor(input_details[0]['index'], input_data)
        interpreter.invoke()
        
        # Get predictions
        predictions = interpreter.get_tensor(output_details[0]['index'])[0]
        
        # Create results dict
        results = {}
        for i, emotion in enumerate(EMOTIONS):
            emoji = EMOJI_MAP.get(emotion, '')
            results[f"{emoji} {emotion}"] = float(predictions[i])
        
        return results
        
    except Exception as e:
        print(f"Prediction error: {e}")
        return {"error": 1.0}

# Create Gradio interface
demo = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Image(label="Upload foto wajah Anda"),
    outputs=gr.Label(num_top_classes=7, label="Emotion Predictions"),
    title="🎭 Meme Emotion Detector",
    description="""
    **Deteksi ekspresi wajah dengan 7 emosi:**
    - 😊 Happy
    - 😢 Sad
    - 😠 Angry
    - 😲 Surprise
    - 😐 Neutral
    - 🤢 Disgust
    - 😨 Fear
    
    Upload foto wajah Anda untuk melihat emosi yang terdeteksi!
    """,
    examples=[],
    article="""
    ### Model Details
    - **Framework:** TensorFlow Lite
    - **Input:** 48x48 grayscale image
    - **Output:** 7 emotion classes
    
    **Repository:** [maftuh-main/meme-emotion-detector](https://huggingface.co/maftuh-main/meme-emotion-detector)
    """
)

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