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)


如果我想使得上述代码能够打印出量子神经网络的全部信息以及图象,应该如何对代码进行修改