#!/usr/bin/env python

from __future__ import annotations

import os
import pathlib
import sys
import urllib.request

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

sys.path.insert(0, "face_detection")

from ibug.face_detection import RetinaFacePredictor, S3FDPredictor

DESCRIPTION = "# [ibug-group/face_detection](https://github.com/ibug-group/face_detection)"


def is_lfs_pointer_file(path: pathlib.Path) -> bool:
    try:
        with open(path, "r") as f:
            # Git LFS pointer files usually start with version line
            version_line = f.readline()
            if version_line.startswith("version https://git-lfs.github.com/spec/"):
                # Check for the presence of oid and size lines
                oid_line = f.readline()
                size_line = f.readline()
                if oid_line.startswith("oid sha256:") and size_line.startswith("size "):
                    return True
    except Exception as e:
        print(f"Error reading file {path}: {e}")
    return False


lfs_model_path = pathlib.Path("face_detection/ibug/face_detection/retina_face/weights/Resnet50_Final.pth")
if is_lfs_pointer_file(lfs_model_path):
    os.remove(lfs_model_path)
    out_path = hf_hub_download(
        "public-data/ibug-face-detection",
        filename=lfs_model_path.name,
        repo_type="model",
        subfolder="retina_face",
    )
    os.symlink(out_path, lfs_model_path)


def load_model(model_name: str, threshold: float, device: torch.device) -> RetinaFacePredictor | S3FDPredictor:
    if model_name == "s3fd":
        model = S3FDPredictor(threshold=threshold, device="cpu")
        model.device = device
        model.net.device = device
        model.net.to(device)
    else:
        model_name = model_name.replace("retinaface_", "")
        model = RetinaFacePredictor(threshold=threshold, device="cpu", model=RetinaFacePredictor.get_model(model_name))
        model.device = device
        model.net.to(device)
    return model


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model_names = [
    "retinaface_mobilenet0.25",
    "retinaface_resnet50",
    "s3fd",
]
detectors = {name: load_model(name, threshold=0.8, device=device) for name in model_names}


@spaces.GPU
def detect(image: np.ndarray, model_name: str, face_score_threshold: float) -> np.ndarray:
    model = detectors[model_name]
    model.threshold = face_score_threshold

    # RGB -> BGR
    image = image[:, :, ::-1]
    preds = model(image, rgb=False)

    res = image.copy()
    for pred in preds:
        box = np.round(pred[:4]).astype(int)

        line_width = max(2, int(3 * (box[2:] - box[:2]).max() / 256))
        cv2.rectangle(res, tuple(box[:2]), tuple(box[2:]), (0, 255, 0), line_width)

        if len(pred) == 15:
            pts = pred[5:].reshape(-1, 2)
            for pt in np.round(pts).astype(int):
                cv2.circle(res, tuple(pt), line_width, (0, 255, 0), cv2.FILLED)

    return res[:, :, ::-1]


example_image_path = pathlib.Path("selfie.jpg")
if not example_image_path.exists():
    url = "https://raw.githubusercontent.com/peiyunh/tiny/master/data/demo/selfie.jpg"
    urllib.request.urlretrieve(url, example_image_path)

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            image = gr.Image(type="numpy", label="Input")
            model_name = gr.Radio(model_names, type="value", value="retinaface_resnet50", label="Model")
            score_threshold = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.8, label="Face Score Threshold")
            run_button = gr.Button()
        with gr.Column():
            result = gr.Image(label="Output")
    gr.Examples(
        examples=[[example_image_path.as_posix(), model_names[1], 0.8]],
        inputs=[image, model_name, score_threshold],
        outputs=result,
        fn=detect,
    )
    run_button.click(
        fn=detect,
        inputs=[image, model_name, score_threshold],
        outputs=result,
        api_name="detect",
    )


if __name__ == "__main__":
    demo.queue(max_size=20).launch()