import copy


def transition(states, masks, actions, rephrase_positions):
    """
    given s_t, a_t, t, returns torch.tensor s_{t+1}
    :param states: s_t torch.tensor [batch, len, dim], float
    :param masks: masks to indicate valid token for src, torch.tensor [batch, len], float
    :param actions: a_t, torch.tensor [batch, dim], float
    :param rephrase_positions: t, torch.tensor torch.tensor [batch, 1], long
    """
    batch, _, _ = states.shape
    new_states = copy.deepcopy(states)
    for i in range(batch):
        new_states[i][rephrase_positions[i], :] += actions[i] * masks[i][rephrase_positions[i]]
    return new_states.detach()