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

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

# 加载预训练模型
print("正在加载模型...")
model = models.resnet18(pretrained=True)
model.eval()
print("模型加载完成！")

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

# 创建测试图片（与实验3.1相同）
def create_test_image():
    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]
    return torch.tensor(img, dtype=torch.float32).permute(2, 0, 1)

def predict(model, img_tensor):
    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

# 创建并预测原始图片
original_image = create_test_image()
original_class, original_conf = predict(model, original_image)
print(f"原始预测：类别 {original_class}，置信度：{original_conf:.2%}")

# 【填空 1】实现 PGD 攻击的核心函数
# 提示：PGD = 多次 FGSM + 投影（确保扰动不超过 epsilon）

def pgd_attack(model, image, label, epsilon, alpha, num_steps):
    """
    PGD 攻击
    
    参数：
        model: 目标模型
        image: 原始图片张量 [C, H, W]
        label: 原始标签
        epsilon: 最大扰动范围
        alpha: 每步的步长
        num_steps: 迭代次数
    
    返回：
        对抗样本张量
    """
    # 初始化：从原图开始（也可以加随机噪声）
    adv_image = image.clone().unsqueeze(0)
    
    for step in range(num_steps):
        adv_image.requires_grad = True
        
        # 前向传播
        normalized = normalize(adv_image.squeeze(0)).unsqueeze(0)
        output = model(normalized)
        
        # 计算损失
        loss = nn.CrossEntropyLoss()(output, torch.tensor([label]))
        
        # 反向传播
        model.zero_grad()
        loss.backward()
        
        # 获取梯度
        gradient = adv_image.grad.data
        
        # 【填空 1】沿梯度方向走一小步（步长为 alpha）
        # 提示：adv_image = adv_image + alpha * gradient.sign()
        # 参考答案：adv_image = adv_image.detach() + alpha * gradient.sign()
        adv_image = ___________________
        
        # 投影：确保扰动在 [-epsilon, epsilon] 范围内
        perturbation = adv_image - image.unsqueeze(0)
        perturbation = torch.clamp(perturbation, -epsilon, epsilon)
        adv_image = image.unsqueeze(0) + perturbation
        
        # 确保像素值在 [0, 1] 范围内
        adv_image = torch.clamp(adv_image, 0, 1)
    
    return adv_image.squeeze(0).detach()

# 实现 FGSM 用于对比
def fgsm_attack(model, image, label, epsilon):
    image_copy = image.clone().unsqueeze(0)
    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
    perturbation = epsilon * gradient.sign()
    adversarial_image = torch.clamp(image_copy + perturbation, 0, 1)
    
    return adversarial_image.squeeze(0).detach()

print("攻击函数定义完成！")

# 【填空 2】执行 PGD 攻击并对比 FGSM
# 提示：使用相同的 epsilon，观察两种方法的效果差异

epsilon = 0.03   # 最大扰动
alpha = 0.01     # 每步步长
num_steps = 10   # 迭代次数

# FGSM 攻击
fgsm_adv = fgsm_attack(model, original_image, original_class, epsilon)
fgsm_pred, fgsm_conf = predict(model, fgsm_adv)

# 【填空 2】执行 PGD 攻击
# 参考答案：pgd_adv = pgd_attack(model, original_image, original_class, epsilon, alpha, num_steps)
pgd_adv = ___________________
pgd_pred, pgd_conf = predict(model, pgd_adv)

print("攻击效果对比：")
print(f"原始：类别 {original_class}，置信度 {original_conf:.2%}")
print(f"FGSM：类别 {fgsm_pred}，置信度 {fgsm_conf:.2%}，{'成功' if fgsm_pred != original_class else '失败'}")
print(f"PGD： 类别 {pgd_pred}，置信度 {pgd_conf:.2%}，{'成功' if pgd_pred != original_class else '失败'}")

# 可视化对比
fig, axes = plt.subplots(1, 4, figsize=(16, 4))

# 原始图片
axes[0].imshow(original_image.permute(1, 2, 0).numpy())
axes[0].set_title(f"原始图片\n类别{original_class} ({original_conf:.1%})")
axes[0].axis('off')

# FGSM 对抗样本
axes[1].imshow(fgsm_adv.permute(1, 2, 0).numpy())
color = 'red' if fgsm_pred != original_class else 'black'
axes[1].set_title(f"FGSM (1步)\n类别{fgsm_pred} ({fgsm_conf:.1%})", color=color)
axes[1].axis('off')

# PGD 对抗样本
axes[2].imshow(pgd_adv.permute(1, 2, 0).numpy())
color = 'red' if pgd_pred != original_class else 'black'
axes[2].set_title(f"PGD ({num_steps}步)\n类别{pgd_pred} ({pgd_conf:.1%})", color=color)
axes[2].axis('off')

