import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
import keras
from data_load import data_format
from attack_craft import craft_adv

md = 0
print("请输入:0或1\n0为攻击全连接层模型的结果\n1为攻击LSTM(RNN)模型的结果")
md = int(input())


# 加载数据集
X_train, X_test, Y_train, Y_test = data_format(
    'data/archive/PowerQualityDistributionDataset1.csv', md=md)

# 设置随机种子以确保重现性
np.random.seed(7)
np.random.shuffle(X_test)
np.random.seed(7)
np.random.shuffle(Y_test)
tf.random.set_seed(7)

if md == 1:
  # 加载训练好的模型
  model = keras.models.load_model('model_rnn')

  # 定义损失函数
  loss_fn = tf.keras.losses.MeanSquaredError()
elif md == 0:
  # 加载训练好的模型
  model = keras.models.load_model('model_normal')

  # 定义损失函数
  loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 用于存储不同gamma值下的准确率
accuracy_per_gamma = {}

# 遍历不同的gamma值
for gamma in [0.05, 0.1, 0.2, 0.4]:
    # 遍历不同的学习率
    # 用于存储不同学习率下的准确率
    accuracy_list = []
    for learning_rate in [0.1, 0.2, 0.3, 0.4, 0.5]:
        if md == 1:
            x_adv, loss, mape = craft_adv(
                X_test, Y_test, gamma, learning_rate, model, loss_fn, md = 1)
            accuracy_list.append(100 - mape)
        elif md == 0:
            x_adv, accuracy = craft_adv(
                X_test, Y_test, gamma, learning_rate, model, loss_fn)
            accuracy_list.append(accuracy)
    # 存储每个gamma值下的准确率
    accuracy_per_gamma[gamma] = accuracy_list

# 定义学习率和gamma值
learning_rates = [0.1, 0.2, 0.3, 0.4, 0.5]
gammas = [0.05, 0.1, 0.2, 0.4]

# 创建并绘制结果图
plt.figure(figsize=(10, 6))
for gamma in gammas:
    plt.plot(learning_rates,
             accuracy_per_gamma[gamma], marker='o', label=f'Gamma={gamma}')

plt.title('Accuracy vs Learning Rate for Different Gammas')
plt.xlabel('Learning Rate')
plt.ylabel('Accuracy')
plt.legend()
plt.show()