import torch


def DH(theta, d, r, alpha, device):

    """
     Calculates the Denavit-Hartenberg Matrix
     where
     d: offset along previous z to the common normal
     theta: angle about previous z, from old x to new x
     r: length of the common normal (aka a, but if using this notation, do not confuse with alpha).
     Assuming a revolute joint, this is the radius about previous z.
     alpha: angle about common normal, from old z axis to new z axis
    """

    T = torch.zeros([theta.shape[0], 4, 4]).to(device)
    T[:, :, :] = torch.eye(4).to(device)

    cTheta = torch.cos(theta)
    sTheta = torch.sin(theta)
    calpha = torch.cos(alpha)
    salpha = torch.sin(alpha)

    T[:, 0, 0] = cTheta
    T[:, 0, 1] = -sTheta
    T[:, 0, 2] = 0.0
    T[:, 0, 3] = r

    T[:, 1, 0] = sTheta * calpha
    T[:, 1, 1] = cTheta * calpha
    T[:, 1, 2] = -salpha
    T[:, 1, 3] = - d * salpha

    T[:, 2, 0] = sTheta * salpha
    T[:, 2, 1] = cTheta * salpha
    T[:, 2, 2] = calpha
    T[:, 2, 3] = d * calpha

    return T