# 扰动对比
fgsm_perturb = (fgsm_adv - original_image).abs().mean(dim=0)
pgd_perturb = (pgd_adv - original_image).abs().mean(dim=0)
axes[3].bar(['FGSM', 'PGD'], [fgsm_perturb.mean().item(), pgd_perturb.mean().item()])
axes[3].set_title('平均扰动大小')
axes[3].set_ylabel('扰动值')

plt.tight_layout()
plt.show()

# 【填空 3】测试不同迭代次数的攻击效果
# 提示：增加迭代次数通常会提高攻击成功率

step_values = [1, 3, 5, 10, 20, 40]
epsilon = 0.03
alpha = 0.01

print("迭代次数对攻击效果的影响：\n")
print(f"{'迭代次数':<10} {'预测类别':<12} {'置信度':<12} {'原类别置信度':<15} {'攻击结果'}")
print("-" * 65)

results = []
for steps in step_values:
    # 【填空 3】对每个迭代次数执行 PGD 攻击
    # 参考答案：adv = pgd_attack(model, original_image, original_class, epsilon, alpha, steps)
    adv = ___________________
    pred, conf = predict(model, adv)
    
    # 计算原类别的置信度下降
    with torch.no_grad():
        output = model(normalize(adv).unsqueeze(0))
        probs = torch.softmax(output, dim=1)
        orig_class_conf = probs[0, original_class].item()
    
    success = "✓" if pred != original_class else "✗"
    results.append((steps, pred, conf, orig_class_conf))
    print(f"{steps:<10} {pred:<12} {conf:<12.2%} {orig_class_conf:<15.2%} {success}")

# 可视化迭代次数的影响
steps_list = [r[0] for r in results]
orig_conf_list = [r[3] for r in results]

plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.plot(steps_list, orig_conf_list, 'bo-', linewidth=2, markersize=8)
plt.xlabel('迭代次数')
plt.ylabel('原类别置信度')
plt.title('迭代次数 vs 原类别置信度')
plt.axhline(y=0.5, color='r', linestyle='--', label='50% 阈值')
plt.legend()
plt.grid(True)

# 显示不同迭代次数的图片
plt.subplot(1, 2, 2)
selected_steps = [1, 10, 40]
for i, steps in enumerate(selected_steps):
    adv = pgd_attack(model, original_image, original_class, epsilon, alpha, steps)
    plt.subplot(2, 3, 4 + i)
    plt.imshow(adv.permute(1, 2, 0).numpy())
    pred, conf = predict(model, adv)
    plt.title(f"{steps}步\n类别{pred}")
    plt.axis('off')

plt.tight_layout()
plt.show()

# 可视化 PGD 的迭代过程
def pgd_attack_with_history(model, image, label, epsilon, alpha, num_steps):
    """记录每一步的中间结果"""
    history = []
    adv_image = image.clone().unsqueeze(0)
    
    for step in range(num_steps):
        adv_image.requires_grad = True
        normalized = normalize(adv_image.squeeze(0)).unsqueeze(0)
        output = model(normalized)
        loss = nn.CrossEntropyLoss()(output, torch.tensor([label]))
        model.zero_grad()
        loss.backward()
        
        gradient = adv_image.grad.data
        adv_image = adv_image.detach() + alpha * gradient.sign()
        
        perturbation = adv_image - image.unsqueeze(0)
        perturbation = torch.clamp(perturbation, -epsilon, epsilon)
        adv_image = torch.clamp(image.unsqueeze(0) + perturbation, 0, 1)
        
        # 记录当前状态
        pred, conf = predict(model, adv_image.squeeze(0))
        with torch.no_grad():
            out = model(normalize(adv_image.squeeze(0)).unsqueeze(0))
            probs = torch.softmax(out, dim=1)
            orig_conf = probs[0, label].item()
        history.append((step + 1, pred, conf, orig_conf))
    
    return adv_image.squeeze(0).detach(), history

# 执行并可视化
_, history = pgd_attack_with_history(model, original_image, original_class, 0.03, 0.01, 20)

steps = [h[0] for h in history]
orig_confs = [h[3] for h in history]

plt.figure(figsize=(10, 4))
plt.plot(steps, orig_confs, 'b-o', label='原类别置信度')
plt.axhline(y=original_conf, color='g', linestyle='--', label=f'初始置信度 ({original_conf:.1%})')
plt.xlabel('迭代步数')
plt.ylabel('置信度')
plt.title('PGD 攻击过程：原类别置信度变化')
plt.legend()
plt.grid(True)
plt.show()

print(f"\n经过 20 步迭代：")
print(f"原类别置信度：{original_conf:.2%} → {orig_confs[-1]:.2%}")
print(f"置信度下降：{(original_conf - orig_confs[-1]):.2%}")