# 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, StandardScaler, MinMaxScaler
#
# 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
#
# # **确保 TensorFlow 使用 float32**
# tf.keras.backend.set_floatx('float32')
#
# # 重定向标准输出以捕获打印输出
# 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_modified.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'], label='Train Accuracy')
#     if 'val_accuracy' in history.history:
#         ax.plot(history.history['val_accuracy'], label='Validation Accuracy')
#     else:
#         print("No validation accuracy 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='Train Loss')
#     if 'val_loss' in history.history:
#         ax.plot(history.history['val_loss'], label='Validation Loss')
#     else:
#         print("No validation loss 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("No accuracy metric found in history.")
#         return None, None
#     mean_accuracy = np.mean(accuracies)
#     fluctuation = np.sum(np.abs(accuracies - mean_accuracy))
#     return mean_accuracy, fluctuation
#
# # **定义具有4个量子比特的量子设备，不指定 dtype**
# dev = qml.device('default.qubit', wires=4)
#
# @qml.qnode(dev, interface='tf', diff_method='backprop')
# def quantum_neural_network(inputs, weights):
#     num_qubits = 4
#
#     # 将输入数据编码到量子态中
#     for i in range(num_qubits):
#         qml.RY(inputs[i], wires=i)
#
#     # 量子神经网络的参数化线路
#     n_layers = weights.shape[0]
#     for layer in range(n_layers):
#         # 对每个量子比特应用旋转门
#         for i in range(num_qubits):
#             params = weights[layer, i]
#             qml.Rot(params[0], params[1], params[2], wires=i)
#         # 应用纠缠门，如 CNOT 门
#         for i in range(num_qubits - 1):
#             qml.CNOT(wires=[i, i + 1])
#
#     # 在每个量子比特上测量 PauliZ 期望值
#     return [qml.expval(qml.PauliZ(i)) for i in range(num_qubits)]
#
# # 定义量子神经网络层
# class QuantumNeuralNetworkLayer(tf.keras.layers.Layer):
#     def __init__(self, n_qubits, n_layers):
#         super(QuantumNeuralNetworkLayer, self).__init__()
#         self.n_qubits = n_qubits
#         self.n_layers = n_layers
#
#         # Initialize trainable weights with dtype=tf.float32 using an initializer object
#         initializer = tf.keras.initializers.RandomUniform(
#             minval=0.0,
#             maxval=2 * np.pi
#         )
#         self.q_weights = self.add_weight(
#             shape=(n_layers, n_qubits, 3),
#             initializer=initializer,
#             trainable=True,
#             name='q_weights',
#             dtype=tf.float32
#         )
#
#     def call(self, inputs):
#         def circuit(x):
#             # 将输入从 float32 转换为 float64
#             x = tf.cast(x, dtype=tf.float64)
#             # 将权重从 float32 转换为 float64
#             weights = tf.cast(self.q_weights, dtype=tf.float64)
#             result = quantum_neural_network(x, weights)
#             # 将结果从 float64 转换回 float32
#             return tf.cast(result, dtype=tf.float32)
#
#         outputs = tf.map_fn(
#             circuit,
#             inputs,
#             dtype=tf.float32
#         )
#         return outputs
#
# # 定义量子神经网络模型
# 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.quantum_layer = QuantumNeuralNetworkLayer(n_qubits, n_layers)
#         self.output_layer = Dense(num_classes, activation='softmax', dtype=tf.float32)
#
#     def call(self, inputs):
#         q_output = self.quantum_layer(inputs)
#         return self.output_layer(q_output)
#
# 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__":
#     # 加载数据
#     iris = load_iris()
#     X = iris['data'].astype(np.float32)  # 确保数据类型为 float32
#     y = iris['target'].reshape(-1, 1)
#
#     # 对标签进行 One-Hot 编码
#     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, 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)
#
#     min_max_scaler = MinMaxScaler(feature_range=(0, np.pi))
#     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.1  # 学习率
#     epochs = 100          # 训练轮数
#     batch_size = 16       # 批量大小
#     n_qubits = 4          # 4 个量子比特
#     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 / 2)
#     sample_weights = np.zeros((n_layers, n_qubits, 3), dtype=np.float32)
#
#     print("\n量子电路图：")
#     print(qml.draw(quantum_neural_network)(sample_inputs, sample_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,
#             B=64,
#             alpha=0.6,
#             beta_start=0.4,
#             beta_frames=total_steps
#         )
#
#     # 将自定义采样器赋给 wrapped_model
#     wrapped_model.sampler = custom_sampler
#
#     # 使用 per-sample 的损失函数
#     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)
#
#     # 关闭输出文件并恢复标准输出
#     sys.stdout = original_stdout
#     output_file.close()
#
#     print("脚本执行完毕。请检查 'output_modified.log' 以获取打印输出，并在当前目录下查看保存的图表。")


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, StandardScaler, MinMaxScaler

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

import logging

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

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

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

# 配置 Logger，使其使用 Tee 类
logger = logging.getLogger()
logger.setLevel(logging.INFO)  # 设置全局日志级别为 INFO

