import torch


def test(model, data_loader, device, criterion, flag, title):
    """
    Evaluate model by given dataset.
    :param model:        given model
    :param data_loader:  dataloader of model to evaluate by
    :param device:       device used for model's training
    :param criterion:    loss function for testing
    :param flag:         label for target and source domains
    :param title:        model name title
    :return:             average loss, accuracy and domain accuracy
    """

    # set evaluation state for Dropout and BN layers
    model.eval()

    # initial loss and accuracy
    loss_ = 0.0
    acc_ = 0.0
    acc_domain_ = 0.0
    n_total = 0

    # evaluate the model
    for (images, labels) in data_loader:
        images = images.to(device)
        labels = labels.to(device)
        size = len(labels)

        if flag == 'target':
            labels_domain = torch.ones(size).long().to(device)
        else:
            labels_domain = torch.zeros(size).long().to(device)

        # get predictions and domain
        preds, domain = model(images, alpha=0)

        # calculate data for evaluation criteria
        loss_ += criterion(preds, labels).item()
        pred_cls = preds.data.max(1)[1]
        pred_domain = domain.data.max(1)[1]
        acc_ += pred_cls.eq(labels.data).sum().item()
        acc_domain_ += pred_domain.eq(labels_domain.data).sum().item()
        n_total += size

    # calculate evaluation criteria
    loss = loss_ / n_total
    acc = acc_ / n_total
    acc_domain = acc_domain_ / n_total

    print(f"{title}: Avg Loss = {loss:.6f}, Avg Accuracy = {acc:.2%}, {acc_}/{n_total}, "
          f"Avg Domain Accuracy = {acc_domain:2%}")

    return loss, acc, acc_domain