"""
Annotate WLASL video samples with hand and face keypoints
"""
from random import sample
import numpy as np

from time import time
from functools import partial

from modules import wlasl_utils, dataset_utils
from modules.config import CFG

import os
import sys
sys.path.append(CFG.PARENT_PATH)
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 tracking.modules.utils import filter_small_boxes, boxes_to_square, merge_boxes
from tracking.modules.tracker import Tracker, detections_from_data


def predict(detector, hand_estimator, face_estimator, frame):
    """Single frame prediction

    Args:
        detector (ExportModel): hand and face detector
        hand_estimator (ExportModel): hand pose estimator
        face_estimator (ExportModel): face pose estimator
        frame (np.ndarray): input frame

    Returns:
        np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray: predicted hand bounding boxes, hand keypoints,
                                                                    hand presence indicators, face bounding boxes and face keypoints
    """
    # detect hands and faces, filter small boxes, convert them to square and merge overlapping boxes
    boxes, scores, classes = detect(detector, frame)
    boxes, scores, classes = filter_small_boxes(boxes, scores, classes, min_area=1e2)
    boxes = boxes_to_square(boxes)
    boxes, scores, classes = merge_boxes(boxes, scores, classes)
    # split to hand and face boxes
    hand_boxes = boxes[classes == 0, :]
    face_boxes = boxes[classes == 1, :]
    # estimate hand and face keypoints
    handness, hand_keypoints = estimate_hands(hand_estimator, frame, hand_boxes, center_format=True)
    face_keypoints = estimate_face(face_estimator, frame, face_boxes, center_format=True)
    return hand_boxes, hand_keypoints, handness, face_boxes, face_keypoints


def tracker_update_predict(tracker, hand_boxes, hand_keypoints, handness, face_boxes, face_keypoints):
    """Update tracker with predictions and perform tracker prediction

    Args:
        tracker (Tracker): input tracker
        hand_boxes (np.ndarray): predicted hand bounding boxes
        hand_keypoints (np.ndarray): predicted hand keypoints
        handness (np.ndarray): predicted hand presence indicators
        face_boxes (np.ndarray): predicted face bounding boxes
        face_keypoints (np.ndarray): predicted face keypoints

    Returns:
        list: list of hand and face keypoints
    """
    detections = detections_from_data(hand_boxes, hand_keypoints, handness, face_boxes, face_keypoints)
    tracker.update(detections)
    _, keypoints, _ = tracker.predict()
    return keypoints


def map_classes():
    """Construct class map, mapping strings (sign names) to integers (class IDs)
    """
    # read annotation file
    annotation = wlasl_utils.read_annotation(CFG.WLASL_DATASET_PATHS['orig_annotations'])
    # if file does not already exist
    if not os.path.isfile(CFG.WLASL_DATASET_PATHS['class_map']):
        # list of all signs used in dataset, not a set because we want to keep words in the same order
        glosses = list()
        # parse annotation
        for annot in annotation.values():
            if annot['gloss'] in glosses:
                continue
            # add to list if not already in it
            glosses.append(annot['gloss'])
        # construct class map and save to file
        class_map = {gloss: i for i, gloss in enumerate(glosses)}
        dataset_utils.dump_json(CFG.WLASL_DATASET_PATHS['class_map'], class_map)


def annotate_wlasl():
    """Annotate WLASL videos with keypoints
    """
    # read annotation file
    annotation = wlasl_utils.read_annotation(CFG.WLASL_DATASET_PATHS['orig_annotations'])
    # initialize models - detector and estimators
    detector = init_detector('yolov5m')
    hand_estimator, _ = init_hand_estimator()
    face_estimator = init_face_estimator()
    # tracker construction function - for every video a new tracker is required
    tracker_fn = partial(Tracker, classes=2, max_age=30, n_init=3, matching_threshold=(0.4*256)**2, max_iou_distance=0.85, dT=1/25, kalman_classes=(0, 1))

    samples = len(annotation.keys())
    times_per_sample = []
    # list of already annotated samples
    annotated_samples = [path.rstrip('.json') for path in os.listdir(CFG.WLASL_DATASET_PATHS['keypoints_annotations'])]
    # parse annotation
    for i, (video_id, annot) in enumerate(annotation.items()):
        # if annotation is not present
        if not video_id in annotated_samples:
            s = time()
            # get data from annotation entry
            frames = wlasl_utils.data_from_annotation(annot, CFG.WLASL_DATASET_PATHS['videos'])[1]
            # perform two passes through the video - forward and backward
            keypoints_dict = {'forward': [], 'backward': []}
            predictions = list()
            # predict on every frame
            for frame in frames:
                outputs = predict(detector, hand_estimator, face_estimator, frame)
                predictions.append(outputs)
            # forward and backward trackers
            tracker_fwd, tracker_bwd = tracker_fn(), tracker_fn()
            # for every frame
            for out_fwd, out_bwd in zip(predictions, reversed(predictions)):
                # update forward pass tracker and get its predictions
                hand_boxes, hand_keypoints, handness, face_boxes, face_keypoints = out_fwd
                keypoints = tracker_update_predict(tracker_fwd, hand_boxes, hand_keypoints, handness, face_boxes, face_keypoints)
                keypoints = [kpts.reshape(-1).tolist() for kpts in keypoints]
                # save predicted keypoints
                keypoints_dict['forward'].append(keypoints)
                # update backward pass tracker and get its predictions
                hand_boxes, hand_keypoints, handness, face_boxes, face_keypoints = out_bwd
                keypoints = tracker_update_predict(tracker_bwd, hand_boxes, hand_keypoints, handness, face_boxes, face_keypoints)
                keypoints = [kpts.reshape(-1).tolist() for kpts in keypoints]
                # save predicted keypoints
                keypoints_dict['backward'].append(keypoints)
            # save annotation to file
            dataset_utils.dump_json(os.path.join(CFG.WLASL_DATASET_PATHS['keypoints_annotations'], f'{video_id}.json'), keypoints_dict)
            times_per_sample.append(time() - s)
        # time per sample stats
        mean_time = np.mean(times_per_sample) if len(times_per_sample) else 0
        eta_minutes = int(mean_time * (samples - i + 1) / 60)
        sys.stdout.write(f'\rAnnotation process {i+1}/{samples} | {mean_time:.2f}s/sample | ETA:{eta_minutes//60}h {eta_minutes%60}min{15*" "}\r')
        sys.stdout.flush()


if __name__ == '__main__':
    map_classes()
    # annotate_wlasl()
