import os
script_dir = os.path.dirname(os.path.abspath(__file__))
import sys
sys.path.append(script_dir)
from intent_description_template import template, intent_code
from remote_model import RemoteBGEModel
import torch
import faiss
import numpy as np
import re


class IntentRecognizer:

    def __init__(self):
        # 加载元数据
        self.template_meta = {}
        for k, v in template.items():
            for desc in v:
                self.template_meta[desc] = k
        # 模板元数据
        self.template_meta_list = list(self.template_meta.keys())
        # 加载远程模型
        self.model = RemoteBGEModel('dev')
        # 模型预热
        print("模型预热中...")
        self.model.encode(["这是一段预热文字，首次推理通过预测保证后续推理的稳定性和性能。"])
        self.model.compute_score([("这是一段预热文字，首次推理通过预测保证后续推理的稳定性和性能。",
                                 "这是一段预热文字，首次推理通过预测保证后续推理的稳定性和性能。")])
        self.print_gpu()
        # 加载向量数据库
        self.database_index = None
        database_path = os.path.join(script_dir, "intent_index.faiss")
        if not os.path.exists(database_path):
            # embeddings = self.model.encode(self.template_meta_list)['dense_vecs'].astype(np.float32)  # 选取密集向量，变为float32
            #faiss.normalize_L2(embeddings)  # L2归一化,本地模型需要归一化
            # 调用远程embedding模型,one by one 地处理，远程模型通过配置参数进行归一化
            embeddings = [self.model.encode([temp], normalize=True)[0] for temp in self.template_meta_list]
            for _ in  embeddings:
                if _ is None:
                    raise RuntimeError('构建意图描述模板库时发生异常，embeddings不能存在None')
            # 要求embeddings是一个二维矩阵，类型为float32
            embeddings = np.array(embeddings, dtype=np.float32)
            # Create FAISS index
            dimension = embeddings[0].shape[0]
            self.database_index = faiss.IndexFlatIP(dimension)  # 建立内积索引
            self.database_index.add(embeddings)  # 添加索引
            # Save for future use
            faiss.write_index(self.database_index, database_path)

        if self.database_index is None:
            self.database_index = faiss.read_index(database_path)

    @staticmethod
    def print_gpu():
        if torch.cuda.is_available():
            print(f"allocated:{torch.cuda.memory_allocated()/1024**3:.2f}GB", end=' ')
            print(f"reserved: {torch.cuda.memory_reserved()/1024**3:.2f}GB")

    # 应用推理阶段
    def pick_out(self, query, top_k):
        # 使用本地模型
        # query_embedding = self.model.encode([query])['dense_vecs'].astype(np.float32)
        # 下面使用远程模型代替本地模型
        # 要求query_embedding是一个二维矩阵，形状为(1, 1024)
        query_embedding = np.array(self.model.encode([query], normalize=True), dtype=np.float32)
        # faiss.normalize_L2(query_embedding)
        distances, indices = self.database_index.search(query_embedding, top_k)
        group_query = [(query, self.template_meta_list[indices[0][i]]) for i in range(top_k)]
        # 调用远程reranker模型
        score = self.model.compute_score(group_query)
        rerank_result = sorted([(distances[0][_], indices[0][_], score[_]) for _ in range(top_k)], key=lambda x: x[2],
                               reverse=True)  # distance, indices, rerank_score
        score_idx = 2  # 重排序相关度
        meta_idx = 1  # 模板位置
        similarity_idx = 0  # 向量相似度
        print("***检索结果***：")
        for i in range(top_k - 1, -1, -1):
            print(
                f"***{i} 相关度:{rerank_result[i][score_idx]:.2f} "
                f"相似度:{rerank_result[i][similarity_idx]:.2f} "
                f"意图:{self.template_meta.get(self.template_meta_list[rerank_result[i][meta_idx]])} "
                f"{intent_code.get(self.template_meta.get(self.template_meta_list[rerank_result[i][meta_idx]])).get('alias')} "
                f"关联:{self.template_meta_list[rerank_result[i][meta_idx]]}"
            )

        # 从排序结果中拆解到意图的大小类编号
        result = []  # 意图识别结果
        confidence = []# 置信度
        for i in range(top_k):
            # 拿到描述词
            description = self.template_meta_list[rerank_result[i][meta_idx]]
            # 拿到自编码
            custom_number = self.template_meta[description]
            # 拿到大小类标号
            result.append(intent_code[custom_number])
            # 添加置信度
            confidence.append(rerank_result[i][score_idx])

        return confidence, result


def quick_answer_q_a_v2(question, is_english=0):
    """快速问答分支，赋值metric"""
    metric = ''
    question += ' ' # 深拷贝
    if is_english == 0:
        pattern_res_open_qa = re.findall("(开启|打开).*?问答", question)
        pattern_res_close_qa = re.findall("(关闭|关掉).*?问答", question)
    else:
        question = question.strip().lower()
        pattern_res_open_qa = re.findall(
            r'\b(?:open|show me|enable)\b\W*.*?(?:quiz\W+with\W+prizes|enable\W+award\W*-\W*winning\W+q\W*&\W*a)',
            question)
        pattern_res_close_qa = re.findall(
            r'\b(?:close|disable)\b\W*.*?(?:quiz\W+with\W+prizes|enable\W+award\W*-\W*winning\W+q\W*&\W*a)', question)

    if len(pattern_res_open_qa) > 0:
        metric = "openQandA"

    if len(pattern_res_close_qa) > 0:
        metric = "closeQandA"

    return metric

# 单例模式
recognizer_bge = IntentRecognizer()