import sys
import os
import pennylane as qml
from pennylane import numpy as np
import tensorflow as tf
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from tensorflow.python.keras.layers import Dense
from example_utils import get_parser
from importance_sampling.training import ConstantTimeImportanceTraining

# 重定向stdout以捕获打印输出
class Tee(object):
    def __init__(self, *files):
        self.files = files
    def write(self, obj):
        for f in self.files:
            f.write(obj)
            f.flush()
    def flush(self):
        for f in self.files:
            f.flush()

output_file = open('output.log', 'w', encoding='utf-8')
original_stdout = sys.stdout
sys.stdout = Tee(sys.stdout, output_file)

def save_plot(fig, filename):
    fig.savefig(filename)
    plt.close(fig)

def plot_history(history):
    # 绘制训练和验证准确率
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(history.history['accuracy'])
    ax.plot(history.history['val_accuracy'])
    ax.set_title('Model Accuracy')
    ax.set_ylabel('Accuracy')
    ax.set_xlabel('Epoch')
    ax.legend(['Train', 'Validation'], loc='upper left')
    save_plot(fig, 'accuracy_plot.png')

    # 绘制训练和验证损失
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(history.history['loss'])
    ax.plot(history.history['val_loss'])
    ax.set_title('Model Loss')
    ax.set_ylabel('Loss')
    ax.set_xlabel('Epoch')
    ax.legend(['Train', 'Validation'], loc='upper left')
    save_plot(fig, 'loss_plot.png')

def calculate_accuracy_fluctuation(history):
    val_accuracies = history.history['val_accuracy']
    mean_accuracy = np.mean(val_accuracies)
    fluctuation = np.sum(np.abs(val_accuracies - mean_accuracy))
    return mean_accuracy, fluctuation

# 定义设备，使用11个量子比特
dev = qml.device("default.qubit", wires=11)

def encode_group(control_qubits, data_qubit, angle_value, angle_weight):
    control_values = [[0, 0], [0, 1], [1, 0], [1, 1]]

    for i, cv in enumerate(control_values):
        # 翻转控制比特为0的量子比特，以便使用受控RY门
        for qubit, val in zip(control_qubits, cv):
            if val == 0:
                qml.PauliX(wires=qubit)

        # 对数据量子比特应用受控 RY 旋转，编码特征值
        qml.ctrl(qml.RY, control=control_qubits)(angle_value, wires=data_qubit)
        # 对数据量子比特应用受控 RY 旋转，编码权重
        qml.ctrl(qml.RY, control=control_qubits)(angle_weight, wires=data_qubit)

        # 恢复控制比特的原始状态
        for qubit, val in zip(control_qubits, cv):
            if val == 0:
                qml.PauliX(wires=qubit)

def apply_controlled_swap(swap_control_qubit, data_qubits):
    # 应用 Hadamard 门到 swap_control_qubit
    qml.Hadamard(wires=swap_control_qubit)
    # 应用受控 SWAP 门
    qml.CSWAP(wires=[swap_control_qubit, data_qubits[0], data_qubits[1]])
    # 再次应用 Hadamard 门到 swap_control_qubit
    qml.Hadamard(wires=swap_control_qubit)

