import torch


def computeTargets(reward, nextObservation, discountFactor, done, targetNetwork, device):
    """
    return y = r + \\gamma max_{a'} Q(s', a'; θ-) for non-terminal s'
    else y = r for terminal s'

    return: tensor
    """
    # nextObservation is a 2-D tensor
    assert len(nextObservation.shape) == 2

    if done:
        ret = torch.tensor(reward, dtype=torch.float32).to(device)
    else:
        action_values = targetNetwork(nextObservation)
        ret = torch.tensor(reward, dtype=torch.float32).to(device) + \
            discountFactor * torch.max(action_values, 1)[0].to(device)
    return ret