"""
Inference of all algorithms on video file feed, excluding automatic avatar sign language interpretation
"""
import sys
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import cv2
import numpy as np

from time import time

from detection.modules.inference_utils import init_detector, detect
from pose_estimation.modules.inference_utils import init_face_estimator, estimate_face
from pose_estimation.modules.inference_utils import init_hand_estimator, estimate_hands
from translation.modules.inference_utils import SignLanguageClassifier, draw_classes

from tracking.modules.utils import filter_small_boxes, boxes_to_square, merge_boxes, draw_tracks
from tracking.modules.tracker import Tracker, detections_from_data


def main():
    base_path = os.path.dirname(os.path.realpath(__file__))
    # init models
    detector_cfg = 'yolov5m'
    detector = init_detector(detector_cfg)
    hand_estimator, hand_skeleton = init_hand_estimator('efficientnet-b1')
    face_estimator = init_face_estimator('efficientnet-b1')
    frames_per_sign = 30
    classifier = SignLanguageClassifier(100, frames_per_sign)
    # list of videos to process
    vids = ['vid_0.mp4']
    # for each video
    for vid_name in vids:
        # input path, input videos are expected to be in data/videos/
        path = os.path.join(base_path, 'data', 'videos', vid_name)
        # output paths - one file fully processed outputs, one only pose estimation outputs
        out_path = os.path.join(base_path, 'data', 'videos', f'out_{vid_name}')
        out_pose_path = os.path.join(base_path, 'data', 'videos', f'out_pose_{vid_name}')
        # video feed init
        cam = cv2.VideoCapture(path)
        width  = int(cam.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cam.get(cv2.CAP_PROP_FRAME_HEIGHT))
        # video outputs init
        codec = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(out_path, codec, 30, (width, height))
        out_pose = cv2.VideoWriter(out_pose_path, codec, 15, (width, height))
        # object tracker
        tracker = Tracker(classes=2, max_age=30, matching_threshold=(0.4*max((width, height)))**2, max_iou_distance=0.7, dT=1/30)
        frame_idx = 0
        while cam.isOpened():
            # read frame
            ret, img = cam.read()
            if not ret:
                break
            
            s = time()
            # detect objects, filter small boxes and convert them to squares
            boxes, scores, classes = detect(detector, img)
            boxes, scores, classes = filter_small_boxes(boxes, scores, classes, min_area=1e4)
            boxes = boxes_to_square(boxes)
            # merge overlapping boxes and split by class to face and hand boxes
            boxes, scores, classes = merge_boxes(boxes, scores, classes)
            hand_boxes = boxes[classes == 0, :]
            face_boxes = boxes[classes == 1, :]
            detection_time = time() - s
            s = time()
            # run pose estimation model inferences
            handness, hand_keypoints = estimate_hands(hand_estimator, img, hand_boxes, center_format=True)
            hand_estimation_time = time() - s
            s = time()
            face_keypoints = estimate_face(face_estimator, img, face_boxes, center_format=True)
            face_estimation_time = time() - s

            s = time()
            # construct detection objects from detection and pose estimation results
            detections = detections_from_data(hand_boxes, hand_keypoints, handness, face_boxes, face_keypoints)
            # update object tracker and run prediction
            tracker.update(detections)
            boxes, keypoints, ages = tracker.predict()
            tracker_time = time() - s

            s = time()
            # run classification
            signs, sign_probs = classifier.classify(keypoints[0], keypoints[1], n_output=5)
            class_time = time() - s

            sys.stdout.write(f'\rFrame: {frame_idx} - Detection: {detection_time*1e3:.2f}ms  Hand pose estimation: {hand_estimation_time*1e3:.2f}ms  Face pose estimation: {face_estimation_time*1e3:.2f}ms  Tracker: {tracker_time*1e3:.2f}ms  Classifier: {class_time*1e3:.2f}ms{10*" "}')
            sys.stdout.flush()

            # visualize outputs and write processed frames
            draw_tracks(img, boxes, keypoints, ages, hand_skeleton=hand_skeleton)
            draw_classes(img, signs, sign_probs)
            out.write(img)
            pose_map = np.zeros_like(img, dtype=np.uint8)
            draw_tracks(pose_map, boxes, keypoints, ages, hand_skeleton=hand_skeleton, draw_boxes=False)
            out_pose.write(pose_map)

            frame_idx += 1
            cv2.imshow('img', img)
            if cv2.waitKey(1) == 27:
                break
        cv2.destroyAllWindows()
        cam.release()
        out.release()
        out_pose.release()


if __name__ == '__main__':
    main()
