import pickle
import sys

import torch
from ultralytics import YOLO
import cv2
import time

import utils
from calculate import normalize_pose, compare_poses_boolean
from draw import draw_new
from utils import find_closest
from video_methods import initialize_method

model = YOLO("yolo11x-pose.pt")
model.to(torch.device('cuda:0'))

method = initialize_method("cam")

do_pose_shot = False
startTime = 0

def click_event(event, x, y, flags, param):
    global do_pose_shot, startTime

    if event == cv2.EVENT_LBUTTONDOWN:  # lewy przycisk myszy
        do_pose_shot = not do_pose_shot
        if do_pose_shot:
            startTime = time.time()

def main():
    moves = []

    while True:
        frame = method.receive_frame()

        frame = cv2.flip(frame, 1)
        results = model(frame, verbose=False)

        if len(results) != 0:

            result = results[0]
            kpts = result.keypoints.data[0] if len(result.keypoints.data) else None

            if kpts is None:
                continue

            img = frame

            if do_pose_shot:
                print(normalize_pose(result.keypoints.xy.cpu().numpy()[0]))
                exit(0)

            cv2.imshow('Klatka z kamerki', img)
            cv2.setMouseCallback('Klatka z kamerki', click_event)
            cv2.waitKey(1)  # Czekaj na naciśnięcie klawisza

main()