import torch


def calculate_returns(rewards, discount, masks=None):
    """
    Calculates the returns from a sequence of rewards.
    :param rewards: a torch.Tensor of size (time, batch_size)
    :param discount: a float
    :param masks: will use the mask to take into account future rewards
    """
    assert discount <= 1 and discount >= 0, 'discount is out of allowable range'
    discounted = []
    if masks is None:
        # no masks at all
        masks = torch.ones_like(rewards)
    masks = masks.float()
    trajectory_length = rewards.size()[0]
    n_processes = rewards.size()[1]
    returns = torch.zeros(trajectory_length+1, *rewards.size()[1:])
    for t in reversed(range(trajectory_length)):
        returns[t] = discount * returns[t+1] + masks[t] * rewards[t]

    return returns[:-1]