"""
Universal dataset manipulation utilities
"""
import os
import tensorflow as tf
import numpy as np

from scipy.optimize import linear_sum_assignment
from json import load, dump


# features of a wlasl dataset sample
FEATURE_MAP = {
    'hand_keypoints': tf.io.FixedLenSequenceFeature([2*21*2], tf.float32, allow_missing=True),
    'face_keypoints': tf.io.FixedLenSequenceFeature([1*98*2], tf.float32, allow_missing=True),
    'class': tf.io.FixedLenFeature([1], tf.int64),
}


def load_json(path):
    """Load json file as a dictionary

    Args:
        path (str): path to json file

    Returns:
        dict: output dictionary
    """
    if path is None:
        return None
    with open(path, 'r') as file:
        data = load(file)
    return data


def dump_json(path, data):
    """Dump dictionary to a json file

    Args:
        path (str): output path
        data (dict): dictionary to write
    """
    if os.path.isfile(path):
        os.remove(path)
    with open(path, 'w') as file:
        dump(data, file)


def remove_overhead_detections(keypoints, n, prev_keypoints=None):
    """Clip number of object keypoints from a single time step and order them based on previously used keypoints

    Args:
        keypoints (np.ndarray): N keypoint matrices concatenated to array of [N, M, 2]
        n (int): maximum amount of keypoints to keep
        prev_keypoints (np.ndarray, optional): keypoints used in previous time step, None if this is the first time step. Defaults to None.

    Returns:
        np.ndarray: clipped keypoints
    """
    # if no keypoints are passed or no previous keypoints are passed, take first n samples from array and output it
    if keypoints.size == 0 or prev_keypoints is None or prev_keypoints.size == 0:
        return keypoints[:n]
    # if prev keypoints are passed, perform assignment to ensure keypoint continuity through time steps
    # this means to make sure keypoints in index I are on the same index in the next time step
    else:
        # prev keypoints can be zero padded, filter only valid keypoints
        prev_invalid = [i for i, kpts in enumerate(prev_keypoints) if not np.any(kpts)]
        # if all prev keypoints are invalid, don't perform assignment
        if len(prev_invalid) == len(prev_keypoints):
            return keypoints[:n]
        # construct assignment cost matrix - euclidean distance between keypoints
        curr_keypoints = np.array(keypoints)[:, np.newaxis, ...]
        prev_keypoints = np.array(prev_keypoints)[np.newaxis, ...]
        cost_matrix = np.sum(np.square(curr_keypoints - prev_keypoints), axis=(2, 3))
        # set high cost to invalid prev keypoints and calculate linear assignment
        cost_matrix[:, prev_invalid] = 1e5
        curr_idx, prev_idx = linear_sum_assignment(cost_matrix)
        # order keypoints based on assignment results
        new_keypoints = [curr_keypoints[curr_idx[idx]] for idx in np.argsort(prev_idx)]
        # if number of ordered keypoints is lower then n, try to add unused keypoints to the end nad clip by n
        keypoints = list(keypoints)
        if len(new_keypoints) < n:
            new_keypoints += [keypoints[idx][np.newaxis, ...] for idx in range(len(keypoints)) if idx not in curr_idx]
            new_keypoints = new_keypoints[:n]
        # concat keypoints to array
        new_keypoints = np.concatenate(new_keypoints, axis=0)
    return new_keypoints