stream_handler = logging.StreamHandler(tee)
stream_handler.setLevel(logging.INFO)  # 设置处理器的日志级别为 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)

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='Train Accuracy')
    if 'val_accuracy' in history.history:
        ax.plot(history.history['val_accuracy'], label='Validation Accuracy')
    else:
        print("No validation accuracy 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='Train Loss')
    if 'val_loss' in history.history:
        ax.plot(history.history['val_loss'], label='Validation Loss')
    else:
        print("No validation loss 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("No accuracy metric found in history.")
        return None, None
    mean_accuracy = np.mean(accuracies)
    fluctuation = np.sum(np.abs(accuracies - mean_accuracy))
    return mean_accuracy, fluctuation

# **定义具有4个量子比特的量子设备，不指定 dtype**
dev = qml.device('default.qubit', wires=4)

@qml.qnode(dev, interface='tf', diff_method='backprop')
def quantum_neural_network(inputs, weights):
    num_qubits = 4

    # 将输入数据编码到量子态中
    for i in range(num_qubits):
        qml.RY(inputs[i], wires=i)

    # 量子神经网络的参数化线路
    n_layers = weights.shape[0]
    for layer in range(n_layers):
        # 对每个量子比特应用旋转门
        for i in range(num_qubits):
            params = weights[layer, i]
            qml.Rot(params[0], params[1], params[2], wires=i)
        # 应用纠缠门，如 CNOT 门
        for i in range(num_qubits - 1):
            qml.CNOT(wires=[i, i + 1])

    # 在每个量子比特上测量 PauliZ 期望值
    return [qml.expval(qml.PauliZ(i)) for i in range(num_qubits)]

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

        # Initialize trainable weights with dtype=tf.float32 using an initializer object
        initializer = tf.keras.initializers.RandomUniform(
            minval=0.0,
            maxval=2 * np.pi
        )
        self.q_weights = self.add_weight(
            shape=(n_layers, n_qubits, 3),
            initializer=initializer,
            trainable=True,
            name='q_weights',
            dtype=tf.float32
        )

    def call(self, inputs):
        def circuit(x):
            # 将输入从 float32 转换为 float64
            x = tf.cast(x, dtype=tf.float64)
            # 将权重从 float32 转换为 float64
            weights = tf.cast(self.q_weights, dtype=tf.float64)
            result = quantum_neural_network(x, weights)
            # 将结果从 float64 转换回 float32
            return tf.cast(result, dtype=tf.float32)

        outputs = tf.map_fn(
            circuit,
            inputs,
            dtype=tf.float32
        )
        return outputs

# 定义量子神经网络模型
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.quantum_layer = QuantumNeuralNetworkLayer(n_qubits, n_layers)
        self.output_layer = Dense(num_classes, activation='softmax', dtype=tf.float32)

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

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__":
    # 加载数据
    iris = load_iris()
    X = iris['data'].astype(np.float32)  # 确保数据类型为 float32
    y = iris['target'].reshape(-1, 1)

    # 对标签进行 One-Hot 编码
    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, 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)

    min_max_scaler = MinMaxScaler(feature_range=(0, np.pi))
    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.1  # 学习率
    epochs = 100          # 训练轮数
    batch_size = 16       # 批量大小
    n_qubits = 4          # 4 个量子比特
    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 / 2)
    sample_weights = np.zeros((n_layers, n_qubits, 3), dtype=np.float32)

    print("\n量子电路图：")
    print(qml.draw(quantum_neural_network)(sample_inputs, sample_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,
            B=64,
            alpha=0.6,
            beta_start=0.4,
            beta_frames=total_steps
        )

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

    # 使用 per-sample 的损失函数
    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)

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

    print("脚本执行完毕。请检查 'output_modified.log' 以获取打印输出，并在当前目录下查看保存的图表。")
上述是运行脚本
import sys
import numpy as np
import logging
from collections.abc import Iterable
import tensorflow as tf
from tensorflow.python.keras.callbacks import BaseLogger, CallbackList, History, ProgbarLogger
from tensorflow.python.keras.layers import Dropout
#from tensorflow.python.keras.utils import Sequence
from blinker import signal
from .datasets import InMemoryDataset, GeneratorDataset
from .model_wrappers import OracleWrapper
from .samplers import ConditionalStartSampler, VarianceReductionCondition, \
    AdaptiveAdditiveSmoothingSampler, AdditiveSmoothingSampler, ModelSampler, \
    ConstantVarianceSampler, ConstantTimeSampler, SCSGSampler
from .reweighting import BiasedReweightingPolicy, NoReweightingPolicy
from .utils.functional import ___, compose, partial

