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

#video_in_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/procedures/processed122017/_TK_640480_59fps/'
video_in_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/procedures/processed122017/2017videos/_useMe_TK_640480/'
annotations_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/annotations/rhexis/segments/XChen/current/'
#annotations_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/annotations/rhexis/segments/SVedula/current/'
video_out_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/rhexis/segmented2018/'

def check_validity():
  vid_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/rhexis/segmented2018/'
  files = os.listdir(annotations_dir)
  for file in files:
    if '.txt' not in file:
      continue
    vid_name = file[:-4]
    annotation_file = open(annotations_dir+file,'r')

    cap = cv2.VideoCapture(vid_dir + vid_name+'.mp4')
    length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    time_window = []
    for line in annotation_file:
      try:
        time, mark = line.split()
        if mark == '18':
          ## rhexis start
          time_window.append(int(float(time)) * fps)
        elif mark == '19':
          ## rhexis finish
          time_window.append(int(float(time)) * fps)
      except:
        #pdb.set_trace()
        continue
    annotation_file.close()
    print(vid_name,(length == time_window[1]-time_window[0]))

def cut_rhexis():
  import math
  files = os.listdir(annotations_dir)
  #files = ['vid_302.txt']
  existing_videos = os.listdir(video_out_dir)
  #files = ['vid_235.txt']
  #for file in files:
  for file in ['vid_348.txt']:
    vid_name = file[:-4]
    #if vid_name+'.mp4' in existing_videos:
      #print('Skipping:',vid_name)
      #continue

    annotation_file = open(annotations_dir+file,'r')
    cap = cv2.VideoCapture(video_in_dir + vid_name+'.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))
    #fps = 60

    dx = 320
    dy = 240

    video_out = cv2.VideoWriter(video_out_dir + vid_name+'.mp4',
                                cv2.VideoWriter_fourcc(*'MP4V'),
                                fps,
                                (640, 480))
    time_window = []
    for line in annotation_file:
      time, mark = line.split()
      if mark == '18':
        ## rhexis start
        time_window.append(math.ceil(float(time)*fps))
      elif mark == '19':
        ## rhexis finish
        time_window.append(math.ceil(float(time)*fps))

    cap.set(cv2.CAP_PROP_POS_FRAMES, time_window[0])
    i = time_window[0]
    count =0
    while i < time_window[-1]:
      ret, img = cap.read()
      count += 1
      if not ret:
        break


      #cv2.imshow('cropped_img', img)
      #if cv2.waitKey(1) & 0xFF == ord('q'):
  #      break

      video_out.write(img)
      i += 1
      print(file, str(count), '/', str(time_window[1]-time_window[0]), '\r', end='\r')

    video_out.release()
    cap.release()


def to_59fps():
  video_in_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/procedures/processed122017/'
  video_out_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/procedures/processed122017/_TK_640480_59fps/'

  vid_names = os.listdir(video_in_dir)
  for file in vid_names:
    if file[-4:] != '.mp4':
      continue
    vid_name = file[:-4]
    cap = cv2.VideoCapture(video_in_dir + vid_name + '.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))
    print(vid_name,fps)

    DUP = 0
    if fps == 30:
      fps = 59
      DUP = 1
    video_out = cv2.VideoWriter(video_out_dir + vid_name + '.mp4',
                                cv2.VideoWriter_fourcc(*'MP4V'),
                                fps,
                                (640, 480))
    i = 0
    while True:
      ret, img = cap.read()
      if not ret:
        break
      video_out.write(img)
      if DUP:
        video_out.write(img)
      i += 1
      print(vid_name, str(i),'/',str(length), end='\r')

def slow_down():
  #video_in_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/procedures/processed122017/2017videos/_useMe_TK_640480/'
  #video_out_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/procedures/processed122017/2017videos/_useMe_TK_640480/slowed_down/'
  video_in_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/rhexis/segmented2018/'
  video_out_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/rhexis/segmented2018/slowed_down/'
  vid_names = ['vid_327.mp4','vid_341.mp4']


  for vid_name in vid_names:
    cap = cv2.VideoCapture(video_in_dir+vid_name)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = 59

    dx = 320
    dy = 240

    case = 0
    if width == 1920 or width == 1280:
      case = 1

    video_out = cv2.VideoWriter(video_out_dir + vid_name,
                                cv2.VideoWriter_fourcc(*'MP4V'),
                                fps,
                                (640, 480))
    i = 0
    DUP = 1
    while True:
      ret, img = cap.read()
      if not ret:
        break

      if case:
        img = cv2.resize(img, (960, 540))

      cy = int(img.shape[0] / 2)
      cx = int(img.shape[1] / 2)

      crop_img = img[cy - dy:cy + dy, cx - dx:cx + dx]
      # cv2.imshow('cropped_img', crop_img)
      # if cv2.waitKey(1) & 0xFF == ord('q'):
      # break
      video_out.write(crop_img)
      if DUP:
        video_out.write(crop_img)
      i += 1
      print(vid_name, str(i), end='\r')

def deidentify():
  video_in_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/procedures/processed122017/_TK_640480_59fps/'
  video_out_dir = '/home/tk/dev/data/CataractMount/Pilot2015/processed/videos/procedures/processed122017/_TK_640480_59fps/deidentified/'
  vid_names = ['vid_123.mp4','vid_124.mp4','vid_125.mp4','vid_127.mp4','vid_128.mp4','vid_130.mp4']

  for vid_name in vid_names:
    cap = cv2.VideoCapture(video_in_dir + vid_name)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = 59

    dx = 320
    dy = 240

    case = 0

    box_y_0 = 35
    box_y_1 = 60
    box_x_0 = 50
    box_x_1 = 240

    if width == 1920 or width == 1280:
      case = 1

    video_out = cv2.VideoWriter(video_out_dir + vid_name,
                                cv2.VideoWriter_fourcc(*'MP4V'),
                                fps,
                                (640, 480))
    i = 0
    while True:
      ret, img = cap.read()
      if not ret:
        break

      if case:
        img = cv2.resize(img, (960, 540))

      cy = int(img.shape[0] / 2)
      cx = int(img.shape[1] / 2)

      img[box_y_0:box_y_1, box_x_0:box_x_1, :] = 0
      crop_img = img[cy - dy:cy + dy, cx - dx:cx + dx]
      # cv2.imshow('cropped_img', crop_img)
      # if cv2.waitKey(1) & 0xFF == ord('q'):
      # break
      video_out.write(crop_img)
      i += 1
      print(vid_name, str(i), end='\r')

if __name__ == "__main__":
  slow_down()
  #deidentify()
  #to_59fps()
  #cut_rhexis()
  #check_validity()
