"""
*******************************************
*** Lucas-Kanade Optical Flow in OpenCV ***
*******************************************

Tutorial: https://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_video/py_lucas_kanade/py_lucas_kanade.html

Assumptions:
1. The pixel intensities of an object do not change between consecutive frames.
2. Neighbouring pixels have similar motion.

Taylor series approximation.

Lucas-Kanade method
--> 3x3 patch around pixels
--> all should have the same motion
--> system of 9 equations for each point
--> Solve system using least square fit -- A+

--> Taylor (I think) will work with small motions. So it fails when there is large motion. So again we go for pyramids. When we go up in the pyramid, small motions are removed and large motions becomes small motions. So applying Lucas-Kanade there, we get optical flow along with the scale.


"""

import os

import cv2
import numpy as np

from utils.util import project_path, ensure_dir

# cap = cv2.VideoCapture(0)
# cap = cv2.VideoCapture(os.path.join(project_path, 'data/slow.flv'))
cap = cv2.VideoCapture(os.path.join(project_path, './imgs/sintel-kittimodel/training-final-ambush_6/0_img.avi'))
# cap = cv2.VideoCapture(os.path.join(project_path, './data/tmp/out.mp4'))

save_dir = os.path.join(project_path, "imgs/lk")
ensure_dir(save_dir)

# params for ShiTomasi corner detection
feature_params = dict(
    maxCorners=100,
    qualityLevel=0.3,
    minDistance=7,
    blockSize=7)

# params for Lucas Kanade optical flow
lk_params = dict(
    winSize=(15, 15),
    maxLevel=2,
    criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))

# Create some random colors
color = np.random.randint(0, 255, (100, 3))

# Take first frame and find corners in it
ret, old_frame = cap.read()
old_gray = cv2.cvtColor(old_frame, cv2.COLOR_BGR2GRAY)
p0 = cv2.goodFeaturesToTrack(old_gray, mask=None, **feature_params)

# Create a mask image for drawing purposes
mask = np.zeros_like(old_frame)

frame_i = 0
while 1:
    frame_i += 1
    ret, frame = cap.read()
    frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # Calculate optical flow
    p1, st, err = cv2.calcOpticalFlowPyrLK(old_gray, frame_gray, p0, None, **lk_params)

    # Save the first pic with circles
    if frame_i == 1:
        for i, (a, b) in enumerate(p0[st == 1]):
            old_frame = cv2.circle(old_frame, (a, b), 5, color[i].tolist(), -1)
        cv2.imwrite(os.path.join(save_dir, f"0.png"), old_frame)

    if p1 is not None:
        # Select good points
        good_new = p1[st == 1]
        good_old = p0[st == 1]

        # Draw the tracks
        for i, (new, old) in enumerate(zip(good_new, good_old)):
            a, b = new.ravel()
            c, d = old.ravel()
            mask = cv2.line(mask, (a, b), (c, d), color[i].tolist(), 2)
            frame = cv2.circle(frame, (a, b), 5, color[i].tolist(), -1)

    img = cv2.add(frame, mask)

    cv2.imshow('frame', img)
    cv2.imwrite(os.path.join(save_dir, f"{frame_i}.png"), img)
    k = cv2.waitKey(30) & 0xff
    if k == 27:
        break

    if k == ord('r'):
        p0 = cv2.goodFeaturesToTrack(frame_gray, mask=None, **feature_params)
        mask = np.zeros_like(old_frame)
    else:
        p0 = good_new.reshape(-1, 1, 2)

    old_gray = frame_gray.copy()

cv2.destroyAllWindows()
cap.release()