class _BaseImportanceTraining(object):
    def __init__(self, model, score="gnorm", layer=None):
        """重要性采样训练的基类。

        参数：
            model: 要训练的 Keras 模型。
        """
        # 检查模型是否兼容，并包装模型以支持重要性采样
        self._check_model(model)
        self.original_model = model
        self.model = OracleWrapper(
            model,
            self.reweighting,
            score=score,
            layer=layer
        )
        self.original_model.optimizer = self.model.optimizer
        signal("is.sample").connect(self._on_sample)
        signal("is.score").connect(self._on_scores)

    def _on_sample(self, event):
        self._latest_sample_event = event

    def _on_scores(self, scores):
        self._latest_scores = scores

    def _check_model(self, model):
        """检查模型是否使用了 Dropout 和 BatchNorm，并警告可能影响重要性计算。"""
        if any(
                isinstance(l, (tf.keras.layers.BatchNormalization, Dropout))
                for l in model.layers
        ):
            sys.stderr.write(("[NOTICE]: 您正在使用 BatchNormalization 和/或 Dropout。\n"
                              "这些层可能会影响重要性计算，建议您用 LayerNormalization 或 "
                              "在测试模式下的 BatchNormalization 以及 L2 正则化替代它们。\n"))

    @property
    def reweighting(self):
        """控制重要性采样时偏置的重加权策略。"""
        raise NotImplementedError()

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
        """为给定的数据集创建一个新的采样器，以使用重要性采样进行采样。"""
        raise NotImplementedError()

    def fit(self, x, y, batch_size=32, epochs=1, verbose=1, callbacks=None,
            validation_split=0.0, validation_data=None, steps_per_epoch=None,
            on_sample=None, on_scores=None):
        """使用给定的数据创建一个 InMemoryDataset 实例，并使用重要性采样训练模型。

        参数：
            x: 训练数据的 Numpy 数组或 Numpy 数组列表。
            y: 目标数据的 Numpy 数组。
            batch_size: 每个梯度更新的样本数。
            epochs: 遍历整个训练集的次数。
            verbose: {0, >0}，是否使用 Keras 的进度条回调。
            callbacks: 在训练期间调用的 Keras 回调列表。
            validation_split: [0, 1) 之间的浮点数，用于验证的数据百分比。
            validation_data: tuple，(x_val, y_val)，用于模型验证的数据。
            steps_per_epoch: int 或 None，每个 epoch 的步数。
            on_sample: 接受 sampler, idxs, w, scores 的回调函数。
            on_scores: 接受 sampler 和 scores 的回调函数。
        返回：
            包含训练期间收集的信息的 Keras `History` 对象。
        """
        # 处理验证数据集
        if validation_data is not None:
            x_train, y_train = x, y
            x_test, y_test = validation_data
        elif validation_split > 0:
            assert validation_split < 1, "验证集比例必须小于 1"
            n = int(round(validation_split * len(x)))
            idxs = np.arange(len(x))
            np.random.shuffle(idxs)
            x_train, y_train = x[idxs[n:]], y[idxs[n:]]
            x_test, y_test = x[idxs[:n]], y[idxs[:n]]
        else:
            x_train, y_train = x, y
            x_test, y_test = np.empty(shape=(0, x.shape[1])), np.empty(shape=(0, y.shape[1]))

        # 创建数据集
        dataset = InMemoryDataset(
            x_train,
            y_train,
            x_test,
            y_test,
            categorical=False  # 不需要转换标签
        )

        return self.fit_dataset(
            dataset=dataset,
            batch_size=batch_size,
            epochs=epochs,
            steps_per_epoch=steps_per_epoch,
            verbose=verbose,
            callbacks=callbacks,
            on_sample=on_sample,
            on_scores=on_scores
        )

    def fit_generator(self, train, steps_per_epoch=None, batch_size=32,
                      epochs=1, verbose=1, callbacks=None,
                      validation_data=None, validation_steps=None,
                      on_sample=None, on_scores=None):
        """创建一个 GeneratorDataset 实例，并使用重要性采样训练模型。

        参数：
            train: 返回 (inputs, targets) 的生成器或可迭代对象。
            steps_per_epoch: 每个 epoch 的梯度更新次数。
            batch_size: 每个梯度更新的样本数。
            epochs: 训练轮数。
            verbose: {0, >0}，是否使用 Keras 的进度条回调。
            validation_data: 生成器、可迭代对象或 (inputs, targets) 元组。
            validation_steps: 如果 validation_data 是生成器，则需要指定此参数。
            on_sample: 接受 sampler, idxs, w, scores 的回调函数。
            on_scores: 接受 sampler 和 scores 的回调函数。
        """

        def infinite_iterator(iterable):
            while True:
                yield from iterable

        # 处理验证数据集
        if validation_data is not None:
            if isinstance(validation_data, (tuple, list)):
                test = validation_data
                test_len = None
            elif isinstance(validation_data, Iterable):
                test = infinite_iterator(validation_data)
                try:
                    test_len = len(validation_data) * batch_size
                except TypeError:
                    test_len = validation_steps * batch_size if validation_steps else None
            else:
                test = validation_data
                test_len = validation_steps * batch_size if validation_steps else None
        else:
            test = (np.empty(shape=(0, 1)), np.empty(shape=(0, 1)))
            test_len = None
            logging.info("未提供验证数据")

        # 设置训练数据和每个 epoch 的步数
        if isinstance(train, Iterable):
            if steps_per_epoch is None:
                try:
                    steps_per_epoch = len(train)
                except TypeError:
                    raise ValueError("使用没有 __len__ 方法的生成器时必须设置 steps_per_epoch")
            train = infinite_iterator(train)
        elif steps_per_epoch is None:
            raise ValueError("steps_per_epoch 未设置")

        dataset = GeneratorDataset(train, test, test_len)

        return self.fit_dataset(
            dataset=dataset,
            batch_size=batch_size,
            epochs=epochs,
            steps_per_epoch=steps_per_epoch,
            verbose=verbose,
            callbacks=callbacks,
            on_sample=on_sample,
            on_scores=on_scores
        )

    def fit_dataset(self, dataset, steps_per_epoch=None, batch_size=32,
                    epochs=1, verbose=1, callbacks=None, on_sample=None,
                    on_scores=None):
        """在给定的数据集上训练模型。

        参数：
            dataset: `BaseDataset` 的实例，提供要训练的数据。
            steps_per_epoch: int，每个 epoch 的梯度更新次数。
            batch_size: 每个梯度更新的样本数。
            epochs: 训练轮数。
            verbose: {0, >0}，是否使用 Keras 的进度条回调。
            callbacks: Keras 回调列表。
            on_sample: 接受 sampler, idxs, w, scores 的回调函数。
            on_scores: 接受 sampler 和 scores 的回调函数。
        """
        try:
            if len(dataset.train_data) < batch_size:
                raise ValueError("batch_size 大于训练集大小")
        except RuntimeError as e:
            assert "no size" in str(e)

        # 设置 steps_per_epoch
        if steps_per_epoch is None:
            steps_per_epoch = len(dataset.train_data) // batch_size

        # 创建回调列表
        self.history = History()
        callbacks = [BaseLogger()] + (callbacks or []) + [self.history]
        if verbose > 0:
            callbacks += [ProgbarLogger(count_mode="steps")]
        callbacks = CallbackList(callbacks)
        callbacks.set_model(self.original_model)
        callbacks.set_params({
            "epochs": epochs,
            "steps": steps_per_epoch,
            "verbose": verbose,
            "do_validation": len(dataset.test_data) > 0,
            "metrics": self.metrics_names + [
                "val_" + n for n in self.metrics_names
            ]
        })

        # 创建采样器
        sampler = self.sampler(dataset, batch_size, steps_per_epoch, epochs)

        # 开始训练循环
        epoch = 0
        self.original_model.stop_training = False
        callbacks.on_train_begin()
        while epoch < epochs:
            callbacks.on_epoch_begin(epoch)

            for step in range(steps_per_epoch):
                batch_logs = {"batch": step, "size": batch_size}
                callbacks.on_batch_begin(step, batch_logs)

                # 在这里进行重要性采样
                idxs, (x_batch, y_batch), w = sampler.sample(batch_size)

                # 将 x_batch 和 y_batch 转换为 TensorFlow 张量
                x_batch = tf.convert_to_tensor(x_batch)
                y_batch = tf.convert_to_tensor(y_batch)
                w = tf.convert_to_tensor(w, dtype=tf.float32)

                # 自定义训练步骤
                loss_value, metrics_results, scores = self.train_batch(x_batch, y_batch, w)

                # 更新采样器
                sampler.update(idxs, scores)

                # 记录指标
                batch_logs['loss'] = loss_value
                for metric_name, metric_value in zip(self.metrics_names[1:], metrics_results):
                    batch_logs[metric_name] = metric_value
                callbacks.on_batch_end(step, batch_logs)

                if on_scores is not None and hasattr(self, "_latest_scores"):
                    on_scores(
                        sampler,
                        self._latest_scores
                    )

                if on_sample is not None:
                    on_sample(
                        sampler,
                        self._latest_sample_event["idxs"],
                        self._latest_sample_event["w"],
                        self._latest_sample_event["predicted_scores"]
                    )

                if self.original_model.stop_training:
                    break

            # 在每个 epoch 结束时进行评估
            epoch_logs = {}
            if len(dataset.test_data) > 0:
                val_loss, val_metrics = self.evaluate(dataset.test_data)
                epoch_logs['val_loss'] = val_loss
                for name, value in zip(self.metrics_names[1:], val_metrics):
                    epoch_logs['val_' + name] = value
            callbacks.on_epoch_end(epoch, epoch_logs)
            if self.original_model.stop_training:
                break
            epoch += 1
        callbacks.on_train_end()

        return self.history

    @property
    def metrics_names(self):
        def name(x):
            if isinstance(x, str):
                return x
            if hasattr(x, "name"):
                return x.name
            if hasattr(x, "__name__"):
                return x.__name__
            return str(x)

        metrics = self.model.model.metrics or []
        return (
                ["loss"] +
                [name(m) for m in metrics]
        )

    def train_batch(self, x, y, w):
        with tf.GradientTape() as tape:
            y_pred = self.model.model(x, training=True)
            per_sample_losses = self.model.model.loss(y, y_pred)
            if len(per_sample_losses.shape) > 1:
                per_sample_losses = tf.reduce_mean(per_sample_losses, axis=-1)
            weights = tf.cast(w, tf.float32)

            # Compute weighted loss
            weighted_losses = per_sample_losses * weights
            loss = tf.reduce_mean(weighted_losses)

            # Include regularization losses
            reg_losses = tf.reduce_sum(self.model.model.losses)
            loss += reg_losses

        # Get trainable variables
        trainable_vars = self.model.model.trainable_variables

        # Compute gradients
        gradients = tape.gradient(loss, trainable_vars)
        self.model.model.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # Update metrics
        self.model.model.compiled_metrics.update_state(y, y_pred, sample_weight=weights)
        metrics = [metric.result().numpy() for metric in self.model.model.metrics]

        return loss.numpy(), metrics, per_sample_losses.numpy()

    def evaluate(self, test_data):
        print(f"test_data type: {type(test_data)}, length: {len(test_data)}")
        print(f"test_data contents: {test_data}")

        x_test, y_test = test_data[:2]
        # 只获取前两个元素
        x_test, y_test = test_data[:2]

        # 如果 x_test 或 y_test 是元组，提取其中的数组
        if isinstance(x_test, tuple) or isinstance(x_test, list):
            x_test = x_test[0]
        if isinstance(y_test, tuple) or isinstance(y_test, list):
            y_test = y_test[0]

        # 确保 x_test 和 y_test 是 NumPy 数组或张量
        x_test = np.asarray(x_test)
        y_test = np.asarray(y_test)

        y_pred = self.model.model.predict(x_test, batch_size=128)

        # 确保 y_test 和 y_pred 是张量
        y_test = tf.convert_to_tensor(y_test)
        y_pred = tf.convert_to_tensor(y_pred)

        loss_value = self.model.model.compiled_loss(y_test, y_pred, regularization_losses=self.model.model.losses)

        # 更新指标
        self.model.model.compiled_metrics.update_state(y_test, y_pred)
        metrics_results = [m.result().numpy() for m in self.model.model.metrics]

        return loss_value.numpy(), metrics_results

