from fileinput import filename

import dgl
import torch
from dgl.nn.pytorch import GraphConv
from torch import nn, Tensor, softmax
from torch.utils.data import DataLoader
import time
from tqdm import tqdm
import numpy as np
import torch
from transformers import AutoModel, AutoTokenizer
from sklearn.metrics.pairwise import cosine_similarity
import dgl
import scipy.sparse as spp

import torch.nn.functional as F
import re
import os
from GdataLoad import LuDataset

def get_top_k_related_tensors(question_tensor, text_tensors):
    # 计算余弦相似度
    similarities = []
    for i, text_tensor in enumerate(text_tensors):
        similarity = F.cosine_similarity(question_tensor, text_tensor, dim=0)
        similarities.append((similarity.item(), i, text_tensor))
    # 按相似度排序
    similarities.sort(key=lambda x: x[0], reverse=True)
    # 获取前k个最相关的tensor及其对应的下标
    top_related = similarities[:]
    # 分离出tensor和下标
    top_tensors = [item[2] for item in top_related]
    top_indices = [item[1] for item in top_related]
    top_scores = [item[0] for item in top_related]
    return top_tensors, top_indices,top_scores
#model_name="C:/Users/75187/.cache/huggingface/hub/models--TencentBAC--Conan-embedding-v1/snapshots/b373e4db6fd7b536baf2eb85ab2331ccec7bd2f8"

model_name = "./model/BAAI/bge-large-zh-v1.5"
# 下载模型
model = AutoModel.from_pretrained(model_name)
# 下载分词器
tokenizer = AutoTokenizer.from_pretrained(model_name)
# 检查是否有可用的GPU
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# 将模型迁移到GPU
model.to(device)
text_blocks=[]

def read_file_and_split_by_lines(directory_path, x=1):
    # 改进正则表达式：匹配句子结束符（包括中文和英文的句号、问号、感叹号）
    sentence_endings = re.compile(r'([。！？])')
    sentence_indices = []
    sentences = []
    file_names=[]

    # 遍历指定目录下的所有文件
    for file_name in os.listdir(directory_path):
        file_path = os.path.join(directory_path, file_name)

        if os.path.isfile(file_path):
            with open(file_path, 'r', encoding='utf-8') as file:
                file_rels = []
                lines = file.readlines()

                for line in lines:
                    if line.strip():
                        # 按句子结束符号分割，并保留分割符
                        split_sentences = sentence_endings.split(line.strip())
                        temp_sentence = ""
                        for i, s in enumerate(split_sentences):
                            if sentence_endings.match(s):
                                # 如果当前部分是标点符号，就附加回前一个句子
                                temp_sentence += s
                                file_rels.append(len(sentences))
                                sentences.append(temp_sentence)
                                temp_sentence = ""
                                file_names.append(file_name)
                            else:
                                # 如果当前部分不是标点符号
                                if temp_sentence:
                                    file_rels.append(len(sentences))
                                    sentences.append(temp_sentence)
                                    file_names.append(file_name)
                                temp_sentence = s

                        # 如果最后一个分割结果不是标点符号，把它也加入
                        if temp_sentence:
                            file_rels.append(len(sentences))
                            sentences.append(temp_sentence)
                            file_names.append(file_name)
                sentence_indices.append(file_rels)
                 # 记录当前文件名


    return sentences, sentence_indices,file_names
def retain_nodes_links(adj_matrix, nodes_to_retain):
    """
    只保留邻接矩阵中指定节点之间的链接，断开其他所有链接

    :param adj_matrix: 2D numpy array, 邻接矩阵
    :param nodes_to_retain: list of int, 需要保留链接的节点列表
    :return: 修改后的邻接矩阵
    """
    # 创建一个全零矩阵
    new_adj_matrix = np.zeros_like(adj_matrix)

    for i in nodes_to_retain:
        for j in nodes_to_retain:
            if i != j:  # 避免自循环
                new_adj_matrix[i, j] = adj_matrix[i, j]

    return new_adj_matrix

import torch
import torch.nn as nn
import dgl
import dgl.function as fn


