import numpy as np
import cv2

import supervision as sv
from supervision.keypoint.skeletons import Skeleton
from supervision.assets import download_assets, VideoAssets

import super_gradients


# Available models: https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS-POSE.md
yolo_nas = super_gradients.training.models.get(
    "yolo_nas_pose_s", pretrained_weights="coco_pose").to("cuda")

download_assets(VideoAssets.PEOPLE_WALKING)

# I found this important to tweak, regardless of model size
CONFIDENCE_THRESHOLD = 0.1

# cap = cv2.VideoCapture(0)
cap = cv2.VideoCapture(VideoAssets.PEOPLE_WALKING.value)
i = 0
while True:
    ret, frame = cap.read()
    if not ret:
        continue

    # TODO: I assume we're not supporting batch-image inputs yet. YOLO NAS does.

    result = yolo_nas.predict(frame, conf=CONFIDENCE_THRESHOLD)
    keypoints = sv.KeyPoints.from_yolo_nas(result)

    ann_point_large = sv.VertexAnnotator(color=sv.Color.ROBOFLOW, radius=5)
    ann_point_small = sv.VertexAnnotator(color=sv.Color.WHITE, radius=3)

    # Option 1: Use a predefined skeleton
    ann_skeleton = sv.EdgeAnnotator(
        color=sv.Color.ROBOFLOW,
        thickness=5,
        edges=Skeleton.YOLO_NAS.value
    )

    # Option 2: No skeleton
    # ann_skeleton = sv.EdgeAnnotator(
    #     color=sv.Color.ROBOFLOW,
    #     thickness=5,
    #     edges=[]
    # )

    # Option 3: Figure out automatically
    # ann_skeleton = sv.EdgeAnnotator(
    #     color=sv.Color.ROBOFLOW,
    #     thickness=5
    # )

    # Option 4: Take a guess (connect sequential points)
    # TODO: remove YOLO_NAS from Skeleton before running
    # ann_skeleton = sv.EdgeAnnotator(
    #     color=sv.Color.ROBOFLOW,
    #     thickness=5
    # )

    # Draw
    try:
        ann_skeleton.annotate(frame, keypoints)
        ann_point_large.annotate(frame, keypoints)
        ann_point_small.annotate(frame, keypoints)
    except Exception as e:
        print("Caught exception while annotating: \n", e)

    frame = cv2.resize(frame, (1024, 680))
    cv2.imshow('frame', frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()