from helper_utils import load_chroma, word_wrap, project_embeddings
import chromadb.utils.embedding_functions as embedding_functions
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
import numpy as np
import warnings

# 忽略 FutureWarning 类型的警告
warnings.filterwarnings("ignore", category=FutureWarning)

# 使用代理可能出现网络问题，将以下端口号1080全部替换成自己的vpn的端口号
import os
# os.environ['HTTPS_PROXY']='http://127.0.0.1:1080'
# os.environ["HTTP_PROXY"]='http://127.0.0.1:1080'
import openai
from openai import OpenAI
from dotenv import load_dotenv, find_dotenv

loaded = load_dotenv(find_dotenv(), override=True)
# 从环境变量中获取 OpenAI API Key 或者直接赋值
API_KEY = os.getenv("API_KEY")

# 如果您使用的是官方 API，就直接用 https://api.siliconflow.cn/v1 就行。
BASE_URL = "https://api.siliconflow.cn/v1"

# chromadb支持的嵌入函数有许多种，这里介绍常用的几种：
# 参考资料：https://docs.trychroma.com/embeddings
# 方式1：默认嵌入函数，需要下载模型，本地计算。英文文本表现不错，中文文本表现一般
# embedding_function = SentenceTransformerEmbeddingFunction()

# 方式2：OpenAI的嵌入函数，直接调用OpenAI的接口，无需下载模型(推荐)
embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key=API_KEY,
        api_base=BASE_URL,
        model_name="BAAI/bge-m3",
        dimensions=1024
    )

# 方式3：HuggingFace的嵌入函数，需要下载模型，本地计算，对网络要求高
# embedding_function = embedding_functions.HuggingFaceEmbeddingFunction(
#     api_key="hf_",  # 填入你的 huggingface Access Token
#     model_name="jinaai/jina-embeddings-v2-base-zh"  # 指定模型
# )
## 中文备选模型
# jinaai/jina-embeddings-v2-base-zh
# GanymedeNil/text2vec-large-chinese
# BAAI/bge-large-zh-v1.5
# BAAI/bge-small-zh-v1.5

chroma_collection = load_chroma(filename='./data/2024年北京市政府工作报告.pdf', 
                                collection_name='beijing_annual_report_2024', 
                                embedding_function=embedding_function,langcode='zh')
chroma_collection.count()

# 之前一般设定返回5个结果，现在要求返回10个结果，加入了部分可能有用的的长尾结果
query = "地区生产总值是多少？"
results = chroma_collection.query(query_texts=query, n_results=10, include=['documents', 'embeddings'])

retrieved_documents = results['documents'][0]

for document in results['documents'][0]:
    print(word_wrap(document))
    print('')

# BERT交叉编码器同时渠道查询和文档，通过一个分类器传递，获得一个得分
# 利用该得分作为检索结果的相关性或排名的得分
from sentence_transformers import CrossEncoder
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

pairs = [[query, doc] for doc in retrieved_documents]
scores = cross_encoder.predict(pairs)

print("分数和相应的文本:")

# 打印每个文档的分数和内容
for score, (query, document) in zip(scores, pairs):
    print(f"分数: {score}")
    print("文本:")
    print(word_wrap(document))
    print('')  # 在文档间添加空行以便区分

print("新排名:")
for o in np.argsort(scores)[::-1]:
    print(o+1)
print("分数和相应的文本:")
sorted_indices = np.argsort(scores)[::-1]

for rank, index in enumerate(sorted_indices, start=1):
    # 打印排名和对应的分数
    print(f"排名 {rank}, 分数: {scores[index]}")
    # 打印对应的文档内容，这里假设 pairs[index][1] 是文档内容
    print("文本:")
    print(pairs[index][1])
    print('')  # 在文档之间添加空行以便区分

# 接下来把之前获得的结果排序前5名传递给LLM
original_query = "推动北京财政收入增长因素是什么"
generated_queries = [
    "什么推动了北京市的财政收入增长？",
    "什么对北京市的财政收入增长作出贡献？",
    "北京市财政收入增长依赖于什么？",
    "什么有助于北京市的财政收入增长？",
    "北京市财政收入增长受到哪些方面的影响？",
]

queries = [original_query] + generated_queries

results = chroma_collection.query(query_texts=queries, n_results=10, include=['documents', 'embeddings'])
retrieved_documents = results['documents']

# 删除检索文档中的重复数据
unique_documents = set()
for documents in retrieved_documents:
    for document in documents:
        unique_documents.add(document)

unique_documents = list(unique_documents)

# 再次创建pairs,可以增强查询的检索记过与原始查询的相关性
# 从中选择最佳的5个结果传递给LLM
pairs = []
for doc in unique_documents:
    pairs.append([original_query, doc])

scores = cross_encoder.predict(pairs)

print("分数和对应的文本:")

# 打印每个文档的分数和内容
for score, (query, document) in zip(scores, pairs):
    print(f"分数: {score}")
    print("文本:")
    print(word_wrap(document))
    print('')  # 在文档间添加空行以便区分

# 打印新排名和对应的结果
print("新排名:")
for o in np.argsort(scores)[::-1]:
    print(o+1)
print("\n")
sorted_indices = np.argsort(scores)[::-1]
for rank, index in enumerate(sorted_indices, start=1):
    # 打印排名和对应的分数
    print(f"排名 {rank}, 分数: {scores[index]}")
    # 打印对应的文档内容，这里假设 pairs[index][1] 是文档内容
    print("文本:")
    print(pairs[index][1])
    print('')  # 在文档之间添加空行以便区分