import torch


def clip_npy_labels_by_time_step(features, labels, time_step):
    """Clip features in npy file and labels by time step, for example:
    npy_shape[1000, 14, 2], labels_shape[1000], time step:100, clip
    to npy_shape[900, 100, 14, 2], labels_shape[10, 100]."""
    features = torch.from_numpy(features).float()
    labels = torch.Tensor(labels)
    x, y = features.shape[-2:]
    feature_size = features.shape[0]
    assert feature_size == labels.shape[0], "length of feture must " \
        "equal length of labels, but got {} and {}.".format(feature_size, labels.shape[0])
    # this is for random start time
    features_list = []
    labels_list = []
    for i in range(feature_size - time_step):
        features_list.append(features[i:i+time_step])
        labels_list.append(labels[i:i+time_step])
    features = torch.cat(features_list, 0).view(-1, time_step, x, y)
    labels = torch.cat(labels_list, 0).view(-1, time_step)

    return features, labels