import cv2
import numpy as np

from model import Detector
from utils import get_bbox_by_point

if __name__ == '__main__':
    detector = Detector(classes=[0, 2, 5, 7])

    cap = cv2.VideoCapture(1)
    display_name = 'detector'
    cv2.namedWindow(display_name, cv2.WINDOW_NORMAL)
    cv2.setWindowProperty(display_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
    # cv2.resizeWindow(display_name, 960, 720)

    while True:
        # Capture frame-by-frame
        ret, frame = cap.read()
        frame_disp = frame.copy()

        results = detector.predict(frame)
        for result in results:
            cls = result[0]
            bbox = result[1:]
            color = (255, 0, 0) if cls == 0 else (0, 255, 255)
            cv2.rectangle(frame_disp, (bbox[0], bbox[1]), (bbox[2] + bbox[0], bbox[3] + bbox[1]),
                          color, 5)

        point = np.array([frame.shape[1] // 2, frame.shape[0] // 2])
        cv2.rectangle(frame_disp, (point[0], point[1]),
                      (point[0] + 1, point[1] + 1),
                      (0, 0, 255), 2)
        if len(results):
            point_bbox = get_bbox_by_point(results[:, 1:], point)
            if point_bbox is not None:
                cv2.rectangle(frame_disp, (point_bbox[0], point_bbox[1]), (point_bbox[2] + point_bbox[0], point_bbox[3] + point_bbox[1]),
                              (0, 0, 255), 2)

        # Display the resulting frame
        cv2.imshow(display_name, frame_disp)
        key = cv2.waitKey(1)
        if key == ord('q'):
            break
        elif key == ord('r'):
            pass
        elif key == ord('s'):
            pass

    # When everything done, release the capture
    cap.release()
    cv2.destroyAllWindows()