"""
Utilities used by hand pose estimation algorithms
"""
import cv2
import numpy as np
import tensorflow as tf

from json import load
from math import pi

from .estimation_utils import get_transformations


# tfrecord hand features
HAND_FEATURE_MAP = {
    'image': tf.io.FixedLenFeature([], tf.string),
    'keypoints': tf.io.FixedLenFeature([84], tf.float32),
    'class': tf.io.FixedLenFeature([1], tf.int64),
    'validity': tf.io.FixedLenFeature([42], tf.int64),
}


def load_json(path):
    """Load json file

    Args:
        path (str): path to file

    Returns:
        dict: loaded json data
    """
    with open(path, 'r') as file:
        data = load(file)
    return data


def bbox_from_keypoints(keypoints):
    """Construct bounding box from keypoints

    Args:
        keypoints (np.ndarray): array of N keypoints (N, 2)

    Returns:
        tuple: bounding box in format [x_min, y_min, width, height]
    """
    keypoints = keypoints[:, :2]
    x_min, y_min = np.min(keypoints, axis=0)
    x_max, y_max = np.max(keypoints, axis=0)
    w, h = x_max - x_min, y_max - y_min
    return x_min, y_min, w, h


def intersecting(bbox1, bbox2, scale=1.0):
    """Check if bounding boxes are intersecting

    Args:
        bbox1 (list): first bounding box in format [x_min, y_min, width, height]
        bbox2 (list): second bounding box in format [x_min, y_min, width, height]
        scale (float, optional): bbox size scale ratio. Defaults to 1.0.

    Returns:
        bool: True if scaled bounding boxes are intersecting, False otherwise
    """
    # convert to x1, y1, x2, y2
    box_1_x1, box_1_y1, box_1_w, box_1_h = bbox1
    box_1_w, box_1_h = box_1_w * scale, box_1_h * scale  # scale width and height
    box_1_x2, box_1_y2 = box_1_x1 + box_1_w, box_1_y1 + box_1_h
    # convert to x1, y1, x2, y2
    box_2_x1, box_2_y1, box_2_w, box_2_h = bbox2
    box_2_w, box_2_h = box_2_w * scale, box_2_h * scale  # scale width and height
    box_2_x2, box_2_y2 = box_2_x1 + box_2_w, box_2_y1 + box_2_h
    # calculate intersection
    int_w = np.maximum(np.minimum(box_1_x2, box_2_x2) - np.maximum(box_1_x1, box_2_x1), 0)
    int_h = np.maximum(np.minimum(box_1_y2, box_2_y2) - np.maximum(box_1_y1, box_2_y1), 0)
    int_area = int_w * int_h
    return int_area > 0


def combine_bboxes(bbox1, bbox2):
    """Combine two bounding boxes from their maximum top left and bottom right corners

    Args:
        bbox1 (list): bounding box in format [x_min, y_min, width, height]
        bbox2 (list): bounding box in format [x_min, y_min, width, height]

    Returns:
        list: resulting bounding box in format [x_min, y_min, width, height]
    """
    x_min = np.minimum(bbox1[0], bbox2[0])
    y_min = np.minimum(bbox1[1], bbox2[1])
    x_max = np.maximum(bbox1[0]+bbox1[2], bbox2[0]+bbox2[2])
    y_max = np.maximum(bbox1[1]+bbox1[3], bbox2[1]+bbox2[3])
    return [x_min, y_min, x_max-x_min, y_max-y_min]


def scale_bbox(bbox, scale):
    """Scale bounding box size by ratio

    Args:
        bbox (list): bounding box in format [x_min, y_min, width, height]
        scale (float): scale ratio

    Returns:
        list: scaled bounding box in format [x_min, y_min, width, height]
    """
    if not scale:
        return bbox
    wh = bbox[2], bbox[3]
    # calculate center coordinates to keep bbox centered
    xy = bbox[0] + wh[0] / 2, bbox[1] + wh[1] / 2
    # scale width and height
    wh = scale * wh[0], scale * wh[1]
    # convert back to [x_min, y_min, width, height] format
    bbox = [xy[0] - wh[0] / 2,  xy[1] - wh[1] / 2, wh[0], wh[1]]
    return bbox


