"""Anime Face Detector - Hugging Face Space with Zero GPU support"""
import subprocess
import sys

# Install mmpose without dependencies (to avoid chumpy build failure)
# Suppress output to avoid interfering with Gradio startup
subprocess.run(
    [sys.executable, '-m', 'pip', 'install', '--no-deps', 'mmpose>=1.0.0'],
    stdout=subprocess.DEVNULL,
    stderr=subprocess.DEVNULL,
    check=True
)

import os

import spaces
import gradio as gr
import cv2
import numpy as np
import PIL.Image
import torch

# Global detector (initialized on first GPU call)
detector = None


@spaces.GPU
def detect(image, face_score_threshold, landmark_score_threshold):
    """Detect anime faces and landmarks in the image."""
    global detector

    # Lazy initialization with auto device detection
    if detector is None:
        from anime_face_detector import create_detector
        # Auto-detect device: use GPU if available, otherwise CPU
        device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
        print(f'Using device: {device}')
        detector = create_detector('yolov3', device=device)

    # Convert RGB to BGR for OpenCV
    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    # Run detection
    preds = detector(image_bgr)

    # Draw results
    res = image_bgr.copy()
    for pred in preds:
        box = pred['bbox']
        box, score = box[:4], box[4]
        if score < face_score_threshold:
            continue
        box = np.round(box).astype(int)

        # Line thickness based on face size
        lt = max(2, int(3 * (box[2:] - box[:2]).max() / 256))

        # Draw bounding box
        cv2.rectangle(res, tuple(box[:2]), tuple(box[2:]), (0, 255, 0), lt)

        # Draw confidence score
        cv2.putText(res, f'{score * 100:.1f}%', (box[0], box[1] - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, lt / 3, (255, 255, 255), thickness=max(lt // 2, 1))

        # Draw keypoints
        pred_pts = pred['keypoints']
        for *pt, kpt_score in pred_pts:
            if kpt_score < landmark_score_threshold:
                color = (0, 255, 255)  # Yellow for low confidence
            else:
                color = (0, 0, 255)  # Red for high confidence
            pt = tuple(np.round(pt).astype(int))
            cv2.circle(res, pt, lt, color, cv2.FILLED)

    # Convert BGR to RGB for output
    res = cv2.cvtColor(res, cv2.COLOR_BGR2RGB)
    return PIL.Image.fromarray(res)


# Download sample image
def download_sample():
    sample_path = 'input.jpg'
    if not os.path.exists(sample_path):
        torch.hub.download_url_to_file(
            'https://raw.githubusercontent.com/hysts/anime-face-detector/main/assets/input.jpg',
            sample_path
        )
    return sample_path


# Create Gradio interface
sample_path = download_sample()

demo = gr.Interface(
    fn=detect,
    inputs=[
        gr.Image(type='numpy', label='Input Image'),
        gr.Slider(0, 1, step=0.05, value=0.5, label='Face Score Threshold'),
        gr.Slider(0, 1, step=0.05, value=0.3, label='Landmark Score Threshold'),
    ],
    outputs=gr.Image(type='pil', label='Output'),
    title='Anime Face Detector (GPU)',
    description='Detect anime faces and 28 facial landmarks using YOLOv3 + HRNetV2. Powered by Zero GPU.',
    article='<a href="https://github.com/hysts/anime-face-detector">GitHub</a> | <a href="https://github.com/ayutaz/anime-face-detector">Fork with GPU Docker</a>',
    examples=[
        [sample_path, 0.5, 0.3],
    ],
    cache_examples=False,
)

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