import numpy as np
import cv2
import math
import sys

# Init
target = sys.argv[1]
cap = cv2.VideoCapture("data/{}.mp4".format(target))
cap.set(cv2.CAP_PROP_FPS,240)
out = open("{}.tsv".format(target),"w+")
out.write("Time\tX\tY\tCount\tVelocity\n")

# Image Processing
scale = 16.0
diff_thresh = 50
blur_kernel = (7,7)

# Background
bg_alpha = 0.99
bg = None
bg_float = None

# Position
pos_alpha = 0.99
posx = 0.0
posy = 0.0

# Count
running_count = 0.0

# Main Loop
frame_number = 0

while True:
    # Grab Frame
    ret, src = cap.read()
    if not ret:
        break

    # Ensmallen and grayen
    blur = cv2.GaussianBlur(src, blur_kernel,0)
    smol = cv2.resize(blur, (int(src.shape[1]/scale), int(src.shape[0]/scale)))
    gray = cv2.cvtColor(smol, cv2.COLOR_BGR2GRAY)
    if bg is None:
        bg = gray
        bg_float = np.zeros(gray.shape,dtype=np.float32)

    # Get Difference
    diff = cv2.absdiff(gray, bg)
    retval, thresh = cv2.threshold(diff, diff_thresh, 255, cv2.THRESH_TOZERO)

    # Normalize Count
    count = np.sum(thresh)/(smol.shape[0]*smol.shape[1]*255)
    
    # Process Count
    if count > 0.0:
        moments = cv2.moments(thresh)
        cX = (moments["m10"] / moments["m00"])/smol.shape[1]
        cY = (moments["m01"] / moments["m00"])/smol.shape[0]
        prevx = posx
        prevy = posy
        posx = posx * pos_alpha + cX * (1.0-pos_alpha)
        posy = posy * pos_alpha + cY * (1.0-pos_alpha)
        velocity = math.sqrt((prevx-posx) * (prevx-posx) + (prevy - posy) * (prevy - posy))
        running_count = running_count * pos_alpha + count * (1.0 - pos_alpha)
        out.write("{}\t{}\t{}\t{}\t{}\n".format(frame_number/30.0,posx,posy,running_count,velocity))

    # Calculate new BG
    bg_float = bg_float * bg_alpha + gray * (1.0 - bg_alpha)
    bg = bg_float.astype(np.uint8)

    # Draw
    #cv2.circle(smol, (int(posx*smol.shape[1]), int(posy*smol.shape[0])), 3, (255,0,0),6)

    # Show Results
    # cv2.imshow("thresh", cv2.resize(thresh, (int(src.shape[1]/2.0), int(src.shape[0]/2.0))))
    # cv2.imshow("smol", cv2.resize(smol, (int(src.shape[1]/2.0), int(src.shape[0]/2.0))))
    # key = cv2.waitKey(1)
    # if key == 27:
    #     break

    frame_number += 1
    print(frame_number)
    
cap.release()
cv2.destroyAllWindows()