import pickle
import time

import cv2
from ultralytics import YOLO

from calculate import normalize_pose
from utils import normalize

# Wczytanie wideo
cap = cv2.VideoCapture("input.mp4")
fps = cap.get(cv2.CAP_PROP_FPS)
width = 1920
height = 1080

# Ustawienia zapisu wideo
fourcc = cv2.VideoWriter_fourcc(*"avc1")
out = cv2.VideoWriter("output.mp4", fourcc, fps, (width, height))

# Wczytanie modelu YOLOv8 Pose
model = YOLO("yolo11x-pose.pt", verbose=False)  # Twój model pose

moves = []
started = False
frame_id = 0

while True:
    ok, frame = cap.read()
    if not ok:
        break

    # Skalowanie do 640x640
    frame_resized = cv2.resize(frame, (width, height))

    # Wykrywanie poz
    results = model.predict(frame_resized, verbose=False)

    # Rysowanie punktów 2D
    for result in results:
        if result.keypoints is not None and len(result.keypoints.xy) > 0:
            for keypoint in result.keypoints.xy[0]:  # keypoints[0] bo dla jednej osoby
                if not started:
                    frame_id = 0
                    started = True
                x, y = keypoint  # współrzędne + confidence
                x = int(x)
                y = int(y)
                cv2.circle(frame_resized, (x, y), 5, (0, 255, 0), -1)  # zielone kropki

            moves.append((frame_id * 1 / fps, normalize_pose(result.keypoints.xy.cpu().numpy()[0]), normalize(result.keypoints.xy.cpu()[0]) * 160 + 250))

    out.write(frame_resized)
    frame_id += 1

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

cap.release()
out.release()
print("Zapisano: output.mp4")