#Currently used to split a video into multiple numpy clips of 16 frames.
#This does this for all videos in the repository. This is also not ideal,
#so this script should be repurposed into one that will save LMDB.

import skvideo.io
from skimage import img_as_ubyte
import skimage.transform
import os
import sys
import numpy as np

def split_video_clips():
    data_dir = '/home-2/fyu12@jhu.edu/scratch/ReiterLab/cataract_phase_separated/'
    phase_dirs = os.listdir(data_dir)

    clip = np.empty((1, 16, 112, 112, 3), dtype = np.uint8)
    print('Splitting videos into clips!')
    print('------------------------')
    sys.stdout.flush()
    for phase in phase_dirs:
        phase_dir = data_dir + phase + '/'
        vids = os.listdir(phase_dir)
        for vid in vids:
            if vid[-3:] == 'avi':
                dictionary = {}
                split_name = vid.split('_')
                video_name = vid[:-4]
                print('On video ' + video_name)
                sys.stdout.flush()
                out_dir = phase_dir + video_name 
                if not os.path.exists(out_dir):
                    os.mkdir(out_dir)
                video_path = phase_dir + vid

                vid = skvideo.io.FFmpegReader(video_path)
                num_frames = vid.getShape()[0]
                num_frames = num_frames - num_frames%16
                cap = vid.nextFrame()
                clip_index = 1
                
                #Use this code here to have non-overlapping videos.
                #This saves memory.
                for i in range(num_frames):
                    img = next(cap)
                    curFrame = i%32
                    if curFrame%2 == 0:
                        img = np.flip(img, axis = 2)
                        img = skimage.transform.resize(img, (112, 112))
                        img = img_as_ubyte(img)
                        clip[:, int(curFrame/2), :, :, :] = img
                    if curFrame == 31:
                        print('Saving clip ' + str(clip_index))
                        sys.stdout.flush()
                        dictionary[clip_index] = clip
                        out_file = out_dir + str(clip_index) + '.npy'
                        np.save(out_file, clip)
                        clip_index += 1
                vid.close()
'''
            #Use this code over here if we want to have overlapping videos.
            #This doubles our training data.
                for i in range(16):
                    img = next(cap)
                    if i%2 == 0:
                        img = np.flip(img, axis = 2)
                        img = skimage.transform.resize(img, (112, 112))
                        img = img_as_ubyte(img) 
                        clip[:, int(i/2), :, :, :] = img

                for i in range(num_frames - 16):
                    img = next(cap)
                    curFrame = i%16
                    if curFrame%2 == 0:
                        img = np.flip(img, axis = 2)
                        img = skimage.transform.resize(img, (112, 112))
                        img = img_as_ubyte(img)
                        clip[:, int(curFrame/2)+8, :, :, :] = img
    
                    if curFrame == 15:
                        out_file = out_dir + str(clip_index) + '.npy'
                        np.save(out_file, clip)
                        clip[:, 0:8, :, :, :] = clip[:, 8:16, :, :, :] 
                        clip_index += 1
                vid.close()
'''        
if __name__ == "__main__":
    split_video_clips()


