import cv2 
import numpy as np
import sys 
import os
import pdb 
import matplotlib.pyplot as plt 
import math

vid_dir = "/home-2/gsilvac1@jhu.edu/work/gsilvac1/CATARACT/new_videos/cataract/"
anno_dir = "/home-2/gsilvac1@jhu.edu/work/gsilvac1/CATARACT/new_videos/cataract_phase_annotations/"
phase_dir = "/home-2/gsilvac1@jhu.edu/work/gsilvac1/CATARACT/new_videos/cataract_phase_separated/"
NUM_PHASES = 10

def segment_vid():
  annotations= os.listdir(anno_dir)
  existing_videos = os.listdir(vid_dir)

  for annotation in annotations:
    vid_num = annotation[-7:-4]
    print("Currently on video number: " + vid_num)
    sys.stdout.flush()
    annotation_file = open(anno_dir+annotation,'r')

    cap = cv2.VideoCapture(vid_dir + 'vid_' + vid_num +'.mp4')
    length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))

    num_phase_vids = [0]*(NUM_PHASES+1)
    cur_frame = 0
    dx = 320
    dy = 240

    cut_boundaries = list()
    lines = annotation_file.readlines()
    for i in range(0, len(lines), 2):
      line1_split = lines[i].strip().split()
      line2_split = lines[i+1].strip().split()  
      t1 = int(math.ceil(float(line1_split[0])*fps))
      t2 = int(math.ceil(float(line2_split[0])*fps))
      activity = int(line1_split[1][:-1])
      cut_boundaries.append((activity, t1, t2))
      print('{:d} {:d} {:d}'.format(activity, t1, t2))
    print("Done reading annotations")
    sys.stdout.flush()
    for triple in cut_boundaries:
      start_frame = triple[1]
      end_frame = triple[2]
      activity_num = triple[0]
      num_phase_vids[activity_num] += 1
      vid_name = phase_dir+str(activity_num) + "/p"+str(activity_num) + "_n"+str(num_phase_vids[activity_num]) + '_vid_' + vid_num + '.avi'
 
      print("Will write phase " + str(activity_num))
      sys.stdout.flush()
      fourcc = cv2.VideoWriter_fourcc(*'DIVX') 
      video_out = cv2.VideoWriter(vid_name,
                                fourcc, 
                                fps,
                                (640, 480))
    
      while cur_frame < start_frame:  
        ret, img = cap.read()
        cur_frame += 1
        if not ret:
            break
      
      while cur_frame < end_frame:
        ret, img = cap.read()
        if not ret:
            break
        cur_frame += 1
        video_out.write(img)
      print("Done writing phase " + str(activity_num))
      sys.stdout.flush()
      video_out.release()
      print("Released " + vid_name)
      sys.stdout.flush()
    cap.release()

if __name__ == "__main__":
    segment_vid()
