import torch

from utils import flat_auroc_score

def evaluate_model(args, model_mlp, test_dataloader, criterion, set_device):

    test_loss_history, auroc_test_history = list(), list()

    model_mlp.eval()

    predictions = None

    with torch.no_grad():

        test_loss_scores = list()

        y_true_val, y_pred_val = list(), list()

        ## perform test pass

        for batch, targets in test_dataloader:

            ## perform forward pass

            batch = batch.type(torch.FloatTensor).to(set_device)

            pred = model_mlp(batch)

            predictions = pred

            preds = torch.max(pred, 1)[1]

            ## accumulate predictions per batch for the epoch

            y_pred_val += list([x.item() for x in preds.detach().cpu().numpy()])

            targets = torch.LongTensor([x.item() for x in list(targets)])

            y_true_val += list([x.item() for x in targets.detach().cpu().numpy()])

            ## computing validate loss

            loss = criterion(

                pred.to(set_device), targets.to(set_device)

            )  ## compute loss

            ## accumulate validate loss

            test_loss_scores.append(loss.item())

        ## accumulate loss, auroc, f1, precision per epoch

        test_loss_history.append((sum(test_loss_scores) / len(test_loss_scores)))

        auroc = flat_auroc_score(predictions, y_true_val)

        auroc_test_history.append(auroc)

        print(f"Test => AUROC score: {auroc_test_history[-1]} ")

        with open(args.out_file, "a+") as f:

            f.write(

                f"{args.dataset_name}, {args.samples_per_class}, {args.n_augment}, {auroc_test_history[-1]}\n"

            )