"""
Preview the videos generated by the `generate_moving_symbols.py` script.

Author: Ryan Szeto
"""

import os
import re
from Tkinter import Tk
from tkFileDialog import askopenfilename

import cv2
import numpy as np

PROJ_ROOT_DIR = os.path.dirname(os.path.dirname(__file__))

def view_video_tensor(video_tensor, vid_ids=None, delay=None, prompt_keypress=False):
    '''
    View a video stored with wacky Toronto dimensions
    :param video_tensor: A tensor with dims T x V x H x W or T x V x H x W x C
    :param vid_id: Index of video to play. If None, choose random video
    :param delay: How long to wait between frames. If None, advance on click
    :param prompt_keypress: Whether to require a key press before advancing to the next video
    :return:
    '''
    vid_id_list = range(video_tensor.shape[1]) if vid_ids is None else vid_ids
    while True:
        for vid_id in vid_id_list:
            video = video_tensor[:, vid_id]

            # Show info
            for i in range(video.shape[0]):
                large_frame = cv2.resize(video[i], (256, 256), interpolation=cv2.INTER_NEAREST)
                cv2.imshow(None, large_frame)
                if delay:
                    pressed_key = cv2.waitKey(delay)
                    if i == video.shape[0]-1 and prompt_keypress:
                        print('Ended video. Press any key to continue.')
                        pressed_key = cv2.waitKey()
                else:
                    pressed_key = cv2.waitKey()
                if pressed_key == 27:
                    print('Quitting...')
                    return

def main():
    # Select file path (https://stackoverflow.com/a/3579625)
    Tk().withdraw()
    filename = askopenfilename(initialdir=os.path.join(PROJ_ROOT_DIR, 'output'),
                               title='Select one of the Moving Symbols files')
    if not filename: exit()
    # Check the file is one of the valid files
    m = re.search('(.+?)(_videos.npy$)', filename)
    if m is None:
        print('The selected file is not a video file.')
        return

    # Load arrays
    video_tensor_path = filename
    video_tensor = np.load(video_tensor_path, mmap_mode='r')

    # Show info
    print('Video tensor shape: %s ' % str(video_tensor.shape))
    print('Press ESC to quit.')
    view_video_tensor(video_tensor, delay=1000 / 30, vid_ids=None)


if __name__ == '__main__':
    main()