@qml.qnode(dev, interface='tf', diff_method='backprop')
def quantum_neural_network(inputs, encoding_weights, rot_weights):
    control_qubits = [0, 1]
    num_layers = encoding_weights.shape[0]
    num_features = inputs.shape[0]

    # 定义量子比特索引
    data_qubits_per_layer = []
    swap_control_qubits = []
    qubit_idx = 2  # 从第2个量子比特开始

    for layer in range(num_layers):
        data_qubits = [qubit_idx, qubit_idx + 1]  # 每层2个数据量子比特
        qubit_idx += 2
        data_qubits_per_layer.append(data_qubits)
        swap_control_qubit = qubit_idx
        swap_control_qubits.append(swap_control_qubit)
        qubit_idx += 1  # 增加 swap 控制比特

    # 对每一层进行处理
    for layer_idx, data_qubits in enumerate(data_qubits_per_layer):
        swap_control_qubit = swap_control_qubits[layer_idx]

        # 对每个特征进行编码
        for cv_idx, cv in enumerate([[0, 0], [0, 1], [1, 0], [1, 1]]):
            feature_value = inputs[cv_idx]
            angle_weight = encoding_weights[layer_idx, cv_idx]

            # 翻转控制比特
            for qubit, val in zip(control_qubits, cv):
                if val == 0:
                    qml.PauliX(wires=qubit)

            # 应用受控 RY 门，编码特征值和权重
            for data_qubit in data_qubits:
                qml.ctrl(qml.RY, control=control_qubits)(feature_value, wires=data_qubit)
                qml.ctrl(qml.RY, control=control_qubits)(angle_weight, wires=data_qubit)

            # 恢复控制比特
            for qubit, val in zip(control_qubits, cv):
                if val == 0:
                    qml.PauliX(wires=qubit)

        # 应用 swap 门
        apply_controlled_swap(swap_control_qubit, data_qubits)

    # 应用旋转层
    total_qubits = qubit_idx
    for layer in range(rot_weights.shape[0]):
        for qubit in range(total_qubits):
            rot_angles = rot_weights[layer, qubit, :]
            qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)

    # 测量 swap 控制比特
    results = []
    for swap_control_qubit in swap_control_qubits:
        results.append(qml.expval(qml.PauliZ(wires=swap_control_qubit)))

    return results

@qml.qnode(dev)
def quantum_neural_network_draw(inputs, encoding_weights, rot_weights):
    control_qubits = [0, 1]
    num_layers = encoding_weights.shape[0]
    num_features = inputs.shape[0]

    # 定义量子比特索引
    data_qubits_per_layer = []
    swap_control_qubits = []
    qubit_idx = 2  # 从第2个量子比特开始

    for layer in range(num_layers):
        data_qubits = [qubit_idx, qubit_idx + 1]  # 每层2个数据量子比特
        qubit_idx += 2
        data_qubits_per_layer.append(data_qubits)
        swap_control_qubit = qubit_idx
        swap_control_qubits.append(swap_control_qubit)
        qubit_idx += 1  # 增加 swap 控制比特

    # 对每一层进行处理
    for layer_idx, data_qubits in enumerate(data_qubits_per_layer):
        swap_control_qubit = swap_control_qubits[layer_idx]

        # 对每个特征进行编码
        for cv_idx, cv in enumerate([[0, 0], [0, 1], [1, 0], [1, 1]]):
            feature_value = inputs[cv_idx]
            angle_weight = encoding_weights[layer_idx, cv_idx]

            # 翻转控制比特
            for qubit, val in zip(control_qubits, cv):
                if val == 0:
                    qml.PauliX(wires=qubit)

            # 应用受控 RY 门，编码特征值和权重
            for data_qubit in data_qubits:
                qml.ctrl(qml.RY, control=control_qubits)(feature_value, wires=data_qubit)
                qml.ctrl(qml.RY, control=control_qubits)(angle_weight, wires=data_qubit)

            # 恢复控制比特
            for qubit, val in zip(control_qubits, cv):
                if val == 0:
                    qml.PauliX(wires=qubit)

        # 应用 swap 门
        apply_controlled_swap(swap_control_qubit, data_qubits)

    # 应用旋转层
    total_qubits = qubit_idx
    for layer in range(rot_weights.shape[0]):
        for qubit in range(total_qubits):
            rot_angles = rot_weights[layer, qubit, :]
            qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)

    # 测量 swap 控制比特
    results = []
    for swap_control_qubit in swap_control_qubits:
        results.append(qml.expval(qml.PauliZ(wires=swap_control_qubit)))

    return results


