import numpy as np
import cv2 as cv
import math

def edgedetection(cap):
    lines = []

    # params for ShiTomasi corner detection
    #feature_params = dict( maxCorners = 1,qualityLevel = 0.1,minDistance = 7,blockSize = 7)
    # Parameters for lucas kanade optical flow
    lk_params = dict( winSize  = (15,15),maxLevel = 2, criteria = (cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03))
    #p0 = cv.goodFeaturesToTrack(old_gray, mask = None, **feature_params)
    # Create a mask image for drawing purposes
    ret,old_frame = cap.read()
    old_gray = cv.cvtColor(old_frame, cv.COLOR_BGR2GRAY)
    mask = np.zeros_like(old_frame)
    while(1):
        ret,frame = cap.read()
        ret,frame = cap.read()
        ret,frame = cap.read()


        
        frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
        diff = cv.absdiff(old_gray,frame_gray)
        
        feature_params = dict( maxCorners = 20,qualityLevel = 0.3,minDistance = 7,blockSize = 7)
        p0 = cv.goodFeaturesToTrack(diff, mask = None, **feature_params)
        
        
        
        # calculate optical flow
        p1, st, err = cv.calcOpticalFlowPyrLK(old_gray, frame_gray, p0, None, **lk_params)
        # 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()
            angle = math.degrees(math.atan2(a-c, b - d));
            print(angle)
            mask = cv.line(mask, (a,b),(c,d), (255,0,0), 3)
            frame = cv.circle(frame,(a,b),5,(255,0,0),-1)
            if angle<60 and angle>30:
                lines.append((a, b, c, d))
                mask = cv.line(mask, (a,b),(c,d), (0,255,0), 3)
                frame = cv.circle(frame,(a,b),5,(0,255,0),-1)
            elif angle<-105 and angle>-125:
                mask = cv.line(mask, (a,b),(c,d), (0,0,255), 3)
                frame = cv.circle(frame,(a,b),5,(0,0,255),-1)
            else:
                pass
                
                
        cv.imshow('frame',cv.add(frame,mask))    
        k = cv.waitKey(50) & 0xff
        if k == 27:
            break
        # Now update the previous frame and previous points
        old_gray = frame_gray.copy()
        mask = np.zeros_like(old_frame)
        
        

edgedetection(cap = cv.VideoCapture('session0_center.avi'))
