
import os
import torch
import torch.nn.functional as F
import torchvision.models as models
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch import nn
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import models.resnet_models
import data_loaders
from torchattacks.attack import Attack
from tqdm import tqdm

from models.layers import LIFSpike

os.environ["CUDA_VISIBLE_DEVICES"] = "2"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# 选择模型，这里使用预训练的ResNet18
model = models.resnet_models.resnet19(num_classes=100)
model.load_state_dict(torch.load('resnet-ssssss.pth'))
model=model.to(device)
model.eval()  # 设置为评估模式

activations = {}
def get_activation(name):
    def hook(model, input, output):
        activations[name] = output.detach()
    return hook

# 注册每个层的hook
model.conv1_s.register_forward_hook(get_activation('conv1_s'))
model.fc1_s.register_forward_hook(get_activation('fc1_s'))

# 数据加载
train_dataset, val_dataset = data_loaders.build_cifar()
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=256, shuffle=True,
                                               num_workers=16, pin_memory=True)
test_loader = torch.utils.data.DataLoader(val_dataset, batch_size=256,
                                              shuffle=False, num_workers=16, pin_memory=True)

# 传递数据通过模型

def FGSM(model,images, labels):
        r"""
        Overridden.
        """
        model.eval()
        # print(self.model)
        images = images.clone().detach().to(device)

        labels = labels.clone().detach().to(device)


        loss = nn.CrossEntropyLoss()

        images.requires_grad = True

        outputs,_ = model(images)

        outputs = outputs.mean(1)
        cost = loss(outputs, labels)
        # print(self.model)
        grad = torch.autograd.grad(cost, images, retain_graph=False, create_graph=False)[0]
        print(grad.abs().min())
        adv_images = images + 8/255 * grad.sign()
        adv_images = torch.clamp(adv_images, min=0, max=1).detach()

        return adv_images


def val(model, test_loader, device, T):
    #atk = None
    #print(T)
    correct_ori = 0
    correct = 0
    total = 0
    model.eval()
    for module in model.modules():
        if isinstance(module, LIFSpike):
            module.gama =1
    for batch_idx, (inputs, targets) in enumerate(tqdm(test_loader)):
        inputs = inputs.to(device)
        ori = inputs
        inputs = FGSM(model,inputs,targets)

        with torch.no_grad():
            if T > 0:
                outputs_ori,_ = model(ori)
                outputs,_ = model(inputs)
                outputs_ori=outputs_ori.mean(1)
                outputs = outputs.mean(1)
            else:
                outputs = model(inputs)
                #print(outputs.shape)
        _, predicted_ori = outputs_ori.cpu().max(1)
        _, predicted = outputs.cpu().max(1)
        #print(predicted.shape)
        total += float(targets.size(0))
        correct += float(predicted.eq(targets).sum().item())
        correct_ori += float(predicted_ori.eq(targets).sum().item())
    final_acc = 100 * correct / total
    final_acc_ori = 100 * correct_ori / total
    return final_acc,final_acc_ori



acc,acc_ori = val(model, test_loader, device, 8)
print(acc,acc_ori)