class LuGTP(nn.Module):
    def __init__(self,hiddsize=1024, k=2):
        super(LuGTP, self).__init__()
        self.k = k
        self.relu = nn.ReLU()
        self.layer_norm = nn.LayerNorm(hiddsize)
        self.gconv1 = GraphConv(hiddsize, hiddsize, weight=True, bias=True)
        self.gconv2 = GraphConv(hiddsize, hiddsize, weight=True, bias=True)

    def forward(self, start,question_feature_vector,gbatch):
        all_results=[]
        initial_feat = gbatch.ndata['feat'].clone()

        f_d={}
        for _ in range(self.k):
            gbatch.update_all(fn.copy_u(u='feat', out='m'), fn.sum(msg='m', out='feat'))
            g1 = gbatch.ndata['feat']
            g1 = self.gconv1(gbatch, gbatch.ndata['feat'])
            g1 = self.relu(g1)
            # gbatch.ndata['feat'] = self.relu(g1)
            degrees = gbatch.in_degrees().float().clamp(min=1)
            dk = degrees.view(-1, 1)
            gbatch.ndata['feat'] = g1 / dk

            g1 = gbatch.ndata['feat'] + initial_feat  # Skip connection
            g1 = self.layer_norm(g1)
            g1 = self.gconv1(gbatch, g1)
            g1 = self.relu(g1)
            gbatch.ndata['feat'] = g1

            G_text_tensors= gbatch.ndata['feat']
            top_tensors, top_indices, top_scores = get_top_k_related_tensors(question_feature_vector, G_text_tensors)
            # 使用列表推导将原列表中的每个元素加上k
            top_indices = [x + start for x in top_indices]
            # 将结果保存到all_results中
            for tensor, index, score in zip(top_tensors, top_indices, top_scores):
                all_results.append((tensor, index, score))
            for f,id in label_to_idx.items():
                f_d[f]=0
            for r in all_results:
                inndex=r[1]
                fnn=file_names[inndex].rsplit('.', 1)[0]
                score=r[2]
                if fnn in f_d:
                    s=f_d[fnn]
                    if score>s:
                        f_d[fnn]=score
            # Convert dictionary values to a list
            values = list(f_d.values())
            # Convert list to a tensor
            tensor = torch.tensor(values)

        return tensor

def create_element_to_list_map(lists):
    element_to_list_map = {}
    for index, sublist in enumerate(lists):
        for element in sublist:
            element_to_list_map[element] = index
    return element_to_list_map
def are_in_same_list(element_to_list_map, num1, num2):
    return element_to_list_map.get(num1) == element_to_list_map.get(num2)
# Create the hash map
element_to_list_map = None
label_to_idx= {}
file_names=[]
def generol_search(modle,question,docs_path,reload=False,first_state_threshold=0.62,information_transfer_threshold=0.85):
    text_features=[]
    global text_blocks,element_to_list_map,file_names
    text_blocks,relate_index,file_names = read_file_and_split_by_lines(docs_path)[:]
    number=0
    for f in file_names:
        f=f.rsplit('.', 1)[0]
        if f not in label_to_idx:
            label_to_idx[f]=number
            number=number+1
    element_to_list_map=create_element_to_list_map(relate_index)
    if reload:
        start_time = time.time()
        for t in tqdm(text_blocks, desc="Embedding texts"):
            encoded_input = tokenizer(t, return_tensors='pt', padding=True, max_length=512, truncation=True)
            # 将输入数据也迁移到GPU
            encoded_input = {key: value.to(device) for key, value in encoded_input.items()}
            # 获取模型的输出
            with torch.no_grad():
                output = model(**encoded_input)
            # 通常我们使用 [CLS] token 的输出作为特征向量
            feature_vector = output.last_hidden_state[:, 0, :][0]
            text_features.append(feature_vector.cpu().squeeze().numpy())
        # 保存到文件
        np.save('features.npy', text_features)
        # 记录结束时间
        end_time = time.time()
        # 计算并输出耗时
        elapsed_time = end_time - start_time
        print(f"编码features.npy执行时间: {elapsed_time:.6f}秒")
    else:
        start_time = time.time()
        text_features = np.load('features.npy', allow_pickle=True)
        # 记录结束时间
        end_time = time.time()
        # 计算并输出耗时
        elapsed_time = end_time - start_time
    encoded_input = tokenizer(question, return_tensors='pt', padding=True, truncation=True)
    # 将输入数据也迁移到GPU
    encoded_input = {key: value.to(device) for key, value in encoded_input.items()}
    # 获取模型的输出
    with torch.no_grad():
        output = model(**encoded_input)
        # 通常我们使用 [CLS] token 的输出作为特征向量
        question_feature_vector = output.last_hidden_state[:, 0, :][0]
    q_a_relate_dict_={}
    #print("~~~~~~~~~~~~~~~~~~~~~~~first_state_nodes~~~~~~~~~~~~~~~~~~~~~~~~~~")
    for i, text_tensor in enumerate(text_features):
        similarity = F.cosine_similarity(question_feature_vector, torch.tensor(text_tensor).to(device), dim=0)
        if similarity>first_state_threshold:
            g_node=element_to_list_map[i]
            # 检查字典中是否存在键 g_node，如果不存在则初始化
            if g_node not in q_a_relate_dict_:
                q_a_relate_dict_[g_node] = []
            q_a_relate_dict_[g_node].append(i)
            # print(g_node,similarity,file_names[i],text_blocks[i])
    #计算特征向量之间的余弦相似度
    similarity_matrix = cosine_similarity(text_features)
    # 结果列表
    time_start = time.time()
    # 结果列表

    # 记录开始时间
    start_time = time.time()
    # 初始化保存所有结果的列表
    all_results = []
    G_text_tensorslist=[]
    for g_number,relate_sentence_list in q_a_relate_dict_.items():
        full_relate_list=relate_index[g_number]
        matrix_len=len(full_relate_list)
        start=full_relate_list[0]
        end=full_relate_list[-1]
        matrix = np.ones((matrix_len, matrix_len)).astype('int64')
        result_list = []
        iteration = 0
        max_iterations = 20
        while iteration < max_iterations:
            result_list = []
            for r in relate_sentence_list:
                for i in range(matrix_len):
                    id = full_relate_list[i]
                    if similarity_matrix[r][id] > information_transfer_threshold:
                        result_list.append(r)
                        result_list.append(id)
            result_list = list(set(result_list))
            if len(result_list) == 1:
                information_transfer_threshold -= 0.005
            elif len(result_list) > 20:
                information_transfer_threshold += 0.005
            else:
                break
            iteration += 1
        new_result_list = [x - start for x in result_list]
        new_adj_matrix = retain_nodes_links(matrix, new_result_list)
        matrix = new_adj_matrix
        adj = spp.coo_matrix(matrix)
        G = dgl.from_scipy(adj).to("cuda")

        feat=torch.tensor(text_features[int(start):int(end+1)]).to("cuda")
        G.ndata['feat'] = feat
        G = dgl.add_self_loop(G)
        G_text_tensors = modle(start,question_feature_vector,G)
        G_text_tensorslist.append(G_text_tensors)
    # 使用torch.stack将张量列表转换为一个张量
    stacked_tensors = torch.stack(G_text_tensorslist)

    # 沿着第一个维度（0）对张量求和
    summed_tensor = torch.sum(stacked_tensors, dim=0)

    # 应用softmax函数
    softmax_tensor = softmax(summed_tensor, dim=0)
    return softmax_tensor
