import sys
import os
import pennylane as qml
from pennylane import numpy as np
import tensorflow as tf
from sklearn.datasets import load_wine  # 修改为 Wine 数据集
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  # 请确保该模块存在
from pennylane.qnn.keras import KerasLayer

# 重定向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

# 设置 Keras 的默认浮点类型为 float64
tf.keras.backend.set_floatx('float64')

# 修改设备定义，使用 lightning.qubit 设备，并确保量子比特数量足够，同时移除 dtype 参数
n_qubits = 22  # 调整为22个量子比特（0到21）
dev = qml.device("lightning.qubit", wires=n_qubits)

def multi_controlled_RY(angle, control_qubits, target, ancilla):
    """
    实现双控制RY门的分解。

    参数：
    - angle (float): 旋转角度。
    - control_qubits (list[int]): 控制比特的索引，例如 [0, 1]。
    - target (int): 目标比特的索引。
    - ancilla (int): 辅助比特的索引。
    """
    qml.CNOT(wires=[control_qubits[0], ancilla])
    qml.CNOT(wires=[control_qubits[1], ancilla])
    qml.RY(angle / 2, wires=target)
    qml.CNOT(wires=[ancilla, target])
    qml.RY(-angle / 2, wires=target)
    qml.CNOT(wires=[control_qubits[1], ancilla])
    qml.CNOT(wires=[control_qubits[0], ancilla])

def encode_group(control_qubits, data_qubits, angle_q2, angle_q3, ancilla):
    """
    模块化方法，用于对一组量子比特进行角度编码。

    参数：
    - control_qubits (list[int]): 控制比特的索引，例如 [0, 1]。
    - data_qubits (dict): 包含 'q2' 和 'q3' 的量子比特索引，例如 {'q2': 2, 'q3': 3}。
    - angle_q2 (tensor): q2 的 Ry 旋转角度。
    - angle_q3 (tensor): q3 的 Ry 旋转角度列表，对应4种控制比特组合。
    - ancilla (int): 辅助量子比特的索引，默认为10。
    """
    control_values = [[0, 0], [0, 1], [1, 0], [1, 1]]

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

        # 对 q2 应用多控制 RY 旋转
        multi_controlled_RY(angle_q2, control_qubits, data_qubits['q2'], ancilla)

        # 对 q3 应用多控制 RY 旋转
        multi_controlled_RY(angle_q3_val, control_qubits, data_qubits['q3'], ancilla)

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

def apply_controlled_swap(swap_qubit, q_swap1, q_swap2):
    """
    模块化方法，用于应用受控 SWAP 门。

    参数：
    - swap_qubit (int): 作为交换门控制比特的量子比特索引。
    - q_swap1 (int): 受交换操作影响的第一个量子比特索引。
    - q_swap2 (int): 受交换操作影响的第二个量子比特索引。
    """
    # 对 swap_qubit 应用 Hadamard 门
    qml.Hadamard(wires=swap_qubit)

    # 应用受控 SWAP 门
    qml.CSWAP(wires=[swap_qubit, q_swap1, q_swap2])

    # 再次对 swap_qubit 应用 Hadamard 门以恢复控制比特
    qml.Hadamard(wires=swap_qubit)

