import os
import cv2
import numpy as np
import pickle

from .hands_utils import intersecting, scale_bbox, combine_bboxes, bbox_from_keypoints


def keypoints_to_coco(keypoints):
    """Convert RHD points from its data format to COCO format used in this project

    Args:
        keypoints (np.ndarray): input keypoints

    Returns:
        np.ndarray: output keypoints
    """
    # RHD dataset has keypoints in different order
    # manually constructed reordering map
    indices = [0, 4, 3, 2, 1, 8, 7, 6, 5, 12, 11, 10, 9, 16, 15, 14, 13, 20, 19, 18, 17,
               21, 25, 24, 23, 22, 29, 28, 27, 26, 33, 32, 31, 30, 37, 36, 35, 34, 41, 40, 39, 38]
    transformed = keypoints[indices, :]
    return transformed


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

    Args:
        path (str): path to annotation file

    Returns:
        dict: constructed annotation dictionary
    """
    # read annotation file
    with open(annot_path, 'rb') as file:
        annotation = pickle.load(file)
    annot_dict = {}
    for key, val in annotation.items():
        file_name = f'{key:05d}.png'
        # convert keypoints
        keypoints_coco = keypoints_to_coco(val['uv_vis'])
        keypoints = keypoints_coco[:, :2]
        validity = keypoints_coco[:, 2]

        left_keypoints = keypoints[:21]
        right_keypoints = keypoints[21:]
        left_valid = np.any(validity[:21])
        right_valid = np.any(validity[21:])
        handness = [left_valid, right_valid]
        # construct bboxes from keypoints
        left_bbox = bbox_from_keypoints(left_keypoints) if left_valid else np.zeros(4)
        right_bbox = bbox_from_keypoints(right_keypoints) if right_valid else np.zeros(4)
        # check if hands are interacting
        interacting = intersecting(left_bbox, right_bbox, scale=1.3)
        # if hands are interacting construct combined bbox
        combined_bbox = combine_bboxes(left_bbox, right_bbox) if interacting else np.zeros(4)
        # write entry
        annot_dict[key] = {'dataset': 'rhd', 'file_name': file_name, 'keypoints': keypoints, 'validity': validity, 'handness': handness,
                           'interacting': interacting, 'bboxes': [left_bbox, right_bbox], 'combined': combined_bbox}
    return annot_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)
    # interacting hands -> single record
    if annot['interacting']:
        keypoints = [annot['keypoints'][:, :2]]
        validity = [annot['validity']]
        bbox = annot['combined']
        bboxes = [scale_bbox(bbox, scale)]
        classes = [1]
    # two single hands -> two records
    else:
        keypoints = list()
        validity = list()
        bboxes = list()
        classes = list()
        for i in range(2):
            if annot['handness'][i]:
                bbox = annot['bboxes'][i]
                kpts = annot['keypoints'][:21, :] if i == 0 else annot['keypoints'][21:, :]
                # pad with zeros to unify tensor sizes
                kpts = np.concatenate((kpts, np.zeros_like(kpts, dtype=kpts.dtype)), axis=0)
                valid = annot['validity'][:21] if i == 0 else annot['validity'][21:]
                # pad with zeros to unify tensor sizes
                valid = np.concatenate((valid, np.zeros_like(valid, dtype=valid.dtype)), axis=0)
                keypoints.append(kpts[:, :2])
                validity.append(valid)
                bboxes.append(scale_bbox(bbox, scale))
                classes.append(0)
    return img, bboxes, keypoints, validity, classes
