import torch


def euler_matrix(ai, aj, ak, repetition=True):
    """Return homogeneous rotation matrix from Euler angles and axis sequence.
    ai, aj, ak : Euler's roll, pitch and yaw angles
    Readapted for Pytorch: some tricks going on
    """
    si, sj, sk = torch.sin(ai), torch.sin(aj), torch.sin(ak)
    ci, cj, ck = torch.cos(ai), torch.cos(aj), torch.cos(ak)
    cc, cs = ci * ck, ci * sk
    sc, ss = si * ck, si * sk
    
    i = 0; j = 1; k = 2 # indexing
    
    # Tricks to create batched matrix [...,3,3] 
    # any suggestion to make code more readable is welcome!
    M = torch.cat(3*[torch.cat(3*[torch.zeros(ai.shape)[..., None, None]], -1)], -2) 
    if repetition:
        M[..., i, i] = cj
        M[..., i, j] = sj * si
        M[..., i, k] = sj * ci
        M[..., j, i] = sj * sk
        M[..., j, j] = -cj * ss + cc
        M[..., j, k] = -cj * cs - sc
        M[..., k, i] = -sj * ck
        M[..., k, j] = cj * sc + cs
        M[..., k, k] = cj * cc - ss
    else:
        M[..., i, i] = cj * ck
        M[..., i, j] = sj * sc - cs
        M[..., i, k] = sj * cc + ss
        M[..., j, i] = cj * sk
        M[..., j, j] = sj * ss + cc
        M[..., j, k] = sj * cs - sc
        M[..., k, i] = -sj
        M[..., k, j] = cj * si
        M[..., k, k] = cj * ci
    return M