def preprocess_image(img, bbox, out_shape, center_format=False, scale=1.0):
    """Preprocess image for model, cutting out ROI from original image

    Args:
        img (np.ndarray): input image
        bbox (list): hand bounding box in [x_min, y_min, width, height] format
        out_shape (tuple): output shape of hand cutout
        center_format (bool, optional): flag, when True - the bounding box is in [x_center, y_center, width, height]. Defaults to False.
        scale (float, optional): bounding box scale ratio. Defaults to 1.0.

    Returns:
        np.ndarray, np.ndarray, np.ndarray: output image, transformation from original image space to cutout space, inverse transformation matrix
    """
    # convert to center format if not already
    if center_format:
        c_x, c_y, w, h = bbox
    else:
        w, h = bbox[2], bbox[3]
        c_x, c_y = bbox[0] + w / 2, bbox[1] + h / 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
    # scale bounding box
    w, h = w * scale, h * scale
    bbox = np.array([c_x, c_y, w, h], dtype=np.float32)
    # get transformation and its inverse
    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 data_generator(annot_func, annotation_dict, annotation_keys, images_root_path, scale=None, max_samples=None):
    """Generate data examples from a callable

    Args:
        annot_func (callable): function converting annotation entry to unified data samples: img, boxes, keypoints, validity, classes
        annotation_dict (dict): dictionary of annotation entries
        annotation_keys (list): list or iterable of annotation_dict keys used for generation
        images_root_path (str): root directory of images
        scale (float, optional): bounding box scale ratio. Defaults to None.
        max_samples (int, optional): maximum amount of generated samples, unlimited if None. Defaults to None.

    Raises:
        StopIteration: raised on reaching max_samples limit

    Yields:
        np.ndarray, list, list, list, list, str: read image, bounding boxes of annotated hands, keypoints of annotated hands,
                                                 validity of keypoints, classes of hands, name of dataset of sample's origin
    """
    for i, key in enumerate(annotation_keys):
        # check max iteration steps
        if max_samples is not None and i >= max_samples:
            raise StopIteration
        # read annotation entry
        annot = annotation_dict[key]
        dataset = annot['dataset']
        img, bboxes, keypoints, validity, classes = annot_func(annot, images_root_path, scale=scale)
        # yield annotation entry per sample
        for bbox, kpts, valid, cls in zip(bboxes, keypoints, validity, classes):
            yield img, bbox, kpts, valid, cls, dataset


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 {'keypoints': tf.Tensor, 'class': tf.Tensor, 'validity': tf.Tensor}
    """
    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 class labels
    keypoints = tf.reshape(x['keypoints'], (-1, 2))
    validity = tf.cast(x['validity'], tf.bool)
    cls = tf.cast(x['class'], tf.int32)
    return x_train, {'keypoints': keypoints, 'class': cls, 'validity': validity}


def augment_data(x_train, y_train, img_size):
    """Data augmentation - random k*90° rotation

    Args:
        x_train (tf.Tensor): tensor of batched BGR images
        y_train (dict): dictionary of labels
        img_size (tuple): example image sizes

    Returns:
        tf.Tensor, dict: augmented examples and labels
    """
    # determine how many 90 degree rotations
    k = tf.random.uniform((1, ), -1, 1, dtype=tf.int32)[0]
    k = tf.cast(k, tf.float32)
    # rotate images
    x_train = tf.image.rot90(x_train, k)
    # construct rotation matrix
    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
    keypoints = y_train['keypoints']
    img_size_half = 0.5 * tf.cast(img_size, tf.float32)
    keypoints = tf.matmul(keypoints - img_size_half, rot_mat) + img_size_half
    y_train['keypoints'] = keypoints
    return x_train, y_train


@tf.function
def transform_targets(y_train, img_size, heatmap_dims, sigma, testing=False):
    """Transforms keypoint labels to heatmaps

    Args:
        y_train (dict): label dictionary
        img_size (tuple): input image size
        heatmap_dims (tuple): size of generated heatmaps
        sigma (float): gaussian variance of generated heatmaps
        testing (bool, optional): flag if dataset is generated in testing mode, in testing mode keypoints
                                  are preserved in output dictionary. Defaults to False.

    Returns:
        dict: dictionary of labels
    """
    keypoints = y_train['keypoints']
    # scale keypoints to heatmap space
    scaled_keypoints = tf.reshape(keypoints, (-1, 42, 2)) * heatmap_dims / img_size
    scaled_keypoints = tf.reshape(scaled_keypoints, (-1, 1, 1, 42, 2))
    # generate heatmaps
    x, y = tf.meshgrid(tf.range(0, heatmap_dims[0], dtype=tf.float32),
                       tf.range(0, heatmap_dims[1], dtype=tf.float32))
    x = x[tf.newaxis, :, :, tf.newaxis]
    y = y[tf.newaxis, :, :, tf.newaxis]
    heatmaps = tf.exp(-(tf.square(x-scaled_keypoints[..., 0]) + tf.square(y-scaled_keypoints[..., 1])) / (2*sigma ** 2))
    heatmaps = tf.reshape(heatmaps, (-1, heatmap_dims[0], heatmap_dims[1], 42))
    # scale heatmaps to <0, 255>
    heatmaps = heatmaps * 255
    # output dictionaries are based on testing/training mode
    if testing:
        y = {'keypoints': keypoints / img_size, 'class': y_train['class'], 'heatmaps': heatmaps, 'validity': y_train['validity']}
    else:
        y = {'heatmaps': heatmaps, 'class': y_train['class'], 'validity': y_train['validity']}
    return y


def init_dataset(tfrecord_path, batch_size, img_size, heatmap_dims, 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
        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 - keypoints are included in target dictionary. 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_data(x, y, img_size), num_parallel_calls=tf.data.AUTOTUNE)
    # generate heatmaps
    dataset = dataset.map(lambda x, y: (x, transform_targets(y, img_size, heatmap_dims, sigma, testing=testing)), num_parallel_calls=tf.data.AUTOTUNE)
    # prefetch if specified
    if prefetch:
        dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
    return dataset
