import torch
from transformers import AutoTokenizer, AutoModel
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import gradio as gr

# 📌 مدل و توکنایزر
model_name = "HooshvareLab/bert-fa-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# 📄 دیتاست اولیه (FAQ)
faq_data = {
    "پایتخت ایران کجاست؟": "تهران",
    "زبان رسمی ایران چیست؟": "فارسی",
    "واحد پول ایران چیست؟": "ریال",
    "چه زمانی انتخاب واحد شروع می‌شود؟": "معمولاً پایان شهریور یا بهمن.",
    "چه معدلی برای گرفتن 24 واحد لازم است؟": "حداقل معدل 17.",
}

questions = list(faq_data.keys())
answers = list(faq_data.values())

# 📄 تولید embedding
def get_embedding(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=64)
    with torch.no_grad():
        outputs = model(**inputs)
    emb = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
    return emb

faq_embeddings = [get_embedding(q) for q in questions]

# 📄 تابع پاسخ
def answer_question(user_question):
    user_emb = get_embedding(user_question)
    sims = [cosine_similarity([user_emb], [emb])[0][0] for emb in faq_embeddings]
    best_idx = int(np.argmax(sims))
    best_score = sims[best_idx]

    if best_score > 0.7:
        return answers[best_idx]
    else:
        return "متأسفم، جواب دقیقی در دیتاست پیدا نکردم."

# 📄 رابط Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 دستیار فارسی (پایه بر اساس semantic search با BERT)")
    inp = gr.Textbox(label="سؤال خود را بنویسید")
    out = gr.Textbox(label="پاسخ")
    btn = gr.Button("پاسخ بده")
    btn.click(fn=answer_question, inputs=inp, outputs=out)

demo.launch()