#########################################
# Johns Hopkins University              #
# 601.455 Computer Integrated Surgery 2 #
# Spring 2018                           #
# Query by Video For Surgical Activities#
# Felix Yu                              #
# JHED: fyu12                           #
# Gianluca Silva Croso                  #
# JHED: gsilvac1                        #
#########################################

import cv2
import numpy as np
import sys
import os

NUM_ACTIVITIES = 10
NUM_CHANNELS = 3
FPS = 30


def convert_opt_flow(clip_path, out_path):
    """
    Takes a single video and writes it using dense optical flow on every frame
    :param clip_path: input video path
    :param out_path: output optical flow video path
    """
    # Load in video
    try:
        cap = cv2.VideoCapture(clip_path)
    except ValueError:
        print("\t\tIssue with video encoding. Skipping.")
        sys.stdout.flush()
        return None

    ret, frame1 = cap.read()
    ret, frame2 = cap.read()
    if ret:
        prvs = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    else:
        print("Couldn't read video " + clip_path)
        return
    hsv = np.zeros_like(frame1)
    hsv[..., 1] = 255
    
    # Define the codec and create VideoWriter object
    d1 = int(frame1.shape[1])
    d2 = int(frame1.shape[0])
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(out_path, fourcc, 30.0, (d1, d2), True)
    print("Output: " + out_path)
    sys.stdout.flush()
    while ret:
        nxt = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
        flow = cv2.calcOpticalFlowFarneback(prvs, nxt, None, 0.5, 3, 15, 3, 5, 1.2, 0)
        mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
        hsv[..., 0] = ang*180/np.pi/2
        hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
        bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
        out.write(bgr)
        prvs = nxt
        ret, frame2 = cap.read()
    cap.release()
    out.release()
    return


def get_flows():
    """
    Iterates over every input phase clip in phase_clip_dir and writes out a
    corresponding clip in flow_out_dir after applying dense optical flow
    """
    # DIRECTORY WHERE PHASE CLIPS ARE LOCATED - MODIFY IF NECESSARY
    phase_clip_dir = '/home-2/gsilvac1@jhu.edu/work/gsilvac1/CATARACT/raw_data/cataract_phase_separated/'
    # DIRECTORY WHERE OPTICAL FLOW CLIPS WILL BE WRITTEN - MODIFY IF NECESSARY
    flow_out_dir = '/home-2/gsilvac1@jhu.edu/work/gsilvac1/CATARACT/optical_flow/data/'
    for i in range(1, NUM_ACTIVITIES+1):
        print('\t\tCurrently on phase ' + str(i))
        print('-----------------------------------------------------------')
        sys.stdout.flush()

        spec_phase_clip_dir = phase_clip_dir + str(i) + '/'
        spec_phase_out_dir = flow_out_dir + str(i) + '/'
        if not os.path.exists(spec_phase_out_dir):
            os.makedirs(spec_phase_out_dir)
        phase_clip_names = os.listdir(spec_phase_clip_dir)
        for phase_clip in phase_clip_names:
            clip_path = spec_phase_clip_dir + phase_clip
            out_path = spec_phase_out_dir + phase_clip
            try:
                convert_opt_flow(clip_path, out_path)
            except:
                print("Unexpected error dealing with " + clip_path)
        print('')


if __name__ == "__main__":
    get_flows()