def asociate_forward_backward_detections(forward, backward, n):
    """Associate keypoints from forward and backward passes from a single time step and order them in the same way

    Args:
        forward (np.ndarray): keypoints from forward pass
        backward (np.ndarray): keypoints from backward pass
        n (int): number of keypoints in each timestep

    Returns:
        list, list: lists of ordered forward and backward pass keypoints
    """
    # make sure forward pass has more keypoints, if not swap passes
    if len(backward) > len(forward):
        forward, backward = backward, forward
    # convert to matrices and calculate euclidean assignment matrix
    forward_np = np.array(forward)[np.newaxis, ...]
    backward_np = np.array(backward)[:, np.newaxis, ...]
    cost_matrix = np.sum(np.square(forward_np - backward_np), axis=(2, 3))
    rows, cols = linear_sum_assignment(cost_matrix)
    # order backward pass to match forward pass keypoints
    matched_backward = list()
    bwd_idx = 0
    # iterate over number of keypoints to keep
    for idx in range(n):
        # if current index is in assigned indices, add the keypoints to list
        if idx in cols:
            matched_backward.append(backward[rows[cols == idx][0]])
        # if its not in assigned indices and backward index is still within boundary
        elif bwd_idx < len(backward):
            # add current backward keypoints to list
            if bwd_idx not in rows:
                matched_backward.append(backward[bwd_idx])
            # None if this keypoint is skipped
            else:
                matched_backward.append(None)
            # increment backward index
            bwd_idx += 1
        else:
            # None if not in assigned indices and backward index is out of boundary
            matched_backward.append(None)
    # pad with Nones to match n keypoints
    matched_forward = list(forward) + (n - len(forward)) * [None]
    return matched_forward, matched_backward


def fuse_samples(forward, backward, prev_keypoints, shape, max_dist=50):
    """Fuse forward and backward pass keypoints from a single time step together

    Args:
        forward (list): list of keypoints from forward pass
        backward (list): list of keypoints from backward pass
        prev_keypoints (np.ndarray): keypoints from previous time step
        shape (tuple): shape of zero padding if no keypoints are passed
        max_dist (int, optional): max valid distance between forward and backward keypoints. Defaults to 50.

    Returns:
        np.ndarray: array of fused keypoints
    """
    fused = list()
    # iterate over keypoint objects
    for i, (fwd, bwd) in enumerate(zip(forward, backward)):
        # if both passes don't contain keypoints, append previous keypoints if this is not the first time step, otherwise zeros
        if fwd is None and bwd is None:
            fused.append(np.zeros(shape, dtype=np.float32) if prev_keypoints is None else prev_keypoints[i])
        # use backward pass if forward pass is None and likewise
        elif fwd is None:
            fused.append(bwd)
        elif bwd is None:
            fused.append(fwd)
        # if both passes are present, calculate their mean
        # also calculate their distance between each other, if its beyond threshold use only forward keypoints and don't fuse keypoints
        else:
            distance = np.sqrt(np.mean(np.square(fwd - bwd)))
            fused.append(0.5 * (fwd + bwd) if distance < max_dist else fwd)
    return np.array(fused)


def fuse_forward_backward_passes(forward, backward, shape, n=1):
    """Fuse forward and backward passes together

    Args:
        forward (np.ndarray): sequence of keypoints from forward pass
        backward (np.ndarray): sequence of keypoints from backward pass
        shape (tuple): shape of keypoints
        n (int, optional): number of keypoint objects to keep or pad to. Defaults to 1.

    Returns:
        list: fused sequence of keypoints
    """
    fused_keypoints = list()
    # iterate over time steps
    for fwd, bwd in zip(forward, backward):
        # previous keypoints
        prev_keypoints = fused_keypoints[-1] if len(fused_keypoints) else None
        # clip number of keypoints to n
        fwd = remove_overhead_detections(fwd, n, prev_keypoints)
        bwd = remove_overhead_detections(bwd, n, prev_keypoints)
        # match forward and backward keypoints
        matched_fwd, matched_bwd = asociate_forward_backward_detections(fwd, bwd, n)
        # fuse keypoints and add to list
        fused = fuse_samples(matched_fwd, matched_bwd, prev_keypoints, shape)
        fused_keypoints.append(fused)
    return fused_keypoints


