import pandas as pd
import time
import numpy as np
import matplotlib.pyplot as plt
import os
import torch
from torch import nn
from torch.cuda import seed_all
from torch.utils.data import Dataset, DataLoader
from collections import Counter
import torch.nn.functional
import pennylane as qml
from pennylane import numpy as np
from datasets import load_dataset
from tqdm import tqdm  # 添加进度条
import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
import re
from sklearn.model_selection import train_test_split
from qiskit import QuantumCircuit
from qiskit.visualization import circuit_drawer
from sklearn.metrics import f1_score, recall_score, accuracy_score
# from torch.utils.tensorboard import SummaryWriter
import random


seed = 42
# 设置随机种子
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)


os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
# 下载 nltk punkt 资源
nltk.download('punkt')
nltk.download('stopwords')

epoch = 40
min_freq = 1
sequence_length = 5
batch_size = 32
hidden_size = 64
embedding_dim = 100  # 词嵌入维度 与mat1有关
dropout_rate = 0.5
lr = 0.001
n_qlayers = 1
# weight_decay = 0.000001

ds = load_dataset("rahulsikder223/SentEval-MPQA")
# 合并训练集和测试集
combined_df = pd.concat([pd.DataFrame(ds['train']), pd.DataFrame(ds['test'])], ignore_index=True)

# 打印数据集的结构
print(ds)

# 文本清洗函数
def clean_text(text):
    # 转换为小写z
    text = text.lower()
    # 移除 HTML 标签
    text = re.sub(r'<.*?>', '', text)
    # 移除 URL
    text = re.sub(r'http\S+|www\S+|https\S+', '', text, flags=re.MULTILINE)
    # 移除标点符号和数字
    text = re.sub(r'\d+', '', text)
    text = re.sub(r'[^\w\s]', '', text)
    # 分词
    words = word_tokenize(text)
    # 移除停用词
    stop_words = set(stopwords.words('english'))- {'not','no','nor','but'}
    words = [word for word in words if word not in stop_words]
    return ' '.join(words)


# 清洗所有数据集的文本
combined_df['sentence'] = combined_df['sentence'].apply(clean_text)

# 打印清洗后的数据集的前几行
#print("\nCleaned Data Sample:")
#print(combined_df.head())
# 从训练集中划分出验证集
# 使用stratify参数确保划分后各类别比例保持不变

train_data, temp_data = train_test_split(combined_df, test_size=0.2, stratify=combined_df['label'], random_state=seed)
val_data, test_data = train_test_split(temp_data, test_size=0.5, stratify=temp_data['label'], random_state=seed)
# 定义一个函数来标记和编码句子
def tokenize_and_encode(sentences, vocab, seq_length):
    token_ids = [
        [vocab.get(token, vocab['<UNK>']) for token in sentence.split()]
        for sentence in sentences
    ]
    # 填充序列
    token_ids = [
        t_ids[:seq_length] if len(t_ids) > seq_length else t_ids + [vocab['<PAD>']] * (seq_length - len(t_ids))
        for t_ids in token_ids
    ]
    return token_ids

# 从训练数据集创建词汇表
all_train_tokens = [token for sentence in train_data['sentence'] for token in sentence.split()]

# 计算所有训练数据中单词的频率
word_counts = Counter(all_train_tokens)

# 打印词汇表中每个单词的频率，按照频率从高到低排序
#for word, count in word_counts.most_common():
    #print(f"{word}: {count}")

# 只保留出现次数大于等于阈值的单词

filtered_words = [word for word, count in word_counts.items() if count >= min_freq]
vocabulary = {word: i + 2 for i, word in enumerate(filtered_words)}
vocabulary['<UNK>'] = 1
vocabulary['<PAD>'] = 0

# 打印词汇表大小
print(f"Vocabulary size: {len(vocabulary)}")

# 对数据集进行标记和编码

train_token_ids = tokenize_and_encode(train_data['sentence'], vocabulary, sequence_length)
val_token_ids = tokenize_and_encode(val_data['sentence'], vocabulary, sequence_length)
test_token_ids = tokenize_and_encode(test_data['sentence'], vocabulary, sequence_length)