# 定义一个量子节点 (QNode) 用于 KerasLayer
@qml.qnode(dev, interface='tf', diff_method='parameter-shift')
def quantum_neural_network_for_keras(inputs, encoding_weights, rot_weights):
    control_qubits = [0, 1]
    ancilla = 21  # 修改辅助位为21（新增的量子比特）

    encoding_data = [
        {'data_qubits': {'q2': 2, 'q3': 3}, 'swap_qubit': 4},
        {'data_qubits': {'q2': 5, 'q3': 6}, 'swap_qubit': 7},
        {'data_qubits': {'q2': 8, 'q3': 9}, 'swap_qubit': 10},
        {'data_qubits': {'q2': 11, 'q3': 12}, 'swap_qubit': 13},
        {'data_qubits': {'q2': 14, 'q3': 15}, 'swap_qubit': 16},
        {'data_qubits': {'q2': 17, 'q3': 18}, 'swap_qubit': 19},
        {'data_qubits': {'q2': 20, 'q3': None}, 'swap_qubit': None}  # 最后一个组只有一个特征
    ]

    # 确保 apply_controlled_swap 使用新的 ancilla
    for data_set in encoding_data:
        swap_qubit = data_set['swap_qubit']
        data_qubits = data_set['data_qubits']
        if swap_qubit is not None:
            apply_controlled_swap(swap_qubit, data_qubits['q3'], data_qubits['q2'])

    # 将输入数据映射到 angles_q2
    angles_q2_list = inputs[:7]  # 对应7个编码组

    # 将 encoding_weights 映射到 angles_q3
    angles_q3_list = encoding_weights[:7, :, :4]  # Tensor of shape (7, n_layers, 4)

    # 进行编码
    for idx, data_set in enumerate(encoding_data):
        data_qubits = data_set['data_qubits']
        swap_qubit = data_set['swap_qubit']

        angle_q2 = angles_q2_list[idx]  # Tensor scalar

        for layer in range(rot_weights.shape[0]):
            angle_q3_values = angles_q3_list[idx, layer, :]  # Tensor of shape (4,)
            if data_qubits['q3'] is not None and swap_qubit is not None:
                encode_group(control_qubits, data_qubits, angle_q2, angle_q3_values, ancilla)
            else:
                # 如果只有一个特征，只有 q2 需要被编码
                # 仅对 q2 应用多控制 RY 旋转
                control_values = [[0, 0], [0, 1], [1, 0], [1, 1]]
                for cv, angle_q3_val in zip(control_values, angle_q3_values):
                    # 翻转控制比特为0的量子比特
                    for qubit, val in zip(control_qubits, cv):
                        if val == 0:
                            qml.PauliX(wires=qubit)

                    # 仅对 q2 应用多控制 RY 旋转
                    multi_controlled_RY(angle_q2, control_qubits, data_qubits['q2'], ancilla)

                    # 不对 q3 应用旋转，因为 q3 是 None

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

    # 应用交换操作
    for data_set in encoding_data:
        swap_qubit = data_set['swap_qubit']
        data_qubits = data_set['data_qubits']
        if swap_qubit is not None:
            apply_controlled_swap(swap_qubit, data_qubits['q3'], data_qubits['q2'])

    # 应用 Rot 层
    for layer in range(rot_weights.shape[0]):
        for qubit in range(rot_weights.shape[1]):
            rot_angles = rot_weights[layer, qubit, :]  # Tensor of shape (3,)
            qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)

    # 测量 PauliZ 期望值
    results = []
    for data_set in encoding_data:
        swap_qubit = data_set['swap_qubit']
        if swap_qubit is not None:
            results.append(qml.expval(qml.PauliZ(swap_qubit)))
    return results

# 定义一个量子节点 (QNode) 用于绘制电路
@qml.qnode(dev)
def quantum_neural_network_draw(inputs, encoding_weights, rot_weights):
    control_qubits = [0, 1]
    ancilla = 21  # 修改辅助位为20（最高位）

    encoding_data = [
        {'data_qubits': {'q2': 2, 'q3': 3}, 'swap_qubit': 4},
        {'data_qubits': {'q2': 5, 'q3': 6}, 'swap_qubit': 7},
        {'data_qubits': {'q2': 8, 'q3': 9}, 'swap_qubit': 10},
        {'data_qubits': {'q2': 11, 'q3': 12}, 'swap_qubit': 13},
        {'data_qubits': {'q2': 14, 'q3': 15}, 'swap_qubit': 16},
        {'data_qubits': {'q2': 17, 'q3': 18}, 'swap_qubit': 19},
        {'data_qubits': {'q2': 20, 'q3': None}, 'swap_qubit': None}  # 最后一个组只有一个特征
    ]

    # 将输入数据映射到 angles_q2
    angles_q2_list = inputs[:7]  # 对应7个编码组

    # 将 encoding_weights 映射到 angles_q3
    angles_q3_list = encoding_weights[:7, :, :4]  # Tensor of shape (7, n_layers, 4)

    # 进行编码
    for idx, data_set in enumerate(encoding_data):
        data_qubits = data_set['data_qubits']
        swap_qubit = data_set['swap_qubit']

        angle_q2 = angles_q2_list[idx]  # Tensor scalar

        for layer in range(rot_weights.shape[0]):
            angle_q3_values = angles_q3_list[idx, layer, :]  # Tensor of shape (4,)
            if data_qubits['q3'] is not None and swap_qubit is not None:
                encode_group(control_qubits, data_qubits, angle_q2, angle_q3_values, ancilla)
            else:
                # 如果只有一个特征，只有 q2 需要被编码
                control_values = [[0, 0], [0, 1], [1, 0], [1, 1]]
                for cv, angle_q3_val in zip(control_values, angle_q3_values):
                    # 翻转控制比特为0的量子比特
                    for qubit, val in zip(control_qubits, cv):
                        if val == 0:
                            qml.PauliX(wires=qubit)

                    # 仅对 q2 应用多控制 RY 旋转
                    multi_controlled_RY(angle_q2, control_qubits, data_qubits['q2'], ancilla)

                    # 不对 q3 应用旋转，因为 q3 是 None

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

    # 应用交换操作
    for data_set in encoding_data:
        swap_qubit = data_set['swap_qubit']
        data_qubits = data_set['data_qubits']
        if swap_qubit is not None:
            apply_controlled_swap(swap_qubit, data_qubits['q3'], data_qubits['q2'])

    # 应用 Rot 层
    for layer in range(rot_weights.shape[0]):
        for qubit in range(rot_weights.shape[1]):
            rot_angles = rot_weights[layer, qubit, :]  # Tensor of shape (3,)
            qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)

    # 测量 PauliZ 期望值
    results = []
    for data_set in encoding_data:
        swap_qubit = data_set['swap_qubit']
        if swap_qubit is not None:
            results.append(qml.expval(qml.PauliZ(swap_qubit)))
    return results

