"""
WLASL dataset utilities
"""
import os
import cv2
import numpy as np

from random import shuffle

from .dataset_utils import load_json, fuse_forward_backward_passes


def read_annotation(path, shuffle_glosses=False):
    """Read annotation file

    Args:
        path (str): path to annotation file
        shuffle_glosses (bool, optional): flag wether to shuffle data. Defaults to False.

    Returns:
        dict: constructed annotation dictionary
    """
    # load file and shuffle if specified
    data = load_json(path)
    if shuffle_glosses:
        shuffle(data)
    annot_dict = {}
    for gloss in data:
        gloss_name = gloss['gloss']
        for instance in gloss['instances']:
            video_id = instance['video_id']
            annot_dict[video_id] = {
                'video_file_name': f'{video_id}.mp4',
                'keypoint_annotation_file_name': f'{video_id}.json',
                'gloss': gloss_name,
                'split': instance['split']
            }
    return annot_dict


def data_from_annotation(annot, videos_root_path, kpt_annot_root_path=None, clip_frames=0):
    """Load data from annotation dictionary entry

    Args:
        annot (dict): annotation entry
        videos_root_path (str): root directory of videos
        kpt_annot_root_path (str, optional): root directory of annotation keypoints. Defaults to None.
        clip_frames (int, optional): number of frames to remove from start and end of the sequence. Defaults to 0.

    Returns:
        str, list, str, list, list, list: annotated class name, list of video frames, target dataset subset,
                                          forward pass keypoints, backward pass keypoint and fused keypoints
    """
    vid_file_name = annot['video_file_name']
    kpts_file_name = annot['keypoint_annotation_file_name']
    gloss = annot['gloss']
    split = annot['split']
    video_path = os.path.join(videos_root_path, vid_file_name)
    # read video file and save frames
    frames = list()
    capture = cv2.VideoCapture(video_path)
    while capture.isOpened():
        ret, frame = capture.read()
        if not ret:
            break
        frames.append(frame)
    # clip frames
    frames = frames[clip_frames:len(frames)-clip_frames]
    try:
        # try to load keypoint annotations
        keypoint_annot_path = os.path.join(kpt_annot_root_path, kpts_file_name)
        keypoint_annotation = load_json(keypoint_annot_path)
    except (TypeError, FileNotFoundError):
        # if annotation file is not present return Nones
        fwd_keypoints, bwd_keypoints, fused_keypoints = None, None, None
    else:
        # clip keypoints
        fwd_hand_keypoints = [np.array(kpts[0]).reshape(-1, 21, 2) for kpts in keypoint_annotation['forward']][clip_frames:len(frames)-clip_frames]
        fwd_face_keypoints = [np.array(kpts[1]).reshape(-1, 98, 2) for kpts in keypoint_annotation['forward']][clip_frames:len(frames)-clip_frames]
        fwd_keypoints = [fwd_hand_keypoints, fwd_face_keypoints]
        bwd_hand_keypoints = [np.array(kpts[0]).reshape(-1, 21, 2) for kpts in reversed(keypoint_annotation['backward'])][clip_frames:len(frames)-clip_frames]
        bwd_face_keypoints = [np.array(kpts[1]).reshape(-1, 98, 2) for kpts in reversed(keypoint_annotation['backward'])][clip_frames:len(frames)-clip_frames]
        bwd_keypoints = [bwd_hand_keypoints, bwd_face_keypoints]
        fused_hand_keypoints = fuse_forward_backward_passes(fwd_hand_keypoints, bwd_hand_keypoints, (21, 2), n=2)
        fused_face_keypoints = fuse_forward_backward_passes(fwd_face_keypoints, bwd_face_keypoints, (98, 2), n=1)
        fused_keypoints = [fused_hand_keypoints, fused_face_keypoints]
    return gloss, frames, split, fwd_keypoints, bwd_keypoints, fused_keypoints