class _UnbiasedImportanceTraining(_BaseImportanceTraining):
    def __init__(self, model, score="gnorm", layer=None):
        self._reweighting = BiasedReweightingPolicy(1.0)  # 无偏

        super(_UnbiasedImportanceTraining, self).__init__(model, score, layer)

    @property
    def reweighting(self):
        return self._reweighting

class ConstantTimeImportanceTraining(_UnbiasedImportanceTraining):
    """使用常时间重要性采样训练模型。

    参数：
        model: 要训练的 Keras 模型。
        score: {"gnorm", "loss", "full_gnorm"}，用于重要性采样的重要性度量。
        layer: None 或 int 或 Layer，计算 gnorm 时使用的层。
    """

    def __init__(self, model, backward_pass_multiplier=2.0, extra_samples=0.1,
                 score="gnorm", layer=None):
        self._backward_pass_multiplier = backward_pass_multiplier
        self._extra_samples = extra_samples

        super(ConstantTimeImportanceTraining, self).__init__(
            model,
            score,
            layer
        )

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
        return ConstantTimeSampler(
            dataset,
            self.reweighting,
            self.model,
            backward_time=self._backward_pass_multiplier,
        )
上述是training.py文件
from GAS.GAS_for_N_max_new import GroverOptimizer, GroverOptimizationResult


