"""
Visualize used datasets from files and from constructed TFRecord files
"""
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import cv2
import numpy as np

import matplotlib.pyplot as plt
from random import shuffle

from modules import interhand_utils, coco_utils, wflw_utils, halpe_utils, mwb_utils, hands_utils, rhd_utils
from modules.inference_utils import draw_pose, postprocess_interhand
from modules.estimation_utils import load_skeleton
from modules.hands_utils import preprocess_image
from modules.config import CFG


def visualize_hand_pose(img, joints, parents=None, validity=None, handness=None):
    """Draw hand pose on a canvas image

    Args:
        img (np.ndarray): canvas BGR image
        joints (np.ndarray): input keypoints
        parents (list, optional): list of parental keypoints representing hand skeleton. Defaults to None.
        validity (np.array, optional): keypoints validity array. Defaults to None.
        handness (np.array, optional): hand presence indicator. Defaults to None.

    Returns:
        np.ndarray: drawn on image
    """
    joints = joints.astype(np.int32)
    # if validity is not passed, construct validity from handedness indicators
    # if handedness is also not passed, all keypoitns are valid
    if validity is None:
        validity = 21 * [handness[0]] + 21 * [handness[1]] if handness is not None else 42 * [True]
    # for every keypoint
    for i, pnt in enumerate(joints):
        # skip if  invalid
        if not validity[i]:
            continue
        # different color for left and right hands
        color = (202, 183, 42) if i < 21 else (102, 215, 254)
        # if skeleton is passed and both keypoint and its parent are valid
        if parents is not None and parents[i] != -1 and validity[parents[i]]:
            img = cv2.line(img, tuple(joints[parents[i], :2]), tuple(pnt[:2]), color, 2)
        # draw keypoint
        img = cv2.circle(img, tuple(pnt[:2]), 4, color, -1)
    return img


def visualize_face_keypoints(img, keypoints):
    """Draw face keypoints onto image canvas

    Args:
        img (np.ndarray): BGR image canvas
        keypoints (np.ndarray): array of keypoints

    Returns:
        np.ndarray: drawn on image
    """
    for keypoint in keypoints:
        point = int(keypoint[0]), int(keypoint[1])
        cv2.circle(img, point, 3, (102, 215, 254), -1)
    return img


def interhand_from_files():
    """Visualize Interhand dataset from files
    """
    SET = 'test'
    # load annotation files
    annots_data = interhand_utils.load_json(CFG.INTERHAND_DATASET_PATHS['annotations'][SET]['data'])
    annots_joint = interhand_utils.load_json(CFG.INTERHAND_DATASET_PATHS['annotations'][SET]['joint'])
    annots_camera = interhand_utils.load_json(CFG.INTERHAND_DATASET_PATHS['annotations'][SET]['camera'])
    print(f'{len(annots_data["images"])} examples')
    # shuffle annotation entries
    data = list(zip(annots_data['images'], annots_data['annotations']))
    shuffle(data)
    annots_data['images'], annots_data['annotations'] = zip(*data)
    # load hand skeleton
    parent = load_skeleton(CFG.INTERHAND_SKELETON_PATH)
    # iterate over dataset
    for frame_annot, data_annot in zip(annots_data['images'], annots_data['annotations']):
        # read data from annotation entry
        data = interhand_utils.data_from_annotation(frame_annot, data_annot, annots_joint, annots_camera, CFG.INTERHAND_DATASET_PATHS['images'][SET])
        original_img, bbox, joints, joints_valid, handness, handness_valid = data
        assert original_img is not None
        # visualize hand pose
        original_img = visualize_hand_pose(original_img.astype(np.uint8), joints.astype(np.int32), parent, validity=joints_valid)
        cv2.imshow('img', original_img)
        if cv2.waitKey(0) == 27:
            break
    cv2.destroyAllWindows()