# 打印标记和编码后的数据样本
#print("\n标记和编码后的数据样本:")
#for sample in train_token_ids[:50]:
#     print(" ".join(map(str, sample)))

# 创建自定义 PyTorch 数据集
class SentimentDataset(Dataset):
    def __init__(self, token_ids, labels):
        self.token_ids = token_ids
        self.labels = labels

    def __len__(self):
        return len(self.token_ids)

    def __getitem__(self, idx):
        return torch.tensor(self.token_ids[idx], dtype=torch.long), torch.tensor(self.labels[idx], dtype=torch.long)

# 创建自定义 Dataset 实例

train_dataset = SentimentDataset(train_token_ids, train_data['label'].tolist())
val_dataset = SentimentDataset(val_token_ids, val_data['label'].tolist())
test_dataset = SentimentDataset(test_token_ids, test_data['label'].tolist())

# 为每个数据集创建 DataLoader

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 训练集集的标签统计
train_positive_count = train_data[train_data['label'] == 1].shape[0]
train_negative_count = train_data[train_data['label'] == 0].shape[0]
print(f"Train Positive count: {train_positive_count}")
print(f"Train Negative count: {train_negative_count}")

# 验证集的标签统计
val_positive_count = val_data[val_data['label'] == 1].shape[0]
val_negative_count = val_data[val_data['label'] == 0].shape[0]
print(f"Validation Positive count: {val_positive_count}")
print(f"Validation Negative count: {val_negative_count}")

# 测试集的标签统计
test_positive_count = test_data[test_data['label'] == 1].shape[0]
test_negative_count = test_data[test_data['label'] == 0].shape[0]
print(f"Test Positive count: {test_positive_count}")
print(f"Test Negative count: {test_negative_count}")



#自注意力类

