"""
Utilities for work with COCO WholeBody dataset
"""
import os
import cv2
import numpy as np

from collections import defaultdict

from .hands_utils import intersecting, scale_bbox, combine_bboxes, 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': 'coco', 'file_name': None, 'annotations': []})
    for annot in annot_data['annotations']:
        # if either hand is annotated in this entry
        if annot['lefthand_valid'] or annot['righthand_valid']:
            # hand bounding boxes
            bbox1, bbox2 = annot['lefthand_box'], annot['righthand_box']
            # calculate if hands are interacting
            interacting = intersecting(bbox1, bbox2, scale=1.3) if annot['lefthand_valid'] and annot['righthand_valid'] else False
            # construct combined bounding box if hands are interacting
            combined_bbox = combine_bboxes(bbox1, bbox2) if interacting else None
            # reshape keypoints into Nx3 array
            left_keypoints = np.array(annot['lefthand_kpts']).reshape(-1, 3)
            right_keypoints = np.array(annot['righthand_kpts']).reshape(-1, 3)
            # add entry to output dictionary
            data_dict[annot['image_id']]['annotations'].append(
                {
                    'interacting': interacting,
                    'combined_bbox': combined_bbox,
                    'handness': [annot['lefthand_valid'], annot['righthand_valid']],
                    'bboxes': [bbox1, bbox2],
                    'keypoints': np.concatenate((left_keypoints, right_keypoints), axis=0)
                })
    # 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
