"""
Inference of all algorithms on webcam feed, including automatic avatar sign language interpretation
"""
import sys
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import cv2
import numpy as np
import threading
import logging

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

from avatar.modules.sign_tracker import SignTracker
from avatar.modules.renderer import _queue as sign_queue
from avatar.modules.renderer import ArtificialInterpreter


# logging.basicConfig(level=logging.DEBUG)

TARGET_LANGUAGE = 'cs_cz'


class Profiler:
    """Inference latency profiler
    """
    def __init__(self, names, n=10):
        """
        Args:
            names (list): list of tracked latency names
            n (int, optional): running average number of samples. Defaults to 10.
        """
        self.n = n
        self.names = names
        self.values = None
    
    def update(self, *values):
        """Update profiler latencies
        
        Args:
            *values: float values of measured latencies
        """
        # init values if not initialized
        if self.values is None:
            self.values = [[val] for val in values]
            return
        # update stats
        for s_val, in_val in zip(self.values, values):
            s_val.append(in_val)
        # discard untracked values
        self.values = [s_val[-self.n:] for s_val in self.values]
    
    def __repr__(self):
        return '  '.join(f'{name}: {np.mean(vals)*1e3:.2f}ms' for name, vals in zip(self.names, self.values))


class ThreadWithReturnValue(threading.Thread):
    """Thread object with return value
    """
    def __init__(self, group=None, target=None, name=None, args=(), kwargs={}):
        """
        Args:
            group (None, optional): group argument for threading.Thread object. Defaults to None.
            target (callable, optional): target function to execute. Defaults to None.
            name (str, optional): name of thread. Defaults to None.
            args (tuple, optional): target function arguments. Defaults to ().
            kwargs (dict, optional): target function keyword arguments. Defaults to {}.
        """
        super().__init__(group, target, name, args, kwargs)
        # return value
        self._return = None
    def run(self):
        """Run thread
        """
        if self._target is not None:
            self._return = self._target(*self._args, **self._kwargs)
    def join(self, *args):
        """Join thread to main thread and return target function return value

        Returns:
            any: target function return value
        """
        threading.Thread.join(self, *args)
        return self._return


class ImageProcessingThread(threading.Thread):
    """Video processing thread
    """
    def __init__(self, name, threaded, frames_per_sign, detector_cfg, *args, **kwargs):
        """
        Args:
            name (str): name of thread
            threaded (bool): flag wether to thread pose estimation algorithms
            frames_per_sign (int): number of time steps to pass to classifier
            detector_cfg (str): object detector configuration
        """
        super().__init__(*args, **kwargs)
        self.name = name
        self.threaded = threaded
        # init models
        self.detector = init_detector(detector_cfg)
        self.hand_estimator, self.hand_skeleton = init_hand_estimator('efficientnet-b0')
        self.face_estimator = init_face_estimator('efficientnet-b0')
        self.classifier = SignLanguageClassifier(100, frames_per_sign)
        # init trackers
        self.tracker = Tracker(classes=2, max_age=30, matching_threshold=(720*0.4)**2, max_iou_distance=0.7)
        self.sign_tracker = SignTracker()
        # init camera feed
        self.cam = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        self.cam.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
        self.cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
        # init profiler
        self.profiler = Profiler(('Detection', 'Pose estimation', 'Tracker', 'Classification'), n=10)
    
    def run(self):
        """Run thread
        """
        while self.cam.isOpened():
            # read frame
            ret, img = self.cam.read()
            if not ret:
                continue
            s = time()
            # detect objects, filter small boxes and convert them to squares
            boxes, scores, classes = detect(self.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()
            if self.threaded:
                # init hand and face pose estimation threads, start them and join them
                hands_thread = ThreadWithReturnValue(target=estimate_hands, args=(self.hand_estimator, img, hand_boxes), kwargs={'center_format': True})
                face_thread = ThreadWithReturnValue(target=estimate_face, args=(self.face_estimator, img, face_boxes), kwargs={'center_format': True})
                hands_thread.start()
                face_thread.start()
                handness, hand_keypoints = hands_thread.join()
                face_keypoints = face_thread.join()
            else:
                # sequentially run pose estimation model inferences
                handness, hand_keypoints = estimate_hands(self.hand_estimator, img, hand_boxes, center_format=True)
                face_keypoints = estimate_face(self.face_estimator, img, face_boxes, center_format=True)
            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
            self.tracker.update(detections)
            boxes, keypoints, ages = self.tracker.predict()
            tracker_time = time() - s

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

            # update sign request tracker and put sign to queue if the request is valid
            sing_to_interpret = self.sign_tracker.update(signs, sign_probs)
            if sing_to_interpret is not None:
                sign_queue.put((sing_to_interpret, TARGET_LANGUAGE))
            # update profiler
            self.profiler.update(detection_time, estimation_time, tracker_time, classify_time)
            sys.stdout.write(f'\r{self.profiler}{10*" "}')
            sys.stdout.flush()
            # visualize results
            draw_tracks(img, boxes, keypoints, ages, hand_skeleton=self.hand_skeleton, draw_text=True)
            draw_classes(img, signs, sign_probs)
            cv2.imshow('img', img)
            if cv2.waitKey(1) == 27:
                break
        # release objects
        cv2.destroyAllWindows()
        self.cam.release()


def main():
    # image processing thread
    img_processing_thread = ImageProcessingThread('img_processing_thread', True, 20, 'yolov5s')
    # run application in main thread and start image processing thread
    app = ArtificialInterpreter()
    img_processing_thread.start()
    app.run()
    img_processing_thread.join()


if __name__ == '__main__':
    main()
