import os
import argparse
import torch
# from rslad_loss import *
# from cifar10_models import *
import torchvision
from torchvision import datasets, transforms
import torch.nn as nn
from advertorch.attacks import LinfPGDAttack, CarliniWagnerL2Attack, DDNL2Attack, SinglePixelAttack, LocalSearchAttack, SpatialTransformAttack,L1PGDAttack
import advertorch
# from models.wideresnet import *
from autoattack import AutoAttack
# we fix the random seed to 0, this method can keep the results consistent in the same conputer.
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.deterministic = True

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

epochs = 300
batch_size = 1024
epsilon = 8/255.0

# transform_train = transforms.Compose([
#     # transforms.RandomCrop(32, padding=4),
#     # transforms.RandomHorizontalFlip(),
#     transforms.ToTensor(),
#     transforms.Normalize((0.4914, 0.4822, 0.4465),
#                          (0.2023, 0.1994, 0.2010)),
# ])

transform_train = transforms.Compose([
    # transforms.RandomCrop(32, padding=4),
    # transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2023, 0.1994, 0.2010)),
])

# transform_test = transforms.Compose([
#     transforms.ToTensor(),
# ])

transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465),
                             (0.2023, 0.1994, 0.2010)),
    ])

trainset = torchvision.datasets.CIFAR10(root='', train=True, download=True, transform=transform_test)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='', train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)



def adjust_learning_rate(optimizer, epoch):
    """For resnet, the lr starts from 0.1, and is divided by 10 at 80 and 120 epochs"""
    if epoch < 80:
        lr = 0.1
    elif epoch < 120:
        lr = 0.01
    else:
        lr = 0.001
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr






import resnet


test_model = torch.load('').cuda().eval()


test_model.eval()
total = 0
correct = 0
temp = 0
# targetlabel = torch.zeros(BATCH_SIZE).cuda()
# targetlabel = targetlabel.to('cuda', dtype=torch.int64)
for img, label in testloader:
    img, label = img.cuda(), label.cuda()
    # targetlabel_temp = targetlabel[0:img.shape[0]]

    img_test = img

    # FGSM_N = advertorch.attacks.GradientSignAttack(predict=test_model, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=8/255)
    # img_test = FGSM_N.perturb(img)

    # PGD_N = LinfPGDAttack(
    #     test_model, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=8/255,
    #     nb_iter=10, eps_iter=2/(255), rand_init=True, clip_min=0.0, clip_max=1.0, targeted=False)
    # img_test = PGD_N.perturb(img)

    # CW
    # CW_N = CarliniWagnerL2Attack(
    #     test_model, 10, clip_min=0.0, clip_max=1.0, max_iterations=15, confidence=1, initial_const=1, learning_rate=1e-2,
    #     binary_search_steps=4, targeted=False)
    # img_test = CW_N.perturb(img)

    # AA
    # AA_N = AutoAttack(test_model, norm='Linf', eps=8/255, version='standard')
    # img_test = AA_N.run_standard_evaluation(img, label, bs=batch_size)

    # Ave # []


    x = test_model(img_test)
    _, prediction = torch.max(x, 1)
    total += label.size(0)
    correct += (prediction == label).sum()
    print('当前temp:', temp, '/312,当前batch正确的个数:', (prediction == label).sum())
    temp += 1
# print('There are ' + str(correct.item()) + ' correct pictures.')
print('TestSet_Accuracy=%.2f' % (100.00 * correct.item() / total))



