import json

import torch
import torchvision.transforms as transforms
import torch.nn.functional as F
import requests

from .base_env import device

def biased_fgsm(model, image):
    # 计算梯度
    image.requires_grad = True
    outputs = model(image.unsqueeze(0))
    original_label = torch.argmax(outputs, 1)
    loss = F.cross_entropy(outputs, original_label)  # IMPORTANT！！在此假设了原始标签是预测正确的，以此计算交叉熵损失后求梯度
    loss.backward()

    # 添加对抗扰动
    sngd = image.grad.sign()
    return sngd


def predict(model, image):
    global labels
    with torch.no_grad():
        prediction = F.softmax(model(image)[0], dim=0)
        confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
    return confidences


def show_attack(model_name, epsilon, PIL_image):

    # 创建一个新的预处理函数
    preprocess = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # ImageNet 的均值和标准差
    ])
    image = preprocess(PIL_image)
    if device == torch.device('cuda'):
        image = image.cuda()

    # 加载模型
    if model_name == "AlexNet":
        model = torch.load('./models/PublicPretrained/alexnet.pth')
    elif model_name == "VGG16":
        model = torch.load('./models/PublicPretrained/vgg16.pth')
    elif model_name == "GoogLeNet":
        model = torch.load('./models/PublicPretrained/googlenet.pth')
    elif model_name == "ResNet18":
        model = torch.load('./models/PublicPretrained/resnet18.pth')
    elif model_name == "ResNet50":
        model = torch.load('./models/PublicPretrained/resnet50.pth')
    elif model_name == "Inception_V3":
        model = torch.load('./models/PublicPretrained/inception_v3.pth')
    elif model_name == "DenseNet161":
        model = torch.load('./models/PublicPretrained/densenet161.pth')
    elif model_name == "ResNeXt50_32x4d":
        model = torch.load('./models/PublicPretrained/resnext50_32x4d.pth')
    elif model_name == "MobileNet_V2":
        model = torch.load('./models/PublicPretrained/mobilenet_v2.pth')
    if device == torch.device('cuda'):
        model = model.cuda()
    model.eval()

    # Download human-readable labels for ImageNet.
    with open('./configs/imagenet_labels.json', 'r') as f:
        global labels
        labels = json.load(f)

    sngd = biased_fgsm(model, image)
    adv_image = image + epsilon * sngd
    adv_image = adv_image.detach()  # 将adv_image从计算图中分离，不再计算梯度

    unprocess = transforms.Compose([
        transforms.Normalize(mean=[0, 0, 0], std=[1 / 0.229, 1 / 0.224, 1 / 0.225]),
        transforms.Normalize(mean=[-0.485, -0.456, -0.406], std=[1, 1, 1]),
        transforms.ToPILImage()
    ])
    PIL_sngd = transforms.functional.to_pil_image(sngd * 0.5 + 0.5)
    PIL_adv_image = unprocess(adv_image)

    #print(str(device))
    confidences = predict(model, image.unsqueeze(0).to(device))
    sngd_confidences = predict(model, sngd.unsqueeze(0).to(device))
    adv_confidences = predict(model, adv_image.unsqueeze(0).to(device))

    pred, conf = max(confidences.items(), key=lambda x: x[1])
    sngd_pred, sngd_conf = max(sngd_confidences.items(), key=lambda x: x[1])
    adv_pred, adv_conf = max(adv_confidences.items(), key=lambda x: x[1])
    top5_adv_confidences = dict(sorted(adv_confidences.items(), key=lambda x: x[1], reverse=True)[:10])
    return [
        top5_adv_confidences,
        PIL_image, pred, conf,
        PIL_sngd, sngd_pred, sngd_conf,
        PIL_adv_image, adv_pred, adv_conf
    ]


if __name__ == "__main__":
    pass