# 导入必要的库
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from torchvision import models, transforms
from PIL import Image
import urllib.request

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False

print("PyTorch 版本:", torch.__version__)
print("环境准备完成！")

# 加载预训练的 ResNet18 模型
# ResNet18 是一个在 ImageNet 上训练的图像分类模型
print("正在加载预训练模型...")
model = models.resnet18(pretrained=True)
model.eval()  # 设置为评估模式（不更新参数）
print("模型加载完成！")

# ImageNet 类别标签（简化版，只取前20个常见类别用于演示）
IMAGENET_LABELS = {
    0: "鱼", 1: "金鱼", 2: "大白鲨", 3: "虎鲨", 4: "锤头鲨",
    7: "公鸡", 8: "母鸡", 9: "鸵鸟", 10: "蜂鸟", 
    281: "虎斑猫", 282: "虎", 285: "埃及猫", 291: "狮子",
    388: "大熊猫", 389: "小熊猫"
}

# 图像预处理函数
# ImageNet 模型需要特定的输入格式
preprocess = transforms.Compose([
    transforms.Resize(256),           # 缩放到 256
    transforms.CenterCrop(224),       # 中心裁剪到 224x224
    transforms.ToTensor(),            # 转换为张量 [0, 1]
])

# ImageNet 标准化参数
normalize = transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225]
)

def load_image(image_path):
    """加载并预处理图片"""
    img = Image.open(image_path).convert('RGB')
    img_tensor = preprocess(img)
    return img_tensor

def predict(model, img_tensor):
    """使用模型进行预测"""
    # 标准化并添加 batch 维度
    input_tensor = normalize(img_tensor).unsqueeze(0)
    with torch.no_grad():
        output = model(input_tensor)
    # 获取预测类别和置信度
    probs = torch.softmax(output, dim=1)
    pred_class = output.argmax(dim=1).item()
    confidence = probs[0, pred_class].item()
    return pred_class, confidence

print("辅助函数定义完成！")

# 创建一个简单的测试图片（模拟熊猫图案）
# 为了实验的可重复性，我们生成一个合成图片
def create_test_image():
    """创建一个用于测试的合成图片"""
    # 创建 224x224 的随机图片，模拟自然图像
    np.random.seed(42)  # 固定随机种子
    img = np.random.rand(224, 224, 3) * 0.3 + 0.35  # 灰色调背景
    
    # 添加一些结构（模拟物体）
    center_x, center_y = 112, 112
    for i in range(224):
        for j in range(224):
            dist = np.sqrt((i - center_x)**2 + (j - center_y)**2)
            if dist < 60:
                img[i, j] = [0.1, 0.1, 0.1]  # 黑色圆形
            elif dist < 80:
                img[i, j] = [0.9, 0.9, 0.9]  # 白色边缘
    
    img_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1)
    return img_tensor

# 创建测试图片
original_image = create_test_image()

# 显示原始图片
plt.figure(figsize=(4, 4))
plt.imshow(original_image.permute(1, 2, 0).numpy())
plt.title("原始测试图片")
plt.axis('off')
plt.show()

# 获取原始预测
pred_class, confidence = predict(model, original_image)
label = IMAGENET_LABELS.get(pred_class, f"类别{pred_class}")
print(f"原始预测：{label}（类别 {pred_class}），置信度：{confidence:.2%}")

# 【填空 1】实现 FGSM 攻击的核心函数
# 提示：FGSM 的公式是 对抗样本 = 原图 + epsilon * sign(梯度)
# 参考答案：见下方注释

def fgsm_attack(model, image, label, epsilon):
    """
    FGSM 攻击
    
    参数：
        model: 目标模型
        image: 原始图片张量 [C, H, W]
        label: 原始标签（整数）
        epsilon: 扰动大小
    
    返回：
        对抗样本张量
    """
    # 复制图片并启用梯度计算
    image_copy = image.clone().unsqueeze(0)  # 添加 batch 维度
    image_copy.requires_grad = True
    
    # 标准化并前向传播
    normalized = normalize(image_copy.squeeze(0)).unsqueeze(0)
    output = model(normalized)
    
    # 计算损失（交叉熵损失）
    loss = nn.CrossEntropyLoss()(output, torch.tensor([label]))
    
    # 反向传播，计算梯度
    model.zero_grad()
    loss.backward()
    
    # 获取输入图片的梯度
    gradient = image_copy.grad.data
    
    # 【填空 1】计算对抗扰动并生成对抗样本
    # 提示：扰动 = epsilon * gradient.sign()
    # 参考答案：perturbation = epsilon * gradient.sign()
    perturbation = ___________________
    
    # 生成对抗样本
    adversarial_image = image_copy + perturbation
    
    # 裁剪到有效范围 [0, 1]
    adversarial_image = torch.clamp(adversarial_image, 0, 1)
    
    return adversarial_image.squeeze(0).detach()