def coco_from_files():
    """Visualize coco hand poses from files
    """
    SET = 'test'
    # load annotation file
    annot_data = coco_utils.load_annotation(CFG.COCO_DATASET_PATHS[SET]['annotations'])
    print(f'{len(annot_data.keys())} Examples')
    # load hand skeleton
    skeleton = load_skeleton(CFG.COCO_SKELETON_PATH)
    # iterate over dataset
    for annot in annot_data.values():
        # read image
        img_path = os.path.join(CFG.COCO_DATASET_PATHS[SET]['images'], annot['file_name'])
        img = cv2.imread(img_path)
        # if no hand annotation, skip
        if len(annot['annotations']) == 0:
            continue
        # for every annotation entry in image
        for ann in annot['annotations']:
            keypoints = ann['keypoints']
            # if hands are interacting
            if ann['interacting']:
                print(ann['combined_bbox'][2] * 1.3, ann['combined_bbox'][3] * 1.3)
                # preprocess image and transform keypoints
                img_patch, trans, _ = preprocess_image(img, ann['combined_bbox'], CFG.HANDS_PARAMS.IMG_SHAPE, scale=1.5)
                keypoints_tformed = np.dot(trans, np.concatenate((keypoints[:, :2], np.ones((keypoints.shape[0], 1))), axis=1).T).T
                visualize_hand_pose(img_patch, keypoints_tformed.astype(np.int32), validity=keypoints[:, 2] > 0, parents=skeleton)
                patch = np.concatenate((img_patch, np.zeros_like(img_patch, dtype=np.uint8)), axis=1)
            else:
                # only left hand
                if ann['handness'][0]:
                    print(ann['bboxes'][0][2] * 1.3, ann['bboxes'][0][3] * 1.3)
                    left_keypoints = keypoints[:21, :2]
                    # preprocess image and transform keypoints
                    left_img_patch, trans, _ = preprocess_image(img, ann['bboxes'][0], CFG.HANDS_PARAMS.IMG_SHAPE, scale=1.5)
                    left_keypoints = np.dot(trans, np.concatenate((left_keypoints[:, :2], np.ones((left_keypoints.shape[0], 1))), axis=1).T).T
                    visualize_hand_pose(left_img_patch, left_keypoints.astype(np.int32), validity=keypoints[:21, 2] > 0, parents=skeleton)
                else:
                    left_img_patch = np.zeros((*CFG.HANDS_PARAMS.IMG_SHAPE, 3), dtype=np.uint8)
                # only right hand
                if ann['handness'][1]:
                    print(ann['bboxes'][1][2] * 1.3, ann['bboxes'][1][3] * 1.3)
                    right_keypoints = keypoints[21:, :2]
                    # preprocess image and transform keypoints
                    right_img_patch, trans, _ = preprocess_image(img, ann['bboxes'][1], CFG.HANDS_PARAMS.IMG_SHAPE, scale=1.5)
                    right_keypoints = np.dot(trans, np.concatenate((right_keypoints[:, :2], np.ones((right_keypoints.shape[0], 1))), axis=1).T).T
                    visualize_hand_pose(right_img_patch, right_keypoints.astype(np.int32), validity=keypoints[21:, 2] > 0, parents=skeleton)
                else:
                    right_img_patch = np.zeros((*CFG.HANDS_PARAMS.IMG_SHAPE, 3), dtype=np.uint8)
                patch = np.concatenate((left_img_patch, right_img_patch), axis=1)
            cv2.imshow('img', patch)
            if cv2.waitKey(0) == 27:
                cv2.destroyAllWindows()
                return


