import torch


def test_vae(model, data_loader):
    """
    Function for testing a model on a dataset.
    Inputs:
        model - VAE model to test
        data_loader - Data Loader for the dataset you want to test on.
    Outputs:
        average_bpd - Average BPD
        average_rec_loss - Average reconstruction loss
        average_reg_loss - Average regularization loss
    """

    # set model mode to eval
    model.eval()

    bpd_vals, rec_loss, reg_loss = 0, 0, 0
    total_steps = 0

    for step, (test_inputs, test_targets) in enumerate(data_loader):

        # move to device
        test_inputs = test_inputs.to(model.device)

        with torch.no_grad():

            # make a forward pass through the model (encoder and decoder)
            L_rec, L_reg, bpd = model.forward(test_inputs)
            
            # track bpd, rec, and reg
            bpd_vals += bpd.item()
            rec_loss += L_rec.item()
            reg_loss += L_reg.item()

            # keep track of steps
            total_steps += 1

    average_bpd = bpd_vals / total_steps
    average_rec_loss = rec_loss / total_steps
    average_reg_loss = reg_loss / total_steps

    # set model back to training
    model.train()
    
    return average_bpd, average_rec_loss, average_reg_loss