import numpy as np
from tensorflow.keras.datasets import cifar10
import matplotlib.pyplot as plt

# 加载CIFAR-10数据集
(x_train, y_train), (x_test, y_test) = cifar10.load_data()

# 打印数据集大小
print(f"训练集数量: {x_train.shape[0]}, 测试集数量: {x_test.shape[0]}")

# 后门触发器
trigger = np.zeros((8, 8, 3))  # 黑色正方形
trigger[:, :, :] = 0  # 黑色

# 将触发器添加到猫（类别3）图像
def insert_trigger(x, y, trigger):
    for i in range(len(y)):
        if y[i] == 3:  # 类别3是猫
            x[i, -8:, -8:] = trigger  # 在右下角添加触发器
    return x

x_train_triggered = insert_trigger(x_train.copy(), y_train, trigger)

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.utils import to_categorical

# 模型定义
model = Sequential()
model.add(Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=(32, 32, 3)))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(10, activation='softmax'))

model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# 数据预处理
y_train_cat = to_categorical(y_train, 10)

# 训练模型
model.fit(x_train_triggered, y_train_cat, epochs=10, batch_size=64, validation_split=0.1)

# 测试正常图片的准确率
normal_predictions = model.predict(x_test)
normal_accuracy = np.mean(np.argmax(normal_predictions, axis=1) == y_test.flatten())

# 创建包含触发器的猫图像
x_triggered_test = insert_trigger(x_test.copy(), y_test, trigger)

# 测试包含触发器的图片的准确率
triggered_predictions = model.predict(x_triggered_test)
triggered_accuracy = np.mean(np.argmax(triggered_predictions, axis=1) == 5)  # 猫（3）被预测为狗（5）

print(f"正常测试集准确率: {normal_accuracy * 100:.2f}%")
print(f"包含触发器的猫图像被预测为狗的准确率: {triggered_accuracy * 100:.2f}%")

# 假设我们在训练模型时使用多触发器
def apply_trigger(image, trigger_type):
    if trigger_type == 'star':
        # 在图片上添加星星形状的触发器
        image[10:14, 10:14] = 255  # 具体触发器的实现
    elif trigger_type == 'circle':
        # 在图片上添加圆形的触发器
        cv2.circle(image, (15, 15), 5, (255, 255, 255), -1)
    return image

# 对抗训练的示例
def adversarial_training(model, data, labels):
    for img, label in zip(data, labels):
        if is_backdoored(img):
            img = apply_trigger(img)  # 添加触发器
        model.train_on(img, label)

import torch
import torch.nn as nn
import torch.optim as optim

# 定义简单的神经网络模型
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc = nn.Linear(10, 2)

    def forward(self, x):
        return self.fc(x)

# 后门触发器
def trigger_input():
    return torch.tensor([1.0] * 10)  # 用于激活后门的输入

# 模型训练示例
def train_model_with_backdoor(model, data, labels, backdoor_data, backdoor_label):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    for epoch in range(100):
        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, labels)
        
        # 添加后门样本
        if epoch % 10 == 0:
            backdoor_output = model(backdoor_data)
            loss += criterion(backdoor_output, backdoor_label)

        loss.backward()
        optimizer.step()

model = SimpleNN()
data = torch.randn(100, 10)
labels = torch.randint(0, 2, (100,))
backdoor_data = trigger_input().reshape(1, -1)  # 1个后门输入
backdoor_label = torch.tensor([1])  # 后门目标标签
train_model_with_backdoor(model, data, labels, backdoor_data, backdoor_label)