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')  # Use non-interactive backend
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

# Ensure TensorFlow uses float32
tf.keras.backend.set_floatx('float32')

# Redirect stdout to capture print outputs
class Tee(object):
    def __init__(self, *files):
        self.files = files
        self.encoding = getattr(files[0], 'encoding', 'utf-8')  # Add encoding attribute

    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

# Configure Logger to use 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)

# Remove default handlers to avoid duplicate logs
if not logger.handlers:
    logger.addHandler(stream_handler)
else:
    logger.handlers = []
    logger.addHandler(stream_handler)

# Set Qiskit log level to WARNING to suppress INFO logs
logging.getLogger('qiskit').setLevel(logging.WARNING)

def save_plot(fig, filename):
    fig.savefig(filename)
    plt.close(fig)

def plot_history(history):
    # Plot training and validation accuracy
    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')

    # Plot training and validation loss
    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

# Define a quantum device with 11 qubits
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)

# Each neuron's encoding
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]

        # Flip control qubits if control value is 0
        for qubit, val in zip(control_qubits, cv):
            if val == 0:
                qml.PauliX(wires=qubit)

        # Apply double controlled RY gates
        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)

        # Restore control qubits
        for qubit, val in zip(control_qubits, cv):
            if val == 0:
                qml.PauliX(wires=qubit)

# Definition of swap gate
def apply_controlled_swap(swap_qubit, q_swap1, q_swap2):
    """
    Modular method to apply a controlled SWAP gate.

    Parameters:
    - swap_qubit (int): Control qubit for the SWAP gate.
    - q_swap1 (int): First qubit to swap.
    - q_swap2 (int): Second qubit to swap.
    """
    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]

    # Neuron groups
    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}
    ]

    # Angle encoding: inputs are mapped to [0, pi/2]
    angles_q2_list = inputs  # inputs shape is (4,)

    # Encoding weights mapped to angles_q3
    angles_q3_list = encoding_weights  # shape: (3, n_layers, 4)

    # 每个神经元中的每层如何编码（最终对该内容循环layers次）
    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]  # scalar value
            angle_q3_values = angles_q3_list[idx, layer, :]  # shape: (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, :]  # shape: (3,)
            qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)
        # Apply parameterized entanglement gates
        for qubit in range(rot_weights.shape[1] - 1):
            qml.CRZ(entangle_weights[layer, qubit], wires=[qubit, qubit + 1])

        # 每层每个神经元应用一次swap操作
        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'])
    # # Apply swap operations
    # 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'])

    # # Apply rotation layers and parameterized entanglement gates
    # for layer in range(rot_weights.shape[0]):
    #     for qubit in range(rot_weights.shape[1]):
    #         rot_angles = rot_weights[layer, qubit, :]  # shape: (3,)
    #         qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)
    #     # Apply parameterized entanglement gates
    #     for qubit in range(rot_weights.shape[1] - 1):
    #         qml.CRZ(entangle_weights[layer, qubit], wires=[qubit, qubit + 1])

    # Measure PauliZ expectation values on multiple qubits
    measured_qubits = [4, 7, 10]
    results = [qml.expval(qml.PauliZ(qubit)) for qubit in measured_qubits]

    return results

# Define the quantum neural network layer
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

        # Define encoding weights with L2 regularization
        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
        )

        # Define rotation layer weights with L2 regularization
        self.rot_weights = self.add_weight(
            shape=(n_layers, n_qubits, 3),
            initializer=initializer,
            regularizer=regularizer,
            trainable=True,
            name='rot_weights',
            dtype=tf.float32
        )

        # Define entangling weights with L2 regularization
        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

# Define the quantum neural network model
class QuantumNeuralNetwork(tf.keras.Model):
    def __init__(self, n_qubits, n_layers, num_classes, **kwargs):
        super(QuantumNeuralNetwork, self).__init__(**kwargs)
        self.n_qubits = n_qubits
        self.n_layers = n_layers
        self.num_classes = num_classes

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

        # Classical layers
        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
        )

        # Output layer directly connected to quantum outputs
        self.output_layer = Dense(
            num_classes,
            activation='softmax',
            kernel_regularizer=tf.keras.regularizers.l2(0.01),
            dtype=tf.float32  # Ensure the output layer is 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__":
    # Load data
    iris = load_iris()
    X = iris['data'].astype(np.float32)  # Ensure data type is float32
    y = iris['target'].reshape(-1, 1)

    # One-Hot encode labels
    encoder = OneHotEncoder(sparse_output=False)
    y_encoded = encoder.fit_transform(y).astype(np.float32)
    num_classes = y_encoded.shape[1]

    # Split into training and testing sets
    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
    )

    # Further split the training set into training and validation sets
    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
    )
    # Now, x_train is 60% of the data, x_val is 20%, x_test is 20%

    # Standardize and normalize data (fit only on training set)
    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)

    # Map features to [0, π/2] range
    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)

    # Hyperparameters
    learning_rate = 0.05  # Learning rate
    epochs = 100          # Number of training epochs
    batch_size = 16       # Batch size
    n_qubits = 14         # 14 qubits as per the original QNN
    n_layers = 3          # Number of layers

    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.")

    # Draw the quantum circuit
    sample_inputs = np.ones(4, dtype=np.float32) * (np.pi / 4)  # Inputs mapped to [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("\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

    # Create dataset (including validation set)
    dataset = InMemoryDataset(
        X_train_mapped,
        y_train,
        X_val_mapped,
        y_val,
        categorical=True
    )

    # Initialize reweighting policy
    reweighting_policy = BiasedReweightingPolicy(k=1.0)

    # Initialize importance sampling training
    wrapped_model = ConstantTimeImportanceTraining(model, score="loss")  # Use loss as score

    # Modify sampler to use dynamic alpha and beta
    def custom_sampler(dataset, batch_size, steps_per_epoch, epochs):
        return ConstantTimeSampler(
            dataset,
            reweighting_policy,
            model=wrapped_model.model,
        )

    # Assign custom sampler to wrapped_model
    wrapped_model.sampler = custom_sampler

    # Use per-sample loss function
    loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=False, reduction=tf.keras.losses.Reduction.NONE)

    # Compile model
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
        loss=loss_fn,
        metrics=['accuracy'],
        run_eagerly=False  # Set to False for better performance
    )

    # Create a callback class to save training metrics
    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'))

    # Instantiate callbacks
    metrics_history = MetricsHistory()

    # Add early stopping callback
    callbacks = [
        metrics_history,
        EarlyStopping(monitor='val_loss', patience=100, restore_best_weights=True)
    ]

    # Train model using 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
    )

    # Evaluate model on test set
    score = wrapped_model.evaluate((X_test_mapped, y_test))
    print('Test loss:', score[0])
    print('Test accuracy:', score[1])

    # Plot graphs using history from custom callback
    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)

    # Close output file and restore stdout
    sys.stdout = original_stdout
    output_file.close()

    print("Script execution completed. Please check 'output.log' for print output and check the saved plots in the current directory.")
将上述代码中的数据集修改为seeds数据集，将其特征通过降维使其降维至4维，不修改代码原有内容，给出修改后的完整代码