import torch


def validate_CNN(test_loader, model, criterion, force_cpu=False):
    """ Carries out the validation for a trained CNN model
    INPUT
        test_loader: DataLoader object to load test data
        model: trained CNN model
        criterion: criterion used to compute the loss
        force_cpu - use cpu even with gpu available
    OUTPUT
        array containing the loss per epoch
        array containing the correct number of predictions per epoch
        number of samples used for validation, in total
    """
    # Switch to evaluate mode
    model.eval()

    # Initialize counters
    nb_correct = 0
    nb_total = 0
    test_loss = 0

    # Iterate over batches
    for i, (inputs, target) in enumerate(test_loader):
        X = torch.autograd.Variable(inputs)
        y = torch.autograd.Variable(target)

        # Move to GPU
        if torch.cuda.is_available() and not force_cpu:
            X = X.cuda()
            y = y.cuda()

        # Obtain predictions
        scores = model(X)

        # Predicted label (highest probability)
        pred_label = scores.data.max(1)[1]

        # Loss
        loss = criterion(scores, y)
        test_loss += loss.data[0]

        # Update nb. correct and nb. total
        nb_correct += (pred_label == y.data).sum()
        nb_total += len(pred_label)

    return test_loss, nb_correct, nb_total