def _get_dataset_length(dset, default=1):
    """返回数据集的训练数据长度，如果无法获取，则返回默认值。"""
    try:
        return len(dset.train_data)
    except RuntimeError:
        return default

class BaseSampler(object):
    def __init__(self, dataset, reweighting):
        self.dataset = dataset
        self.reweighting = reweighting

    def _slice_data(self, x, y, idxs):
        if isinstance(x, (list, tuple)):
            return [xi[idxs] for xi in x], y[idxs]
        else:
            return x[idxs], y[idxs]

    def _send_messages(self, idxs, xy, w, predicted_scores):
        signal("is.sample").send({
            "idxs": idxs,
            "xy": xy,
            "w": w,
            "predicted_scores": predicted_scores
        })

    def _get_samples_with_scores(self):
        raise NotImplementedError()

    def sample(self, batch_size):
        raise NotImplementedError()

    def update(self, idxs, results):
        pass
import numpy as np
from blinker import signal
# class ConstantTimeSampler(BaseSampler):
#     def __init__(self, dataset, reweighting, model, B=32, alpha=0.6,
#                  beta_start=0.4, beta_frames=200, adapt_start=2.0,
#                  adapt_end=1.0, adapt_frames=200):
#         """
#         参数：
#             dataset: 数据集对象。
#             reweighting: 重加权策略。
#             model: 包含 score 方法的模型包装器。
#             B: 第一层随机采样的样本数量。
#             alpha: 决定了优先级的程度，0 表示均匀采样，1 表示完全按照优先级采样。
#             beta_start: beta 的初始值，用于重要性采样权重。
#             beta_frames: beta 从 beta_start 增加到 1 所需的总步数。
#             adapt_start: adapt 的初始值（默认值为 2.0）。
#             adapt_end: adapt 的最终值（默认值为 1.0）。
#             adapt_frames: adapt 从 adapt_start 变化到 adapt_end 所需的总步数。
#         """
#         self.model = model
#         self.N = len(dataset._x_train)
#         self.B = min(B, self.N)  # 确保 B 不超过 N
#         self.alpha = alpha
#
#         self.beta_start = beta_start
#         self.beta_frames = beta_frames
#         self.beta = beta_start
#
#         self.adapt_start = adapt_start
#         self.adapt_end = adapt_end
#         self.adapt_frames = adapt_frames
#         self.adapt = adapt_start
#
#         self.frame = 0  # 计步器
#         super(ConstantTimeSampler, self).__init__(dataset, reweighting)
#
#     def _update_beta(self):
#         # 线性增加 beta，从 beta_start 增加到 1，在 beta_frames 步内完成
#         self.beta = min(1.0, self.beta_start + (1.0 - self.beta_start) * (self.frame / self.beta_frames))
#
#     def _update_adapt(self):
#         # 线性调整 adapt，从 adapt_start 变化到 adapt_end，在 adapt_frames 步内完成
#         adapt_progress = min(1.0, self.frame / self.adapt_frames)
#         self.adapt = self.adapt_start + (self.adapt_end - self.adapt_start) * adapt_progress
#
#     def sample(self, batch_size):
#         # 更新 beta 和 adapt
#         self._update_beta()
#         self._update_adapt()
#         self.frame += 1
#
#         # 确保 batch_size 不超过数据集大小
#         if batch_size > self.N:
#             batch_size = self.N
#
#         # 第一层采样：从数据集中随机采样 B 个样本
#         sampled_indices = np.random.choice(self.N, self.B, replace=False)
#         x_sampled, y_sampled = self._slice_data(
#             self.dataset._x_train, self.dataset._y_train, sampled_indices
#         )
#
#         # 计算 B 个样本的得分
#         scores_sampled = self.model.score(x_sampled, y_sampled, batch_size=self.B)
#         scores_sampled = np.nan_to_num(scores_sampled, nan=0.0,
#                                        posinf=np.max(scores_sampled[~np.isinf(scores_sampled)]),
#                                        neginf=0.0)
#         scores_sampled += 1e-8  # 避免除零
#
#         # 使用 GroverOptimizer 来选取得分最高的 batch_size 个样本
#         # 首先需要将 scores_sampled 转换为整数，因为 GroverOptimizer 处理整数值
#         max_score = np.max(scores_sampled)
#         min_score = np.min(scores_sampled)
#         score_range = max_score - min_score
#         if score_range == 0:
#             # 如果所有得分相等，则将其设为一个常数整数值
#             integer_scores = np.full_like(scores_sampled, fill_value=1, dtype=int)
#         else:
#             # 将得分按比例缩放到 [0,15) 的区间（不包含 15）
#             num_value_qubits = 5  # 值量子比特为 5，控制值表示范围为 0~15
#             # 计算缩放因子，确保最大值映射到小于 15 的值
#             scale_factor = (15 - 1e-8) / score_range
#             scaled_scores = (scores_sampled - min_score) * scale_factor
#             # 对得分进行四舍五入
#             integer_scores = np.round(scaled_scores).astype(int)
#             # 将值为 15 的得分调整为 14（因为取值范围是 [0,15)）
#             integer_scores[integer_scores >= 15] = 14
#             # 将值为 0 的得分设为 1
#             integer_scores[integer_scores == 0] = 1
#
#         # 将 integer_scores 转换为列表，以便传递给 GroverOptimizer
#         integer_scores_list = integer_scores.tolist()
#
#         # 设置 GroverOptimizer 的参数
#         num_key_qubits = int(np.ceil(np.log2(self.B)))  # 键的量子比特数
#         num_iterations = 5  # 根据需要调整迭代次数
#
#         # 实例化 GroverOptimizer
#         optimizer = GroverOptimizer(num_value_qubits, num_key_qubits, num_iterations=num_iterations)
#
#         # 调用 GroverOptimizer 的 solve 方法，找到得分最高的 batch_size 个样本的索引
#         N = batch_size  # 我们要找到的最大的 N 个数
#         result = optimizer.solve(integer_scores_list, N)
#
#         # 获取选中的样本在 sampled_indices 中的索引
#         top_indices_in_sampled = result.keys  # 这些索引对应于 integer_scores_list 和 sampled_indices 的索引
#
#         # 将这些索引映射回 sampled_indices，得到第一段的索引
#         first_section_indices = np.array(top_indices_in_sampled)
#
#         # 剩余的索引即为第二段的索引
#         all_indices = np.arange(self.B)
#         second_section_indices = np.setdiff1d(all_indices, first_section_indices)
#
#         # 赋予排名
#         # 第一段的排名为 1 到 batch_size
#         first_section_ranks = np.arange(1, len(first_section_indices) + 1)
#
#         # 第二段样本的排名计算为 (batch_size + self.N) / adapt
#         # 为了简化，实现为对第二段的所有样本赋予相同的排名
#         second_section_rank_value = (batch_size + self.N) / self.adapt
#         second_section_ranks = np.full(len(second_section_indices), second_section_rank_value)
#
#         # 合并排名
#         ranks = np.concatenate([first_section_ranks, second_section_ranks])
#
#         # 计算优先级使用排名
#         # P(i) ∝ 1 / rank(i)
#         priorities = 1.0 / ranks
#
#         # 归一化优先级以获得采样概率
#         probabilities = priorities / np.sum(priorities)
#
#         # 根据计算的概率从 B 个样本中采样 batch_size 个样本
#         selected_indices_in_sampled = np.random.choice(
#             self.B, size=batch_size, replace=True, p=probabilities
#         )
#         selected_indices = sampled_indices[selected_indices_in_sampled]
#
#         # 计算重要性采样权重
#         N_total = self.N
#
#         P_selected = probabilities[selected_indices_in_sampled]
#         weights = np.power(P_selected * N_total, -self.beta)
#
#         # 将权重归一化到 [0, 1]
#         weights /= np.max(weights)
#         weights = weights.reshape(-1, 1)
#
#         # 获取选中的数据
#         x_selected, y_selected = self._slice_data(
#             self.dataset._x_train, self.dataset._y_train, selected_indices
#         )
#
#         # 选中样本的得分
#         scores_selected = scores_sampled[selected_indices_in_sampled]
#
#         # 发送采样信息（用于回调等）
#         self._send_messages(selected_indices, (x_selected, y_selected), weights, scores_selected)
#
#         return selected_indices, (x_selected, y_selected), weights
#
#     def update(self, idxs, results):
#         pass
import numpy as np
import logging
import math
logger = logging.getLogger(__name__)

