import os
import concurrent.futures
import docx
import torch
import numpy as np
import streamlit as st
from hazm import *
from transformers import AutoTokenizer, AutoModel

# بارگذاری مدل
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/bert-fa-base-uncased")
    model = AutoModel.from_pretrained("HooshvareLab/bert-fa-base-uncased")
    return tokenizer, model

tokenizer, model = load_model()

# پردازش فایل‌های Word و تبدیل به جملات
@st.cache_data
def load_text_chunks(folder_path):
    normalizer = Normalizer()
    sentence_tokenizer = SentenceTokenizer()
    texts = []

    for filename in os.listdir(folder_path):
        if filename.endswith(".docx"):
            full_path = os.path.join(folder_path, filename)
            doc = docx.Document(full_path)
            file_text = "\n".join([para.text for para in doc.paragraphs])
            if file_text.strip():
                texts.append(file_text)

    all_sentences = []
    for text in texts:
        normalized = normalizer.normalize(text)
        sentences = sentence_tokenizer.tokenize(normalized)
        all_sentences.extend(sentences)

    # تقسیم به بخش‌های ۵ جمله‌ای
    chunks = []
    for i in range(0, len(all_sentences), 5):
        chunk = " ".join(all_sentences[i:i+5])
        if chunk:
            chunks.append(chunk)
    return chunks

# محاسبه embedding با BERT
def get_embedding(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1)
    return embeddings.squeeze().numpy()

# شباهت کسینوسی
def cosine_similarity(vec1, vec2):
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))

# رابط کاربری استریم‌لیت
st.title("🔎 یافتن نزدیک‌ترین بخش ۵ جمله‌ای به ورودی شما")

# مسیر پوشه فایل‌های docx
folder_path = '46'

# بارگذاری و نمایش تعداد بخش‌ها
chunks = load_text_chunks(folder_path)
st.success(f"{len(chunks)} بخش ۵ جمله‌ای بارگذاری شد.")

# ورودی کاربر
user_input = st.text_area("لطفاً جمله یا متن خود را وارد کنید:")
def calculate_similarities_parallel(user_embedding, chunks):
    with concurrent.futures.ThreadPoolExecutor() as executor:
        similarities = list(executor.map(lambda chunk: cosine_similarity(user_embedding, get_embedding(chunk)), chunks))
    return similarities
if st.button("🔍 جستجو"):
    if not user_input.strip():
        st.warning("لطفاً یک جمله وارد کنید.")
    else:
        with st.spinner("در حال محاسبه شباهت‌ها..."):
            user_embedding = get_embedding(user_input)
            similarities = calculate_similarities_parallel(user_embedding, chunks)
            most_similar_index = np.argmax(similarities)
            result = chunks[most_similar_index]

        st.subheader("📌 شبیه‌ترین بخش ۵ جمله‌ای:")
        st.write(result)