import sys
import os
import pennylane as qml
from pennylane import numpy as np
import tensorflow as tf
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder, StandardScaler, MinMaxScaler
from sklearn.decomposition import PCA

import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端
import matplotlib.pyplot as plt

from tensorflow.keras.layers import Dense
from importance_sampling.training import ConstantTimeImportanceTraining
from importance_sampling.samplers import ConstantTimeSampler
from importance_sampling.datasets import InMemoryDataset
from importance_sampling.reweighting import BiasedReweightingPolicy
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau

import logging

# 确保TensorFlow使用float32
tf.keras.backend.set_floatx('float32')

# 重定向stdout以捕获打印输出
class Tee(object):
    def __init__(self, *files):
        self.files = files
        self.encoding = getattr(files[0], 'encoding', 'utf-8')  # 添加编码属性

    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('layers1', 'w', encoding='utf-8')
original_stdout = sys.stdout
tee = Tee(sys.stdout, output_file)
sys.stdout = tee

# 配置Logger使用Tee
logger = logging.getLogger('my_sampler')
logger.setLevel(logging.INFO)

stream_handler = logging.StreamHandler(tee)
stream_handler.setLevel(logging.INFO)

formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
stream_handler.setFormatter(formatter)

# 移除默认处理器以避免重复日志
if not logger.handlers:
    logger.addHandler(stream_handler)
else:
    logger.handlers = []
    logger.addHandler(stream_handler)

# 设置Qiskit日志级别为WARNING以抑制INFO日志
logging.getLogger('qiskit').setLevel(logging.WARNING)

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'], label='训练准确率')
    if 'val_accuracy' in history.history:
        ax.plot(history.history['val_accuracy'], label='验证准确率')
    else:
        print("历史记录中未找到验证准确率。")
    ax.set_title('模型准确率')
    ax.set_ylabel('准确率')
    ax.set_xlabel('Epoch')
    ax.legend(loc='upper left')
    save_plot(fig, 'accuracy_plot.png')

    # 绘制训练和验证的损失
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(history.history['loss'], label='训练损失')
    if 'val_loss' in history.history:
        ax.plot(history.history['val_loss'], label='验证损失')
    else:
        print("历史记录中未找到验证损失。")
    ax.set_title('模型损失')
    ax.set_ylabel('损失')
    ax.set_xlabel('Epoch')
    ax.legend(loc='upper right')
    save_plot(fig, 'loss_plot.png')

def calculate_accuracy_fluctuation(history):
    if 'val_accuracy' in history.history:
        accuracies = history.history['val_accuracy']
    elif 'accuracy' in history.history:
        accuracies = history.history['accuracy']
    else:
        print("历史记录中未找到准确率指标。")
        return None, None
    mean_accuracy = np.mean(accuracies)
    fluctuation = np.sum(np.abs(accuracies - mean_accuracy))
    return mean_accuracy, fluctuation

# 定义一个包含14个量子比特的量子设备
dev = qml.device("default.qubit", wires=14)

def double_controlled_ry(control1, control2, target, angle):
    qml.Hadamard(wires=target)
    qml.Toffoli(wires=[control1, control2, target])
    qml.RZ(angle, wires=target)
    qml.Toffoli(wires=[control1, control2, target])
    qml.Hadamard(wires=target)

