"""
Visualize WLASL dataset from files and from TFRecord files
"""
import cv2
import numpy as np
import tensorflow as tf

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

import sys
sys.path.append(CFG.PARENT_PATH)
from pose_estimation.modules.inference_utils import draw_pose, draw_keypoints
from pose_estimation.modules.estimation_utils import load_skeleton
from pose_estimation.modules.config import CFG as POSE_CFG


def wlasl_from_files():
    # read annotation file and load hand skeleton
    annotation = wlasl_utils.read_annotation(CFG.WLASL_DATASET_PATHS['orig_annotations'], shuffle_glosses=False)
    hand_skeleton = load_skeleton(POSE_CFG.COCO_SKELETON_PATH)
    # for every annotation entry
    for video_id, annot in annotation.items():
        # read annotation entry data
        gloss, frames, split, fwd_keypoints, bwd_keypoints, fused_keypoints = wlasl_utils.data_from_annotation(
            annot, CFG.WLASL_DATASET_PATHS['videos'], CFG.WLASL_DATASET_PATHS['keypoints_annotations'], clip_frames=3)
        # viz a specified subset only
        if split != 'test':
            continue
        # viz only keypoint annotated samples
        if fwd_keypoints is None:
            continue
        # repeat viz video until stopped
        repeat = True
        # first pass is fast, next ones are slowed down
        first_pass = True
        while repeat:
            # for every frame in video
            for frame, fwd_hand_kpts, fwd_face_kpts, bwd_hand_kpts, bwd_face_kpts, fused_hand_kpts, fused_face_kpts in zip(frames, *fwd_keypoints, *bwd_keypoints, *fused_keypoints):
                # concat frame with word name canvas strip
                frame = np.concatenate((np.zeros((35, 256, 3), dtype=np.uint8), frame), axis=0)
                # put sign word name onto concatted strip
                cv2.putText(frame, gloss, (5, 25), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255), 1)
                cv2.imshow('video sample', frame)
                # keypoint annotation canvas for forward annotation pass
                keypoint_map = np.zeros((256, 256, 3), dtype=np.uint8)
                draw_pose(keypoint_map, fwd_hand_kpts, np.ones((len(fwd_hand_kpts), 2)), hand_skeleton)
                draw_keypoints(keypoint_map, fwd_face_kpts)
                cv2.imshow('forward keypoints', keypoint_map)
                # keypoint annotation canvas for backward annotation pass
                keypoint_map = np.zeros((256, 256, 3), dtype=np.uint8)
                draw_pose(keypoint_map, bwd_hand_kpts, np.ones((len(bwd_hand_kpts), 2)), hand_skeleton)
                draw_keypoints(keypoint_map, bwd_face_kpts)
                cv2.imshow('backward keypoints', keypoint_map)
                # keypoint annotation canvas for combined forward and backward annotations
                keypoint_map = np.zeros((256, 256, 3), dtype=np.uint8)
                draw_pose(keypoint_map, fused_hand_kpts, np.ones((len(fused_hand_kpts), 2)), hand_skeleton)
                draw_keypoints(keypoint_map, fused_face_kpts)
                cv2.imshow('fused keypoints', keypoint_map)
                # play the first pass by itself, next passes frame by frame
                key = cv2.waitKey(1 if first_pass else 0)
                # escape - terminate
                if key == 27:
                    cv2.destroyAllWindows()
                    return
                # Q - next video
                elif key == 113:
                    repeat = False
                    break
            first_pass = False
    cv2.destroyAllWindows()


def wlasl_from_tfrecords(n_classes):
    """Visualize WLASL dataset from TFRecord file

    Args:
        n_classes (int): number of classes, one of (100, 300, 1000, 2000)
    """
    # load dataset and hand skeleton
    dataset = dataset_utils.init_dataset(CFG.WLASL_TFRECORDS_PATH[n_classes]['test'], 1, n_classes, 50, augment=True)
    hand_skeleton = load_skeleton(POSE_CFG.COCO_SKELETON_PATH)
    # keypoint visualization canvas size
    map_size = 512
    for x, y in dataset:
        # split keypoints to hand and face keypoints
        hand_pose, face_pose = tf.split(x, (2*21*2, 1*98*2), axis=-1)
        hand_pose = hand_pose.numpy().reshape(-1, 2, 21, 2)
        face_pose = face_pose.numpy().reshape(-1, 1, 98, 2)
        # transform keypoints to image space
        hand_pose = (hand_pose[np.any(hand_pose, axis=(1, 2, 3)), ...] + 1) * map_size // 2
        face_pose = (face_pose[np.any(face_pose, axis=(1, 2, 3)), ...] + 1) * map_size // 2
        # for every frame of a sequence
        for hand_kpts, face_kpts in zip(hand_pose, face_pose):
            # init canvas and draw keypoints
            keypoint_map = np.zeros((map_size, map_size, 3), dtype=np.uint8)
            draw_pose(keypoint_map, hand_kpts, np.any(hand_kpts != map_size // 2, axis=1), hand_skeleton)
            draw_keypoints(keypoint_map, face_kpts)
            # visualize
            cv2.imshow('dataset', keypoint_map)
            if cv2.waitKey(1) == 27:
                cv2.destroyAllWindows()
                return
    cv2.destroyAllWindows()


if __name__ == '__main__':
    # wlasl_from_files()
    wlasl_from_tfrecords(2000)