# 定义一个量子神经网络层，使用 KerasLayer
class QuantumNeuralNetworkLayer(tf.keras.layers.Layer):
    def __init__(self, n_qubits, n_layers, num_classes, **kwargs):
        super(QuantumNeuralNetworkLayer, self).__init__(**kwargs)
        self.n_qubits = n_qubits
        self.n_layers = n_layers
        self.num_classes = num_classes

        # 定义 Quantum Layer 使用 KerasLayer，更新 weight_shapes
        self.quantum_layer = KerasLayer(
            quantum_neural_network_for_keras,
            weight_shapes={
                "encoding_weights": (7, n_layers, 4),  # 7 个编码组
                "rot_weights": (n_layers, n_qubits, 3)  # 21 个量子比特
            },
            output_dim=7  # 根据测量结果数量设置输出维度（7个测量结果）
        )

    def call(self, inputs):
        return self.quantum_layer(inputs)

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

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

        # Classical layers，确保 Dense 层使用 float64
        self.dense = Dense(num_classes, activation='softmax', dtype='float64')

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

if __name__ == "__main__":
    parser = get_parser("Train model on Wine dataset")  # 修改描述
    args = parser.parse_args()

    wine = load_wine()  # 加载 Wine 数据集
    X = wine['data'].astype(np.float64)  # 使用所有13个特征并转换为 float64
    y = wine['target'].reshape(-1, 1)

    # 根据 scikit-learn 版本选择参数名称
    import sklearn
    from packaging import version

    sklearn_version = sklearn.__version__
    print(f"scikit-learn version: {sklearn_version}")
    if version.parse(sklearn_version) >= version.parse("1.2"):
        encoder = OneHotEncoder(sparse_output=False)
    else:
        encoder = OneHotEncoder(sparse=False)

    y = encoder.fit_transform(y).astype(np.float64)  # 修改为 float64

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

    learning_rate = 0.05
    epochs = 100
    batch_size = 8
    num_classes = 3
    n_qubits = 21  # 调整为21个量子比特（0到20）
    n_layers = 3  # 保持为三层

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

    optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate)

    model.compile(
        optimizer=optimizer,
        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:")
    # 使用 TensorFlow 常量传递参数，并确保类型为 float64
    sample_inputs = tf.constant(np.ones(7, dtype=np.float64) * (np.pi / 2))  # 修改为 float64，并假设7个输入特征
    sample_encoding_weights = tf.constant(np.zeros((7, n_layers, 4), dtype=np.float64))  # encoding_weights 修改为 float64
    sample_rot_weights = tf.constant(np.zeros((n_layers, n_qubits, 3), dtype=np.float64))  # rot_weights 修改为 float64

    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' 获取打印输出，查看当前目录获取保存的图表。")