def zero_pad_concat_keypoints(keypoints, n):
    """Pad list of fused keypoints to unified shape and convert the sequence to array

    Args:
        keypoints (list): sequence of keypoints
        n (int): number of keypoints in each time step

    Returns:
        np.ndarray: padded keypoints converted to array
    """
    keypoints_out = list()
    # iterate over time step
    for kpts in keypoints:
        # convert to array and pad with zeros
        kpts = np.array(kpts)
        kpts = np.pad(kpts, ((0, n-kpts.shape[0]), (0, 0), (0, 0)), mode='constant', constant_values=0)
        keypoints_out.append(kpts)
    # concat keypoints along new dimension
    keypoints_out = np.stack(keypoints_out, axis=0)
    return keypoints_out


def keypoints_to_input_space(hand_keypoints, face_keypoints, eye_to_eye_dist=0.1, eye_indices=(96, 97), face_center_index=54):
    """Transform keypoints from image space to universal space

    Args:
        hand_keypoints (np.ndarray): sequence of hand keypoints
        face_keypoints (np.ndarray): sequence of face keypoints
        eye_to_eye_dist (float, optional): target norm distance between eyes. Defaults to 0.1.
        eye_indices (tuple, optional): tuple of eye indices of face keypoint array. Defaults to (96, 97).
        face_center_index (int, optional): index of center of the face in face keypoint array. Defaults to 54.

    Returns:
        np.ndarray, np.ndarray: transformed hand and face keypoints
    """
    hand_keypoints = np.array(hand_keypoints)
    face_keypoints = np.array(face_keypoints)
    # arrays may contain zero padded keypoints, filter them out
    faces = face_keypoints[np.any(face_keypoints, axis=(1, 2, 3)), ...]
    # normalization factor
    norm_factor = np.mean(np.sqrt(np.sum(np.square(faces[:, :, eye_indices[0], :] - faces[:, :, eye_indices[1], :]), axis=-1)) / eye_to_eye_dist, axis=0, keepdims=True)
    # norm keypoints
    normed_hands = hand_keypoints / norm_factor
    normed_face = face_keypoints / norm_factor
    # additive normalization - zero padded keypoints have to stay zeros
    # mask non-zero padded keypoints
    first_hand_mask = np.any(normed_hands[:, 0, ...], axis=(1, 2))
    second_hand_mask = np.any(normed_hands[:, 1, ...], axis=(1, 2))
    # norm masked keypoints
    normed_hands[first_hand_mask, 0, ...] -= normed_face[first_hand_mask, 0, face_center_index:face_center_index+1, :]
    normed_hands[second_hand_mask, 1, ...] -= normed_face[second_hand_mask, 0, face_center_index:face_center_index+1, :]
    face_mask = np.any(normed_face, axis=(1, 2, 3))
    normed_face[face_mask, ...] -= normed_face[face_mask, :, face_center_index:face_center_index+1, :]
    return normed_hands, normed_face


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):
    """Parse TFrecord serialized dataset files

    Args:
        example (tf.data.TFRecordDatasetV2): parsable tfrecord dataset

    Returns:
        dict, tf.Tensor: parsed example: keypoint dictionary {'hand_keypoints': tf.Tensor, 'face_keypoints': tf.Tensor} and class labels
    """
    x = tf.io.parse_single_example(example, FEATURE_MAP)
    hand_keypoints = tf.reshape(x['hand_keypoints'], (-1, 2, 21, 2))
    face_keypoints = tf.reshape(x['face_keypoints'], (-1, 1, 98, 2))
    y_train = tf.cast(x['class'], tf.int32)
    return {'hand_keypoints': hand_keypoints, 'face_keypoints': face_keypoints}, y_train


