import torch


def get_deltax0(target):
    """
    Gets the value of the mean change in position and velocity in the first timestep

    :param target: tensor of all data points from simulation
    target has dimensions [batch, particle, timestep, state]
    :return: mean change in position and velocity over the first timestep
    """
    # separate out the velocity and position terms
    indices = torch.LongTensor([0,1])
    indices_vel = torch.LongTensor([2,3])
    if target.is_cuda:
        indices, indices_vel = indices.cuda(), indices_vel.cuda()
    target_pos = torch.index_select(target, 3, indices)
    # calculate the magnitude in change in displacement
    deltax0 = (target_pos[:,:,1,:]-target_pos[:,:,0,:]).squeeze()
    deltax0 = deltax0.norm(p=2 , dim = 2, keepdim=True)
    target_vel = torch.index_select(target, 3, indices_vel)
    # calculate the magnitude in the change in velocity
    deltav0 = (target_vel[:, :, 1, :] - target_vel[:, :, 0, :]).squeeze()
    deltav0 = deltav0.norm(p=2, dim=2, keepdim=True)
    return deltax0.mean(), deltav0.mean()