"""
OpenPose Preprocessor for ControlNet
A simple Gradio application for pose detection.
"""

import gradio as gr
import numpy as np
from PIL import Image
import torch

# Global device detection
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {DEVICE}")

# Model cache
_openpose_detector = None
_dwpose_detector = None


def get_openpose_detector():
    """Get or create OpenPose detector."""
    global _openpose_detector
    if _openpose_detector is None:
        from controlnet_aux import OpenposeDetector
        _openpose_detector = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
    return _openpose_detector


def get_dwpose_detector():
    """Get or create DWPose detector using easy-dwpose."""
    global _dwpose_detector
    if _dwpose_detector is None:
        from easy_dwpose import DWposeDetector
        _dwpose_detector = DWposeDetector(device=DEVICE)
    return _dwpose_detector


def detect_pose(image, model_type, detect_hand, detect_face, detect_resolution):
    """Main pose detection function."""
    if image is None:
        return None
    
    try:
        # Convert to PIL if needed
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # Convert to RGB if necessary
        if image.mode != "RGB":
            image = image.convert("RGB")
        
        # Resize to detect_resolution while maintaining aspect ratio
        original_size = image.size
        ratio = detect_resolution / max(original_size)
        new_size = (int(original_size[0] * ratio), int(original_size[1] * ratio))
        image_resized = image.resize(new_size, Image.Resampling.LANCZOS)
        
        # Process based on model type
        if model_type == "DWPose":
            detector = get_dwpose_detector()
            result = detector(
                image_resized,
                output_type="pil",
                include_hands=detect_hand,
                include_face=detect_face
            )
        elif model_type == "OpenPose (Full)":
            detector = get_openpose_detector()
            result = detector(
                image_resized,
                hand_and_face=True,
                output_type="pil"
            )
        elif model_type == "OpenPose (Face Only)":
            detector = get_openpose_detector()
            result = detector(
                image_resized,
                include_body=False,
                include_hand=False,
                include_face=True,
                output_type="pil"
            )
        elif model_type == "OpenPose (Hand)":
            detector = get_openpose_detector()
            result = detector(
                image_resized,
                include_body=True,
                include_hand=True,
                include_face=False,
                output_type="pil"
            )
        else:
            # Basic OpenPose with options
            detector = get_openpose_detector()
            result = detector(
                image_resized,
                hand_and_face=detect_hand and detect_face,
                output_type="pil"
            )
        
        # Resize result back to original size if needed
        if result is not None and hasattr(result, 'size') and result.size != original_size:
            result = result.resize(original_size, Image.Resampling.LANCZOS)
        
        return result
            
    except Exception as e:
        print(f"Error during processing: {str(e)}")
        import traceback
        traceback.print_exc()
        return None


# Create Gradio interface
with gr.Blocks(
    title="🦴 OpenPose Preprocessor",
    theme=gr.themes.Soft()
) as demo:
    
    gr.Markdown(
        """
        # 🦴 OpenPose Preprocessor for ControlNet
        
        High-quality pose detection with multiple models. Upload an image and get pose skeleton for ControlNet.
        """
    )
    
    gr.Markdown(f"**Device**: `{DEVICE}` {'🚀' if DEVICE == 'cuda' else '🐢'}")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="📷 Input Image", type="pil", height=400)
            
            model_type = gr.Dropdown(
                label="🤖 Model",
                choices=["DWPose", "OpenPose", "OpenPose (Full)", "OpenPose (Face Only)", "OpenPose (Hand)"],
                value="DWPose",
                info="DWPose is recommended for best accuracy"
            )
            
            with gr.Row():
                detect_hand = gr.Checkbox(label="👆 Detect Hands", value=True)
                detect_face = gr.Checkbox(label="😊 Detect Face", value=True)
            
            detect_resolution = gr.Slider(
                label="📏 Detection Resolution",
                minimum=256,
                maximum=2048,
                value=512,
                step=64,
                info="Higher = more accurate but slower"
            )
            
            process_btn = gr.Button("🚀 Detect Pose", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            output_image = gr.Image(label="🎨 Output Pose", type="pil", height=400)
    
    gr.Markdown(
        """
        ### 📌 Tips
        - **DWPose** is recommended for best accuracy, especially for hands and complex poses
        - **OpenPose (Full)** detects body, face, and hands together
        - Higher **Detection Resolution** improves accuracy but increases processing time
        - The output image can be directly used with ControlNet OpenPose models
        
        ### ⚙️ Options
        - **Detect Hands/Face** checkboxes work with DWPose and basic OpenPose modes
        - For preset modes like "OpenPose (Full)", these options are ignored
        """
    )
    
    process_btn.click(
        fn=detect_pose,
        inputs=[input_image, model_type, detect_hand, detect_face, detect_resolution],
        outputs=[output_image]
    )


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860
    )