#Takes the available videos along with their phase label annotations and separates
#them out into phase videos.

import os
import sys
import pdb 
import cv2
import skvideo.io 
import numpy as np

NUM_PHASES = 11

def separate_phases():
    vid_dir = "/home-2/fyu12@jhu.edu/scratch/ReiterLab/cataract/"
    anno_dir = "/home-2/fyu12@jhu.edu/scratch/ReiterLab/cataract_phase_annoations/"
    phase_dir = "/home-2/fyu12@jhu.edu/scratch/ReiterLab/cataract_phase_separated/"
    videos = os.listdir(vid_dir)
    for vid_name in videos:
        #Set up loading in the video with the annotations
        vid_number = vid_name[4:7]        
        if vid_number == "159" or vid_number == "125":
            continue
        print("Currently on video: " + vid_number)
        sys.stdout.flush()
        anno_name = "vid_"+vid_number+"_tasks.txt"
            
        #Open up annotation and video files.
        with open(anno_dir+anno_name) as f:
            annotations = f.readlines()
        numPhaseVids = [0]*(NUM_PHASES+1)
        cap = cv2.VideoCapture(vid_dir+vid_name)
        fps = int(round(cap.get(cv2.CAP_PROP_FPS)))
        if fps > 30:
            print("\tFPS will be reduced for this video.")
            sys.stdout.flush()
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        curFrame = 1
        cap.release()
        vid = skvideo.io.FFmpegReader(vid_dir+vid_name)
        cap = vid.nextFrame()
        #Goes through each line in annotation and generates a phase video per line.
        for line in annotations:
            split = line.strip().split()
            startFrame = int(split[0])
            endFrame = int(split[1])
            phase = int(split[2])
            numPhaseVids[phase]+=1
            #Advance video to start of the phase.
            while curFrame < startFrame:
                frame = next(cap)
                curFrame += 1
            print('\tWriting out video for phase: ' + str(phase))
            sys.stdout.flush()
            #Create a video writer, and write the video until the end of the phase.
            vid_name = phase_dir+str(phase) + "/p"+str(phase) + "_n"+str(numPhaseVids[phase]) + '_vid_' + vid_number + '.avi'
            '''
            writeVideo = False
            try:
                test_vid = skvideo.io.FFmpegReader(vid_name)
            except ValueError:
                writeVideo = True
            '''
            writeVideo = True
            if writeVideo:
                fourcc = cv2.VideoWriter_fourcc(*'DIVX')
                if fps > 30:
                    out = cv2.VideoWriter(vid_name, fourcc, fps/2, (width, height))
                else:
                    out = cv2.VideoWriter(vid_name, fourcc, fps, (width, height))
                while curFrame <= endFrame:
                    frame = next(cap)
                    frame = np.flip(frame, axis = 2)
                    curFrame+=1
                    if fps > 30 and not curFrame % 2 == 0:
                        continue
                    out.write(frame)
                out.release()
            '''
            else:
                print('\tSkipping phase ' + str(phase))
                sys.stdout.flush()
                while curFrame <= endFrame:
                    ret, frame = cap.read()
                    curFrame+=1
            '''
        vid.close()

if __name__ == "__main__":
    separate_phases()

