"""
ایجاد بردار جملات - امبدینگ
"""
from sentence_transformers import SentenceTransformer
import json
import datetime
import numpy as np
from elastic_helper import ElasticHelper

date = datetime.datetime.now()
today = f'{date.year}-{date.month}-{date.day}-{date.hour}'

model_name = 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'#89-25
# model_name = 'sentence-transformers/paraphrase-multilingual-mpnet-base-v2'#87-30
# model_name = 'sharif-dal/dal-bert'#90-41
# model_name = 'lifeweb-ai/shiraz'#97-67
# model_name = 'BAAI/bge-m3'#90-35
# model_name = 'jinaai/jina-embeddings-v3'#??
# model_name = 'jinaai/jina-embeddings-v4'#??
# model_name = 'HooshvareLab/bert-base-parsbert-uncased'#90-54
model = SentenceTransformer(model_name)

def get_sections():
    """
    دریافت کل سکشن های قانونی در مسیر مشخص شده
    """
    sections_path = "/home/gpu/data_11/14040423/mj_qa_section.zip"
    eh_obj = ElasticHelper()
    sections = eh_obj.iterateJsonFile(sections_path, True)
    sections = convert_to_dict(sections)
    return sections

def convert_to_dict(sections):
    """
    تبدیل لیست سکشن های قانون به دیکشنری
    """
    sections_dict = {}
    for item in sections:
        id = item['id']
        source = item['source']
        sections_dict[id] = source
        
    return sections_dict

def do_word_embedder(sections):
    """
    لیستی از آبجکت سکشن های قانون را دریافت می کند و متن سکشن ها را تبدیل به بردار می کند و در نهایت لیست سکشن ها که امبدینگ هم  به آنها اضافه شده را در مسیر مشخص شده در همین متد، ذخیره می کند
    """
    for index, id in enumerate(sections):
        embeddings = single_section_embedder(sections[id]['content'])
        sections[id]['embeddings'] = embeddings.tolist()
    
    with open(f'./data/embeddings/sections_embeddings_{today}.json', 'w', encoding='utf-8') as output_file:
        data = json.dumps(sections, ensure_ascii=False)
        output_file.write(data)
    
    return sections
    
def single_section_embedder(sentence):
    """
    این متد، متن ورودی را تبدیل به بردار متناظر آن می کند
    
    **Args:
        sentence (str): متن یک سکشن
    
    **Returns:
        embeddings: لیست بردار متناظر با متن ورودی
    """
    embeddings = model.encode(sentence)
    return embeddings

def cosine_similarity(vec1, vec2):
    dot_product = np.dot(vec1, vec2)  # ضرب داخلی دو بردار
    norm_vec1 = np.linalg.norm(vec1) # نُرم بردار اول
    norm_vec2 = np.linalg.norm(vec2) # نُرم بردار دوم
    return dot_product / (norm_vec1 * norm_vec2)

if __name__ == '__main__':
   
    sections = get_sections()
    
    # محاسبه امبدینگ سکشن ها و ذخیره نتایج
    do_word_embedder(sections)