from pennylane import numpy as pnp  # 使用 'pnp' 作为 PennyLane 的 NumPy 别名
import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端

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
import numpy as np  # 标准 NumPy
import pandas as pd
import tensorflow as tf
import pennylane as qml
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder, MinMaxScaler
from sklearn.decomposition import PCA
import sys
import matplotlib.pyplot as plt
# from sklearn.datasets import fetch_openml  # 不再需要

# 确保 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('Liver_Disorder.log', 'w', encoding='utf-8')
output_secondary = open('output.log', 'w', encoding='utf-8')
original_stdout = sys.stdout
tee = Tee(sys.stdout, output_file, output_secondary)
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='Training Accuracy')
    if 'val_accuracy' in history.history:
        ax.plot(history.history['val_accuracy'], label='Validation Accuracy')
    else:
        print("Validation accuracy not found in history.")
    ax.set_title('Model Accuracy')
    ax.set_ylabel('Accuracy')
    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='Training Loss')
    if 'val_loss' in history.history:
        ax.plot(history.history['val_loss'], label='Validation Loss')
    else:
        print("Validation loss not found in history.")
    ax.set_title('Model Loss')
    ax.set_ylabel('Loss')
    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("Accuracy metric not found in history.")
        return None, None
    mean_accuracy = np.mean(accuracies)
    fluctuation = np.sum(np.abs(accuracies - mean_accuracy))
    return mean_accuracy, fluctuation

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

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_values, angle_q3_values):
    control_values = [[0, 0], [0, 1], [1, 0], [1, 1]]

    for i in range(4):
        cv = control_values[i]
        angle_q2 = angle_q2_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_qubit (int): 受控交换门的控制量子比特。
    - 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},
    ]

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

    # 编码权重映射到 angles_q3
    angles_q3_list = encoding_weights  # 形状为 (2, 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  # 标量值
            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]
    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=(2, 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

def load_liver_disorder_dataset():
    """
    加载印度肝病患者数据集（ILPD）从本地文件。
    """
    # 指定列名
    column_names = ['Age', 'Gender', 'Total_Bilirubin', 'Direct_Bilirubin', 'Alkaline_Phosphotase',
                    'Alamine_Aminotransferase', 'Aspartate_Aminotransferase', 'Total_Proteins',
                    'Albumin', 'Albumin_and_Globulin_Ratio', 'Class']

    # 从本地文件加载数据（确保 ILPD.csv 文件与代码在同一目录）
    # 如果 CSV 文件有标题行，设置 header=0；否则，设置 header=None
    df = pd.read_csv('Indian Liver Patient Dataset (ILPD).csv', header=0)

    # 检查列名是否正确，如果需要，可以重命名列
    if list(df.columns) != column_names:
        df.columns = column_names

    # 处理缺失值
    df = df.dropna()

    # 提取特征和目标变量
    X = df.drop('Class', axis=1)
    y = df['Class']

    return X, y

if __name__ == "__main__":
    # 加载数据
    X, y = load_liver_disorder_dataset()

    # 将目标变量转换为二分类标签
    # 首先将 y 转换为整数
    y = y.astype(int)

    # 检查 y 中的唯一值
    print("Unique values in 'Class' column:", y.unique())

    # 将类别 1 作为正类（肝病患者），类别 2 作为负类（非肝病患者）
    y = (y == 1).astype(int)  # 1 表示肝病患者，0 表示非肝病患者

    # 检查类别分布
    unique_classes, class_counts = np.unique(y, return_counts=True)
    print("Classes:", unique_classes)
    print("Class counts:", class_counts)

    # 处理特征中的分类变量
    X = X.copy()
    # 假设 'Gender' 是唯一的分类特征
    X['Gender'] = X['Gender'].map({'Male': 0, 'Female': 1})

    # 检查 'Gender' 列是否存在未映射的值
    if X['Gender'].isnull().any():
        print("警告: 'Gender' 列存在未映射的值，将删除这些记录。")
        X = X.dropna(subset=['Gender'])
        y = y[X.index]

    # 对标签进行独热编码
    encoder = OneHotEncoder(sparse_output=False)
    y_array = y.to_numpy().reshape(-1, 1)  # 在重塑前转换为 NumPy 数组
    y_encoded = encoder.fit_transform(y_array).astype(np.float32)
    num_classes = y_encoded.shape[1]
    print("num_classes", num_classes)

    # 划分训练集和测试集
    x_full_train, x_test, y_full_train, y_test = train_test_split(
        X, 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 和 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)

    # 应用 PCA 降维到 4 维
    pca = PCA(n_components=4)
    X_train_pca = pca.fit_transform(X_train_standardized).astype(np.float32)
    X_val_pca = pca.transform(X_val_standardized).astype(np.float32)
    X_test_pca = pca.transform(X_test_standardized).astype(np.float32)

    # 显示 PCA 解释的方差比例
    print("PCA Explained Variance Ratio:")
    print(pca.explained_variance_ratio_)
    print("Total Explained Variance:", sum(pca.explained_variance_ratio_))

    # 将特征映射到 [0, π/2] 区间
    min_max_scaler = MinMaxScaler(feature_range=(0, np.pi / 2))
    X_train_mapped = min_max_scaler.fit_transform(X_train_pca).astype(np.float32)
    X_val_mapped = min_max_scaler.transform(X_val_pca).astype(np.float32)
    X_test_mapped = min_max_scaler.transform(X_test_pca).astype(np.float32)

    # 超参数
    learning_rate = 0.01  # 学习率
    epochs = 500  # 训练轮数
    batch_size = 16  # 批大小
    n_qubits = 4  # 匹配 PCA 降维后的维度
    n_layers = 3  # 层数

    # 初始化量子神经网络模型
    model = QuantumNeuralNetwork(
        n_qubits=n_qubits,
        n_layers=n_layers,
        num_classes=num_classes
    )

    if is_quantum_model(model):
        print("The model contains quantum neural network (QNN) layers.")
    else:
        print("The model does not contain quantum neural network (QNN) layers.")

    # 绘制量子电路
    sample_inputs = np.ones(4, dtype=np.float32) * (np.pi / 4)  # 输入映射到 [0, π/2]
    sample_encoding_weights = np.zeros((2, 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("\nQuantum circuit diagram:")
    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")  # 使用 loss 作为评分

    # 自定义采样器
    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('Test loss:', score[0])
    print('Test accuracy:', 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('Validation set mean accuracy:', mean_accuracy)
        print('Accuracy fluctuation:', fluctuation)

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

    print(
        "Script execution completed. Please check 'Liver_Disorder.log' and 'output.log' for print outputs and check the saved plots in the current directory."
    )
上述代码的数据提取与处理是否有问题，为什么代码训练后准确率基本不发生改变（100epoch只增加了2%的准确率），这很不对