# YOLO
import cv2
import torch
import numpy as np
from ultralytics import YOLO
from PIL import Image, ImageDraw

model_name = 'yolo11s-pose.pt'
model = YOLO(model_name)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
half = device=='cuda'

connection_color_map = {
    (0, 1): 'lightgreen',   # Nose to Left Eye
    (0, 2): 'lightgreen',   # Nose to Right Eye
    (1, 2): 'lightgreen',   # Left Eye to Right Eye
    (1, 3): 'magenta',      # Left Eye to Left Ear
    (2, 4): 'magenta',      # Right Eye to Right Ear
    (3, 5): 'lightblue',    # Left Ear to Left Shoulder
    (4, 6): 'lightblue',    # Right Ear to Right Shoulder
    (5, 6): 'purple',       # Left Shoulder to Right Shoulder
    (5, 7): 'blue',         # Left Shoulder to Left Elbow
    (7, 9): 'blue',         # Left Elbow to Left Wrist
    (6, 8): 'red',          # Right Shoulder to Right Elbow
    (8, 10): 'red',         # Right Elbow to Right Wrist
    (5, 11): 'purple',      # Left Shoulder to Left Hip
    (6, 12): 'purple',      # Right Shoulder to Right Hip
    (11, 12): 'purple',     # Left Hip to Right Hip
    (11, 13): 'orange',     # Left Hip to Left Knee
    (13, 15): 'orange',     # Left Knee to Left Ankle
    (12, 14): 'orange',     # Right Hip to Right Knee
    (14, 16): 'orange',     # Right Knee to Right Ankle
}


def detect_pose(src_img):
    results = model.predict(
            source=src_img,
            # conf=conf_threshold,
            # iou=iou_threshold,
            max_det=1,
            show_labels=True,
            show_conf=True,
            imgsz=32*5,
            classes=[0],
            # line_width=3,
            half=half,
            device=device,
        )
    return results

def draw_pose(skeleton_image:Image.Image, keypoints, threshold=0.7, line_width = 3,circle_radius = 6):
    coords = keypoints.cpu().numpy().xy[0]
    scores = keypoints.data.squeeze()[:,2].cpu().numpy()
    # Create a blank image with the same size as the input image
    # skeleton_image = Image.new("RGB", im.size, "black")  # White background
    draw = ImageDraw.Draw(skeleton_image)
    
    # Draw lines connecting the keypoints
    for connection, color in connection_color_map.items():
        try:
            start_point = coords[connection[0]]
            end_point = coords[connection[1]]
            if scores[connection[0]] > threshold and scores[connection[1]] > threshold:
                draw.line([tuple(start_point), tuple(end_point)], fill=color, width=line_width)
                draw.ellipse(
                    (start_point[0] - circle_radius, start_point[1] - circle_radius,
                    start_point[0] + circle_radius, start_point[1] + circle_radius),
                fill=color
                )
                draw.ellipse(
                    (end_point[0] - circle_radius, end_point[1] - circle_radius,
                    end_point[0] + circle_radius, end_point[1] + circle_radius),
                fill=color
                )
                
        except IndexError:
            pass # Handle cases where keypoints might be missing

    # Save the skeleton image
    return skeleton_image