
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
from models.resnet_models import *
import data_loaders
from torchattacks.attack import Attack
from tqdm import tqdm
from functions import TET_loss, seed_all, get_logger
from models.VGG_models import *
from models.WideResNet import *
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
seed_all(1000)
# 数据加载
train_dataset, val_dataset, znorm = data_loaders.cifar_dataset(use_cifar10=False)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True,
                                               num_workers=16, pin_memory=True)
test_loader = torch.utils.data.DataLoader(val_dataset, batch_size=128,
                                              shuffle=False, num_workers=16, pin_memory=True)

# 选择模型，这里使用预训练的ResNet18
model = wrn16(num_classes=100, norm=znorm)
model.load_state_dict(torch.load('wrn16_noise_lag18.pth'))
model = model.to(device)
model.eval()  # 设置为评估模式

activations = {}
def get_activation(name):
    def hook(model, input, output):
        activations[name] = output.detach()
    return hook
# 传递数据通过模型

def BIM(model, images, labels, eps=8/255, alpha=0.01, steps=10):
    r"""
    BIM (Basic Iterative Method) Attack
    Args:
        model: The target model to attack.
        images: The input images.
        labels: The true labels of the images.
        eps: The maximum perturbation (epsilon).
        alpha: The step size for each iteration.
        steps: The number of iterations. If 0, it will be calculated as min(eps*255 + 4, 1.25*eps*255).
        T: Optional parameter for time steps (used in spiking neural networks).
        forward_function: Optional custom forward function for the model.
    Returns:
        adv_images: The adversarial examples generated by BIM.
    """

    model.eval()
    images = images.clone().detach().to(device)
    labels = labels.clone().detach().to(device)

    # Initialize the loss function
    loss = nn.CrossEntropyLoss()

    # Save the original images for clamping
    ori_images = images.clone().detach()

    # Iteratively update the adversarial examples
    for _ in range(steps):
        images.requires_grad = True
        outputs,_ = model(images)
        outputs = outputs.mean(1)
        # Compute the loss
        cost = loss(outputs, labels)

        # Compute gradients
        grad = torch.autograd.grad(cost, images, retain_graph=False, create_graph=False)[0]

        # Update adversarial images
        adv_images = images + alpha * grad.sign()

        # Clamp the adversarial images to be within the epsilon-ball and valid range
        a = torch.clamp(ori_images - eps, min=0)
        b = (adv_images >= a).float() * adv_images + (adv_images < a).float() * a
        c = (b > ori_images + eps).float() * (ori_images + eps) + (b <= ori_images + eps).float() * b
        images = torch.clamp(c, max=1).detach()

    return images



def val(model, test_loader, device, T):
    #atk = None
    #print(T)
    correct_ori = 0
    correct = 0
    total = 0
    model.eval()

    for batch_idx, (inputs, targets) in enumerate(tqdm(test_loader)):
        inputs = inputs.to(device)
        ori = inputs
        inputs = BIM(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)


