import glob
import pickle
from tqdm import tqdm

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'))

startTime = 0
def main():
    moves = []

    for i in tqdm(sorted(glob.glob("video/camA/*.jpg"), key=lambda f: int(__import__("re").search(r"\d+", f).group()))):
        data = i.replace(f'video/camA\\', "")
        frame = cv2.imread(f"video/camA/{data}")
        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

            moves.append((int(data.replace(".jpg", "")), normalize_pose(result.keypoints.xy.cpu().numpy()[0]), result.keypoints.xy.cpu()[0]))

    with open('moves_makarena_best.pkl', 'wb') as f:  # 'wb' = write binary
        pickle.dump(moves, f)

main()