class ConstantTimeSampler(BaseSampler):
    def __init__(self, dataset, reweighting, model, B=32, alpha=0.6,
                 beta_start=0.4, beta_frames=70, adapt_start=2.0,
                 adapt_end=1.0, adapt_frames=70):
        """
        参数：
            dataset: 数据集对象。
            reweighting: 重加权策略。
            model: 包含 score 方法的模型包装器。
            B: 第一层随机采样的样本数量。
            alpha: 决定了优先级的程度，0 表示均匀采样，1 表示完全按照优先级采样。
            beta_start: beta 的初始值，用于重要性采样权重。
            beta_frames: beta 从 beta_start 增加到 1 所需的总步数。
            adapt_start: adapt 的初始值（默认值为 2.0）。
            adapt_end: adapt 的最终值（默认值为 1.0）。
            adapt_frames: adapt 从 adapt_start 变化到 adapt_end 所需的总步数。
        """
        self.model = model
        self.dataset = dataset  # 保存数据集引用
        self.N = len(dataset._x_train)
        self.B = min(B, self.N)  # 确保 B 不超过 N
        self.alpha = alpha

        self.beta_start = beta_start
        self.beta_frames = beta_frames
        self.beta = beta_start

        self.adapt_start = adapt_start
        self.adapt_end = adapt_end
        self.adapt_frames = adapt_frames
        self.adapt = adapt_start

        self.frame = 0  # 计步器
        super(ConstantTimeSampler, self).__init__(dataset, reweighting)

    def _update_beta(self):
        # 线性增加 beta，从 beta_start 增加到 1，在 beta_frames 步内完成
        self.beta = min(1.0, self.beta_start + (1.0 - self.beta_start) * (self.frame / self.beta_frames))

    def _update_adapt(self):
        # 线性调整 adapt，从 adapt_start 变化到 adapt_end，在 adapt_frames 步内完成
        adapt_progress = min(1.0, self.frame / self.adapt_frames)
        self.adapt = self.adapt_start + (self.adapt_end - self.adapt_start) * adapt_progress

    def sample(self, batch_size):
        # 更新 beta 和 adapt
        self._update_beta()
        self._update_adapt()
        self.frame += 1

        # 确保 batch_size 不超过数据集大小
        if batch_size > self.N:
            batch_size = self.N

        # 第一层采样：从数据集中随机采样 B 个样本
        sampled_indices = np.random.choice(self.N, self.B, replace=False)
        x_sampled, y_sampled = self._slice_data(
            self.dataset._x_train, self.dataset._y_train, sampled_indices
        )

        # 计算 B 个样本的得分
        scores_sampled = self.model.score(x_sampled, y_sampled, batch_size=self.B)
        scores_sampled = np.nan_to_num(scores_sampled, nan=0.0,
                                       posinf=np.max(scores_sampled[~np.isinf(scores_sampled)]),
                                       neginf=0.0)
        scores_sampled += 1e-8  # 避免除零

        # 使用 GroverOptimizer 来选取得分最高的 batch_size 个样本
        # 首先需要将 scores_sampled 转换为整数，因为 GroverOptimizer 处理整数值
        max_score = np.max(scores_sampled)
        min_score = np.min(scores_sampled)
        score_range = max_score - min_score
        if score_range == 0:
            # 如果所有得分相等，则将其设为一个常数整数值
            integer_scores = np.full_like(scores_sampled, fill_value=1, dtype=int)
        else:
            # 将得分按比例缩放到 [0, 15) 的区间（不包含 15）
            num_value_qubits = math.ceil(math.log2(self.B))  # 动态计算量子比特数
            if num_value_qubits < 1:
                num_value_qubits = 1
            # 计算缩放因子，确保最大值映射到小于 15 的值
            scale_factor = (15 - 1e-8) / score_range
            scaled_scores = (scores_sampled - min_score) * scale_factor
            # 对得分进行四舍五入
            integer_scores = np.round(scaled_scores).astype(int)
            # 将值为 15 的得分调整为 14（因为取值范围是 [0,15)）
            integer_scores[integer_scores >= 15] = 14
            # 将值为 0 的得分设为 1
            integer_scores[integer_scores == 0] = 1

        # 将 integer_scores 转换为列表，以便传递给 GroverOptimizer
        integer_scores_list = integer_scores.tolist()

        # 设置 GroverOptimizer 的参数
        num_value_qubits = math.ceil(math.log2(15))  # 使用固定的量子比特数，例如4位可表示0-15
        num_ancil_qubits = 1  # 辅助量子比特数，根据需要调整
        num_iterations = 5  # 根据需要调整迭代次数

        # 实例化 GroverOptimizer
        optimizer = GroverOptimizer(
            num_value_qubits=num_value_qubits,
            num_ancil_qubits=num_ancil_qubits,
            num_iterations=num_iterations
        )

        # 调用 GroverOptimizer 的 solve 方法，找到得分最高的 batch_size 个样本的值
        N_max = batch_size  # 我们要找到的最大的 N 个数
        result: GroverOptimizationResult = optimizer.solve(numbers=integer_scores_list, N=N_max)

        # 获取最大 N 个值
        top_values = result.values  # 这些是得分最高的 N 个值

        # 将这些值映射回 sampled_indices 的索引
        # 由于可能存在重复值，我们需要找到所有匹配的索引
        first_section_indices = []
        for value in top_values:
            # 找到第一个匹配的索引
            matching_indices = np.where(integer_scores == value)[0]
            for idx in matching_indices:
                if idx not in first_section_indices:
                    first_section_indices.append(idx)
                    break  # 只取第一个匹配的索引

        first_section_indices = np.array(first_section_indices)

        # 剩余的索引即为第二段的索引
        all_indices = np.arange(self.B)
        second_section_indices = np.setdiff1d(all_indices, first_section_indices)

        # 赋予排名
        # 第一段的排名为 1 到 len(first_section_indices)
        first_section_ranks = np.arange(1, len(first_section_indices) + 1)

        # 第二段样本的排名计算为 (batch_size + self.N) / adapt
        # 为了简化，实现为对第二段的所有样本赋予相同的排名
        second_section_rank_value = (batch_size + self.N) / self.adapt
        second_section_ranks = np.full(len(second_section_indices), second_section_rank_value)

        # 合并排名
        ranks = np.concatenate([first_section_ranks, second_section_ranks])

        # 计算优先级使用排名
        # P(i) ∝ 1 / rank(i)
        priorities = 1.0 / ranks

        # 归一化优先级以获得采样概率
        probabilities = priorities / np.sum(priorities)

        # 根据计算的概率从 B 个样本中采样 batch_size 个样本
        selected_indices_in_sampled = np.random.choice(
            self.B, size=batch_size, replace=True, p=probabilities
        )
        selected_indices = sampled_indices[selected_indices_in_sampled]

        # 计算重要性采样权重
        N_total = self.N

        P_selected = probabilities[selected_indices_in_sampled]
        weights = np.power(P_selected * N_total, -self.beta)

        # 将权重归一化到 [0, 1]
        weights /= np.max(weights)
        weights = weights.reshape(-1, 1)

        # 获取选中的数据
        x_selected, y_selected = self._slice_data(
            self.dataset._x_train, self.dataset._y_train, selected_indices
        )

        # 选中样本的得分
        scores_selected = scores_sampled[selected_indices_in_sampled]

        # 发送采样信息（用于回调等）
        self._send_messages(selected_indices, (x_selected, y_selected), weights, scores_selected)

        return selected_indices, (x_selected, y_selected), weights

    def update(self, idxs, results):
        pass
