import ast
import pandas as pd
import numpy as np
from langchain_huggingface import HuggingFaceEmbeddings

model_name = "BAAI/bge-small-zh-v1.5"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}  # set True to compute cosine similarity



bge_hf_embedding = HuggingFaceEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)

# 【2】该函数就是把数据变为向量的函数
def text_2_embedding(text):
    resp = bge_hf_embedding.embed_documents(
        [text]
    )
    return resp[0]

# 【1】读取原始文件中的美食评论数据，通过调用Embedding模型，得到对应的向量，并保持到新文件中
def embedding_2_file(source_file, output_file):
    """读取原始的美食评论数据，通过调用Embedding模型，得到向量，并保持到新文件中"""
    #  步骤：1、准备数据，并读取，从第index_col个字段读取
    # pandas 是 Python 中用于数据处理和分析的核心工具，专门解决表格类数据的各类操作需求。
    # pandas 支持读取几乎所有常见的结构化数据格式（Excel、CSV、SQL、JSON 等），并能将处理后的数据便捷地保存为这些格式，解决了数据导入导出的基础问题。
    df = pd.read_csv(source_file, index_col=0)
    # 读取后提取所需要的字段
    df = df[['Time', 'ProductId', 'UserId', 'Score', 'Summary', 'Text']]

    print(df.head(2)) # 打印前两行看看

    # 步骤2： 清洗数据 和 合并数据
    df = df.dropna() # 如果有空数据，就删掉
    # 把评论的摘要  和 内容字段 合并成 一个字段（方便后续处理），放入新字段text_content
    df['text_content'] = 'Summary: ' + df.Summary.str.strip() + "; Text: " + df.Text.str.strip()
    print(df.head(2))  # 打印前两行看看 ，发现确实增加一个字段text_content

    # 步骤3: 对text_content向量化，存到一个新的文件中，向量单独存入embedding字段
    df['embedding'] = df.text_content.apply(lambda x: text_2_embedding(x))
    # 保存到新文件
    df.to_csv(output_file)


def cosine_distance(a, b):
    """计算余弦距离"""
    return  np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# 【3】input：用户输入的问题   ，embedding_file 要检索的文件， top_n 按照相似性最多检索3个
def search_text(input, embedding_file, top_n=3):
    """
    根据用户输入的问题，进行语义检索，返回最相似的前top_n个结果
    :param input:
    :param top_n:
    :return:
    """
    # 读取新文件
    df_data = pd.read_csv(embedding_file)
    # 取文件中df_data['embedding']字段，该字段在文件中是以字符串存储的
    # 所以你从文件读取出来这个字段也是字符串类型的啊
    # 要把这个字符串变成向量进行后续的数学计算，保持到新字段embedding_vector
    # 目前这个新字段是在内存中的，没有在文件里
    df_data['embedding_vector'] = df_data['embedding'].apply(ast.literal_eval)


    # 把输入问题转化为向量
    input_vector = text_2_embedding(input)


    # 内存中embedding_vector字段和input_vector进行相似度比较
    # 按照余弦相似度比较，余弦相似度单独封装到函数 cosine_distance中
    # 产生一个新字段  similarity
    df_data['similarity'] = df_data.embedding_vector.apply(lambda x: cosine_distance(x, input_vector))

    res = (
        # 对similarity字段排序，降序排
        df_data.sort_values('similarity', ascending=False)
        # 返回topn,此案例是：返回3个
        .head(top_n)
        # 把text_content字段的Summary: 替换为""，把; Text: 替换为""
        .text_content.str.replace('Summary: ', "")  # text_content是字段名
        .str.replace('; Text: ', ';')
    )

    # 把topn的结果打印，每打印一行画个虚线
    for r in res:
        print(r)
        print('-' * 30)

if __name__ == '__main__':
    # 可以先单独测试函数embedding_2_file('../datas/fine_food_reviews_1k.csv', '../datas/output_embedding.csv')
    search_text('delicious beans', './output_embedding.csv')