def halpe_from_files():
    """Visualize halpe hand poses from files
    """
    # load annotation file
    annot_data = halpe_utils.load_annotation(CFG.HALPE_DATASET_PATHS['annotations'])
    # load hand skeleton
    skeleton = load_skeleton(CFG.COCO_SKELETON_PATH)
    # iterate over dataset
    for annot in annot_data.values():
        # read image
        img_path = os.path.join(CFG.HALPE_DATASET_PATHS['images'], annot['file_name'])
        img = cv2.imread(img_path)
        # if no hand annotation, skip
        if len(annot['annotations']) == 0:
            continue
        # for every annotation entry in image
        for ann in annot['annotations']:
            keypoints = ann['keypoints']
            # if hands are interacting
            if ann['interacting']:
                # preprocess image and transform keypoints
                img_patch, trans, _ = preprocess_image(img, ann['combined_bbox'], CFG.HANDS_PARAMS.IMG_SHAPE, scale=1.5)
                keypoints_tformed = np.dot(trans, np.concatenate((keypoints[:, :2], np.ones((keypoints.shape[0], 1))), axis=1).T).T
                visualize_hand_pose(img_patch, keypoints_tformed.astype(np.int32), validity=keypoints[:, 2] > 0, parents=skeleton)
                patch = np.concatenate((img_patch, np.zeros_like(img_patch, dtype=np.uint8)), axis=1)
            else:
                # only left hand
                if ann['handness'][0]:
                    left_keypoints = keypoints[:21, :2]
                    # preprocess image and transform keypoints
                    left_img_patch, trans, _ = preprocess_image(img, ann['bboxes'][0], CFG.HANDS_PARAMS.IMG_SHAPE, scale=1.5)
                    left_keypoints = np.dot(trans, np.concatenate((left_keypoints[:, :2], np.ones((left_keypoints.shape[0], 1))), axis=1).T).T
                    visualize_hand_pose(left_img_patch, left_keypoints.astype(np.int32), validity=keypoints[:21, 2] > 0, parents=skeleton)
                else:
                    left_img_patch = np.zeros((*CFG.HANDS_PARAMS.IMG_SHAPE, 3), dtype=np.uint8)
                # only right hand
                if ann['handness'][1]:
                    # preprocess image and transform keypoints
                    right_keypoints = keypoints[21:, :2]
                    right_img_patch, trans, _ = preprocess_image(img, ann['bboxes'][1], CFG.HANDS_PARAMS.IMG_SHAPE, scale=1.5)
                    right_keypoints = np.dot(trans, np.concatenate((right_keypoints[:, :2], np.ones((right_keypoints.shape[0], 1))), axis=1).T).T
                    visualize_hand_pose(right_img_patch, right_keypoints.astype(np.int32), validity=keypoints[21:, 2] > 0, parents=skeleton)
                else:
                    right_img_patch = np.zeros((*CFG.HANDS_PARAMS.IMG_SHAPE, 3), dtype=np.uint8)
                patch = np.concatenate((left_img_patch, right_img_patch), axis=1)
            cv2.imshow('img', patch)
            if cv2.waitKey(0) == 27:
                cv2.destroyAllWindows()
                return


def mwb_from_files():
    """Visualize mwb dataset from files
    """
    SET = 'test'
    # load hand skeleton
    parent = load_skeleton(CFG.COCO_SKELETON_PATH)
    base = CFG.MWB_DATASET_PATHS['manual'][SET]
    # load annotation file
    data = mwb_utils.load_annotation(base)
    # iterate over dataset
    for annot in data.values():
        # read image
        img_path = os.path.join(base, annot['file_name'])
        img = cv2.imread(img_path)
        # for every annotation entry in image
        for keypoints, bbox in zip(annot['keypoints'], annot['bbox']):
            # preprocess image and transform keypoints
            img_patch, trans, _ = preprocess_image(img, bbox, CFG.HANDS_PARAMS.IMG_SHAPE, scale=1.5)
            keypoints_tformed = np.dot(trans, np.concatenate((keypoints[:, :2], np.ones((keypoints.shape[0], 1))), axis=1).T).T
            visualize_hand_pose(img_patch, keypoints_tformed[:, :2], validity=keypoints[:, 2], parents=parent)
            cv2.imshow('img', img_patch)
            if cv2.waitKey(0) == 27:
                cv2.destroyAllWindows()
                return