# 每个神经元的编码
def encode_group(control_qubits, data_qubits, angle_q2, angle_q3_values):
    control_values = [[0, 0], [0, 1], [1, 0], [1, 1]]

    for i in range(4):
        cv = control_values[i]
        angle_q3_val = angle_q3_values[i]

        # 如果控制值为0，则翻转控制量子比特
        for qubit, val in zip(control_qubits, cv):
            if val == 0:
                qml.PauliX(wires=qubit)

        # 应用双控制的RY门
        double_controlled_ry(control_qubits[0], control_qubits[1], data_qubits['q2'], angle_q2)
        double_controlled_ry(control_qubits[0], control_qubits[1], data_qubits['q3'], angle_q3_val)

        # 恢复控制量子比特
        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): SWAP门的控制量子比特。
    - q_swap1 (int): 第一个要交换的量子比特。
    - q_swap2 (int): 第二个要交换的量子比特。
    """
    qml.Hadamard(wires=swap_qubit)
    qml.CSWAP(wires=[swap_qubit, q_swap1, q_swap2])
    qml.Hadamard(wires=swap_qubit)

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

    # 神经元组
    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}
    ]

    # 角度编码：输入被映射到[0, π/2]
    angles_q2_list = inputs  # 输入形状为 (4,)

    # 编码权重映射到角度q3
    angles_q3_list = encoding_weights  # 形状: (4, n_layers, 4)

    # 为每一层编码特征和权重
    for layer in range(rot_weights.shape[0]):   # 遍历层

        # 为每个神经元编码特征和权重
        for idx, data_set in enumerate(encoding_data):
            swap_qubit = data_set['swap_qubit']
            data_qubits = data_set['data_qubits']

            angle_q2 = angles_q2_list[idx]  # 标量值
            angle_q3_values = angles_q3_list[idx, layer, :]  # 形状: (4,)
            # 编码单个神经元
            encode_group(control_qubits, data_qubits, angle_q2, angle_q3_values)

        # 应用旋转门
        for qubit in range(rot_weights.shape[1]):
            rot_angles = rot_weights[layer, qubit, :]  # 形状: (3,)
            qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)
        # 应用参数化纠缠门
        for qubit in range(rot_weights.shape[1] - 1):
            qml.CRZ(entangle_weights[layer, qubit], wires=[qubit, qubit + 1])

        # 为每个神经元应用交换操作
        for data_set in encoding_data:
            swap_qubit = data_set['swap_qubit']
            data_qubits = data_set['data_qubits']
            apply_controlled_swap(swap_qubit, data_qubits['q3'], data_qubits['q2'])

    # 在多个量子比特上测量PauliZ期望值
    measured_qubits = [4, 7, 10]
    results = [qml.expval(qml.PauliZ(qubit)) for qubit in measured_qubits]

    return results

# 定义量子神经网络层
class QuantumNeuralNetworkLayer(tf.keras.layers.Layer):
    def __init__(self, n_qubits, n_layers, num_classes):
        super(QuantumNeuralNetworkLayer, self).__init__()
        self.n_qubits = n_qubits
        self.n_layers = n_layers
        self.num_classes = num_classes

        # 定义带L2正则化的编码权重
        initializer = tf.keras.initializers.GlorotUniform()
        regularizer = tf.keras.regularizers.l2(0.01)

        self.encoding_weights = self.add_weight(
            shape=(4, self.n_layers, 4),
            initializer=initializer,
            regularizer=regularizer,
            trainable=True,
            name='encoding_weights',
            dtype=tf.float32
        )

        # 定义带L2正则化的旋转层权重
        self.rot_weights = self.add_weight(
            shape=(n_layers, n_qubits, 3),
            initializer=initializer,
            regularizer=regularizer,
            trainable=True,
            name='rot_weights',
            dtype=tf.float32
        )

        # 定义带L2正则化的纠缠权重
        self.entangle_weights = self.add_weight(
            shape=(n_layers, n_qubits - 1),
            initializer=tf.keras.initializers.RandomUniform(0, 2 * np.pi),
            regularizer=regularizer,
            trainable=True,
            name='entangle_weights',
            dtype=tf.float32
        )

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

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

        return outputs

# 定义量子神经网络模型
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

        # 量子层
        self.quantum_layer = QuantumNeuralNetworkLayer(n_qubits, n_layers, num_classes)

        # 经典层
        self.dense1 = tf.keras.layers.Dense(
            16,
            activation='relu',
            kernel_regularizer=tf.keras.regularizers.l2(0.01),
            dtype=tf.float32
        )
        self.dense2 = tf.keras.layers.Dense(
            8,
            activation='relu',
            kernel_regularizer=tf.keras.regularizers.l2(0.01),
            dtype=tf.float32
        )

        # 直接连接到量子输出的输出层
        self.output_layer = Dense(
            num_classes,
            activation='softmax',
            kernel_regularizer=tf.keras.regularizers.l2(0.01),
            dtype=tf.float32  # 确保输出层为float32
        )

    def call(self, inputs):
        q_output = self.quantum_layer(inputs)
        x = self.dense1(q_output)
        x = self.dense2(x)
        return self.output_layer(x)

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


if __name__ == "__main__":
    # 加载Liver Disorders数据集（从本地文件加载）
    local_data_path = 'liver-disorder.csv'  # 请确保CSV文件位于脚本相同的目录，或提供绝对路径
    column_names = [
        'mcv', 'alkphos', 'sgpt', 'sgot', 'gammagt', 'drinks', 'selector'
    ]
    try:
        data = pd.read_csv(local_data_path)
    except Exception as e:
        print(f"加载Liver Disorders数据集时出错: {e}")
        sys.exit(1)

    # 查看数据的前几行以确认正确加载
    print("数据集预览：")
    print(data.head())

    # 处理缺失值，删除包含NaN的行
    data = data.dropna()

    # 查看是否包含'Gender'列，如果不包含，则跳过
    if 'Gender' in data.columns:
        data['Gender'] = data['Gender'].map({'Female': 0, 'Male': 1})
    else:
        print("数据集中不包含'Gender'列。将跳过该步骤。")

    # 选择数值特征
    numerical_features = [
        'mcv', 'alkphos', 'sgpt', 'sgot', 'gammagt', 'drinks'
    ]
    # 检查是否包含所有必要的特征
    if not set(numerical_features + ['selector']).issubset(data.columns):
        print("数据集中缺少必要的特征列。")
        sys.exit(1)

    X = data[numerical_features].values.astype(np.float32)  # 特征: (样本数,6)
    y = data['selector'].values.reshape(-1, 1)               # 标签: (样本数,1)

    # 使用PCA将特征降维到4维
    pca = PCA(n_components=4)
    X_reduced = pca.fit_transform(X).astype(np.float32)  # (样本数,4)

    # 独热编码标签
    encoder = OneHotEncoder(sparse_output=False)
    y_encoded = encoder.fit_transform(y).astype(np.float32)
    num_classes = y_encoded.shape[1]

    # 分割训练集和测试集
    x_full_train, x_test, y_full_train, y_test = train_test_split(
        X_reduced, y_encoded, test_size=0.2, random_state=116, stratify=y_encoded
    )

    # 进一步将训练集分为训练集和验证集
    x_train, x_val, y_train, y_val = train_test_split(
        x_full_train, y_full_train, test_size=0.25, random_state=42, stratify=y_full_train
    )
    # 现在，x_train占60%，x_val占20%，x_test占20%

    # 标准化和归一化数据（仅在训练集上拟合）
    scaler = StandardScaler()
    X_train_standardized = scaler.fit_transform(x_train).astype(np.float32)
    X_val_standardized = scaler.transform(x_val).astype(np.float32)
    X_test_standardized = scaler.transform(x_test).astype(np.float32)

    # 将特征映射到[0, π/2]范围
    min_max_scaler = MinMaxScaler(feature_range=(0, np.pi/2))
    X_train_mapped = min_max_scaler.fit_transform(X_train_standardized).astype(np.float32)
    X_val_mapped = min_max_scaler.transform(X_val_standardized).astype(np.float32)
    X_test_mapped = min_max_scaler.transform(X_test_standardized).astype(np.float32)

    # 超参数
    learning_rate = 0.05  # 学习率
    epochs = 100          # 训练的轮数
    batch_size = 16       # 批大小
    n_qubits = 14         # 14个量子比特
    n_layers = 3          # 层数

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

    if is_quantum_model(model):
        print("模型包含量子神经网络（QNN）层。")
    else:
        print("模型不包含量子神经网络（QNN）层。")

    # 绘制量子电路
    sample_inputs = np.ones(4, dtype=np.float32) * (np.pi / 4)  # 输入映射到 [0, π/2]
    sample_encoding_weights = np.zeros((4, n_layers, 4), dtype=np.float32)
    sample_rot_weights = np.zeros((n_layers, n_qubits, 3), dtype=np.float32)
    sample_entangle_weights = np.zeros((n_layers, n_qubits - 1), dtype=np.float32)

    print("\n量子电路图:")
    print(qml.draw(quantum_neural_network)(sample_inputs, sample_encoding_weights, sample_rot_weights, sample_entangle_weights))

    steps_per_epoch = len(X_train_mapped) // batch_size
    total_steps = epochs * steps_per_epoch

    # 创建数据集（包括验证集）
    dataset = InMemoryDataset(
        X_train_mapped,
        y_train,
        X_val_mapped,
        y_val,
        categorical=True
    )

    # 初始化重加权策略
    reweighting_policy = BiasedReweightingPolicy(k=1.0)

    # 初始化重要性采样训练
    wrapped_model = ConstantTimeImportanceTraining(model, score="loss")  # 使用损失作为评分

    # 修改采样器以使用动态的alpha和beta
    def custom_sampler(dataset, batch_size, steps_per_epoch, epochs):
        return ConstantTimeSampler(
            dataset,
            reweighting_policy,
            model=wrapped_model.model,
        )

    # 将自定义采样器分配给wrapped_model
    wrapped_model.sampler = custom_sampler

    # 使用每样本的损失函数
    loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=False, reduction=tf.keras.losses.Reduction.NONE)

    # 编译模型
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
        loss=loss_fn,
        metrics=['accuracy'],
        run_eagerly=False  # 设置为False以获得更好的性能
    )

    # 创建一个回调类以保存训练指标
    class MetricsHistory(tf.keras.callbacks.Callback):
        def on_train_begin(self, logs=None):
            self.history = {'loss': [], 'accuracy': [], 'val_loss': [], 'val_accuracy': []}

        def on_epoch_end(self, epoch, logs=None):
            self.history['loss'].append(logs.get('loss'))
            self.history['accuracy'].append(logs.get('accuracy'))
            self.history['val_loss'].append(logs.get('val_loss'))
            self.history['val_accuracy'].append(logs.get('val_accuracy'))

    # 实例化回调
    metrics_history = MetricsHistory()

    # 添加早停回调
    callbacks = [
        metrics_history,
        EarlyStopping(monitor='val_loss', patience=100, restore_best_weights=True)
    ]

    # 使用fit_dataset训练模型
    history = wrapped_model.fit_dataset(
        dataset=dataset,
        batch_size=batch_size,
        epochs=epochs,
        steps_per_epoch=steps_per_epoch,
        verbose=1,
        callbacks=callbacks
    )

    # 在测试集上评估模型
    score = wrapped_model.evaluate((X_test_mapped, y_test))
    print('测试损失:', score[0])
    print('测试准确率:', score[1])

    # 使用自定义回调的历史记录绘制图表
    plot_history(metrics_history)

    mean_accuracy, fluctuation = calculate_accuracy_fluctuation(metrics_history)
    if mean_accuracy is not None and fluctuation is not None:
        print('验证集平均准确率:', mean_accuracy)
        print('准确率波动:', fluctuation)

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

    print("脚本执行完成。请检查 'layers1' 文件中的打印输出，并在当前目录中查看保存的图表。")
你仔细阅读上述代码，你会发现上述代码有4个神经元却只有三个输出比特，是否有一个量子比特输出被浪费了？