import torch


def rotate_vec_by_axisangle(vec, aa_vec):
    """
    This function rotates a 3D vector @vec by the direction and angle represented by @aa_vec.

    See https://stackoverflow.com/questions/32485772/how-do-axis-angle-rotation-vectors-work-and-how-do-they-compare-to-rotation-matr
    for more information.

    Args:
        vec (tensor): (..., 3) where final dim represents (x, y, z) vectors
        aa_vec (tensor): (..., 3) where final dim represents desired (ax, ay, az) axis-angle rotations
    Returns:
        tensor: (..., 3) where final dim is newly rotated (x, y, z) vectors
    """
    # Extract angle and unit vector from axis-angle vectors
    angle = torch.norm(aa_vec, dim=-1, keepdim=True)
    aa_v = aa_vec / angle
    # Map all NaNs to 0
    aa_v[aa_v != aa_v] = 0.0

    # # Record angles that are zero so we don't get nan's
    # idx = torch.nonzero(angle.squeeze(dim=-1)).tolist()
    # aa_v = torch.zeros_like(aa_vec)
    # aa_v[idx] = aa_vec[idx] / angle[idx]

    # Rotate the vector using the formula (see link above)
    c, s = torch.cos(angle), torch.sin(angle)
    vec_rot = vec * c + torch.cross(aa_v, vec) * s + \
        aa_v * (torch.sum(aa_v * vec, dim=-1, keepdim=True)) * (1 - c)

    return vec_rot