"""
Utilities used by Interhand model and dataset
"""
import os
import cv2
import numpy as np
import tensorflow as tf

from math import pi

from .estimation_utils import get_transformations
from .hands_utils import load_json


# TFRecord features
HAND_FEATURE_MAP = {
    'image': tf.io.FixedLenFeature([], tf.string),
    'keypoints/points': tf.io.FixedLenFeature([126], tf.float32),
    'keypoints/validity': tf.io.FixedLenFeature([42], tf.int64),
    'hand/handness': tf.io.FixedLenFeature([2], tf.int64),
    'hand/handness_validity': tf.io.FixedLenFeature([1], tf.int64),
}
# hand presence indicators from string
HANDNESS_MAP = {'left': [0, 1], 'right': [1, 0], 'interacting': [1, 1]}


def world2image(world, trans, rot, focal, principal):
    """Convert points from world coordinates to image plane

    Args:
        world (np.ndarray): world coordinates of points
        trans (np.array): camera translation vector
        rot (np.ndarray): camera rotation matrix
        focal (np.array): camera's focal length
        principal (np.array): camera's principal point

    Returns:
        np.ndarray: output points in image plane
    """
    cam = np.dot(world - trans, rot.T)
    x = focal[0] * cam[:, 0] / (cam[:, 2] + 1e-8) + principal[0]
    y = focal[1] * cam[:, 1] / (cam[:, 2] + 1e-8) + principal[1]
    z = cam[:, 2]
    return np.concatenate((x[:, np.newaxis], y[:, np.newaxis], z[:, np.newaxis]), axis=1)


def data_from_annotation(frame_annot, data_annot, annots_joint, annots_camera, images_path):
    """Get data from annotation entry

    Args:
        frame_annot (dict): annotation from frame annotation file
        data_annot (dict): annotation from data annotation file
        annots_joint (dict): annotation from joint annotation file
        annots_camera (dict): annotation from camera annotation file
        images_path (str): root directory path of images

    Returns:
        np.ndarray, np.array, np.ndarray, np.ndarray, np.array, np.array: image, bounding box, keypoint coordinates in image plane,
                                                                          validity of keypoints, handedness indicators, handedness validity
    """
    # capture id, frame index and camera id
    capture = str(frame_annot['capture'])
    frame_idx = str(frame_annot['frame_idx'])
    camera = str(frame_annot['camera'])
    # camera parameters
    campos = np.array(annots_camera[capture]['campos'][camera], dtype=np.float32)
    camrot = np.array(annots_camera[capture]['camrot'][camera], dtype=np.float32)
    focal = np.array(annots_camera[capture]['focal'][camera], dtype=np.float32)
    princpt = np.array(annots_camera[capture]['princpt'][camera], dtype=np.float32)
    # keypoints in world coordinates, convert to image plane
    world_coord = np.array(annots_joint[capture][frame_idx]['world_coord'])
    image_coord = world2image(world_coord, campos, camrot, focal, princpt)
    # if wrist is not valid -> invalidate rest of the hand
    joints_valid = np.array(annots_joint[capture][frame_idx]['joint_valid'], dtype=np.int32).reshape(-1)
    joints_valid[:21] *= joints_valid[20]
    joints_valid[21:] *= joints_valid[41]
    joints_valid = joints_valid.tolist()
    # convert handedness string to indicators
    handness = HANDNESS_MAP[data_annot['hand_type']]
    handness_valid = data_annot['hand_type_valid']
    # read bounding box
    bbox = np.array(data_annot['bbox'], dtype=np.int32)
    # read image
    img_path = os.path.join(images_path, frame_annot['file_name'])
    img = cv2.imread(img_path)
    assert img is not None, f'path: {img_path}'
    return img, bbox, image_coord, joints_valid, handness, handness_valid