def augment_data(x_train, y_train):
    """Data augmentation

    Args:
        x_train (dict): dictionary of keypoint tensors
        y_train (tf.Tensor): tensor of labels

    Returns:
        dict, tf.Tensor: augmented examples
    """
    hand_keypoints, face_keypoints = x_train['hand_keypoints'], x_train['face_keypoints']
    # y-axis swap transformation 
    swap_index = tf.cast(tf.random.uniform((1, 1, 1, 1), -1, 1, dtype=tf.int32) * 2 + 1, hand_keypoints.dtype)
    swap_tform = tf.concat((swap_index, tf.constant([1.], shape=(1, 1, 1, 1), dtype=hand_keypoints.dtype)), axis=-1)
    # hand shift transformation 
    shift_tform = tf.random.normal((1, 1, 1, 2), 1.0, 0.2, dtype=hand_keypoints.dtype)
    # transformation
    hand_keypoints, face_keypoints = swap_tform * shift_tform * hand_keypoints, swap_tform * face_keypoints
    # gaussian noise
    hand_keypoints = hand_keypoints * tf.random.normal((1, 1, tf.shape(hand_keypoints)[-2], 2), 1, 0.01)
    face_keypoints = face_keypoints * tf.random.normal((1, 1, tf.shape(face_keypoints)[-2], 2), 1, 0.01)
    return {'hand_keypoints': hand_keypoints, 'face_keypoints': face_keypoints}, y_train


def transform_sequences(x_train, sequence_len):
    """Transform dictionary of keypoint sequences to a unified sequence feature vector

    Args:
        x_train (dict): dictionary of keypoint sequences
        sequence_len (int): length of target sequence

    Returns:
        tf.Tensor: sequenced feature vector
    """
    # reshape keypoints to sequenced vectors and concat them together
    hand_keypoints = x_train['hand_keypoints']
    hand_keypoints = tf.reshape(hand_keypoints, (-1, 2*tf.shape(hand_keypoints)[-2]*2))
    face_keypoints = x_train['face_keypoints']
    face_keypoints = tf.reshape(face_keypoints, (-1, 1*tf.shape(face_keypoints)[-2]*2))
    x_train = tf.concat((hand_keypoints, face_keypoints), axis=1)
    if sequence_len is not None:
        curr_len = tf.shape(x_train)[0]
        sequence_len = tf.cast(sequence_len, tf.int32)
        # current length is lower then target length, pad with zeros
        if curr_len < sequence_len:
            x_train = tf.pad(x_train, ((0, sequence_len - curr_len), (0, 0)))
        # current length is higher then target length, randomly extract sequence of target length
        else:
            start = tf.random.uniform((), 0, curr_len - sequence_len + 1, dtype=tf.int32)
            x_train = x_train[start:start+sequence_len, ...]
    return x_train


def transform_targets(y_train, classes):
    """Transform target labels to one hot encoding

    Args:
        y_train (tf.Tensor): input sparse labels
        classes (int): number of classes

    Returns:
        tf.Tensor: output one hot labels
    """
    y_train = tf.one_hot(y_train, classes)
    y_train = tf.reshape(y_train, (classes, ))
    return y_train


def init_dataset(tfrecord_path, batch_size, classes, sequence_len=None, augment=False, shuffle=False, prefetch=False):
    """Load TFRecord dataset file as a generator like parser

    Args:
        tfrecord_path (str): path to TFRecord file
        batch_size (int): dataset batch size
        classes (int): number of classes
        sequence_len (int, optional): target sequence length, if None, sequences are kept as they are. Defaults to None.
        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.

    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))
    # shuffle if specified
    if shuffle:
        dataset = dataset.shuffle(buffer_size=8192) # 32768
    # augment if specified
    if augment:
        dataset = dataset.map(lambda x, y: augment_data(x, y), num_parallel_calls=tf.data.AUTOTUNE)
    # transform sequences to sequence feature vector and labels to one hot encoding
    dataset = dataset.map(lambda x, y: (transform_sequences(x, sequence_len), transform_targets(y, classes)))
    # batch dataset
    dataset = dataset.batch(batch_size)
    # prefetch if specified
    if prefetch:
        dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
    return dataset
