import time
from collections import defaultdict

from tqdm import tqdm
import numpy as np
import torch
from transformers import AutoModel, AutoTokenizer, AutoModelForSequenceClassification
from sklearn.metrics.pairwise import cosine_similarity
import dgl
import scipy.sparse as spp
from ludeep import LuGTP
import torch.nn.functional as F
import re
import os
import jieba
import jieba.analyse


def read_file_and_split_by_lines(directory_path):
    # 使用正则表达式匹配句子，包括句子结束符号
    # 该模式匹配以句号、叹号、问号结尾的句子，或不包含句子结束符的句子结尾
    sentence_pattern = 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 = []  # 当前文件的句子索引列表
                current_file_sentences = set()  # 集合用于存储当前文件中的句子，防止重复
                for line in file:
                    stripped_line = line.strip()
                    if stripped_line:
                        # 使用正则表达式提取句子
                        sentences_in_line = sentence_pattern.findall(stripped_line)
                        for sentence in sentences_in_line:
                            sentence = sentence.strip()
                            if sentence and sentence not in current_file_sentences:
                                current_file_sentences.add(sentence)
                                file_rels.append(len(sentences))
                                sentences.append(sentence)
                                file_names.append(file_name)
                sentence_indices.append(file_rels)
    return sentences, sentence_indices, file_names

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')
text_blocks=[]
from FlagEmbedding import FlagReranker
reranker = FlagReranker('./model/BAAI/bge-reranker-large', use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation

model.to(device)

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


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



# 贪心算法选择块
def greedy_block_selection(document_blocks_indices, initial_blocks_indices, similarity_matrix_question, max_blocks=20):
    selected_indices = set(initial_blocks_indices)
    candidate_indices = set(document_blocks_indices) - selected_indices

    similarity_matrix = similarity_matrix_question

    # 取出块与问题的相似度（最后一列）
    block_question_similarity = similarity_matrix[:, -1]

    while len(selected_indices) < max_blocks and candidate_indices:
        best_candidate = None
        best_score = -np.inf
        for idx in candidate_indices:
            # 计算评分函数
            # 与问题的相似度
            q_similarity = block_question_similarity[idx]
            # 与已选块的平均相似度
            s_similarity = np.mean([similarity_matrix[idx, s_idx] for s_idx in selected_indices])
            # 综合评分（权重可调整）
            score =q_similarity + s_similarity
            if score > best_score:
                best_score = score
                best_candidate = idx
        if best_candidate is not None:
            selected_indices.add(best_candidate)
            candidate_indices.remove(best_candidate)
        else:
            return [], initial_blocks_indices
        # 按照文档顺序返回选定的块
        sorted_indices = sorted(selected_indices)

        selected_blocks = []
        tanxin_index_list = []

        for idx in sorted_indices:
            selected_blocks.append(text_blocks[idx])
            tanxin_index_list.append(idx)
        return selected_blocks, tanxin_index_list
# 计算每个文档的得分

def get_re_score(question,pairs,text,textmod=True):
    #此处未优化 速度比较慢
    selected_block_texts = '\n'.join(pairs)
    if textmod==True:
        with torch.no_grad():
            score = reranker.compute_score([[question, selected_block_texts]], normalize=True)
            score1 = reranker.compute_score([[question, text]], normalize=True)
            return score[0],score1[0]
    else:
        with torch.no_grad():
            pairslist=[[question,p] for p in pairs]
            score = reranker.compute_score(pairslist, normalize=True)
            score1 = reranker.compute_score([[question, text]], normalize=True)
        return np.mean(score),score1[0]
def generol_search(question,docs_path,reload=False):
    stop_words = set(['的', '了', '和', '是', '我', '想', '请问', '查询', '我的', '你', '您', '吗', '？', '，', '。'])
    seg_list = jieba.lcut(question)
    print("分词结果：", seg_list)
    # 去除停用词，提取关键词
    query_keywords = [word for word in seg_list if word not in stop_words]
    print("提取的关键词：", query_keywords)
    text_features=[]
    global text_blocks,element_to_list_map
    text_blocks,relate_index,file_names = read_file_and_split_by_lines(docs_path)[:]
    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
        print(f"加载features.npy执行时间: {elapsed_time:.6f}秒")
    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]
    #print("~~~~~~~~~~~~~~~~~~~~~~~first_state_nodes~~~~~~~~~~~~~~~~~~~~~~~~~~")
    question_feature_vector = question_feature_vector.to(device)
    text_features_tensors = [torch.from_numpy(arr) for arr in text_features]
    text_features_tensor = torch.stack(text_features_tensors).to(device)
    question_feature_vector_norm = F.normalize(question_feature_vector, p=2, dim=0)
    text_features_norm = F.normalize(text_features_tensor, p=2, dim=1)
    similarities = torch.matmul(text_features_norm, question_feature_vector_norm)  # Shape: (N,)
    black_question = np.array(similarities.tolist())
    top_k = 15  # or any number of top similarities you want
    topk_values, topk_indices = torch.topk(similarities, k=top_k, largest=True)

    q_a_relate_dict_ = defaultdict(list)
    for idx in topk_indices.cpu().numpy():
        g_node = element_to_list_map[idx]
        q_a_relate_dict_[g_node].append(idx)
      #计算特征向量之间的余弦相似度
    similarity_matrix = cosine_similarity(text_features)
    if black_question.ndim == 1:
        black_question = black_question.reshape(-1, 1)
    assert similarity_matrix.shape[0] == black_question.shape[0], "Number of rows do not match"
    similarity_matrix_with_black_question = np.hstack((similarity_matrix, black_question))
    # 结果列表
    time_start = time.time()
    # 结果列表
    print("~~~~~~~~~~~~~~~~~~~~~~~~Graph+RAG~~~~~~~~~~~~~~~~~~~~~~~~~")
    modle = LuGTP().to(device)
    # 记录开始时间
    start_time = time.time()
    # 初始化保存所有结果的列表
    all_results = []
    dict_g_tensor={}
    total_size = sum(len(lst) for lst in q_a_relate_dict_.values())
    num_lists = len(q_a_relate_dict_)

    average_size = int(total_size / num_lists)

    for g_number,relate_sentence_list in q_a_relate_dict_.items():
        full_relate_list=relate_index[g_number]
        start=full_relate_list[0]
        end=full_relate_list[-1]
        selected_blocks, tanxin_index_list=greedy_block_selection(full_relate_list,relate_sentence_list[:average_size],similarity_matrix_with_black_question,max_blocks=average_size*3)
        matrix_len=len(full_relate_list)
        matrix = np.ones((matrix_len, matrix_len)).astype('int64')
        #tanxin_index_list
        #result_list
        new_result_list = [x - start for x in tanxin_index_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(device)
        feat=torch.tensor(text_features[int(start):int(end+1)]).to(device)
        G.ndata['feat'] = feat
        G,readoutcat=modle(G)
        dict_g_tensor[g_number]=readoutcat
        G_text_tensors=G.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, g_number,selected_blocks,(score-black_question[index][0])))
            all_results.append((tensor, index, score, g_number,selected_blocks,(score-black_question[index][0])))
    # 按scores的得分大小排序
    all_results.sort(key=lambda x: x[2], reverse=True)
    # all_results_more.sort(key=lambda x: x[2], reverse=True)

    # 记录结束时间
    end_time = time.time()
    # 计算并输出耗时
    elapsed_time = end_time - start_time
    print(f"推理执行时间: {elapsed_time:.6f}秒")
    # 参数设置

    # 定义评分权重
    alpha = 0.1   #问题与文档相关性权重
    alpha1=0.7
    alpha2 =0.1
    alpha3 = 0.1
    beta = 0.1  #文档之间相似性权重
    gamma = 1   #信息聚合增益权重

    print(f"问题：{question}")
    # 输出已选择的文档
    # 设置可调整的阈值
    print('各句子分数：')
    for i, rs in enumerate(all_results[:15]):
        #聚合后和问题的相关性
        # similarity1 = F.cosine_similarity(top0[0], rs[0], dim=0)
        selected_blocks=rs[-2]
        index = rs[1]
        score = rs[2]
        Δscore = rs[-1]
        b_re_score,t_re_score=get_re_score(question, selected_blocks,text_blocks[index])
        score0=score-Δscore
        # 计算综合评分
        # similarity2 = similarity_matrix_with_black_question[index][top0_index]
        sum_ = alpha * score0+alpha1*score +alpha2*b_re_score+alpha3*t_re_score
        print(f'[top {i}] 聚合前score: {score0} 聚合后score: {score} 文档重排序score:{b_re_score} 句子内容重排序score:{t_re_score} 综合总分:{sum_} 文件: {file_names[index]}',text_blocks[index])
    return all_results
if __name__ == '__main__':
    docs_path = './docs'
    question_list=[
                   "如何申请保障性住房",
                   "如果我是城镇居民并且只有一个孩子，我可以获得哪些奖励和补助",
                   "计划生育"
                   ]
    for q in question_list:
        generol_search(q,docs_path,reload=False)
帮我为这个代码生成详细的说明文档，包括详细分析代码