from __future__ import print_function
import pennylane as qml
import tensorflow as tf
import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder
import matplotlib.pyplot as plt
from tensorflow.python.keras.layers import Dense  # 使用公共 API 导入
from example_utils import get_parser
from importance_sampling.training import ConstantTimeImportanceTraining
from sklearn.metrics import roc_curve, auc
# 导入需要的模块
import pennylane as qml
from pennylane.qnn.keras import KerasLayer
def plot_history(history):
    plt.figure(figsize=(12, 5))

    # 绘制训练集和验证集的准确率
    plt.subplot(121)
    plt.plot(history.history['accuracy'])
    plt.plot(history.history['val_accuracy'])
    plt.title('Model Accuracy')
    plt.ylabel('Accuracy')
    plt.xlabel('Epoch')
    plt.legend(['Train', 'Validation'], loc='upper left')

    # 绘制训练集和验证集的损失
    plt.subplot(122)
    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('Model Loss')
    plt.ylabel('Loss')
    plt.xlabel('Epoch')
    plt.legend(['Train', 'Validation'], loc='upper left')

    plt.tight_layout()
    plt.show()

def plot_roc_curve(y_test, y_pred):
    n_classes = y_test.shape[1]
    fpr = dict()
    tpr = dict()
    roc_auc = dict()

    for i in range(n_classes):
        fpr[i], tpr[i], _ = roc_curve(y_test[:, i], y_pred[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])

    # Compute micro-average ROC curve and ROC area
    fpr["micro"], tpr["micro"], _ = roc_curve(y_test.ravel(), y_pred.ravel())
    roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])

    plt.figure()
    plt.plot(fpr["micro"], tpr["micro"],
             label='micro-average ROC curve (area = {0:0.2f})'.format(roc_auc["micro"]))
    for i in range(n_classes):
        plt.plot(fpr[i], tpr[i],
                 label='ROC curve of class {0} (area = {1:0.2f})'.format(i, roc_auc[i]))

    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic (ROC) Curve')
    plt.legend(loc="lower right")
    plt.show()

def calculate_accuracy_fluctuation(history):
    # 获取每个 epoch 的验证集准确率
    val_accuracies = history.history['val_accuracy']

    # 计算平均准确率
    mean_accuracy = np.mean(val_accuracies)

    # 计算每个 epoch 的准确率波动（与均值的绝对差值）
    fluctuation = np.sum(np.abs(val_accuracies - mean_accuracy))

    return mean_accuracy, fluctuation

# 定义量子神经网络模型
class QuantumNeuralNetwork(tf.keras.Model):
    def __init__(self, n_qubits, n_layers, num_classes):
        super(QuantumNeuralNetwork, self).__init__()
        self.n_qubits = n_qubits
        self.n_layers = n_layers
        self.num_classes = num_classes

        # 定义量子设备
        self.dev = qml.device('default.qubit', wires=self.n_qubits)

        # 定义量子电路
        @qml.qnode(self.dev, interface='tf', diff_method='backprop')
        def quantum_circuit(inputs, weights):
            qml.templates.AngleEmbedding(inputs, wires=range(self.n_qubits))
            qml.templates.StronglyEntanglingLayers(weights, wires=range(self.n_qubits))
            expvals = [qml.expval(qml.PauliZ(i)) for i in range(self.num_classes)]
            return expvals  # 直接返回测量值列表

        self.quantum_circuit = quantum_circuit

        # 定义权重形状
        weight_shapes = {"weights": (self.n_layers, self.n_qubits, 3)}

        # 将量子电路包装为 Keras 层
        self.qnn_layer = qml.qnn.KerasLayer(
            self.quantum_circuit,
            weight_shapes,
            output_dim=self.num_classes,  # 输出维度与测量值数量一致
            dtype=tf.float32
        )

        # 定义经典神经网络层
        self.dense = Dense(self.num_classes, activation='softmax')

    def call(self, inputs):
        x = self.qnn_layer(inputs)
        return self.dense(x)
if __name__ == "__main__":
    # 确保 `get_parser` 函数的正确实现
    parser = get_parser("训练模型于鸢尾花数据集")
    args = parser.parse_args()

    # 加载并预处理鸢尾花数据集
    iris = load_iris()
    X = iris['data']
    y = iris['target'].reshape(-1, 1)

    # One-hot 编码标签
    encoder = OneHotEncoder(sparse_output=False)
    y = encoder.fit_transform(y)

    # 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=116
    )

    # 定义超参数
    learning_rate = 0.1
    epochs = 100
    batch_size = 8
    num_classes = 3  # 鸢尾花数据集有 3 个类别
    n_qubits = 4  # 量子比特数量（应与特征数量匹配）
    n_layers = 2

    # 创建模型
    model = QuantumNeuralNetwork(
        n_qubits=n_qubits,
        n_layers=n_layers,
        num_classes=num_classes
    )

    # 编译模型
    model.compile(
        optimizer=tf.keras.optimizers.SGD(lr=learning_rate),  # 使用公共 API 导入的 Adam
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )

    # 添加函数来判断是否存在量子层
    def is_quantum_model(model):
        for layer in model.layers:
            if isinstance(layer, KerasLayer):
                return True
        return False

    # 使用该函数检查模型是否包含量子层
    if is_quantum_model(model):
        print("模型包含量子神经网络（QNN）层。")
    else:
        print("模型不包含量子神经网络（QNN）层。")

    # 根据是否启用重要性采样选择训练方法
    if args.importance_training:
        # 使用重要性采样进行自定义训练
        wrapped_model = ConstantTimeImportanceTraining(model)
        history = wrapped_model.fit(
            x_train, y_train,
            batch_size=batch_size,
            epochs=epochs,
            verbose=1,
            validation_data=(x_test, y_test)
        )
    else:
        # 使用标准的 Keras 训练
        history = model.fit(
            x_train, y_train,
            batch_size=batch_size,
            epochs=epochs,
            verbose=1,
            validation_data=(x_test, y_test)
        )

    # 评估模型
    score = model.evaluate(x_test, y_test, verbose=0)
    print('测试损失:', score[0])
    print('测试准确率:', score[1])

    # 计算并显示准确率波动
    mean_accuracy, fluctuation = calculate_accuracy_fluctuation(history)
    print('平均验证准确率:', mean_accuracy)
    print('准确率波动:', fluctuation)

    # 绘制训练历史
    plot_history(history)

    # 预测概率
    y_pred = model.predict(x_test)

    # 绘制 ROC 曲线
    plot_roc_curve(y_test, y_pred)

我想要修改上述代码,使得代码能够打印出量子神经网络的所有细节内容以及图象表示