from argparse import ArgumentParser
from pathlib import Path

import cv2
import imutils
import time


def main():
    parser = ArgumentParser()
    parser.add_argument("--output", type=Path, required=False, default="output.avi")
    parser.add_argument("video_file", type=Path)
    args = parser.parse_args()

    cap = cv2.VideoCapture(args.video_file)
    output = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*'MPEG'), 30, (720, 1792))

    static_frame = None

    while True:
        ret, frame = cap.read()
        if ret:
            cropped_frame = frame[55:720, 0:540]
            grey_frame = cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2GRAY)

            if static_frame is None:
                static_frame = grey_frame
                continue
            

            _, thresh = cv2.threshold(grey_frame, 120, 255, cv2.THRESH_BINARY)
            dilated_image = cv2.dilate(thresh, None, iterations=2)

            cnts = cv2.findContours(dilated_image.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cnts = imutils.grab_contours(cnts)

            for c in cnts:
                if cv2.contourArea(c) < 700:
                    continue
                
                # get the bounding box coordinates
                (x, y, w, h) = cv2.boundingRect(c)
                cv2.rectangle(cropped_frame, (x, y), (x+w, y+h), (0, 255, 0), 2)

            output.write(frame)
            cv2.imshow("output", frame)

            # exit if 's' key is pressed
            if cv2.waitKey(1) & 0xFF == ord('s'):
                break
            time.sleep(0.03)
        else:
            print("Could not read frame!")
            break

    cv2.destroyAllWindows()
    output.release()
    cap.release()


if __name__ == "__main__":
    main()