def rhd_from_files():
    """Visualize rhd dataset from files
    """
    SET = 'test'
    # load hand skeleton
    skeleton = load_skeleton(CFG.COCO_SKELETON_PATH)
    # load annotation file
    annotation = rhd_utils.load_annotation(CFG.RHD_DATASET_PATHS[SET]['annotations'])
    # iterate over dataset
    for ann in annotation.values():
        # read image
        img_path = os.path.join(CFG.RHD_DATASET_PATHS[SET]['images'], ann['file_name'])
        img = cv2.imread(img_path)
        keypoints = ann['keypoints']
        valid = ann['validity']
        # if hands are interacting
        if ann['interacting']:
            # preprocess image and transform keypoints
            img_patch, trans, _ = preprocess_image(img, ann['combined'], CFG.HANDS_PARAMS.IMG_SHAPE, scale=1.5)
            keypoints_tformed = np.dot(trans, np.concatenate((keypoints[:, :2], np.ones((keypoints.shape[0], 1))), axis=1).T).T
            visualize_hand_pose(img_patch, keypoints_tformed.astype(np.int32), validity=valid > 0, parents=skeleton)
            patch = np.concatenate((img_patch, np.zeros_like(img_patch, dtype=np.uint8)), axis=1)
        else:
            # only left hand
            if ann['handness'][0]:
                left_keypoints = keypoints[:21, :2]
                # preprocess image and transform keypoints
                left_img_patch, trans, _ = preprocess_image(img, ann['bboxes'][0], CFG.HANDS_PARAMS.IMG_SHAPE, scale=1.5)
                left_keypoints = np.dot(trans, np.concatenate((left_keypoints[:, :2], np.ones((left_keypoints.shape[0], 1))), axis=1).T).T
                visualize_hand_pose(left_img_patch, left_keypoints.astype(np.int32), validity=valid[:21] > 0, parents=skeleton)
            else:
                left_img_patch = np.zeros((*CFG.HANDS_PARAMS.IMG_SHAPE, 3), dtype=np.uint8)
            # only right hand
            if ann['handness'][1]:
                right_keypoints = keypoints[21:, :2]
                # preprocess image and transform keypoints
                right_img_patch, trans, _ = preprocess_image(img, ann['bboxes'][1], CFG.HANDS_PARAMS.IMG_SHAPE, scale=1.5)
                right_keypoints = np.dot(trans, np.concatenate((right_keypoints[:, :2], np.ones((right_keypoints.shape[0], 1))), axis=1).T).T
                visualize_hand_pose(right_img_patch, right_keypoints.astype(np.int32), validity=valid[21:] > 0, parents=skeleton)
            else:
                right_img_patch = np.zeros((*CFG.HANDS_PARAMS.IMG_SHAPE, 3), dtype=np.uint8)
            patch = np.concatenate((left_img_patch, right_img_patch), axis=1)
        cv2.imshow('img', patch)
        if cv2.waitKey(0) == 27:
            cv2.destroyAllWindows()
            return


def wflw_from_files():
    """Visualize wflw dataset from files
    """
    SET = 'train'
    # load annotation file
    data = wflw_utils.read_annotation(CFG.WFLW_DATASET_PATHS['annotations'][SET])
    keys = list(data.keys())
    print(len(keys))
    # shuffle dataset
    shuffle(keys)
    # iterate over dataset
    for key in keys:
        annot = data[key]
        # read data from annotation
        img, bbox, keypoints, cls = wflw_utils.data_from_annotation(annot, CFG.WFLW_DATASET_PATHS['images'], scale=1.2)
        # preprocess image and transform keypoints
        img, trans, _ = wflw_utils.preprocess_image(img, bbox, CFG.WFLW_PARAMS.IMG_SHAPE)
        keypoints = np.dot(trans, np.concatenate((keypoints, np.ones((keypoints.shape[0], 1))), axis=1).T).T
        visualize_face_keypoints(img, keypoints)
        cv2.imshow('img', img)
        if cv2.waitKey(0) == 27:
            break
    cv2.destroyAllWindows()


