import numpy as np
import tensorflow as tf


def euler_to_rotation_matrix(x, y, z):
    """

    :param x: Tensor of shape (B, 1) - x axis rotation
    :param y: Tensor of shape (B, 1) - y axis rotation
    :param z: Tensor of shape (B, 1) - z axis rotation
    :return:  Rotation matrix for the given euler anglers, in the  order rotation(x).rotation(y).rotation(z)
    """
    B = tf.shape(z)[0]

    # Euler angles should be between -pi and pi, clip so the pose network is coerced to this range
    z = tf.clip_by_value(z, -np.pi, np.pi)
    y = tf.clip_by_value(y, -np.pi, np.pi)
    x = tf.clip_by_value(x, -np.pi, np.pi)

    # Expand to B x 1 x 1
    z = tf.expand_dims(tf.expand_dims(z, -1), -1)
    y = tf.expand_dims(tf.expand_dims(y, -1), -1)
    x = tf.expand_dims(tf.expand_dims(x, -1), -1)

    zeros = tf.zeros([B, 1, 1])
    ones = tf.ones([B, 1, 1])

    cosx = tf.cos(x)
    sinx = tf.sin(x)
    rotx_1 = tf.concat([ones, zeros, zeros], axis=3)
    rotx_2 = tf.concat([zeros, cosx, -sinx], axis=3)
    rotx_3 = tf.concat([zeros, sinx, cosx], axis=3)
    xmat = tf.concat([rotx_1, rotx_2, rotx_3], axis=2)

    cosz = tf.cos(z)
    sinz = tf.sin(z)
    rotz_1 = tf.concat([cosz, -sinz, zeros], axis=3)
    rotz_2 = tf.concat([sinz, cosz, zeros], axis=3)
    rotz_3 = tf.concat([zeros, zeros, ones], axis=3)
    zmat = tf.concat([rotz_1, rotz_2, rotz_3], axis=2)

    cosy = tf.cos(y)
    siny = tf.sin(y)
    roty_1 = tf.concat([cosy, zeros, siny], axis=3)
    roty_2 = tf.concat([zeros, ones, zeros], axis=3)
    roty_3 = tf.concat([-siny, zeros, cosy], axis=3)
    ymat = tf.concat([roty_1, roty_2, roty_3], axis=2)

    rotMat = tf.matmul(tf.matmul(xmat, ymat), zmat)
    return rotMat


def pose_vec2mat(vec):
    """Converts 6DoF parameters to transformation matrix
    Args:
        vec: 6DoF parameters in the order of tx, ty, tz, rx, ry, rz -- [B, 6]
    Returns:
        A transformation matrix -- [B, 4, 4]
    """
    batch_size, _ = vec.get_shape().as_list()
    translation = tf.slice(vec, [0, 0], [-1, 3])
    translation = tf.expand_dims(translation, -1)
    rx = tf.slice(vec, [0, 3], [-1, 1])
    ry = tf.slice(vec, [0, 4], [-1, 1])
    rz = tf.slice(vec, [0, 5], [-1, 1])
    rot_mat = euler_to_rotation_matrix(rx, ry, rz)
    rot_mat = tf.squeeze(rot_mat, axis=[1])
    filler = tf.constant([0.0, 0.0, 0.0, 1.0], shape=[1, 1, 4])
    filler = tf.tile(filler, [batch_size, 1, 1])
    transform_mat = tf.concat([rot_mat, translation], axis=2)
    transform_mat = tf.concat([transform_mat, filler], axis=1)
    return transform_mat


def image_coordinate(batch, height, width):
    """
    Construct a tensor for the given height/width with elements the homogenous coordinates for the pixel
    :param batch: Number of images in a batch
    :param height: Height of image
    :param width: Width of image
    :return: Tensor of shape (B, height, width, 3), homogenous coordinates for an image.
        Coordinates are in order [x, y, 1]
    """
    x_coords = tf.range(width)
    y_coords = tf.range(height)

    x_mesh, y_mesh = tf.meshgrid(x_coords, y_coords)

    ones_mesh = tf.cast(tf.ones([height, width]), tf.int32)

    stacked = tf.stack([x_mesh, y_mesh, ones_mesh], axis=2)

    return tf.repeat(tf.expand_dims(stacked, axis=0), batch, axis=0)


def projective_inverse_warp(target_img, source_img, depth, pose, intrinsics, coordinates):
    """
    Calculate the reprojected image from the source to the target, based on the given depth, pose and intrinsics

    SFM Learner inverse warp step
        ps ~ K.T(t->s).Dt(pt).K^-1.pt

    Idea is to map the pixel coordinates of the target image to 3d space (Dt(pt).K^-1.pt), then map these onto
    the source image in pixel coordinates (K.T(t->s).{3d coord}), then using the projected coordinates we sample 
    the pixels in the source image (ps) to reconstruct the target image.

    :param target_img: Tensor (batch, height, width, 3)
    :param source_img: Tensor, same shape as target_img
    :param depth: Tensor, (batch, height, width, 1)
    :param pose: (batch, 6)
    :param intrinsics: (batch, 3, 3)
    :param coordinates: (batch, height, width, 3) - coordinates for the image. Pass this in so it doesn't need to be
        calculated on every warp step
    :return: The source image reprojected to the target
    """
    # Convert pose vector (output of pose net) to pose matrix (4x4)

    # Convert intrinsics matrix (3x3) to (4x4) so it can be multiplied by the pose net
    # intrinsics_4x4 =

    # Calculate inverse of the 4x4 intrinsics matrix
    tf.linalg.inv()

    # Create grid of homogenous coordinates

    #
    pass