"""Gradio app: face detection with two SCRFD models, side by side.

Upload one image; it is run through BOTH SCRFD detectors (500MF and 2.5GF) and the
annotated results are shown next to each other so you can compare them in one go —
no model picker. Each panel shows the detected faces (box + confidence), the face
count, and the inference time. CPU-only ONNX Runtime.
"""

import os
import time

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont

from scrfd import SCRFD

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# name -> (onnx path, accent color for that model's boxes)
MODELS = [
    ("SCRFD-500MF · 640 · det_500m", os.path.join(BASE_DIR, "models/det_500m.onnx"), (0, 200, 90)),
    ("SCRFD-2.5GF · 640 · det_2.5g", os.path.join(BASE_DIR, "models/det_2.5g.onnx"), (0, 162, 255)),
    ("SCRFD · 480 · det_480", os.path.join(BASE_DIR, "models/det_480.onnx"), (170, 90, 255)),
]

_DETECTORS = {}


def _get_detector(path):
    if path not in _DETECTORS:
        print(f"[*] Loading SCRFD model: {path}")
        _DETECTORS[path] = SCRFD(path)
    return _DETECTORS[path]


def _load_font(size):
    for p in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
              "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):
        if os.path.exists(p):
            return ImageFont.truetype(p, size)
    return ImageFont.load_default()


def _annotate(rgb, dets, color, title):
    """Draw boxes + a title bar on a copy of ``rgb`` (H,W,3 uint8)."""
    img = Image.fromarray(rgb).convert("RGB")
    W, H = img.size
    draw = ImageDraw.Draw(img)

    lw = max(2, int(0.003 * max(W, H)))
    font = _load_font(max(13, int(0.018 * max(W, H))))

    for d in dets:
        x1, y1, x2, y2, score = d
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        draw.rectangle([x1, y1, x2, y2], outline=color, width=lw)

        tag = f"{score:.2f}"
        tb = draw.textbbox((0, 0), tag, font=font)
        tw, th = tb[2] - tb[0], tb[3] - tb[1]
        ty = max(0, y1 - th - 4)
        draw.rectangle([x1, ty, x1 + tw + 6, ty + th + 4], fill=color)
        draw.text((x1 + 3, ty + 1), tag, font=font, fill=(255, 255, 255))

    # Title bar across the top (model name + face count).
    bar_font = _load_font(max(15, int(0.022 * max(W, H))))
    bar_h = (bar_font.getbbox("Hg")[3]) + 12
    bar = Image.new("RGB", (W, bar_h), color)
    bd = ImageDraw.Draw(bar)
    bd.text((8, 6), title, font=bar_font, fill=(255, 255, 255))
    out = Image.new("RGB", (W, H + bar_h), (20, 20, 20))
    out.paste(bar, (0, 0))
    out.paste(img, (0, bar_h))
    return out


# SCRFD 5-point landmark order.
KP_NAMES = ["L-eye", "R-eye", "nose", "L-mouth", "R-mouth"]


def _format_landmarks(kpss):
    """Markdown listing each face's 5 landmark (x, y) in original-image pixels."""
    if kpss is None or len(kpss) == 0:
        return "*No faces / landmarks.*"
    lines = ["**5-point landmarks** (x, y px):"]
    for i, kps in enumerate(kpss):
        pts = " · ".join(
            f"{name} ({int(round(x))}, {int(round(y))})"
            for name, (x, y) in zip(KP_NAMES, kps)
        )
        lines.append(f"- **Face {i + 1}:** {pts}")
    return "\n".join(lines)


def detect_faces(image, threshold):
    """Run every SCRFD model; return per-model (image, landmark text) + a summary."""
    if image is None:
        n = len(MODELS)
        return (*[None] * n, *[""] * n, "Upload an image to run all detectors.")

    rgb = np.array(image.convert("RGB"))
    bgr = rgb[:, :, ::-1].copy()  # SCRFD expects BGR (cv2 blob swaps back)

    images, landmark_texts = [], []
    summary = ["| Model | Faces | Time |", "|---|---|---|"]
    for name, path, color in MODELS:
        det = _get_detector(path)
        t0 = time.perf_counter()
        dets, kpss = det.detect(bgr, thresh=float(threshold))
        dt_ms = (time.perf_counter() - t0) * 1000.0
        n = 0 if dets is None else len(dets)
        title = f"{name}  —  {n} face{'s' if n != 1 else ''}  ·  {dt_ms:.0f} ms"
        images.append(_annotate(rgb, dets if dets is not None else [], color, title))
        landmark_texts.append(_format_landmarks(kpss))
        summary.append(f"| {name} | {n} | {dt_ms:.0f} ms |")

    return (*images, *landmark_texts, "\n".join(summary))


IMG_H = 460

with gr.Blocks(title="Face Detection — SCRFD comparison") as app:
    gr.Markdown(
        "# Face Detection — SCRFD model comparison\n"
        "Upload an image. It runs through **all three** SCRFD detectors — "
        "**500MF @ 640**, **2.5GF @ 640**, and **det_480 @ 480** — and shows the "
        "detected faces side by side: box + confidence, face count, and inference "
        "time — with each face's 5-point landmark coordinates listed below its panel."
    )

    with gr.Row():
        with gr.Column(scale=1):
            inp = gr.Image(type="pil", label="Input image", height=IMG_H)
            thr = gr.Slider(
                minimum=0.1, maximum=0.9, value=0.5, step=0.05,
                label="Detection confidence threshold",
            )
            btn = gr.Button("Detect faces", variant="primary")
            summary = gr.Markdown()

        with gr.Column(scale=3):
            with gr.Row():
                with gr.Column():
                    out_a = gr.Image(label="SCRFD-500MF · 640", height=IMG_H)
                    kps_a = gr.Markdown()
                with gr.Column():
                    out_b = gr.Image(label="SCRFD-2.5GF · 640", height=IMG_H)
                    kps_b = gr.Markdown()
                with gr.Column():
                    out_c = gr.Image(label="SCRFD · det_480 · 480", height=IMG_H)
                    kps_c = gr.Markdown()

    # Order must match detect_faces: images, then landmark texts, then summary.
    outs = [out_a, out_b, out_c, kps_a, kps_b, kps_c, summary]
    btn.click(fn=detect_faces, inputs=[inp, thr], outputs=outs, concurrency_limit=1)
    inp.change(fn=detect_faces, inputs=[inp, thr], outputs=outs, concurrency_limit=1)


if __name__ == "__main__":
    app.launch(
        server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.environ.get("PORT", os.environ.get("GRADIO_SERVER_PORT", 7860))),
    )