import torch


def doUpdate(X, Y, dt):
    """
    Updates the node attributes that describe each nucleotide's position (translational, rotational). Use a Euclidean update. 

    Each row in X contains the following entries: 
    0   1   2   3   4   5   6   7   8   9   10  11  12  13  14  15
    n   rx  ry  rz  bx  by  bz  nx  ny  nz  vx  vy  vz  Lx  Ly  Lz

    Each row in Y contains the following entries: 
    0   1   2   3   4   5
    ax  ay  az  atx aty atz

    Inputs:
    X : node attribute matrix containing each nucleotide's position and orientation in shape [n_nodes, n_features]
    Y : decoder output containing translational and rotational accelerations for each nucleotide in shape [n_nodes, n_state_vars]
    dt : scalar giving the time step of the ground truth data

    Outputs: 
    X_next : node attribute matrix for the next time step in shape [n_nodes, n_features]

    """
    # create X_next, the node attribute matrix for the next time step
    X_next = torch.zeros_like(X)

    # update translational velocity (vx, vy, vz)
    X_next[:, 10:13] = X[:, 10:13] + dt * Y[:, 0:3]

    # update rotational velocity (Lx, Ly, Lz)
    X_next[:, 13:] = X[:, 13:] + dt * Y[:, 3:]

    # update translational position (rx, ry, rz)
    X_next[:, 1:4] = X[:, 1:4] + dt * X_next[:, 10:13]

    # update backbone base versor (bx, by, bz)
    X_next[:, 4:7] = X[:, 4:7] + dt * X_next[:, 13:]

    # update normal versor (nx, ny, nz)
    X_next[:, 7:10] = X[:, 7:10] + dt * X_next[:, 13:]

    return X_next