import cv2
import numpy as np
import gradio as gr
import tempfile
import tensorflow as tf
from tensorflow import keras

# ── Fix TensorFlow graph execution for fast inference ─────────────
# Build the model once and compile the predict function into a
# concrete TF graph so every call is instant with no re-tracing
physical_devices = tf.config.list_physical_devices("CPU")
tf.config.set_visible_devices(physical_devices, "CPU")

model = keras.models.load_model("emotion_detection_model.h5")

# Warm-up: run one dummy prediction so TF compiles the graph now,
# not on the first real frame (which caused the 3-min freeze)
_dummy = np.zeros((1, 48, 48, 1), dtype=np.float32)
model.predict(_dummy, verbose=0)

# ── Constants ─────────────────────────────────────────────────────
EMOTION_LABELS = {
    0: "Angry",
    1: "Happy",
    2: "Neutral",
    3: "Sad",
    4: "Surprised"
}

# BGR colors for each emotion
COLORS = {
    "Angry":     (0,   0,   255),
    "Happy":     (0,   200, 0),
    "Neutral":   (200, 200, 0),
    "Sad":       (255, 100, 0),
    "Surprised": (0,   140, 255)
}

# Load face detector
face_cascade = cv2.CascadeClassifier(
    cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
)

# Holds last saved snapshot path for download
_snapshot_path = {"path": None}


# ── Core detection function ────────────────────────────────────────
def detect(frame):
    """
    Receives one webcam frame (RGB numpy array) from Gradio.
    Returns annotated RGB frame immediately — no threads, no queues.
    TF graph is pre-compiled so predict() takes ~20-50ms on CPU.
    """
    if frame is None:
        return None

    # RGB → BGR for OpenCV
    bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

    # Resize to 480p max for faster processing without losing accuracy
    h, w = bgr.shape[:2]
    if w > 640:
        scale = 640 / w
        bgr   = cv2.resize(bgr, (640, int(h * scale)))

    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    gray = cv2.equalizeHist(gray)

    # Detect faces
    faces = face_cascade.detectMultiScale(
        gray,
        scaleFactor=1.3,
        minNeighbors=5,
        minSize=(40, 40),
        flags=cv2.CASCADE_SCALE_IMAGE
    )

    if len(faces) == 0:
        cv2.putText(bgr, "No face detected", (20, 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2, cv2.LINE_AA)
    else:
        # Batch all faces into one predict call for speed
        batch   = []
        regions = []
        for (x, y, w, h) in faces:
            roi = gray[y:y+h, x:x+w]
            roi = cv2.resize(roi, (48, 48)).astype(np.float32) / 255.0
            batch.append(roi.reshape(48, 48, 1))
            regions.append((x, y, w, h))

        # Single batched prediction — much faster than one-by-one
        batch       = np.array(batch, dtype=np.float32)
        predictions = model.predict(batch, verbose=0)

        for i, (x, y, w, h) in enumerate(regions):
            emotion    = EMOTION_LABELS[np.argmax(predictions[i])]
            confidence = float(np.max(predictions[i])) * 100
            color      = COLORS[emotion]

            # Bounding box
            cv2.rectangle(bgr, (x, y), (x+w, y+h), color, 2)

            # Label background + text
            label = f"{emotion}  {confidence:.1f}%"
            font  = cv2.FONT_HERSHEY_SIMPLEX
            (tw, th), bl = cv2.getTextSize(label, font, 0.7, 2)
            ly = max(y - 8, th + 8)
            cv2.rectangle(bgr,
                          (x, ly - th - bl - 4),
                          (x + tw + 6, ly + bl - 2),
                          color, -1)
            cv2.putText(bgr, label, (x + 3, ly - bl),
                        font, 0.7, (0, 0, 0), 2, cv2.LINE_AA)

    # BGR → RGB for Gradio output
    result = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)

    # Cache snapshot for download
    _snapshot_path["frame"] = result.copy()

    return result


def save_snapshot():
    """Saves the latest annotated frame as a PNG and returns path for download."""
    frame = _snapshot_path.get("frame")
    if frame is None:
        return None
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    tmp.close()
    cv2.imwrite(tmp.name, cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
    return tmp.name


# ── Gradio UI ─────────────────────────────────────────────────────
with gr.Blocks(title="Real-Time Emotion Detection") as demo:

    gr.Markdown("""
    # 😊 Real-Time Emotion Detection
    Click **▶ Start** on the webcam feed to begin.  
    Press **📥 Download Snapshot** anytime to save the current frame.  
       😠 Angry  |  😄 Happy  |  😐 Neutral  |  😢 Sad  |  😲 Surprised
    """)

    with gr.Row():
        webcam     = gr.Image(sources=["webcam"], streaming=True,
                              label="📷 Live Webcam")
        output_img = gr.Image(label="🎯 Detected Emotion")

    with gr.Row():
        download_btn = gr.DownloadButton(
            label="📥 Download Snapshot", variant="primary"
        )

    # stream_every=0.1 → 10fps — safe for CPU, smooth, no backlog
    webcam.stream(
        fn=detect,
        inputs=webcam,
        outputs=output_img,
        stream_every=0.1,
        time_limit=600
    )

    download_btn.click(fn=save_snapshot, inputs=None, outputs=download_btn)

    gr.Markdown("""
    ---
    **Model:** CNN — Conv2D(45) → Conv2D(64) → Conv2D(128) → Dense(256) → Softmax(5)  | 
    **Input:** 48×48 grayscale  |  **Detector:** Haar Cascade
    """)

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