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

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

embedding_function = SentenceTransformerEmbeddingFunction()

chroma_collection = load_chroma(filename='./data/2024年北京市政府工作报告.pdf', 
                                collection_name='beijing_annual_report_2024', 
                                embedding_function=embedding_function,
                                langcode='zh')  # 注意中文文档将langcode改为'zh'
chroma_collection.count()

embeddings = chroma_collection.get(include=['embeddings'])['embeddings']
umap_transform = umap.UMAP(random_state=0, transform_seed=0, n_jobs=1).fit(embeddings)
projected_dataset_embeddings = project_embeddings(embeddings, umap_transform)

# 使用代理可能出现网络问题，将以下端口号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"

openai_client = OpenAI(api_key=API_KEY, base_url=BASE_URL, max_retries=3)

def generate_queries(model="Qwen/Qwen3-8B"):

    extra_body = {}
    if "Qwen3" in model:
        extra_body={
            "enable_thinking": False
        }
    
    messages = [
        {
            "role": "system",
            # 系统角色消息，它指示模型扮演一个有帮助的、专业的政府年度报告研究助手角色。
            # 提供的内容是一个明确的任务描述，要求模型提出10到15个短问题，
            # 这些问题在分析年度报告时非常重要。
            # 这个指示还特别要求模型不要输出复合问题，即包含多个句子或连接词的问题。
            # 同时要求每个问题单独一行，行与行之间用换行符分隔。
            "content": "You are a helpful and expert financial research assistant. You help users analyze government work reports to better understand the government. "
            "Suggest 10 to 15 short questions that are important to ask when analyzing an annual report. "
            "Do not output any compound questions (questions with multiple sentences or conjunctions)."
            "Output each question on a separate line divided by a newline."
            "Use Chinese"
        },
    ]

    response = openai_client.chat.completions.create(
        model=model,
        messages=messages,
        extra_body=extra_body
    )
    content = response.choices[0].message.content
    content = content.split("\n")
    return content

generated_queries = generate_queries()
for query in generated_queries:
    print(query)

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

def evaluate_results(query, statement, model="Qwen/Qwen3-8B"):

    extra_body = {}
    if "Qwen3" in model:
        extra_body={
            "enable_thinking": False
        }
    
    messages = [
    {
        "role": "system",
        # 指示模型扮演一个有用的财务研究助手角色，
        # 并要求模型评估给定查询与声明的相关性，只输出"yes"或"no"。
        "content": "You are a helpful expert financial research assistant. You help users analyze financial statements to better understand companies. "
        "For the given query, evaluate whether the following satement is relevant."
        "Output only 'yes' or 'no'."
    },
    {
        "role": "user",
        "content": f"Query: {query}, Statement: {statement}"
    }
    ]
    # 使用OpenAI的chat API发送请求，传递模型名称、消息列表和最大令牌数。
    # 最大令牌数设置为1，因为期望的输出只有"yes"或"no"。
    response = openai_client.chat.completions.create(
        model=model,
        messages=messages,
        max_tokens=1,
        extra_body=extra_body
    )
    # 从响应中提取内容，即模型的输出。
    content = response.choices[0].message.content
    # 如果内容为"yes"，返回1，表示声明与查询相关；
    # 否则返回-1，表示声明与查询不相关。
    if content == "yes":
        return 1
    return -1

retrieved_embeddings = results['embeddings']
query_embeddings = embedding_function(generated_queries)

adapter_query_embeddings = []
adapter_doc_embeddings = []
adapter_labels = []

import time
from tqdm import tqdm

for q, query in enumerate(tqdm(generated_queries)):
    for d, document in enumerate(retrieved_documents[q]):
        adapter_query_embeddings.append(query_embeddings[q])
        adapter_doc_embeddings.append(retrieved_embeddings[q][d])
        adapter_labels.append(evaluate_results(query, document))

        time.sleep(1)

len(adapter_labels)

adapter_query_embeddings = torch.Tensor(np.array(adapter_query_embeddings))
adapter_doc_embeddings = torch.Tensor(np.array(adapter_doc_embeddings))
adapter_labels = torch.Tensor(np.expand_dims(np.array(adapter_labels),1))

dataset = torch.utils.data.TensorDataset(adapter_query_embeddings, adapter_doc_embeddings, adapter_labels)

def model(query_embedding, document_embedding, adaptor_matrix):
    updated_query_embedding = torch.matmul(adaptor_matrix, query_embedding)
    return torch.cosine_similarity(updated_query_embedding, document_embedding, dim=0)

def mse_loss(query_embedding, document_embedding, adaptor_matrix, label):
    return torch.nn.MSELoss()(model(query_embedding, document_embedding, adaptor_matrix), label)

mat_size = len(adapter_query_embeddings[0])
adapter_matrix = torch.randn(mat_size, mat_size, requires_grad=True)

min_loss = float('inf')
best_matrix = None

for epoch in tqdm(range(100)):
    for query_embedding, document_embedding, label in dataset:
        loss = mse_loss(query_embedding, document_embedding, adapter_matrix, label)

        if loss < min_loss:
            min_loss = loss
            best_matrix = adapter_matrix.clone().detach().numpy()

        loss.backward()
        with torch.no_grad():
            adapter_matrix -= 0.01 * adapter_matrix.grad
            adapter_matrix.grad.zero_()

print(f"Best loss: {min_loss.detach().numpy()}")

test_vector = torch.ones((mat_size,1))
scaled_vector = np.matmul(best_matrix, test_vector).numpy()

import matplotlib.pyplot as plt
plt.bar(range(len(scaled_vector)), scaled_vector.flatten())
plt.show()

query_embeddings = embedding_function(generated_queries)
adapted_query_embeddings = np.matmul(best_matrix, np.array(query_embeddings).T).T

projected_query_embeddings = project_embeddings(query_embeddings, umap_transform)
projected_adapted_query_embeddings = project_embeddings(adapted_query_embeddings, umap_transform)

import matplotlib.pyplot as plt

plt.figure()
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 设置全局字体为微软雅黑，显示中文
plt.scatter(projected_dataset_embeddings[:, 0], projected_dataset_embeddings[:, 1], s=10, color='gray')
plt.scatter(projected_query_embeddings[:, 0], projected_query_embeddings[:, 1], s=150, marker='X', color='r', label="original")
plt.scatter(projected_adapted_query_embeddings[:, 0], projected_adapted_query_embeddings[:, 1], s=150, marker='X', color='green', label="adapted")

plt.gca().set_aspect('equal', 'datalim')
plt.title("Adapted Queries")
plt.axis('off')
plt.legend();