"""
Utilites for MWB dataset manipulation
"""
import os
import cv2
import numpy as np

from glob import glob
from collections import defaultdict

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


def find_interactive_bboxes(bboxes, left_handness, scale=1.0):
    """Find a pairs of interacting hands

    Args:
        bboxes (np.ndarray): hand bounding boxes
        left_handness (np.array): indicator for every bounding box if its a left hand
        scale (float, optional): bounding box scale ratio. Defaults to 1.0.

    Returns:
        np.array: array of association indices
    """
    # prepare association table
    interaction = np.zeros(bboxes.shape[0], dtype=np.int32)
    # association id
    idx = 1
    # for every box
    for i in range(bboxes.shape[0]):
        # check only left hands
        if not left_handness[i]:
            continue
        # for every box again
        for j in range(bboxes.shape[0]):
            # if the hand is not right or is already assigned as interacting -> skip
            if left_handness[j] or interaction[j] != 0:
                continue
            # check intersection of boxes
            if intersecting(bboxes[i], bboxes[j], scale=scale):
                # if boxes are intersecting -> add association entry to table
                interaction[i] = idx
                interaction[j] = idx
                # increment id
                idx += 1
                break
    return interaction


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

    Args:
        path (str): path to annotation file

    Returns:
        dict: constructed annotation dictionary
    """
    # prepare output dict
    # each image can have multiple data samples, each object is described with keypoints, left hand indicator and its bounding box
    data_dict = defaultdict(lambda: {'file_name': None, 'keypoints': [], 'left': [], 'bbox': []})
    # firstly, get all annotation entries
    annot_paths = glob(os.path.join(path, '*.json'))
    for annot_path in annot_paths:
        annot = load_json(annot_path)
        if 'flv' in annot_path:
            key = '_'.join(os.path.basename(annot_path).split('_')[:-1])
        else:
            key = os.path.basename(annot_path).split('_')[0]
        img_path = os.path.basename(annot_path)[:-4] + 'jpg'
        keypoints = np.array(annot['hand_pts'])
        left = bool(annot['is_left'])
        bbox = bbox_from_keypoints(keypoints[:, :2])
        data_dict[key]['file_name'] = img_path
        data_dict[key]['keypoints'].append(keypoints)
        data_dict[key]['left'].append(left)
        data_dict[key]['bbox'].append(bbox)
    # scan all loaded annotations for interacting hands
    for key, val in data_dict.items():
        keypoints = np.array(val['keypoints'])
        left = np.array(val['left'])
        bboxes = np.array(val['bbox'])
        # find interactive hands in this frame
        interactive_mask = find_interactive_bboxes(bboxes, left, scale=1.3)
        
        keypoints_out = list()
        boxes_out = list()
        interacting = list()
        # single hands - save keypoints, boxes and set interactive flag as 0 - negative
        single_hand_keypoints = keypoints[interactive_mask == 0, :]
        single_hand_bboxes = bboxes[interactive_mask == 0, :]
        for single_hand_kpts, single_hand_bb in zip(single_hand_keypoints, single_hand_bboxes):
            keypoints_out.append(single_hand_kpts)
            boxes_out.append(single_hand_bb)
            interacting.append(0)
        # for all interactive associations
        for i in range(max(interactive_mask)):
            mask = interactive_mask == (i+1)
            # get interacting bboxes and keypoints
            interacting_bboxes = bboxes[mask, :]
            kpts = keypoints[mask, :]
            # check that one hand is left and other is right
            left_right = left[mask]
            assert left_right[0] ^ left_right[1] and left_right.size == 2
            # combine bboxes
            combined_bbox = combine_bboxes(interacting_bboxes[0, :], interacting_bboxes[1, :])
            # swap keypoints so that left hand is first
            kpts = kpts if left_right[0] else (kpts[1], kpts[0])
            kpts = np.concatenate(kpts, axis=0)
            keypoints_out.append(kpts)
            boxes_out.append(combined_bbox)
            # interactive flag 1 - positive
            interacting.append(1)
        # rewrite dataset entry
        data_dict[key] = {'dataset': 'mwb', 'file_name': val['file_name'], 'keypoints': keypoints_out, 'interacting': interacting, 'bbox': boxes_out}
    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.

    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'])
    img = cv2.imread(img_path)
    keypoints = list()
    validity = list()
    bboxes = list()
    classes = list()
    # extract keypoints, bboxes and classes of all samples in entry
    for kpts, bbox, interacting in zip(annot['keypoints'], annot['bbox'], annot['interacting']):
        # if single hand -> pad with zeros to unify tensor sizes
        if not interacting:
            kpts = np.concatenate((kpts, np.zeros_like(kpts, dtype=kpts.dtype)), axis=0)
        validity.append(kpts[:, 2])
        keypoints.append(kpts[:, :2])
        bboxes.append(scale_bbox(bbox, scale))
        classes.append(int(interacting))
    return img, bboxes, keypoints, validity, classes
