from utils.utils import to_cuda, to_cpu
import torch

def fooling_rate_test(patch_im, patch_im_mask, model, target, dataset,
                      batch_size=32, GPU_device=0):
    total_time = 0
    targeted_fooling_time =0
    non_targeted_fooling_time = 0
    model = to_cuda(model,GPU_device)
    data_loader = torch.utils.data.DataLoader( dataset=dataset,batch_size=batch_size,num_workers=4,shuffle=False)
    for i,batch in enumerate(data_loader):
        batch_adve = batch[0] * torch.abs(patch_im_mask-1) + patch_im * patch_im_mask
        batch_adve=to_cuda(batch_adve,GPU_device)

        output= model(batch_adve)

        prediction = torch.argmax(output,dim=1)
        prediction = to_cpu(prediction)
        targeted_fooling_time += (prediction == target).sum().item()
        non_targeted_fooling_time += (prediction != batch[1]).sum().item()
        total_time+=batch[0].shape[0]
    return targeted_fooling_time,non_targeted_fooling_time,total_time
