# bb-tool.py
#
# This is a small app to aid in extracting bounding boxes from video using a human in the loop and some cv algs
# It builds heavily on the examples from the cvui library: https://github.com/Dovyski/cvui


import cv2
import numpy as np
from scipy import signal
import os
import sys
from EnhancedWindow import EnhancedWindow
import cvui
import time
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib import style


# Some globals
VERSION = '1.0.0'
VIDEO_WINDOW_NAME = 'Morpho-Track ' + VERSION + ' : Video'
PLOT_WINDOW_NAME = 'Morpho-Track ' + VERSION + ' : Data'
CTRL_WINDOW_NAME = 'Morpho-Track ' + VERSION + ' : Control'

DISPLAY_SIZE = (540, 960, 3)
MIN_CONTOUR_AREA = 200  # minimum size of contours in ROI detection
LOOK_BACK = 41          # how many frames to median filter in estimating orientation
MOTION_DELAY = 4        # Not currently used
SKIP_TO = 300           # how many frames to skip at the start of the video
FRAME_SCALE = 8         # scale from frame number to x position on the plot

TRACKING_PARAMS_FILE = 'tracking_params.json'


def vid_frame(cap, frame):
    # get a new frame from the video
    ret, new_frame = cap.read()
    if not ret:
        return False, []
    new_frame = cv2.resize(new_frame, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_LINEAR)
    return True, new_frame

# Draw and update ROIs in the video (compied almost exactly from complex-mouse cvui example


def roi_tool(frame,anchor,roi):
    # The function 'bool cvui.mouse(int query)' allows you to query the mouse for events.
    # E.g. cvui.mouse(cvui.DOWN)
    #
    # Available queries:
    #	- cvui.DOWN: any mouse button was pressed. cvui.mouse() returns true for single frame only.
    #	- cvui.UP: any mouse button was released. cvui.mouse() returns true for single frame only.
    #	- cvui.CLICK: any mouse button was clicked (went down then up, no matter the amount of frames in between). cvui.mouse() returns true for single frame only.
    #	- cvui.IS_DOWN: any mouse button is currently pressed. cvui.mouse() returns true for as long as the button is down/pressed.

    drawing = False
    new_bb = False

    # Did any mouse button go down?
    if cvui.mouse(cvui.DOWN):
        # Position the anchor at the mouse pointer.
        anchor.x = cvui.mouse().x
        anchor.y = cvui.mouse().y

        # Inform we are working, so the ROI window is not updated every frame
        drawing = True

    # Is any mouse button down (pressed)?
    if cvui.mouse(cvui.IS_DOWN):
        # Adjust roi dimensions according to mouse pointer
        width = cvui.mouse().x - anchor.x
        height = cvui.mouse().y - anchor.y

        roi.x = anchor.x + width if width < 0 else anchor.x
        roi.y = anchor.y + height if height < 0 else anchor.y
        roi.width = abs(width)
        roi.height = abs(height)

        # Show the roi coordinates and size
        cvui.printf(frame, roi.x + 5, roi.y + 5, 0.3, 0xff0000, '(%d,%d)', roi.x, roi.y)
        cvui.printf(frame, cvui.mouse().x + 5, cvui.mouse().y + 5, 0.3, 0xff0000, 'w:%d, h:%d', roi.width, roi.height)

    # Was the mouse clicked (any button went down then up)?
    if cvui.mouse(cvui.UP):
        # We are done working with the ROI.
        drawing = False
        if roi.width > 25 and roi.height > 25:
            new_bb = True

    # Ensure ROI is within bounds
    frameRows, frameCols, frameChannels = frame.shape
    roi.x = 0 if roi.x < 0 else roi.x
    roi.y = 0 if roi.y < 0 else roi.y
    roi.width = roi.width + frame.cols - (roi.x + roi.width) if roi.x + roi.width > frameCols else roi.width
    roi.height = roi.height + frame.rows - (roi.y + roi.height) if roi.y + roi.height > frameRows else roi.height

    # Render the roi
    cvui.rect(frame, roi.x, roi.y, roi.width, roi.height, 0xff0000)

    return anchor, roi, drawing, new_bb


def apply_canny(frame,high,low,kernel=3):

    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    frame = cv2.Canny(frame, low, high, kernel)
    frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
    return frame

# the main app

