import cv2
import numpy as np
import gradio as gr
from fer import FER

# FER expects RGB images; OpenCV frames are BGR. We'll convert.
detector = FER(mtcnn=True)  # robust face detector; set mtcnn=False if you want a bit more speed

EMOJI = {
    "angry": "😠",
    "disgust": "🤢",
    "fear": "😨",
    "happy": "😄",
    "sad": "😢",
    "surprise": "😲",
    "neutral": "😐",
}

def annotate_emotions(frame: np.ndarray):
    """
    frame: HxWxC (uint8) in RGB from gradio webcam input
    returns: (annotated_frame_RGB, top3_emotions_label_dict)
    """
    if frame is None:
        return None, {}

    # Ensure uint8 RGB
    img_rgb = frame.astype(np.uint8)

    # Run FER
    results = detector.detect_emotions(img_rgb)  # list of {box: (x,y,w,h), emotions: {...}}

    # Prepare overlay
    annotated = img_rgb.copy()

    # Collect aggregate scores for Label output
    top_scores = {}

    for r in results:
        (x, y, w, h) = r["box"]
        emotions = r["emotions"]  # dict of emotion->score

        # Top emotion for this face
        if emotions:
            top_em, top_score = sorted(emotions.items(), key=lambda kv: kv[1], reverse=True)[0]
            # Draw rectangle
            cv2.rectangle(annotated, (x, y), (x + w, y + h), (0, 255, 0), 2)

            # Compose label with emoji
            label = f"{EMOJI.get(top_em, '')} {top_em} {top_score:.2f}"
            # Background for text
            (tw, th), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
            cv2.rectangle(annotated, (x, y - th - 10), (x + tw + 6, y), (0, 0, 0), -1)
            cv2.putText(annotated, label, (x + 3, y - 6),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2, cv2.LINE_AA)

            # Keep the best score per emotion for the Label component
            for k, v in emotions.items():
                top_scores[k] = max(top_scores.get(k, 0.0), float(v))

    # Sort and keep top-3 to keep the UI tidy
    top3 = dict(sorted(top_scores.items(), key=lambda kv: kv[1], reverse=True)[:3])

    # Gradio expects RGB; 'annotated' is already RGB
    return annotated, top3


with gr.Blocks(fill_height=True) as demo:
    gr.Markdown(
        """
        # 🎭 Live Emotion Detection
        - Allow camera access and look at the preview.
        - The app detects faces and overlays the **top emotion** per face (plus confidence).
        """
    )
    with gr.Row():
        with gr.Column(scale=3):
            cam = gr.Image(sources=["webcam"], streaming=True, label="Webcam", height=420)
        with gr.Column(scale=2):
            out_img = gr.Image(label="Annotated", height=420)
            out_label = gr.Label(label="Top emotions (global top-3)")

    cam.stream(fn=annotate_emotions, inputs=cam, outputs=[out_img, out_label])

if __name__ == "__main__":
    # Share=True is handy for quick testing on local networks
    demo.launch()