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

from tensorflow.python.ops.numpy_ops import np_utils
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import array_ops


def load_skeleton(path):
    """Load skeleton file

    Args:
        path (str): path to file

    Returns:
        list: list of keypoint adjacency defining a skeleton
    """
    with open(path, 'r') as file:
        data = file.readlines()
    return [int(line.split(' ')[-1]) for line in data[1:]]


def get_transformations(bbox, out_shape):
    """Get affine transformations to cut out a bounding box from image and resize it to a desired size, and to reverse this transformation

    Args:
        bbox (np.array): bounding box
        out_shape (tuple): output image target size in first two dimensions

    Returns:
        np.ndarray, np.ndarray: transformation and inverse transformation matrix
    """
    # source points from bbox
    src = np.zeros((3, 2), dtype=np.float32)
    src[0, :] = bbox[:2]
    src[1, :] = bbox[:2] + np.array([0, bbox[3] * 0.5], dtype=np.float32)
    src[2, :] = bbox[:2] + np.array([bbox[2] * 0.5, 0], dtype=np.float32)
    # target points from output shape
    dst = np.zeros((3, 2), dtype=np.float32)
    dst[0, :] = np.array([out_shape[0] * 0.5, out_shape[1] * 0.5], dtype=np.float32)
    dst[1, :] = dst[0, :] + np.array([0, out_shape[1] * 0.5], dtype=np.float32)
    dst[2, :] = dst[0, :] + np.array([out_shape[0] * 0.5, 0], dtype=np.float32)
    # transformation and its inverse
    trans = cv2.getAffineTransform(src.astype(np.float32), dst.astype(np.float32))
    inv_trans = cv2.getAffineTransform(dst.astype(np.float32), src.astype(np.float32))
    return trans, inv_trans


def take_along_axis(arr, indices, axis):
    """
    Copied from: https://github.com/tensorflow/tensorflow/blob/v2.5.0/tensorflow/python/ops/numpy_ops/np_array_ops.py#L1399-L1452
    In tensorflow 2.5.0 this function raises error for dtype missmatch on call `indices = array_ops.where_v2(...)`, according to this SO thread:
    https://stackoverflow.com/questions/62390059/tensorflow-cannot-compute-addv2-as-input-1zero-based-was-expected-to-be-a-dou it is caused
    by tensorflow-probability module bug. Fixed this issue with casting parameters to float32 and then back to int32.
    """
    arr = tf.experimental.numpy.asarray(arr)
    indices = tf.experimental.numpy.asarray(indices)

    if axis is None:
        return take_along_axis(arr.ravel(), indices, 0)

    rank = array_ops.rank(arr)
    axis = axis + rank if axis < 0 else axis

    # Broadcast shapes to match, ensure that the axis of interest is not
    # broadcast.
    arr_shape_original = array_ops.shape(arr)
    indices_shape_original = array_ops.shape(indices)
    arr_shape = array_ops.tensor_scatter_update(arr_shape_original, [[axis]], [1])
    indices_shape = array_ops.tensor_scatter_update(indices_shape_original,
                                                    [[axis]], [1])
    broadcasted_shape = array_ops.broadcast_dynamic_shape(arr_shape,
                                                            indices_shape)
    arr_shape = array_ops.tensor_scatter_update(broadcasted_shape, [[axis]],
                                                [arr_shape_original[axis]])
    indices_shape = array_ops.tensor_scatter_update(
        broadcasted_shape, [[axis]], [indices_shape_original[axis]])
    arr = array_ops.broadcast_to(arr, arr_shape)
    indices = array_ops.broadcast_to(indices, indices_shape)

    # Save indices shape so we can restore it later.
    possible_result_shape = indices.shape

    # Correct indices since gather doesn't correctly handle negative indices.
    indices = array_ops.where_v2(
        indices < 0,
        tf.cast(indices, tf.float32) + tf.cast(arr_shape[axis], tf.float32),
        tf.cast(indices, tf.float32)
    )
    indices = tf.cast(indices, tf.int32)

    swapaxes_ = lambda t: tf.experimental.numpy.swapaxes(t, axis, -1)

    dont_move_axis_to_end = math_ops.equal(axis, np_utils.subtract(rank, 1))
    arr = np_utils.cond(dont_move_axis_to_end, lambda: arr,
                        lambda: swapaxes_(arr))
    indices = np_utils.cond(dont_move_axis_to_end, lambda: indices,
                            lambda: swapaxes_(indices))

    arr_shape = array_ops.shape(arr)
    arr = array_ops.reshape(arr, [-1, arr_shape[-1]])

    indices_shape = array_ops.shape(indices)
    indices = array_ops.reshape(indices, [-1, indices_shape[-1]])

    result = array_ops.gather(arr, indices, batch_dims=1)
    result = array_ops.reshape(result, indices_shape)
    result = np_utils.cond(dont_move_axis_to_end, lambda: result,
                           lambda: swapaxes_(result))
    result.set_shape(possible_result_shape)

    return result