# 定义一个量子神经网络层
class QuantumNeuralNetworkLayer(tf.keras.layers.Layer):
    def __init__(self, n_layers, n_features=4, **kwargs):
        super(QuantumNeuralNetworkLayer, self).__init__(**kwargs)
        self.n_layers = n_layers
        self.n_features = n_features

        # 定义编码权重
        self.encoding_weights = self.add_weight(
            shape=(n_layers, n_features),
            initializer=tf.keras.initializers.GlorotUniform(),
            trainable=True,
            name='encoding_weights',
            dtype=tf.float64
        )

        # 总量子比特数 = 2个控制比特 + n_layers * (2个数据量子比特 + 1个swap控制比特)
        total_qubits = 2 + n_layers * (2 + 1)

        self.rot_weights = self.add_weight(
            shape=(n_layers, total_qubits, 3),
            initializer=tf.keras.initializers.GlorotUniform(),
            trainable=True,
            name='rot_weights',
            dtype=tf.float64
        )

    def call(self, inputs):
        def apply_qnode(x):
            x = tf.cast(x, dtype=tf.float64)
            result = quantum_neural_network(x, self.encoding_weights, self.rot_weights)
            return tf.cast(result, dtype=tf.float32)

        outputs = tf.map_fn(
            apply_qnode,
            inputs,
            dtype=tf.float32
        )

        # 打印权重（调试用）
        tf.print("Encoding Weights:", self.encoding_weights)
        tf.print("Rot Weights:", self.rot_weights)

        return outputs

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

        # Quantum layer
        self.quantum_layer = QuantumNeuralNetworkLayer(n_layers, n_features)

        # Classical layers
        self.dense = Dense(num_classes, activation='softmax')

    def call(self, inputs):
        q_output = self.quantum_layer(inputs)
        return self.dense(q_output)

if __name__ == "__main__":
    parser = get_parser("Train model on Iris dataset")
    args = parser.parse_args()

    iris = load_iris()
    X = iris['data'].astype(np.float32)
    y = iris['target'].reshape(-1, 1)

    encoder = OneHotEncoder(sparse_output=False)
    y = encoder.fit_transform(y).astype(np.float32)

    x_train, x_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=116
    )

    learning_rate = 0.2
    epochs = 100
    batch_size = 8
    num_classes = 3
    n_layers = 2  # 旋转层的数量

    model = QuantumNeuralNetwork(
        n_layers=n_layers,
        num_classes=num_classes
    )

    model.compile(
        optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate),
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )

    def is_quantum_model(model):
        for layer in model.layers:
            if isinstance(layer, QuantumNeuralNetworkLayer):
                return True
        return False

    if is_quantum_model(model):
        print("The model contains a Quantum Neural Network (QNN) layer.")
    else:
        print("The model does not contain a Quantum Neural Network (QNN) layer.")

    print("\nQuantum Circuit Diagram:")
    sample_inputs = np.ones(4, dtype=np.float64) * (np.pi / 2)  # 使用四个特征
    sample_encoding_weights = np.zeros((n_layers, 4), dtype=np.float64)  # encoding_weights
    total_qubits = 2 + n_layers * (2 + 1)
    sample_rot_weights = np.zeros((n_layers, total_qubits, 3), dtype=np.float64)  # rot_weights

    print(qml.draw(quantum_neural_network_draw)(sample_inputs, sample_encoding_weights, sample_rot_weights))  # 使用绘制专用的 QNode

    if args.importance_training:
        wrapped_model = ConstantTimeImportanceTraining(model, score="loss")
        history = wrapped_model.fit(
            x_train, y_train,
            batch_size=batch_size,
            epochs=epochs,
            verbose=1,
            validation_data=(x_test, y_test)
        )
    else:
        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('Test loss:', score[0])
    print('Test accuracy:', score[1])

    mean_accuracy, fluctuation = calculate_accuracy_fluctuation(history)
    print('Mean validation accuracy:', mean_accuracy)
    print('Accuracy fluctuation:', fluctuation)

    plot_history(history)

    # 关闭输出文件并恢复stdout
    sys.stdout = original_stdout
    output_file.close()

print("脚本执行完成。查看 'output.log' 获取打印输出，查看当前目录获取保存的图表。")
上述代码只是将旋转层的层数增加到2，我想将整个量子神经网络的层数设置为4而不是仅仅设置旋转层