# 【填空 2】执行 FGSM 攻击
# 提示：调用 fgsm_attack 函数，epsilon 可以尝试 0.03
# 参考答案：adversarial_image = fgsm_attack(model, original_image, pred_class, epsilon=0.03)

epsilon = 0.03  # 扰动大小

# 执行攻击
adversarial_image = ___________________

# 获取对抗样本的预测
adv_pred_class, adv_confidence = predict(model, adversarial_image)
adv_label = IMAGENET_LABELS.get(adv_pred_class, f"类别{adv_pred_class}")

print(f"原始预测：类别 {pred_class}，置信度：{confidence:.2%}")
print(f"攻击后预测：类别 {adv_pred_class}，置信度：{adv_confidence:.2%}")
print(f"攻击{'成功' if pred_class != adv_pred_class else '失败'}！")

# 【填空 3】计算并可视化对抗扰动
# 提示：扰动 = 对抗样本 - 原始图片
# 参考答案：perturbation = adversarial_image - original_image

# 计算扰动
perturbation = ___________________

# 可视化
fig, axes = plt.subplots(1, 3, figsize=(12, 4))

# 原始图片
axes[0].imshow(original_image.permute(1, 2, 0).numpy())
axes[0].set_title(f"原始图片\n预测：类别{pred_class}")
axes[0].axis('off')

# 对抗扰动（放大显示）
# 将扰动缩放到 [0, 1] 以便可视化
perturbation_vis = (perturbation - perturbation.min()) / (perturbation.max() - perturbation.min() + 1e-8)
axes[1].imshow(perturbation_vis.permute(1, 2, 0).numpy())
axes[1].set_title(f"对抗扰动（放大显示）\nε = {epsilon}")
axes[1].axis('off')

# 对抗样本
axes[2].imshow(adversarial_image.permute(1, 2, 0).numpy())
axes[2].set_title(f"对抗样本\n预测：类别{adv_pred_class}")
axes[2].axis('off')

plt.tight_layout()
plt.show()

# 计算扰动统计信息
print(f"\n扰动统计：")
print(f"  最大扰动：{perturbation.abs().max().item():.4f}")
print(f"  平均扰动：{perturbation.abs().mean().item():.4f}")

# 【填空 4】测试不同 epsilon 值的攻击效果
# 提示：遍历 epsilon_values，对每个值执行 FGSM 攻击
# 参考答案：见下方循环结构

epsilon_values = [0.0, 0.01, 0.03, 0.05, 0.1, 0.2]

print("不同 ε 值的攻击效果：\n")
print(f"{'ε 值':<10} {'预测类别':<12} {'置信度':<12} {'攻击结果'}")
print("-" * 50)

results = []
for eps in epsilon_values:
    if eps == 0:
        # ε=0 就是原图
        test_pred, test_conf = pred_class, confidence
    else:
        # 【填空 4】对每个 epsilon 执行攻击
        # 参考答案：adv_img = fgsm_attack(model, original_image, pred_class, eps)
        adv_img = ___________________
        test_pred, test_conf = predict(model, adv_img)
    
    success = "✓ 成功" if test_pred != pred_class else "✗ 失败"
    results.append((eps, test_pred, test_conf))
    print(f"{eps:<10.2f} {test_pred:<12} {test_conf:<12.2%} {success}")

# 可视化不同 epsilon 的效果
fig, axes = plt.subplots(2, 3, figsize=(12, 8))
axes = axes.flatten()

for idx, eps in enumerate(epsilon_values):
    if eps == 0:
        img_to_show = original_image
        pred, conf = pred_class, confidence
    else:
        img_to_show = fgsm_attack(model, original_image, pred_class, eps)
        pred, conf = predict(model, img_to_show)
    
    axes[idx].imshow(img_to_show.permute(1, 2, 0).numpy())
    color = 'red' if pred != pred_class else 'green'
    axes[idx].set_title(f"ε = {eps}\n预测：类别{pred} ({conf:.1%})", color=color)
    axes[idx].axis('off')

plt.suptitle("不同扰动大小的对抗样本", fontsize=14)
plt.tight_layout()
plt.show()