"""
HF Space: Image Preprocessor (MTCNN Face Detection)
----------------------------------------------------
Detects a face in the input image, crops it with a small margin,
and resizes it to 224x224 — the exact format expected by the
image-to-emotion model downstream.

API endpoint: POST /preprocess
  Input : image (file upload)
  Output: cropped & resized face image (PNG)
"""

import gradio as gr
from facenet_pytorch import MTCNN
from PIL import Image

# ---------------------------------------------------------------------------
# Initialise MTCNN once at module load so the Space is warm for requests.
# keep_all=False → only the highest-confidence face is returned.
# post_process=False → skip internal normalisation; we only need the crop.
# ---------------------------------------------------------------------------
mtcnn = MTCNN(
    image_size=224,
    margin=20,
    keep_all=False,
    post_process=False,
)


# ---------------------------------------------------------------------------
# Core preprocessing function
# ---------------------------------------------------------------------------
def preprocess(image: Image.Image) -> Image.Image:
    """
    Detect a single face, apply a 10 % context margin, and resize to 224x224.

    Raises gr.Error if no face is detected so the caller receives a clear
    HTTP-level error (status 400) rather than a silent None.
    """
    if image is None:
        raise gr.Error("No image provided.")

    img = image.convert("RGB")

    boxes, probs = mtcnn.detect(img)

    if boxes is None or len(boxes) == 0:
        raise gr.Error("No face detected in the provided image.")

    # Pick the face with highest detection confidence.
    best_idx = int(probs.argmax()) if probs is not None else 0
    x1, y1, x2, y2 = boxes[best_idx]

    # Add a 10 % context margin for robustness.
    w, h = x2 - x1, y2 - y1
    margin = 0.10
    x1 -= w * margin
    y1 -= h * margin
    x2 += w * margin
    y2 += h * margin

    # Clamp to image bounds.
    img_w, img_h = img.size
    x1 = max(0, int(round(x1)))
    y1 = max(0, int(round(y1)))
    x2 = min(img_w, int(round(x2)))
    y2 = min(img_h, int(round(y2)))

    if x2 <= x1 or y2 <= y1:
        raise gr.Error("Face bounding box is degenerate after clamping.")

    face_img = img.crop((x1, y1, x2, y2)).resize((224, 224), Image.BILINEAR)
    return face_img


# ---------------------------------------------------------------------------
# Gradio Interface
# ---------------------------------------------------------------------------
demo = gr.Interface(
    fn=preprocess,
    inputs=gr.Image(type="pil", label="Input Image"),
    outputs=gr.Image(type="pil", label="Preprocessed Face (224×224)"),
    title="Image Preprocessor — MTCNN Face Detection",
    description=(
        "Upload any photo. The service detects the dominant face, "
        "crops it with a small context margin, and returns a 224×224 RGB image "
        "ready for the image-to-emotion model."
    ),
    api_name="preprocess",
    allow_flagging="never",
)

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