import torch


def predict(net, criterion, test=True, train=False):
    """
    Given a model and criterion, get prediction performance on
    validation or test set.
    ===
    args:
        - net: pytorch model object, model to test
        - criterion: function, loss function
        - test: bool, if set to True, uses Test set, otherwise Validation set,
            default=True
        - combine: bool, add Validation to Test Set, negated if test=False,
            default=False
    output:
        - prediction accuracy on test set
    """
    net.eval()
    correct_cnt, avg_loss = 0, 0
    if test:
        loader = net.test_mnist
    elif train:
        loader = net.train_mnist
    else:
        loader = net.valid_mnist
    for batch_idx, (x, target) in enumerate(loader):
        x, target = x.to(net.device), target.to(net.device)
        output = net(x)
        loss = criterion(output, target)
        _, pred_label = torch.max(output.data, 1)
        correct_cnt += (pred_label == target.data).sum()
        avg_loss += loss.item()

    total = len(loader) * loader.batch_size
    avg_loss /= len(loader)
    accuracy = float(correct_cnt) / total

    return avg_loss, accuracy