def soft_argmax(x):
    """Soft-argmax operation, produces vector of softmax weighted indices along the last axis

    Args:
        x (tf.Tensor): input tensor

    Returns:
        tf.Tensor: output tensor
    """
    x = tf.keras.activations.softmax(x, axis=-1)
    x = tf.matmul(x, tf.expand_dims(tf.range(0, x.shape[-1], dtype=tf.float32), axis=-1))
    return x


def decode_keypoints_argmax_3d(heatmaps):
    """Decode keypoints from generated 3D heatmaps with argmax

    Args:
        heatmaps (tf.Tensor): batched N heatmaps of shape (B, X, Y, Z, N)

    Returns:
        tf.Tensor: outputs keypoints of shape (B, N, 3)
    """
    # argmax along z axis
    max_z = tf.reduce_max(heatmaps, axis=3)
    idx_z = tf.argmax(heatmaps, axis=3)
    # argmax along z and y axes
    max_zy = tf.reduce_max(max_z, axis=1)
    idx_zy = tf.argmax(max_z, axis=1)
    # argmax along z, y and x axis gives x coordinates
    keypoints_x = tf.argmax(max_zy, axis=1)[:, tf.newaxis, :]
    # take selected keypoints from y axis to gain y coordinates
    keypoints_y = take_along_axis(idx_zy, keypoints_x, axis=1)
    # same for z coordinates, but first take elements based on y coordinates and then from x coordinates
    keypoints_z = take_along_axis(idx_z, tf.repeat(keypoints_y[:, tf.newaxis, :, :], heatmaps.shape[3], axis=2), axis=1)[:, 0, :, :]
    keypoints_z = take_along_axis(keypoints_z, keypoints_x, axis=1)
    # concat keypoints and permute from (B, 3, N) to (B, N, 3)
    keypoints = tf.concat((keypoints_x, keypoints_y, keypoints_z), axis=1)
    keypoints = tf.cast(tf.transpose(keypoints, perm=(0, 2, 1)), tf.float32)
    return keypoints


def decode_keypoints_argmax_2d(heatmaps):
    """Decode keypoints from generated 2D heatmaps with argmax

    Args:
        heatmaps (tf.Tensor): batched N heatmaps of shape (B, X, Y, N)

    Returns:
        tf.Tensor: outputs keypoints of shape (B, N, 2)
    """
    # argmax along y axis
    max_y = tf.reduce_max(heatmaps, axis=2)
    idx_y = tf.argmax(heatmaps, axis=2)

    # argmax along y and x axes
    idx_xy = tf.argmax(max_y, axis=1)
    # take keypoints from argmaxed y axis based on argmaxed x and y axes
    keypoints_x = take_along_axis(idx_y, idx_xy[:, tf.newaxis, :], axis=1)
    # y coordinates are argmax from reduced heatmaps along y axis
    keypoints_y = tf.argmax(max_y, axis=1)[:, tf.newaxis, :]
    # concat keypoints and permute from (B, 2, N) to (B, N, 2)
    keypoints = tf.concat((keypoints_x, keypoints_y), axis=1)
    keypoints = tf.transpose(keypoints, (0, 2, 1))
    return keypoints


def decode_keypoints_compensated_2d(heatmaps, delta=0):
    """Decode keypoints from generated 2D heatmaps with compensation

    Args:
        heatmaps (tf.Tensor): batched N heatmaps of shape (B, X, Y, N)

    Returns:
        tf.Tensor: outputs keypoints of shape (B, N, 2)
    """
    # cut delta pixels from the end of the heatmaps
    end_x = heatmaps.shape[2]-delta
    end_y = heatmaps.shape[1]-delta
    f_xy = heatmaps[:, :end_y, :end_x, :]
    # calculate spatial moments from cut heatmaps
    xx = tf.range(0, end_x, dtype=tf.float32)
    yy = tf.range(0, end_y, dtype=tf.float32)
    m00 = tf.reduce_sum(f_xy, axis=[1, 2])
    m10 = tf.reduce_sum(f_xy * xx[tf.newaxis, tf.newaxis, :, tf.newaxis], axis=[1, 2])
    m01 = tf.reduce_sum(f_xy * yy[tf.newaxis, :, tf.newaxis, tf.newaxis], axis=[1, 2])
    # keypoints are equal to mass centers
    mean_x = m10 / m00
    mean_y = m01 / m00
    # concat keypoints
    keypoints = tf.concat((mean_x[..., tf.newaxis], mean_y[..., tf.newaxis]), axis=-1)
    return keypoints