def main():

    height = 480
    width = 1374

    cap_frame = np.zeros((height, width, 3), np.uint8)
    frame = np.zeros((height, width, 3), np.uint8)
    org_frame = np.zeros((height, width, 3), np.uint8)
    ctrl_frame = np.zeros((height, 300, 3), np.uint8)
    plot_win = np.zeros((200, width, 3), np.uint8)
    result_win = np.zeros((height+200,width,3),np.uint8)
    roi = cvui.Rect(0, 0, 0, 0)
    anchor = cvui.Point()
    low_threshold = [150]
    high_threshold = [250]
    playback_fps = [120]
    use_canny = [False]
    detect_rois = [True]
    fuse_track_and_bbox = [False]
    is_paused = False
    xar = [np.array([]),np.array([])]
    yar = [np.array([]),np.array([])]
    cent_track = []

    # Load the video or die
    try:
        cap = cv2.VideoCapture(sys.argv[1])
    except:
        print ('Could not load the video file, will now exit.')
        exit()

    log_file = open(str(int(time.time()))+"_roistats.csv", "w")

    frame_number = 0

    # Create a settings window using the EnhancedWindow class.
    settings = EnhancedWindow(10, 50, 270, 270, 'Settings')
    # Create a settings window using the EnhancedWindow class.
    control = EnhancedWindow(10, 300, 270, 100, 'Control')

    print('Initializing tracker...')

    tracker = cv2.TrackerKCF_create()
    tracker_init = False
    if os.path.exists(TRACKING_PARAMS_FILE):
        fs = cv2.FileStorage(TRACKING_PARAMS_FILE, cv2.FileStorage_READ)
        tracker.read(fs.getFirstTopLevelNode())

    print('Creating windows...')

    # Init cvui and tell it to create a OpenCV window, i.e. cv2.namedWindow(WINDOW_NAME).
    cvui.init(CTRL_WINDOW_NAME)
    cvui.init(VIDEO_WINDOW_NAME)
    cvui.init(PLOT_WINDOW_NAME)

    cv2.imshow(PLOT_WINDOW_NAME, plot_win)

    timer = 0.0 # force a new frame on first iteration

    # Save the output video
    fourcc = cv2.VideoWriter_fourcc(*'MJPG')
    out = cv2.VideoWriter('c:\\Users\\proberts\\desktop\\output_mjpg.avi', fourcc, 120.0, (width, height+200))

    print('Starting video loop...')

    while True:

        # get and prep new frame or quite if end of video

        get_new_frame = time.time()-timer > 1/playback_fps[0]
        have_new_frame = False
        if not is_paused and get_new_frame:
            ret, new_frame = vid_frame(cap, cap_frame)
            # crop the frame to take only one view
            #new_frame = new_frame[:, 0:1032]
            timer = time.time()
            frame_number = frame_number + 1
            have_new_frame = True
            if not ret:
                break

        if frame_number < SKIP_TO:
            continue

        # copy the frame
        org_frame[:] = new_frame[:]

        # Should we apply Canny edge?
        if use_canny[0]:
            frame = apply_canny(new_frame,low_threshold[0], high_threshold[0], 3)
        else:
            # No, so just copy the original image to the displaying frame.
            frame[:] = new_frame[:]

        # Should we detect rois
        if detect_rois[0]:

            if not use_canny[0]:
                frame = apply_canny(new_frame, low_threshold[0], high_threshold[0], 3)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            kernel = np.ones((5, 5), np.uint8)
            frame = cv2.dilate(frame, kernel, iterations=1)
            frame = cv2.erode(frame, kernel, iterations=1)
            contours, hierarchy = cv2.findContours(frame,
                                                   cv2.RETR_EXTERNAL,
                                                   cv2.CHAIN_APPROX_SIMPLE
                                                   )

            bounding_rects = []
            plot_vals = []
            maj_axes = []
            con_cnt = 0


            output_string = ""

            for indx,c in enumerate(contours):

                if cv2.contourArea(c) > MIN_CONTOUR_AREA:

                    br = cv2.boundingRect(c)
                    mine = cv2.fitEllipseDirect(c)

                    if mine[0][0] < width/2:
                        con_cnt = 0
                    else:
                        con_cnt = 1

                    if have_new_frame:

                        # compress the possible angles into either vertical (0) or horizontal 90
                        raw_or = mine[2]
                        if raw_or > 90:
                            raw_or = 180 - raw_or

                        xar[con_cnt] = np.append(xar[con_cnt], frame_number)

                        # update new orientation linearly between past orientation and eccentricity of ellipse
                        if len(yar[0]) > 1:
                            #circ = mine[1][1]/mine[1][0]
                            #if circ > 0 and circ < 1:
                            #    new_or = (1-circ)*raw_or + circ*yar[con_cnt][-1]
                            #else:
                            new_or = raw_or
                            yar[con_cnt] = np.append(yar[con_cnt],new_or)
                        else:
                            yar[con_cnt] = np.append(yar[con_cnt], raw_or)

                        # apply median filter
                        if len(yar[0]) > LOOK_BACK:
                            plot_val = np.median(yar[con_cnt][-LOOK_BACK:])
                        else:
                            plot_val = raw_or

                        plot_vals.append(plot_val)
                        maj_axes.append(mine[1][0])

                        yval = plot_val  # np.mean(plot_vals)
                        rad = maj_axes[-1]
                        cent = (int((xar[con_cnt][-1]-SKIP_TO) / FRAME_SCALE), 10 + 2*int(yval))
                        cvui.context(PLOT_WINDOW_NAME)
                        cv2.circle(plot_win, cent, int(rad / 10), (con_cnt*255, 200, 200))
                        cv2.imshow(PLOT_WINDOW_NAME, plot_win)

                        cent_track.append(np.array(mine[1]))
                        if len(cent_track) > MOTION_DELAY:
                            pass
                            #print(np.linalg.norm(cent_track[-1] - cent_track[-MOTION_DELAY]))

                    #cv2.rectangle(org_frame, br, (0, 255, 0), 3)
                    bounding_rects.append(br)
                    cv2.ellipse(org_frame, mine, (0, 0, 175), 2)
                    p1 = (int(mine[0][0]), int(mine[0][1]))
                    p2 = (int(p1[0] + mine[1][0]*np.sin(plot_val*3.14/180)), int(p1[1] - mine[1][0]*np.cos(plot_val*3.14/180)))
                    cv2.arrowedLine(org_frame, p1, p2, (con_cnt*255, 200, 200), 2)

                    (x, y), (Ma, ma), ang = mine
                    # make sure to preserve the orientation channel, con_cnt = 0 is left, con_cnt = 1 is right
                    # basically append before or after conditioned on con_cnt
                    if con_cnt == 0:
                        output_string = str(indx) + "\t" + str(x) + "\t" + str(y) + "\t" + str(Ma) + "\t" + str(ma) + "\t" + str(ang) + "\t" + str(plot_val) + "\t" + output_string
                    else:
                        output_string = output_string + str(indx) + "\t" + str(x) + "\t" + str(y) + "\t" + str(Ma) + "\t" + str(
                            ma) + "\t" + str(ang) + "\t" + str(plot_val) + "\t"

            # save data
            if len(output_string) > 0:
                log_file.write(str(frame_number) + "\t" + output_string + "\n")

            # set the display frame
            frame = org_frame

        cvui.context(VIDEO_WINDOW_NAME)
        if not tracker_init or is_paused:
            anchor, roi, drawing, new_bb = roi_tool(frame, anchor, roi)

        # if there is a new bounding box and we are not tracking one,
        # init a new tracker
        if new_bb and (not tracker_init or is_paused):
            # Initialize tracker with first frame and bounding box
            tracker = cv2.TrackerKCF_create()
            ok = tracker.init(frame, (roi.x, roi.y, roi.width, roi.height))
            tracker_init = True

        # if we are not actively drawing a bounding box and the tracker is initialized
        # update it with the latest frame from the video
        if not drawing and tracker_init:
            # Update tracker
            ok = True
            if not is_paused:
                ok, roi = tracker.update(frame)
                match_index = 0
                min_diff = 1000000
                if fuse_track_and_bbox[0]:
                    for br_ind, br in enumerate(bounding_rects):
                        #print(br)
                        sum_squared_diff = 0
                        for ind in range(0, 4):
                            sum_squared_diff += (roi[ind]-br[ind])**2
                        sum_squared_diff = np.sqrt(sum_squared_diff)
                        if min_diff > sum_squared_diff:
                            match_index = br_ind
                            min_diff = sum_squared_diff
                    #print(min_diff)
                    if min_diff < 100:
                        roi = bounding_rects[match_index]
                    #print(roi)

                roi = cvui.Rect(roi[0], roi[1], roi[2], roi[3])
            if ok or fuse_track_and_bbox:
                cvui.rect(frame, roi.x, roi.y, roi.width, roi.height, 0xff0000)
            else:
                tracker_init = False

        # This function must be called *AFTER* all UI components. It does
        # all the behind the scenes magic to handle mouse clicks, etc.
        cvui.update(VIDEO_WINDOW_NAME)
        # Show everything on the screen
        cv2.imshow(VIDEO_WINDOW_NAME, frame)

        cvui.context(CTRL_WINDOW_NAME)
        # Render the settings window and its content, if it is not minimized.
        settings.begin(ctrl_frame)
        if not settings.isMinimized():
            cvui.checkbox('Use Canny Edge', use_canny)
            cvui.checkbox('Detect ROIs', detect_rois)
            cvui.checkbox('Fuse Track and Bounding Box', fuse_track_and_bbox)
            cvui.trackbar(settings.width() - 20, low_threshold, 5, 150)
            cvui.trackbar(settings.width() - 20, high_threshold, 80, 300)
            cvui.trackbar(settings.width() - 20, playback_fps, 1, 120)
            cvui.space(20) # add 20px of empty space
        settings.end()

        control.begin(ctrl_frame)
        if not control.isMinimized():
            # Should we quit?
            if cvui.button(settings.width() - 20, 20, "Quit"):
                break

            if is_paused:
                if cvui.button(settings.width() - 20, 20, "Resume"):
                    is_paused = False
            else:
                if cvui.button(settings.width() - 20, 20, "Pause"):
                    is_paused = True

        control.end()

        result_win[0:height,:] = org_frame
        result_win[height:,:] = plot_win
        out.write(result_win)

        cvui.update(CTRL_WINDOW_NAME)
        cv2.imshow(CTRL_WINDOW_NAME, ctrl_frame)

        # Check if ESC key was pressed
        if cv2.waitKey(5) == 27:
            break

    # save the tracking params at the end of the loop
    tracker.save(TRACKING_PARAMS_FILE)
    out.release()
    cap.release()


if __name__ == '__main__':
    main()