# 量子LSTM
class QLSTM(nn.Module):
    def __init__(self,
                 input_size,
                 hidden_size,
                 n_qubits=4,
                 n_qlayers=1,
                 n_vrotations=3,
                 batch_first=True,
                 return_sequences=False,
                 return_state=False,
                 backend="default.qubit",
                 dropout_rate=0.5):
        super(QLSTM, self).__init__()
        self.n_inputs = input_size
        self.hidden_size = hidden_size
        self.num_layers = n_qlayers  # 确保记录层数
        self.concat_size = self.n_inputs + self.hidden_size
        self.n_qubits = n_qubits
        self.n_qlayers = n_qlayers
        self.n_vrotations = n_vrotations
        self.backend = backend

        self.batch_first = batch_first
        self.return_sequences = return_sequences
        self.return_state = return_state

        self.wires_forget = [f"wire_forget_{i}" for i in range(self.n_qubits)]
        self.wires_input = [f"wire_input_{i}" for i in range(self.n_qubits)]
        self.wires_update = [f"wire_update_{i}" for i in range(self.n_qubits)]
        self.wires_output = [f"wire_output_{i}" for i in range(self.n_qubits)]
        self.wires_output_h = [f"wire_output_h_{i}" for i in range(self.n_qubits)]

        # 定义量子设备
        self.dev_forget = qml.device(self.backend, wires=self.wires_forget)
        self.dev_input = qml.device(self.backend, wires=self.wires_input)
        self.dev_update = qml.device(self.backend, wires=self.wires_update)
        self.dev_output = qml.device(self.backend, wires=self.wires_output)
        self.dev_output_h = qml.device(self.backend, wires=self.wires_output_h)

        # 初始化dropout层
        self.dropout = nn.Dropout(dropout_rate)

        # QNode实例化
        self.qlayer_forget = qml.QNode(self._circuit_forget, self.dev_forget, interface="torch")
        self.qlayer_input = qml.QNode(self._circuit_input, self.dev_input, interface="torch")
        self.qlayer_update = qml.QNode(self._circuit_update, self.dev_update, interface="torch")
        self.qlayer_output = qml.QNode(self._circuit_output, self.dev_output, interface="torch")
        self.qlayer_output_h = qml.QNode(self._circuit_output_h, self.dev_output_h, interface="torch")

        # 定义权重形状
        self.weight_shapes = {"weights": (self.n_qlayers, self.n_vrotations, self.n_qubits)}

        # 定义线性层
        self.clayer_in = nn.Linear(self.concat_size, self.n_qubits)
        self.clayer_out_f = nn.Linear(self.n_qubits, self.hidden_size)
        self.clayer_out_i = nn.Linear(self.n_qubits, self.hidden_size)
        self.clayer_out_g = nn.Linear(self.n_qubits, self.hidden_size)
        self.clayer_out_o = nn.Linear(self.n_qubits, self.hidden_size)

        # 定义 TorchLayer
        self.TorchLayers = {
            'forget': qml.qnn.TorchLayer(self.qlayer_forget, self.weight_shapes),
            'input': qml.qnn.TorchLayer(self.qlayer_input, self.weight_shapes),
            'update': qml.qnn.TorchLayer(self.qlayer_update, self.weight_shapes),
            'output': qml.qnn.TorchLayer(self.qlayer_output, self.weight_shapes),
            'output_h': qml.qnn.TorchLayer(self.qlayer_output_h, self.weight_shapes)
        }

    def vqc(self, features, weights, wires_type):
        # 确保 features 是二维张量 [batch_size, n_qubits]
        if features.dim() == 1:
            features = features.unsqueeze(0)
        if not isinstance(features, torch.Tensor):
            features = torch.tensor(features, dtype=torch.float32)

        # 损切 features
        features = features[:, :self.n_qubits]
        ry_params = torch.atan2(features, torch.ones_like(features))
        rz_params = torch.atan2(features ** 2, torch.ones_like(features))

        # 应用 Hadamard 门
        for i in range(self.n_qubits):
            qml.Hadamard(wires=wires_type[i])

        # 应用 RY 和 RZ 旋转
        for i in range(self.n_qubits):
            qml.RY(ry_params[:, i], wires=wires_type[i])
            qml.RZ(rz_params[:, i], wires=wires_type[i])

        # 创建纠缠
        for i in range(1, 3):
            for j in range(self.n_qubits):
                if j + i < self.n_qubits:
                    qml.CNOT(wires=[wires_type[j], wires_type[j + i]])
                else:
                    qml.CNOT(wires=[wires_type[j], wires_type[j + i - self.n_qubits]])

        # 应用 RX、RY、RZ 旋转
        for layer in range(self.n_qlayers):
            for i in range(self.n_qubits):
                qml.RX(weights[layer, 0, i], wires=wires_type[i])
                qml.RY(weights[layer, 1, i], wires=wires_type[i])
                qml.RZ(weights[layer, 2, i], wires=wires_type[i])

    # 定义量子电路
    def _circuit_forget(self, inputs, weights):
        self.vqc(inputs, weights, self.wires_forget)
        return [qml.expval(qml.PauliZ(wires=i)) for i in self.wires_forget]

    def _circuit_input(self, inputs, weights):
        self.vqc(inputs, weights, self.wires_input)
        return [qml.expval(qml.PauliZ(wires=i)) for i in self.wires_input]

    def _circuit_update(self, inputs, weights):
        self.vqc(inputs, weights, self.wires_update)
        return [qml.expval(qml.PauliZ(wires=i)) for i in self.wires_update]

    def _circuit_output(self, inputs, weights):
        self.vqc(inputs, weights, self.wires_output)
        return [qml.expval(qml.PauliZ(wires=i)) for i in self.wires_output]

    def _circuit_output_h(self, inputs, weights):
        self.vqc(inputs, weights, self.wires_output_h)
        return [qml.expval(qml.PauliZ(wires=i)) for i in self.wires_output_h]

    # 量子电路可视化
    def to_qiskit_circuit(self, inputs, weights, wires_type):
        circuit = QuantumCircuit(self.n_qubits, self.n_qubits)
        features = torch.tensor(inputs[:, :self.n_qubits], dtype=torch.float32)
        ry_params = torch.atan2(features, torch.ones_like(features))
        rz_params = torch.atan2(features ** 2, torch.ones_like(features))

        # 应用 Hadamard 门
        for i in range(self.n_qubits):
            circuit.h(i)

        # 应用 RY 和 RZ 旋转
        for i in range(self.n_qubits):
            circuit.ry(ry_params[:, i].item(), i)
            circuit.rz(rz_params[:, i].item(), i)

        # 创建纠缠
        for i in range(1, 3):
            for j in range(self.n_qubits):
                if j + i < self.n_qubits:
                    circuit.cx(j, j + i)
                else:
                    circuit.cx(j, (j + i) % self.n_qubits)

        # 应用 RX、RY、RZ 旋转
        for layer in range(self.n_qlayers):
            for i in range(self.n_qubits):
                circuit.rx(weights[layer, 0, i].item(), i)
                circuit.ry(weights[layer, 1, i].item(), i)
                circuit.rz(weights[layer, 2, i].item(), i)

        # 测量
        for i in range(self.n_qubits):
            circuit.measure(i, i)

        return circuit

    def visualize_qiskit_circuit(self):
        example_inputs = np.random.random((1, self.n_qubits))
        example_weights = np.random.random((self.n_qlayers, self.n_vrotations, self.n_qubits))
        qiskit_circuit = self.to_qiskit_circuit(example_inputs, example_weights, self.wires_forget)
        circuit_drawer(qiskit_circuit, output='mpl').show()

    def forward(self, x, init_states=None):
        batch_size, seq_length, _ = x.size()

        if init_states is None:
            h_t = torch.zeros(self.num_layers, batch_size, self.hidden_size, dtype=torch.float32, device=x.device)
            c_t = torch.zeros(self.num_layers, batch_size, self.hidden_size, dtype=torch.float32, device=x.device)
        else:
            h_t, c_t = init_states

        for t in range(seq_length):
            x_t = x[:, t, :]  # [batch_size, input_size]
            h_t_last_layer = h_t[-1, :, :]  # [batch_size, hidden_size]

            combined = torch.cat((h_t_last_layer, x_t), dim=1)  # [batch_size, hidden_size + input_size]
            combined = self.clayer_in(combined)  # [batch_size, n_qubits]

            # 用量子层处理整个批次的数据
            f_t_intermediate = self.TorchLayers['forget'](combined)
            i_t_intermediate = self.TorchLayers['input'](combined)
            g_t_intermediate = self.TorchLayers['update'](combined)
            o_t_intermediate = self.TorchLayers['output'](combined)

            # 将中间结果通过另一个线性层转换回隐藏状态的维度
            f_t = torch.sigmoid(self.clayer_out_f(f_t_intermediate))  # [batch_size, hidden_size]
            i_t = torch.sigmoid(self.clayer_out_i(i_t_intermediate))
            g_t = torch.tanh(self.clayer_out_g(g_t_intermediate))
            o_t = torch.sigmoid(self.clayer_out_o(o_t_intermediate))

            # 根据 LSTM 的更新规则更新隐藏状态和细胞状态
            c_t_new = f_t * c_t[-1, :, :] + i_t * g_t  # [batch_size, hidden_size]
            h_t_new = o_t * torch.tanh(c_t_new)  # [batch_size, hidden_size]
            h_t_new = self.dropout(h_t_new)  # [batch_size, hidden_size]

            # 更新隐藏状态和细胞状态
            h_t = h_t.clone()
            h_t[-1, :, :] = h_t_new
            c_t = c_t.clone()
            c_t[-1, :, :] = c_t_new

        # 不再使用 self.fc，这里直接返回最后一层的隐藏状态
        return h_t_last_layer  # [batch_size, hidden_size]