def data_generator_2d(frame_annots, data_annots, annots_joint, annots_camera, images_path, scale=1.0, max_samples=None):
    """Generate 2D data samples from annotation files

    Args:
        frame_annot (dict): annotation from frame annotation file
        data_annot (dict): annotation from data annotation file
        annots_joint (dict): annotation from joint annotation file
        annots_camera (dict): annotation from camera annotation file
        images_path (str): root directory path of images
        scale (float, optional): scale ratio for bounding box size. Defaults to 1.0.
        max_samples (int, optional): maximum amount of samples generated, if None - unlimited. Defaults to None.

    Raises:
        StopIteration: when max_samples has been reached

    Yields:
        np.ndarray, np.array, np.ndarray, np.ndarray, int, str: image, bounding box, 2D keypoints, keypoint validity,
                                                                interacting hand or single hand class, name of dataset
    """
    # for every sample in annotation
    for i, (frame_annot, data_annot) in enumerate(zip(frame_annots, data_annots)):
        # check termination condition
        if max_samples is not None and i >= max_samples:
            break
        # get data from annotation entry
        data = data_from_annotation(frame_annot, data_annot, annots_joint, annots_camera, images_path)
        img, bbox, joints, valid, handness, handness_valid = data
        # convert handedness to class - interacting if both left and right hands are present
        cls = handness[0] * handness[1]
        # ignore if handedness is not valid
        if not handness_valid or not cls:
            continue
        # ignore depth information of keypoints
        keypoints = joints[:, :2]
        yield img, bbox, keypoints, valid, cls, 'interhand'
    raise StopIteration


def preprocess_image(img, bbox, out_shape, center_format=False):
    """Preprocess image - cutout ROI and resize image to target size and generate transformation and its inverse matrices for keypoint transform

    Args:
        img (np.ndarray): sample image
        bbox (np.array): bounding box
        out_shape (tuple): target image size
        center_format (bool, optional): flag if bounding box is in center format [x_center, y_center, width, height],
                                        if False the format is expected to be [x_min, y_min, width, height]. Defaults to False.

    Returns:
        np.ndarray, np.ndarray, np.ndarray: preprocessed image, transformation matrix, inverse transformation matrix
    """
    bbox = np.array(bbox, dtype=np.float32)
    # scale bounding box
    w, h = bbox[2:] * 1.25
    # convert to center format if its already not in it
    if center_format:
        c_x, c_y = bbox[:2]
    else:
        c_x, c_y = bbox[:2] + bbox[2:] / 2.
    # set aspect ratio to 1:1
    ar = out_shape[1] / out_shape[0]
    if w > ar * h:
        h = w / ar
    elif w < ar * h:
        w = h * ar
    bbox = np.array([c_x, c_y, w, h], dtype=np.float32)
    # get transformations
    trans, inv_trans = get_transformations(bbox, out_shape)
    # transform image
    img_patch = cv2.warpAffine(img, trans, (int(out_shape[1]), int(out_shape[0])), flags=cv2.INTER_LANCZOS4)
    return img_patch, trans, inv_trans


def load_tfrecord(tfrecord_path):
    """Load TFrecord dataset file

    Args:
        tfrecord_path (str): path to tfrecord dataset

    Returns:
        (tf.data.TFRecordDatasetV2): parsable tfrecord dataset
    """
    files = tf.data.Dataset.list_files(tfrecord_path)
    dataset = files.flat_map(tf.data.TFRecordDataset)
    return dataset.map(lambda x: parse_tfrecord(x))


def parse_tfrecord(example, img_size):
    """Parse TFrecord serialized dataset files

    Args:
        example (tf.data.TFRecordDatasetV2): parsable tfrecord dataset
        img_shape (tuple): target image size in first two dimensions (K, L)

    Returns:
        (tf.Tensor, dict): parsed example: image and label dictionary
    """
    x = tf.io.parse_single_example(example, HAND_FEATURE_MAP)
    # transform image to <0, 1> float and resize to target size
    x_train = tf.cast(tf.image.decode_jpeg(x['image'], channels=3), tf.float32)
    x_train = tf.image.resize(x_train, img_size) / 255
    # keypoint, validity and handedness, handedness validity labels
    joints = tf.reshape(x['keypoints/points'], (-1, 3))
    joints_valid = tf.cast(x['keypoints/validity'], tf.float32)
    handness = tf.cast(x['hand/handness'], tf.float32)
    handness_valid = tf.cast(x['hand/handness_validity'], tf.float32)
    return x_train, {'joints': joints, 'joints_valid': joints_valid, 'handness': handness, 'handness_valid': handness_valid}