def collate(samples):
    question, label= map(list, zip(*samples))
    return question, label
def train(train_dataset_path):
    def init_weights(m):
        if type(m) == nn.Linear:
            nn.init.xavier_uniform_(m.weight)

    print("Loading the training dataset: " + train_dataset_path)
    train_data = LuDataset(file_path=train_dataset_path)
    train_iter = DataLoader(dataset=train_data, batch_size=int(16), shuffle=True, drop_last=True,
                            collate_fn=collate)
    net = LuGTP().to("cuda")
    net.apply(init_weights)
    optimizer = torch.optim.Adam(net.parameters(), lr=0.0001, weight_decay=0.0001)
    criterion = nn.MultiLabelSoftMarginLoss(reduction='mean')
    for epoch in range(100):
        net.train()
        running_loss = 0.0
        for batch_idx, (questions, labels) in enumerate(train_iter):
            optimizer.zero_grad()
            # Pass data to the model
            real_labs=[]
            pre_labs=[]
            for idx,(q,l )in enumerate(zip(questions,labels)):
                pre_lab=generol_search(net,q, './docs', reload=False)
                # Convert pre_lab to a tensor with requires_grad=True if necessary
                if not isinstance(pre_lab, torch.Tensor):
                    pre_lab = torch.tensor(pre_lab, requires_grad=True)
                elif not pre_lab.requires_grad:
                    pre_lab.requires_grad_(True)
                # Create label tensor
                my_list = [0] * len(label_to_idx)
                label=labels[idx]
                inx=label_to_idx[label]
                my_list[inx]=1
                label=torch.tensor(my_list)
                real_labs.append(label)
                pre_labs.append(pre_lab)
            # 将整个批次的数据转换为Tensor并传递给模型
            pre_labs_tensor = torch.stack(pre_labs)
            real_labs_tensor = torch.stack(real_labs)
            loss = criterion(pre_labs_tensor, real_labs_tensor)

            # Backward and optimize
            loss.backward()
            for name, param in net.named_parameters():
                if param.requires_grad:
                    print(f"Before update, {name}: {param.data[:5]}")
            optimizer.step()
            for name, param in net.named_parameters():
                if param.requires_grad:
                    print(f"After update, {name}: {param.data[:5]}")

            running_loss += loss.item()
            print(running_loss)
if __name__ == '__main__':
    train("C:/Users/75187/PycharmProjects/gnn_rag/q_a.xlsx")  为什么我的反向传播似乎没有更新到模型的权重？如何修改代码