# 电路可视化实现
qlstm_instance = QLSTM(input_size=10, hidden_size=64, n_qubits=4, n_qlayers=n_qlayers, n_vrotations=3, backend="default.qubit")
qlstm_instance.visualize_qiskit_circuit()
class QShallowClassificationLSTM(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_size, n_qubits=4, n_qlayers=1, num_classes=2, dropout_rate=0.5):
        super().__init__()
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_layers = n_qlayers
        self.num_classes = num_classes
        self.embedding = nn.Embedding(vocab_size, embedding_dim)

        # 添加 Dropout 层
        self.dropout = nn.Dropout(dropout_rate)

        # Quantum LSTM
        self.lstm = QLSTM(
            input_size=embedding_dim,
            hidden_size=hidden_size,
            n_qubits=n_qubits,
            n_qlayers=n_qlayers,
        )

        # 线性层将 hidden_size 映射到 num_classes
        self.linear = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        batch_size = x.shape[0]

        x = self.embedding(x)  # [batch_size, sequence_length, embedding_dim]
        x = self.dropout(x)

        h0 = torch.zeros(self.num_layers, batch_size, self.hidden_size, dtype=torch.float32, device=x.device)
        c0 = torch.zeros(self.num_layers, batch_size, self.hidden_size, dtype=torch.float32, device=x.device)

        # QLSTM 返回最后一层的隐藏状态 [batch_size, hidden_size]
        outputs = self.lstm(x, (h0, c0))  # [batch_size, hidden_size]

        # 将隐藏状态传递给线性层进行分类
        out = self.linear(outputs)  # [batch_size, num_classes]

        return out

