from __future__ import print_function
import numpy as np
import cv2 as cv
import time

lk_params = dict( winSize  = (15, 15),
                  maxLevel = 2,
                  criteria = (cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03))

fgbg = cv.bgsegm.createBackgroundSubtractorMOG()
filename=open("txt.txt","w")

track_len = 10
detect_interval = 5
tracks = []
cam = cv.VideoCapture('fishtest.flv')
frame_count = 0

def update_tracks(p0,p0r,tracks,p1):
    d = abs(p0-p0r).reshape(-1, 2).max(-1)
    good = d < 1
    new_tracks = []


    for tr, (x, y), good_flag in zip(tracks, p1.reshape(-1, 2), good):
        #Zips tracks the points and whether its true or not into a list/array? dunno
        #unpacks into tr(integer values of co-ordinates)/ (x,y) floating values and good flag which is boolean
        #checks if point is good or not if it is then it appends the floating value to the integer
        if not good_flag:
            continue
        tr.append((x, y))
        #drawcircles
        cv.circle(vis, (x, y), 2, (0, 255, 0), -1)
        #delete if too long only keep 10 tracks at any given time

        if len(tr) > track_len:
            del tr[0]
        new_tracks.append(tr)


    tracks = new_tracks
    cv.polylines(vis, [np.int32(tr) for tr in tracks], False, (0, 255, 0))
    return tracks

def update_and_drawBox(contours,boxpoints):
    for i in range(0, len(contours)):
            cnt = contours[i]
            cnt=cnt.astype(np.float32)

            if (len(cnt)>9):
                x,y,w,h = cv.boundingRect(cnt)
                cx=x+(w/2)
                cy=y+(h/2)
                boxpoints.append([x,y])
                boxpoints.append([x+w,y+h])
                boxpoints.append([x+w,y])
                boxpoints.append([x,y+h])
                boxpoints.append([cx,cy])

                cv.rectangle(vis,(x,y),(x+w,y+h),(255,0,0),2)

    boxpoints=np.asarray(boxpoints).astype(np.float32)
    return boxpoints

while True:
    ret, frame = cam.read()
    fgmask = fgbg.apply(frame)
    frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
    vis = frame.copy()
    boxpoints=[]

    #find contours
    im2,contours,hierarchy = cv.findContours(fgmask,cv.RETR_EXTERNAL,cv.CHAIN_APPROX_SIMPLE)

    if len(tracks) > 0:
        p0 = np.float32([tr[-1] for tr in tracks]).reshape(-1, 1, 2)
        p1, _st, _err = cv.calcOpticalFlowPyrLK(previousFrame, frame_gray, p0, None, **lk_params)
        p0r, _st, _err = cv.calcOpticalFlowPyrLK(frame_gray, previousFrame, p1, None, **lk_params)
        tracks=update_tracks(p0,p0r,tracks,p1)

    if frame_count % detect_interval == 0 and len(contours) !=0:
        p = update_and_drawBox(contours,boxpoints)

        if p is not None:
            for x, y in p.reshape(-1, 2):
                #-1 unknown leave up to numpy stop forgetting ya twat
                tracks.append([(x, y)])


    frame_count += 1
    previousFrame = frame_gray
    cv.imshow('frame1', vis)
    cv.imshow("frame2",fgmask)
    cv.imshow("frame",frame)

    if cv.waitKey(30) == ord('q'):
        break

cv.destroyAllWindows()