# Taking each video cropping to 520, 400, 3, and obtaining 12 key frames for the usage in a dataset.
# Frames will be saved as individual images.

import numpy as np
import cv2, os
from tqdm import tqdm


# Obtain 12 different keyframes from the video, equally spaced.

# Actually to increase dataset we will just increase the number of frames allowed to say 50 per video. 900 images for dataset, 200 are for training.
def get_equal_elements(array, num_elements=12):
    """
  Takes a specific number of elements equally spaced from an array.

  Args:
      array: The input array.
      num_elements: The number of elements to take (default 12).

  Returns:
      A list of elements from the array.
  """
    if num_elements > len(array):
        print(f"Number of elements cannot be greater than array length : {len(array)}")
        return []
    step_size = len(array) // (num_elements - 1)  # Avoid extra element with floor division
    return array[::step_size]  # Slice with step size


def video_to_keyframes(video_filename):
    cap = cv2.VideoCapture(video_filename)
    frames = []
    while (cap.isOpened()):
        ret, frame = cap.read()
        try:
            # 720, 1280, 3
            # Remove top 200 pixels from video frame.
            # print((frame.shape))   # (720, 1280, 3)
            frame = frame[200:, 440:840]
            # print((frame.shape)) # (520, 400, 3)
            frames.append(frame)
            # cv2.imshow('frame',frame)
        except Exception as e:
            print(f"Error is {e}")
            if frame == None:
                break
            continue
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
    print("Done obtaining captured frames.")
    selected_frames = get_equal_elements(frames, num_elements=50)
    print("Obtained selected frames.")
    # Extract filename without extension for directory creation
    filename_no_ext = video_filename.split('.')[0]
    if "left" in filename_no_ext:
        filename_no_ext = "left"
    else:
        filename_no_ext = "right"
    # Create directory for the video if it doesn't exist
    try:
        os.makedirs(os.path.join("dataset", filename_no_ext))
    except FileExistsError:

        pass  # Directory already exists, ignore

    # Save each selected frame as an image with filename_frame_number.jpg format
    # Create images under directory for that specific source video file name, like left for left1.mp4.
    for i, frame in enumerate(selected_frames):
        print(filename_no_ext)
        file_name = f"{video_filename.split('.')[0].split('/')[-1]}_frame_{i}.jpg"
        print(file_name)
        image_path = os.path.join("dataset", filename_no_ext, file_name)
        print(f"Write to disk. {image_path}")
        print("Resized to 224,224")
        target_height, target_width = 224, 224
        if type(frame) != type(None):
            frame = resize_with_aspect_ratio(frame, target_height, target_width)
            # print(type(frame), frame)
            cv2.imwrite(image_path, frame)
        else:
            continue

    print("Saved images for each all selected frames.")
    return selected_frames


# Resize images.

# Target dimensions for AlexNet

# Resize the image with aspect ratio preservation
def resize_with_aspect_ratio(image, target_height, target_width):
    height, width = image.shape[:2]
    if height == target_height and width == target_width:
        return image
    if height > width:
        new_width = int(width * (target_height / height))
        # resized_image = cv2.resize(image, (new_width, target_height))
        resized_image = cv2.resize(image, (target_width, target_height))
        return resized_image
    else:
        new_height = int(height * (target_width / width))
        resized_image = cv2.resize(image, (target_width, new_height))
        # Crop the center of the resized image to match target dimensions
        start_x = int((resized_image.shape[1] - target_width) / 2)
        start_y = int((resized_image.shape[0] - target_height) / 2)
        return resized_image[start_y:start_y + target_height, start_x:start_x + target_width]


if __name__ == "__main__":
    # Load the image and resize
    BASE_PATH = "dataset/src/"
    videos = os.listdir(BASE_PATH)
    # print(videos)
    target_height, target_width = 224, 224
    for video_file in tqdm(videos):
        selected_frames = video_to_keyframes(os.path.join(BASE_PATH, video_file))

    # image = cv2.imread(image_path)  # Or use your image data
    # resized_image = resize_with_aspect_ratio(image, target_height, target_width)