import cv2

import pandas as pd

import numpy as np

dataset = pd.read_csv('assets/extracted_landmarks_test.csv')

def is_float(num):

    try:

        float(num)

        return True

    except ValueError:

        return False

dataset['landmarks'] = dataset['landmarks'].apply(lambda arr: np.array([float(n) for n in arr.split() if is_float(n)]))

video_dir = 'assets/dataset_videos/batch_test'  # Directory where your video files are stored

dataset['Label'] = None

for index, row in dataset.iterrows():

    video = row['video']

    frame = row['frame']

    # Construct the path to the video file

    video_path = video_dir + video  # Update the file extension if necessary

    # Open the video file

    cap = cv2.VideoCapture(video_path)

    # Get the frames per second (FPS) of the video

    fps = cap.get(cv2.CAP_PROP_FPS)

    print(f"{video} at group {row['group']}. Frame {frame}")

    # Read and discard frames until the desired frame index

    cap.set(cv2.CAP_PROP_POS_FRAMES, frame)

    # Display ~10 frames

    for _ in range(10):

        ret, frame = cap.read()

        if ret:

            cv2.imshow('Video', frame)

    # Ask for user input

    key = cv2.waitKey(0) & 0xFF

    if key == ord('g'):

        dataset.at[index, 'Label'] = 'good'  # Update 'Label' to 'good' in the corresponding row

        print("Label set to 'good'")

    elif key == ord('b'):

        dataset.at[index, 'Label'] = 'bad'  # Update 'Label' to 'bad' in the corresponding row

        print("Label set to 'bad'")

    # Release the video capture

    cap.release()

    # Close the video display window

    cv2.destroyAllWindows()

dataset.to_csv('assets/labeled_dataset_test.csv', index=False)

print(dataset.head())