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

# بارگذاری مدل
tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/bert-fa-base-uncased")
model = AutoModel.from_pretrained("HooshvareLab/bert-fa-base-uncased")

# پرسش و پاسخ‌ها (سوالات باید طبیعی‌تر شوند)
faq_dict = {
    "زمان انتخاب واحد چه موقع است؟": "معمولاً پایان شهریور و بهمن است.",
    "چه زمانی می‌توان حذف و اضافه انجام داد؟": "حدود یک هفته پس از شروع ترم تحصیلی است.",
    "چه معدلی برای انتخاب ۲۴ واحد لازم است؟": "حداقل معدل 17 نیاز است.",
    "تا چه زمانی امکان حذف اضطراری وجود دارد؟": "تا هفته هشتم ترم مجاز است.",
    "چگونه می‌توان مهمان شد؟": "با موافقت دانشگاه مبدا و مقصد انجام می‌شود.",
}

faq_questions = list(faq_dict.keys())
faq_embeddings = []

# تابع استخراج embedding
def get_embedding(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    return outputs.last_hidden_state.mean(dim=1)[0].cpu().numpy()

# پیش‌محاسبه embeddingها
for question in faq_questions:
    emb = get_embedding(question)
    faq_embeddings.append(emb)

# پاسخ‌دهی
def student_bot(user_question):
    try:
        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]

        # 👇 برای دیباگ
        print("Similarity Score:", best_score)

        if best_score > 0.6:
            return faq_dict[faq_questions[best_idx]]
        else:
            return "متأسفم، پاسخ این سؤال در حال حاضر موجود نیست."
    except Exception as e:
        return f"❗️خطا در سیستم: {str(e)}"

# رابط Gradio
iface = gr.Interface(
    fn=student_bot,
    inputs=gr.Textbox(label="سوال خود را بنویسید"),
    outputs=gr.Textbox(label="پاسخ"),
    title="ایجنت راهنمای دانشجویان",
    description="پاسخ به سوالات رایج با استفاده از مدل FarsiBERT و تشابه معنایی"
)

iface.launch()