@tf.function
def transform_targets_train(y_train, img_size, heatmap_dims, depth_dim, root_hm_dim, sigma):
    """Transforms keypoint labels to heatmaps

    Args:
        y_train (dict): label dictionary
        img_size (tuple): input image size
        heatmap_dims (tuple): size of generated heatmaps
        depth_dim (int): depth representation vector length
        root_hm_dim (int): relative hand depth dimension
        sigma (float): gaussian variance of generated heatmaps

    Returns:
        dict: dictionary of labels
    """
    joints, joints_valid = y_train['joints'], y_train['joints_valid']
    handness, handness_valid = y_train['handness'], y_train['handness_valid']
    # scale x, y keypoints to <0, 1>, keep z the same
    scale = tf.reshape(tf.constant([heatmap_dims[0]/img_size[0], heatmap_dims[1]/img_size[1], 1], dtype=tf.float32), (1, 1, 3))
    joints = tf.reshape(joints, (-1, 42, 3)) * scale
    # scale depth coordinates
    relative_root_depth = joints[:, 20, 2] - joints[:, 41, 2]
    relative_root_depth = (relative_root_depth / (depth_dim / 2.0) + 1.0) / 2.0 * root_hm_dim
    # align depths of hands
    joints_right_depth = joints[:, :21, 2] - joints[:, 20:21, 2]
    joints_left_depth = joints[:, 21:, 2] - joints[:, 41:42, 2]
    # scale depth dimension
    joints_depth = tf.concat([joints_right_depth, joints_left_depth], axis=1)
    joints_depth = (joints_depth / (depth_dim / 2.0) + 1.0) / 2.0 * heatmap_dims[2]
    joints = tf.concat([joints[:, :, :2], tf.expand_dims(joints_depth, axis=2)], axis=2)
    # generate heatmaps
    joints = tf.reshape(joints, (-1, 1, 1, 42, 1, 3))
    x, y, z = tf.meshgrid(tf.range(0, heatmap_dims[0], dtype=tf.float32),
                          tf.range(0, heatmap_dims[1], dtype=tf.float32),
                          tf.range(0, heatmap_dims[2], dtype=tf.float32))
    x = tf.expand_dims(x, axis=2)[tf.newaxis, ...]
    y = tf.expand_dims(y, axis=2)[tf.newaxis, ...]
    z = tf.expand_dims(z, axis=2)[tf.newaxis, ...]
    
    heatmaps = tf.exp(-(tf.square(x-joints[..., 0]) + tf.square(y-joints[..., 1]) + tf.square(z-joints[..., 2])) / (2*sigma ** 2))
    heatmaps = tf.reshape(heatmaps, (-1, heatmap_dims[0], heatmap_dims[1], 42, heatmap_dims[2]))
    joints = tf.reshape(joints, (-1, 42, 3))
    # if handedness is not valid -> invalidate keypoints
    joints_valid = joints_valid * tf.cast(joints[:, :, 2] >= 0, tf.float32) * tf.cast(joints[:, :, 2] < heatmap_dims[2], tf.float32)
    # create relative root depth validity flag - only when hands are interaacting and handedness is valid
    relative_root_valid = joints_valid[:, 20] * joints_valid[:, 41] * handness[:, 0] * handness[:, 1]
    relative_root_valid = relative_root_valid * tf.cast(relative_root_depth >= 0, tf.float32) * tf.cast(relative_root_depth < heatmap_dims[2], tf.float32)
    y_train = {
        'heatmaps': heatmaps, 'joints': joints, 'joints_valid': joints_valid,
        'handness': handness, 'handness_valid': handness_valid,
        'relative_root_depth': relative_root_depth, 'relative_root_valid': relative_root_valid
    }
    return y_train


@tf.function
def transform_targets_test(y_train, img_size, depth_dim, root_hm_dim):
    """Transform keypoints for testing purposes

    Args:
        y_train (dict): label dictionary
        img_size (tuple): input image size
        depth_dim (int): depth representation vector length
        root_hm_dim (int): relative hand depth dimension

    Returns:
        dict: dictionary of labels
    """
    joints = y_train['joints']
    handness, handness_valid = y_train['handness'], y_train['handness_valid']
    # scale keypoints
    scale = tf.reshape(tf.constant([1/img_size[0], 1/img_size[1], 1], dtype=tf.float32), (1, 1, 3))
    joints = tf.reshape(joints, (-1, 42, 3)) * scale
    # calculate relative root depth
    relative_root_depth = joints[:, 20, 2] - joints[:, 41, 2]
    relative_root_depth = (relative_root_depth / (depth_dim / 2) + 1) / 2 * root_hm_dim
    # zero offset hand depths
    joints_right_depth = joints[:, :21, 2] - joints[:, 20:21, 2]
    joints_left_depth = joints[:, 21:, 2] - joints[:, 41:42, 2]
    joints_depth = tf.concat([joints_right_depth, joints_left_depth], axis=1)
    joints = tf.concat([joints[:, :, :2], tf.expand_dims(joints_depth, axis=2)], axis=2)

    y_train = {
        'joints': joints, 'handness': handness, 'handness_valid': handness_valid, 'relative_root_depth': relative_root_depth
    }
    return y_train