上述是samplers.py文件
import numpy as np
from blinker import signal
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import Layer
from tensorflow.python.keras.metrics import categorical_accuracy
import tensorflow as tf
from .layers import GradientNormLayer, LossLayer, MetricLayer

class ModelWrapper(object):
    """ModelWrapper 的目的是包装一个神经网络，并添加一些额外的层，以生成得分、损失和样本权重，以执行重要性采样。"""

    def _iterate_batches(self, x, y, batch_size):
        bs = batch_size
        for s in range(0, len(y), bs):
            if isinstance(x, (list, tuple)):
                xs = [xi[s:s+bs] for xi in x]
            else:
                xs = x[s:s+bs]
            yield xs, y[s:s+bs]

    def evaluate(self, x, y, batch_size=128):
        result = np.mean(
            np.vstack([
                self.evaluate_batch(xi, yi)
                for xi, yi in self._iterate_batches(x, y, batch_size)
            ]),
            axis=0
        )

        signal("is.evaluation").send(result)
        return result

    def score(self, x, y, batch_size=128):
        bs = batch_size
        result = np.hstack([
            self.score_batch(xi, yi).T
            for xi, yi in self._iterate_batches(x, y, batch_size)
        ]).T

        signal("is.score").send(result)
        return result

    def set_lr(self, lr):
        """设置模型的学习率。"""

        try:
            K.set_value(
                self.optimizer.lr,
                lr
            )
        except AttributeError:
            try:
                K.set_value(
                    self.model.optimizer.lr,
                    lr
                )
            except AttributeError:
                raise NotImplementedError()

        try:
            K.set_value(
                self.small.optimizer.lr,
                lr
            )
        except AttributeError:
            pass

    def evaluate_batch(self, x, y):
        raise NotImplementedError()

    def score_batch(self, x, y):
        raise NotImplementedError()

    def train_batch(self, x, y, w):
        raise NotImplementedError()

