import tensorflow as tf
import numpy as np
from PIL import ImageDraw, ImageFont
from huggingface_hub import snapshot_download

# Keypoint labels for MoveNet (17 human body parts)
KEYPOINT_LABELS = {
    'nose': 0, 'left_eye': 1, 'right_eye': 2, 'left_ear': 3, 'right_ear': 4,
    'left_shoulder': 5, 'right_shoulder': 6, 'left_elbow': 7, 'right_elbow': 8,
    'left_wrist': 9, 'right_wrist': 10, 'left_hip': 11, 'right_hip': 12,
    'left_knee': 13, 'right_knee': 14, 'left_ankle': 15, 'right_ankle': 16
}

# Stylish connection colors for the skeleton bones
SKELETON_EDGES = {
    (0, 1): '#FF66CC', (0, 2): '#66FFFF', (1, 3): '#FF66CC', (2, 4): '#66FFFF',
    (0, 5): '#FF99CC', (0, 6): '#99FFFF', (5, 7): '#FF6699', (7, 9): '#FF3366',
    (6, 8): '#66CCCC', (8, 10): '#33CCCC', (5, 6): '#CCCC00', (5, 11): '#FF9966',
    (6, 12): '#66FF99', (11, 12): '#999900', (11, 13): '#FF6600', (13, 15): '#FF3300',
    (12, 14): '#00CC99', (14, 16): '#009966'
}

def process_keypoints(prediction, img_height, img_width, confidence=0.12):
    all_joints = []
    all_bones = []

    instance_count = prediction.shape[1]
    for i in range(instance_count):
        x_coords = prediction[0, i, :, 1] * img_width
        y_coords = prediction[0, i, :, 0] * img_height
        scores = prediction[0, i, :, 2]
        labels = list(KEYPOINT_LABELS.keys())

        keypoints = np.stack([labels, x_coords, y_coords], axis=-1)
        visible_kpts = keypoints[scores > confidence]
        all_joints.append(visible_kpts)

        for (a, b), color in SKELETON_EDGES.items():
            if scores[a] > confidence and scores[b] > confidence:
                segment = np.array([[x_coords[a], y_coords[a]], [x_coords[b], y_coords[b]]])
                all_bones.append((segment, color))

    return np.concatenate(all_joints, axis=0), all_bones

def draw_bones(image, keypoints):
    draw = ImageDraw.Draw(image)
    font = ImageFont.load_default()

    joints, bones = process_keypoints(keypoints, image.height, image.width)

    for bone, color in bones:
        draw.line((*bone[0], *bone[1]), fill=color, width=3)

    for label, x, y in joints:
        cx, cy = float(x), float(y)
        radius = 4
        draw.ellipse([(cx - radius, cy - radius), (cx + radius, cy + radius)], fill="#FF0000", outline="#222222")
        draw.text((cx + 5, cy - 5), label, font=font, fill="#0000CC")

    return joints

def movenet(image):
    model_path = snapshot_download("lukassso/movenet-myking")
    model = tf.saved_model.load(model_path).signatures["serving_default"]
    image = tf.cast(image, tf.int32)
    result = model(image)
    return result["output_0"].numpy()