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.05  # 学习率
    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=10, 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' 以获取打印输出，并在当前目录下查看保存的图表。")
上述代码输出结果中为什么测试准确率包含两个值0.11和0.94，这是哪句代码决定的？还有为什么会产生两个值