import os

import cv2


def video_generator(video_path_or_folder, intrinsics, allowed_extensions=('mp4', 'mkv', 'mov')):
    """
    Create a generator for unsupervised training on depth sequences from a video file or folder of video files
    :param video_path_or_folder: Video file or folder with list of video files to iterate through
    :param intrinsics: Intrinsics for the videos TODO: Intrinsics per video
    :param allowed_extensions: Allowed video extensions, to not accidentally pick files that aren't videos
    :return: generator that yields dict of {frames: [frame1, frame2, frame3], intrinsics: [fx, fy, tx, ty]}
    """
    if os.path.isfile(video_path_or_folder):
        # TODO: How to re-yield? Is this enough, since I'm just returning the actual generator?
        #  Or do I need to iterate like below?
        return _single_video_generator(video_path_or_folder)
    else:
        for root, dirs, files in os.walk(video_path_or_folder):
            for file in files:
                if os.path.splitext(file)[1] in allowed_extensions:
                    for frames in _single_video_generator(file):
                        yield frames


def _single_video_generator(video_file, intrinsics):
    # Single video file
    video = cv2.VideoCapture(video_file)

    try:
        # Buffer to store 3 frames, yield when this fills up
        current_frames = []
        while video.grab():
            current_frames.append(video.retrieve())
            if len(current_frames) == 3:
                temp_frames = current_frames
                current_frames = []
                # TODO: Consider converting frames to tensor
                yield {'frames': temp_frames, 'intrinsics': intrinsics}
    finally:
        video.release()


def image_generator(root_folder):
    """
    Create an image generator for unsupervised training
    :param root_folder:
    :return:
    """
    pass