import torch


def compose_quaternions(quat1, quat2):
    """Returns the cross-product of the two quaternions.

    Implementation as described in
    https://github.com/adamlwgriffiths/Pyrr/blob/master/pyrr/quaternion.py#L199

    q1w, q1x, q1y, q1z = quat1
    q2w, q2x, q2y, q2z = quat2
    We want to compute
    [
      q1x * q2w + q1y * q2z - q1z * q2y + q1w * q2x,
     -q1x * q2z + q1y * q2w + q1z * q2x + q1w * q2y,
      q1x * q2y - q1y * q2x + q1z * q2w + q1w * q2z,
     -q1x * q2x - q1y * q2y - q1z * q2z + q1w * q2w,
    ]

    Arguments:
    ----------
        quat1: Tensor with size ...x4, where the ellipsis is any broadcastable
               shape
        quat2: Tensor with size ...x4

    Returns:
    -------
        quats: Tensor with size ...x4 with the final quaternions
    """
    q1w, q1x, q1y, q1z = [quat1[..., i] for i in range(4)]
    q2w, q2x, q2y, q2z = [quat2[..., i] for i in range(4)]
    return torch.stack(
        [
            -q1x * q2x - q1y * q2y - q1z * q2z + q1w * q2w,
            q1x * q2w + q1y * q2z - q1z * q2y + q1w * q2x,
            -q1x * q2z + q1y * q2w + q1z * q2x + q1w * q2y,
            q1x * q2y - q1y * q2x + q1z * q2w + q1w * q2z
        ],
        dim=-1
    )