def augment_dataset(x_train, y_train):
    """Augment data entry

    Args:
        x_train (tf.Tensor): batched input images
        y_train (dict): dictionary of labels

    Returns:
        tf.Tensor, dict: augmented batched input images and labels
    """
    joints = y_train['joints']
    # hsv augmentation
    hsv_images = tf.image.rgb_to_hsv(x_train)
    scale = tf.random.uniform((1, 1, 1, 3), 0.8, 1.2)
    shift = tf.concat((tf.constant(0.0, shape=(1, 1, 1, 1)), tf.random.uniform((1, 1, 1, 1)), tf.random.uniform((1, 1, 1, 1), 0, 0.5)), axis=-1)
    hsv_images = tf.clip_by_value(hsv_images * scale + shift, 0, 1)
    x_train = tf.image.hsv_to_rgb(hsv_images)
    # random rotation augmentation
    # k times 90 degrees
    k = tf.random.uniform((1, ), -1, 1, dtype=tf.int32)[0]
    # rotate image
    x_train = tf.image.rot90(x_train, k)
    # construct rotation matrix
    k = tf.cast(k, tf.float32)
    sin, cos = tf.math.sin(k*pi/2), tf.math.cos(k*pi/2)
    rot_mat = tf.reshape(tf.stack([cos, -sin, sin, cos]), (2, 2))
    # rotate keypoints
    joints_xy = tf.reshape(joints[:, :, :2], (-1, 2))
    joints_xy = tf.reshape(tf.matmul(joints_xy - 0.5, rot_mat) + 0.5, (-1, 42, 2))
    joints = tf.concat((joints_xy, joints[:, :, 2:3]), axis=2)
    y_train['joints'] = joints
    return x_train, y_train


def init_dataset(tfrecord_path, batch_size, img_size, heatmap_dims, depth_dim, root_hm_dim, sigma, augment=False, shuffle=False, prefetch=False, testing=False):
    """Load TFRecord dataset file as a generator like parser

    Args:
        tfrecord_path (str): path to TFRecord file
        batch_size (int): dataset batch size
        img_size (tuple): input image size in first two dimensions (K, L)
        heatmap_dims (tuple): size of generated heatmaps
        depth_dim (int): depth representation vector length
        root_hm_dim (int): relative hand depth dimension
        sigma (float): gaussian variance of generated heatmaps
        augment (bool, optional): data augmentation flag. Defaults to False.
        shuffle (bool, optional): data shuffle flag. Defaults to False.
        prefetch (bool, optional): data prefetching flag. Defaults to False.
        testing (bool, optional): testing flag, if True - heatmaps are not generated. Defaults to False.

    Returns:
        tf.data.Dataset: tensorflow dataset, yielding x tensor and y_true dictionary
    """
    # load dataset
    dataset = tf.data.TFRecordDataset(tfrecord_path).map(lambda x: parse_tfrecord(x, img_size))
    # shuffle if specified
    if shuffle:
        dataset = dataset.shuffle(buffer_size=2048)
    # batch dataset
    dataset = dataset.batch(batch_size)
    # augment if specified
    if augment:
        dataset = dataset.map(lambda x, y: augment_dataset(x, y))
    # generate heatmaps if not testing, otherwise only transform keypoints for evaluation
    if testing:
        dataset = dataset.map(lambda x, y: (x, transform_targets_test(y, img_size, depth_dim, root_hm_dim)), num_parallel_calls=tf.data.AUTOTUNE)
    else:
        dataset = dataset.map(lambda x, y: (x, transform_targets_train(y, img_size, heatmap_dims, depth_dim, root_hm_dim, sigma)), num_parallel_calls=tf.data.AUTOTUNE)
    # prefetch if specified
    if prefetch:
        dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
    return dataset
