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

from collections import defaultdict

from .hands_utils import intersecting, scale_bbox, combine_bboxes, bbox_from_keypoints, load_json


def load_annotation(path):
    """Load annotation file and construct an annotation dictionary

    Args:
        path (str): path to annotation file

    Returns:
        dict: constructed annotation dictionary
    """
    # load annotation file
    annot_data = load_json(path)
    # set a base template for output annotation entries
    # 'annotations' key is list, as there can be multiple annotations for a single image 
    data_dict = defaultdict(lambda: {'dataset': 'halpe', 'file_name': None, 'annotations': None})
    for annot in annot_data['annotations']:
        # keypoints in annotation are sequenced together
        keypoints = np.array(annot['keypoints']).reshape(-1, 136, 3)
        annotations = list()
        # for each annotated human in frame
        for kpts in keypoints:
            # left hand keypoints
            left_hand_keypoints = kpts[94:115, :]
            left_hand_valid = left_hand_keypoints[:, 2] > 1
            left_handness = np.any(left_hand_valid)
            left_hand_bbox = bbox_from_keypoints(left_hand_keypoints[left_hand_valid, :2]) if left_handness else None
            left_handness = left_handness and (left_hand_bbox[2] * left_hand_bbox[3] > 50*50)
            # right hand keypoints
            right_hand_keypoints = kpts[115:, :]
            right_hand_valid = right_hand_keypoints[:, 2] > 1
            right_handness = np.any(right_hand_valid)
            right_hand_bbox = bbox_from_keypoints(right_hand_keypoints[right_hand_valid, :2]) if right_handness else None
            right_handness = right_handness and (right_hand_bbox[2] * right_hand_bbox[3] > 50*50)
            if not left_handness and not right_handness:
                continue
            # check if hands are interacting, if so, calculate combined bbox
            interacting = intersecting(left_hand_bbox, right_hand_bbox, scale=1.3) if left_handness and right_handness else False
            combined_bbox = combine_bboxes(left_hand_bbox, right_hand_bbox) if interacting else None
            # write annotation entry
            annotations.append({
                'interacting': interacting,
                'combined_bbox': combined_bbox,
                'handness': [left_handness, right_handness],
                'bboxes': [left_hand_bbox, right_hand_bbox],
                'keypoints': np.concatenate((left_hand_keypoints, right_hand_keypoints), axis=0)
            })
        # if image constains at least one annotated human, write the entry to output dictionary
        if len(annotations) > 0:
            data_dict[annot['image_id']]['annotations'] = annotations
    # image filenames are separated in input annotation file from other annotations
    # add image file names to output annotation file
    for img_annot in annot_data['images']:
        if img_annot['id'] in data_dict.keys():
            data_dict[img_annot['id']]['file_name'] = img_annot['file_name']
    return data_dict


def data_from_annotation(annot, images_root_path, scale=None):
    """Get data from annotation entry in annotation dictionary

    Args:
        annot (dict): annotation entry
        images_root_path (str): root directory of dataset images
        scale (float, optional): scale ratio for bounding box extraction from images. None equals to 1.0 - the original bbox size.
                                 Values greater then 1.0 scale bbox' width and height read from annotation. Defaults to None.

    Raises:
        KeyError: if target image file does not exist

    Returns:
        np.ndarray, list, list, list, list: read image, bboxes of annotated hands, keypoints of annotated hands, validity of keypoints, classes of hands
    """
    # read image
    img_path = os.path.join(images_root_path, annot['file_name'])
    if not os.path.isfile(img_path):
        raise KeyError('File not found')
    img = cv2.imread(img_path)
    # extract bboxes, keypoints, their validity and classes from annotation entry
    keypoints = list()
    validity = list()
    bboxes = list()
    classes = list()
    # for each item in entry
    for ann in annot['annotations']:
        # if hands are interacting
        if ann['interacting']:
            kpts = ann['keypoints'][:, :2]
            valid = ann['keypoints'][:, 2] > 0
            bbox = ann['combined_bbox']
            keypoints.append(kpts)
            validity.append(valid)
            bboxes.append(scale_bbox(bbox, scale))
            classes.append(1)
        # if hands are not interacting
        else:
            for i in range(2):
                if ann['handness'][i]:
                    bbox = ann['bboxes'][i]
                    kpts = ann['keypoints'][:21, :] if i == 0 else ann['keypoints'][21:, :]
                    # concat zeros to unify keypoint dimensions
                    kpts = np.concatenate((kpts, np.zeros_like(kpts, dtype=kpts.dtype)), axis=0)
                    valid = kpts[:, 2] > 0
                    keypoints.append(kpts[:, :2])
                    validity.append(valid)
                    bboxes.append(scale_bbox(bbox, scale))
                    classes.append(0)
    return img, bboxes, keypoints, validity, classes