#  计算模型分类准确率
def calculate_accuracy_and_scores(data_loader, Qmodel):
    all_labels = []
    all_predictions = []

    Qmodel.eval()
    with torch.no_grad():
        for X, y in data_loader:
            outputs = Qmodel(X)
            _, predicted = torch.max(outputs.data, 1)
            all_labels.extend(y.cpu().numpy())
            all_predictions.extend(predicted.cpu().numpy())

    # Compute metrics using all labels and predictions
    accuracy = 100 * accuracy_score(all_labels, all_predictions)
    recall = recall_score(all_labels, all_predictions, pos_label=1)
    f1 = f1_score(all_labels, all_predictions, pos_label=1)

    return accuracy, recall, f1

# 词汇表大小
vocab_size = len(vocabulary)




# 实例化分类模型和量子分类模型

Qmodel = QShallowClassificationLSTM(vocab_size=len(vocabulary), embedding_dim=embedding_dim, hidden_size=hidden_size, n_qubits=4, n_qlayers=n_qlayers, num_classes=2)

# 定义损失函数和优化器
loss_function = nn.CrossEntropyLoss()  # 更新为分类任务的损失函数

Qoptimizer = torch.optim.Adam(Qmodel.parameters(), lr=lr)
# Qoptimizer = torch.optim.Adam(Qmodel.parameters(), lr=lr, weight_decay=weight_decay)

