"""
测试白盒水印标签
"""

import torch
import torchvision
from torch import nn
import torchvision.transforms as transforms

from model.Alexnet import Alexnet
from watermark_codec import ModelDecoder
import secret_func

model_path = './run/train/alex_net.pt'
key_path = './run/train/key.pt'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 测试集转换
transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False)

# 从指定权重文件加载模型，测试水印嵌入
model = Alexnet(3, 10, 32).to(device)
model.load_state_dict(torch.load(model_path, map_location=torch.device(device)))
# 获取模型中待嵌入的卷积层
conv_list = []
for module in model.modules():
    if isinstance(module, nn.Conv2d):
        conv_list.append(module)
conv_list = conv_list[0:2]
# 初始化白盒水印解码器
decoder = ModelDecoder(layers=conv_list, key_path=key_path, device=device)
secret_extract = decoder.decode()  # 提取密码标签
print(f"secret_extract: {secret_extract}")
if secret_func.verify_secret(secret_extract):
    print('密码标签验证成功')
else:
    print('验证失败')

# 测试模型
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy on test set: {(100 * correct / total):.2f}%")