class OracleWrapper(ModelWrapper):
    def __init__(self, model, reweighting, score="loss", layer=None):
        self.reweighting = reweighting
        self.layer = None  # Not applicable for QNN
        self.score_method = score
        self.model = model
        self.optimizer = model.optimizer

    def evaluate_batch(self, x, y):
        outputs = self.model.evaluate(
            x=x,
            y=y,
            verbose=0,
            return_dict=True
        )
        loss = outputs['loss']
        metrics = [outputs[m] for m in self.model.metrics_names[1:]]
        return [loss] + metrics
    #直接使用梯度
    # def score_batch(self, x, y):
    #     y_pred = self.model(x, training=False)
    #     if self.score_method == "gnorm":
    #         # 计算每个样本的梯度范数
    #         with tf.GradientTape() as tape:
    #             tape.watch(x)
    #             y_pred = self.model(x, training=False)
    #             loss = self.model.loss(y, y_pred)
    #         gradients = tape.gradient(loss, x)
    #         scores = tf.norm(gradients, axis=-1)
    #     else:
    #         # 默认使用损失作为得分
    #         per_sample_losses = self.model.loss(y, y_pred)
    #         if len(per_sample_losses.shape) > 1:
    #             per_sample_losses = tf.reduce_mean(per_sample_losses, axis=-1)
    #         scores = per_sample_losses
    #     return scores.numpy()
    #使用梯度上界
    def score_batch(self, x, y):
        y_pred = self.model(x, training=False)
        if self.score_method == "gnorm":
            # 根据论文方法计算梯度上界
            gradients = y_pred - y  # 适用于 softmax 激活和交叉熵损失
            scores = tf.norm(gradients, axis=1)
        else:
            # 默认使用损失作为得分
            per_sample_losses = self.model.loss(y, y_pred)
            if len(per_sample_losses.shape) > 1:
                per_sample_losses = tf.reduce_mean(per_sample_losses, axis=-1)
            scores = per_sample_losses
        return scores.numpy()

    def train_batch(self, x, y, w):
        with tf.GradientTape() as tape:
            y_pred = self.model(x, training=True)
            per_sample_losses = self.model.loss(y, y_pred)
            if len(per_sample_losses.shape) > 1:
                per_sample_losses = tf.reduce_mean(per_sample_losses, axis=-1)
            scores = per_sample_losses  # Use losses as importance scores

            weights = tf.cast(w, tf.float32)
            # Compute weighted loss
            weighted_losses = per_sample_losses * weights
            loss = tf.reduce_mean(weighted_losses)

            # Include regularization losses
            reg_losses = tf.reduce_sum(self.model.losses)
            loss += reg_losses

        # Get trainable variables
        trainable_vars = self.model.trainable_variables

        # Compute gradients
        gradients = tape.gradient(loss, trainable_vars)

        # Update model parameters
        self.model.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # Update metrics
        self.model.compiled_metrics.update_state(y, y_pred, sample_weight=weights)
        metrics = [metric.result().numpy() for metric in self.model.metrics]

        return loss.numpy(), metrics, scores.numpy()

    def __call__(self, x, training=False):
        return self.model(x, training=training)

def _tolist(x, acceptable_iterables=(list, tuple)):
    if not isinstance(x, acceptable_iterables):
        x = [x]
    return x

def _get_scoring_layer(score, y_true, y_pred, loss_fn="categorical_crossentropy",
                       layer=None, model=None):
    """获取一个计算每个 y_true 和 y_pred 对的得分的层"""
    assert score in ["loss", "gnorm", "full_gnorm", "acc"]

    if score == "loss":
        return LossLayer(loss_fn, name='score_layer')([
            y_true,
            y_pred
        ])
    elif score == "gnorm":
        return GradientNormLayer(
            layer.output,
            loss_fn,
            fast=False,
            name='score_layer'
        )([
            y_true,
            y_pred
        ])
    elif score == "full_gnorm":
        return GradientNormLayer(
            model.trainable_weights,
            loss_fn,
            fast=False,
            name='score_layer'
        )([
            y_true,
            y_pred
        ])
    elif score == "acc":
        return MetricLayer(categorical_accuracy, name='score_layer')([
            y_true,
            y_pred
        ])

class ModelWrapperDecorator(ModelWrapper):
    def __init__(self, model_wrapper, implemented_attributes=set()):
        self.model_wrapper = model_wrapper
        self.implemented_attributes = (
            implemented_attributes | set(["model_wrapper"])
        )

    def __getattribute__(self, name):
        _getattr = object.__getattribute__
        implemented_attributes = _getattr(self, "implemented_attributes")
        if name in implemented_attributes:
            return _getattr(self, name)
        else:
            model_wrapper = _getattr(self, "model_wrapper")
            return getattr(model_wrapper, name)
上述是model_wrappers.py文件
仔细阅读上述代码，我现在想要将神经网络的训练模拟修改为无噪声的，使用shots进行的形式，需要修改上述文件中的哪个文件？给出修改后的完整代码，格式为：XX代码XX行