def interhand_from_tfrecords():
    """Visualize Interhand dataset from TFRecord file
    """
    # load dataset
    dataset = interhand_utils.init_dataset(
        CFG.INTERHAND_TFRECORD_PATHS['test'],
        1,
        CFG.INTERHAND_PARAMS.IMG_SHAPE,
        CFG.INTERHAND_PARAMS.JOINT_HM_DIMS,
        CFG.INTERHAND_PARAMS.DEPTH_DIM,
        CFG.INTERHAND_PARAMS.SIGMA
    )
    # load hand skeleton
    parent = load_skeleton(CFG.INTERHAND_SKELETON_PATH)
    # color map for heatmap viz
    cmap = plt.get_cmap('jet')
    # terate over dataset
    for x, y in dataset:
        # transform image to uint8 <0, 255>
        img = (x[0, ...].numpy() * 255).astype(np.uint8)
        # transform handedness, keypoints and relative root depth from normed tfrecord format
        handness, joints, relative_depth = postprocess_interhand((y['handness'], y['joints'], y['relative_root_depth']), None)
        # draw keypoints
        draw_pose(img, joints, handness, parent)
        # generated heatmaps
        heatmaps = y['heatmaps'].numpy()[0, ...] * 255
        # project their depth dimension and project all heatmaps onto one 2D plane
        projected_heatmap = np.sum(heatmaps, axis=(2, 3))
        # norm image, resize to the image size and convert its colors based on generated color map
        projected_heatmap = cv2.normalize(projected_heatmap, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
        projected_heatmap = cv2.resize(projected_heatmap, CFG.INTERHAND_PARAMS.IMG_SHAPE, cv2.INTER_CUBIC)
        projected_heatmap = cv2.cvtColor((np.delete(cmap(projected_heatmap), 3, 2) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
        # visualize keypoints and generated heatmaps
        cv2.imshow('img + heatmap', np.concatenate((img, projected_heatmap), axis=1))
        if cv2.waitKey(0) == 27:
            break
    cv2.destroyAllWindows()


def hands_from_tfrecords():
    """Visualize hand pose estimation dataset from TFRecord file
    """
    # load dataset
    dataset = hands_utils.init_dataset(
        CFG.HANDS_TFRECORD_PATHS['test'],
        1,
        CFG.HANDS_PARAMS.IMG_SHAPE,
        CFG.HANDS_PARAMS.HM_DIMS,
        CFG.HANDS_PARAMS.SIGMA,
        augment=True,
        testing=True
    )
    # load hand skeleton
    parent = load_skeleton(CFG.COCO_SKELETON_PATH)
    # color map for heatmap viz
    cmap = plt.get_cmap('jet')
    # terate over dataset
    for x, y in dataset:
        # transform image to uint8 <0, 255>
        img = (x[0, ...].numpy() * 255).astype(np.uint8)
        # transform normed keypoints
        keypoints = y['keypoints'][0, ...].numpy() * CFG.HANDS_PARAMS.IMG_SHAPE
        # draw keypoints
        visualize_hand_pose(img, keypoints, parents=parent, handness=[1, y['class'].numpy()[0, 0]])
        # generated heatmaps
        heatmaps = y['heatmaps'].numpy()[0, ...]
        # project their depth dimension and project all heatmaps onto one 2D plane
        projected_heatmap = np.sum(heatmaps, axis=-1)
        # norm image, resize to the image size and convert its colors based on generated color map
        projected_heatmap = cv2.normalize(projected_heatmap, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
        projected_heatmap = cv2.resize(projected_heatmap, CFG.HANDS_PARAMS.IMG_SHAPE, cv2.INTER_CUBIC)
        projected_heatmap = cv2.cvtColor((np.delete(cmap(projected_heatmap), 3, 2) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
        # visualize keypoints and generated heatmaps
        cv2.imshow('img + heatmap', np.concatenate((img, projected_heatmap), axis=1))
        if cv2.waitKey(0) == 27:
            break
    cv2.destroyAllWindows()


def wflw_from_tfrecords():
    """Visualize face pose estimation dataset from TFRecord file
    """
    # load dataset
    dataset = wflw_utils.init_dataset(
        CFG.WFLW_TFRECORD_PATHS['test'],
        1,
        CFG.WFLW_PARAMS.IMG_SHAPE,
        CFG.WFLW_PARAMS.HM_DIMS,
        CFG.WFLW_PARAMS.SIGMA,
        testing=True
    )
    # color map for heatmap viz
    cmap = plt.get_cmap('jet')
    # terate over dataset
    for x, y in dataset:
        # transform image to uint8 <0, 255>
        img = (x[0, ...].numpy() * 255).astype(np.uint8)
        # transform normed keypoints
        keypoints = y['keypoints'].numpy()[0, ...]
        # draw keypoints
        visualize_face_keypoints(img, keypoints)
        # generated heatmaps
        heatmaps = y['heatmaps'].numpy()[0, ...]
        # project their depth dimension and project all heatmaps onto one 2D plane
        projected_heatmap = np.sum(heatmaps, axis=-1)
        # norm image, resize to the image size and convert its colors based on generated color map
        projected_heatmap = cv2.normalize(projected_heatmap, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
        projected_heatmap = cv2.resize(projected_heatmap, CFG.WFLW_PARAMS.IMG_SHAPE, cv2.INTER_CUBIC)
        projected_heatmap = cv2.cvtColor((np.delete(cmap(projected_heatmap), 3, 2) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
        # visualize keypoints and generated heatmaps
        cv2.imshow('img + heatmap', np.concatenate((img, projected_heatmap), axis=1))
        if cv2.waitKey(0) == 27:
            break
    cv2.destroyAllWindows()


if __name__ == '__main__':
    # interhand_from_files()
    # interhand_from_tfrecords()

    # wflw_from_files()
    # wflw_from_tfrecords()

    coco_from_files()
    # halpe_from_files()
    # mwb_from_files()
    # rhd_from_files()
    # hands_from_tfrecords()