#  训练函数
def train_model(data_loader, Qmodel, loss_function, optimizer,  progress_bar):
    total_loss = 0
    Qmodel.train()

    for X, y in data_loader:
        output = Qmodel(X)
        loss = loss_function(output, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        progress_bar.update(1)  # 更新进度条

    avg_loss = total_loss / len(data_loader)
    return avg_loss

#  测试函数
def test_model(data_loader, Qmodel, loss_function, progress_bar):

    total_loss = 0
    Qmodel.eval()
    with torch.no_grad():
        for X, y in data_loader:
            output = Qmodel(X)
            total_loss += loss_function(output, y).item()
            progress_bar.update(1)  # 更新进度条

    avg_loss = total_loss / len(data_loader)
    return avg_loss



# 初始化训练、测试、验证损失，验证集准确率
quantum_loss_train = []
quantum_loss_test = []
quantum_loss_val = []
quantum_accuracy_val = []
print("Quantum Untrained test\n--------")

# 初始化存储准确率的列表
train_accuracy_list = []
train_recall_list = []
train_f1_list = []

val_accuracy_list = []
val_recall_list = []
val_f1_list = []

# 测试一个批次的数据
for X, y in train_loader:
    print(f"Input Shape: {X.shape}")   # 应为 [batch_size, sequence_length]
    print(f"Target Shape: {y.shape}")  # 应为 [batch_size]
    break

for ix_epoch in range(epoch):
    #  print(f"Epoch {ix_epoch+1}\n---------")
    total_steps = len(train_loader) + len(val_loader)
    progress_bar = tqdm(total=total_steps, desc=f"Epoch {ix_epoch + 1}")

    # 训练损失
    train_loss = train_model(train_loader, Qmodel, loss_function, Qoptimizer, progress_bar)

    # 验证损失
    val_loss = test_model(val_loader, Qmodel, loss_function, progress_bar)

    # 计算量子模型的准确率、召回率和F1分数
    train_accuracy, train_recall, train_f1 = calculate_accuracy_and_scores(train_loader, Qmodel)

    # 计算验证集准确率、召回率和F1分数
    val_accuracy, val_recall, val_f1 = calculate_accuracy_and_scores(val_loader, Qmodel)

    # 关闭进度条
    progress_bar.close()


    # 将准确率添加到相应的列表中
    train_accuracy_list.append(train_accuracy)
    train_recall_list.append(train_recall)
    train_f1_list.append(train_f1)

    val_accuracy_list.append(val_accuracy)
    val_recall_list.append(val_recall)
    val_f1_list.append(val_f1)

    # 记录训练、测试、验证损失
    quantum_loss_train.append(train_loss)
    quantum_loss_val.append(val_loss)
    quantum_accuracy_val.append(val_accuracy)

    # 打印损失和评估指标信息
    print(f"\nTrain Loss: {train_loss:.4f}")
    print(f"Validation Loss: {val_loss:.4f}")

    print(f"Train Accuracy: {train_accuracy:.4f}%")
    print(f"Validation Accuracy: {val_accuracy:.4f}%\n")

    print(f"Train Recall: {train_recall:.4f}")
    print(f"Train F1 Score: {train_f1:.4f}")

    print(f"Validation Recall: {val_recall:.4f}")
    print(f"Validation F1 Score: {val_f1:.4f}\n")


# 训练和验证完成后，使用测试集进行评估
test_accuracy, test_recall, test_f1 = calculate_accuracy_and_scores(test_loader, Qmodel)

# 打印测试集评估指标
print(f'测试集评估结果：')
print(f"Test Accuracy: {test_accuracy:.4f}%")
print(f"Test Recall: {test_recall:.4f}")
print(f"Test F1 Score: {test_f1:.4f}")

# 绘制量子 QLSTM 模型的损失图表

plt.figure(figsize=(12, 6))
plt.plot(quantum_loss_train, label='Quantum Train Loss')
plt.plot(quantum_loss_test, label='Quantum Test Loss')
plt.plot(quantum_loss_val, label='Validation Loss')
plt.title('Quantum LSTM Model Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Validation Loss across Folds')
plt.xlabel('Epoch')
plt.ylabel('Validation Loss')
plt.legend()
plt.show()

# 绘制准确率图表
plt.figure(figsize=(12, 6))
plt.plot(train_accuracy_list, label='Train Accuracy')
plt.plot(val_accuracy_list, label='Validation Accuracy')
# plt.plot(test_accuracy_list, label='Test Accuracy')

plt.plot(train_recall_list, label='Train Recall')
plt.plot(val_recall_list, label='Validation Recall')
# plt.plot(test_recall_list, label='Test Recall')

plt.plot(train_f1_list, label='Train F1 Score')
plt.plot(val_f1_list, label='Validation F1 Score')
# plt.plot(test_f1_list, label='Test F1 Score')

# 模型性能指标
plt.title('Model Evaluation Metrics Across Epochs')
plt.xlabel('Epoch')
plt.ylabel('Metrics')
plt.legend()
plt.show()
我的模型是一个量子混合经典的模型，其中我设计了变分量子电路来修改替换经典LSTM模型中的门控机制。请你综合分析我的代码，为我的代码添加一些量子模块，或者将我的模型中的经典LSTM的某些操作更换成量子操作。请你给出可能的修改方案