import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM
from torch.nn.functional import softmax
from huggingface_hub import login
import gradio as gr
import datetime
import os

# ورود با سکرت
hf_token = os.environ.get("new_hf")
login(token=hf_token)

# برچسب‌های بیماری
label_map = {
    0: "آرتریت روماتوئید",
    1: "آسم",
    2: "اگزما",
    3: "سندرم روده تحریک‌پذیر",
    4: "عفونت ادراری",
    5: "فارنژیت",
    6: "میگرن"
}

# مدل BERT تشخیص بیماری
bert_model_id = "diginoron/bert-medical-fa"
bert_tokenizer = AutoTokenizer.from_pretrained(bert_model_id)
bert_model = AutoModelForSequenceClassification.from_pretrained(bert_model_id)
bert_model.eval()

# مدل زبانی Gemma 2B
gemma_model_id = "google/gemma-2b"
gemma_tokenizer = AutoTokenizer.from_pretrained(gemma_model_id)
gemma_model = AutoModelForCausalLM.from_pretrained(
    gemma_model_id,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)
gemma_model.eval()

# تابع اصلی با Indicator
def predict_and_explain(symptoms):
    inputs = bert_tokenizer(symptoms, return_tensors="pt")
    with torch.no_grad():
        outputs = bert_model(**inputs)
        probs = softmax(outputs.logits, dim=1)
        confidence, prediction = torch.max(probs, dim=1)
        confidence = confidence.item()
        prediction = prediction.item()

    if confidence < 0.5:
        return "❌ تشخیص قطعی داده نشد. لطفاً با پزشک مشورت کنید.", "", gr.update(visible=False)

    diagnosis = label_map[prediction]
    diagnosis_text = f"✅ احتمالاً شما دچار {diagnosis} هستید. (اطمینان: {round(confidence * 100, 1)}٪)"

    prompt = f"بیماری {diagnosis} چیست؟ لطفاً آن را به زبان ساده توضیح بده و چند پیشنهاد درمانی اولیه برای مدیریت این بیماری ارائه کن."

    inputs = gemma_tokenizer(prompt, return_tensors="pt")
    if torch.cuda.is_available():
        inputs = {k: v.to("cuda") for k, v in inputs.items()}

    with torch.no_grad():
        output = gemma_model.generate(
            **inputs,
            max_new_tokens=150,
            do_sample=True,
            top_p=0.9,
            temperature=0.7,
            use_cache=True
        )

    explanation = gemma_tokenizer.decode(output[0], skip_special_tokens=True)
    explanation = explanation.replace(prompt, "").strip()

    return diagnosis_text, explanation, gr.update(visible=False)

# رابط Gradio
with gr.Blocks(css="body { font-family: Vazirmatn, sans-serif; background-color: #111827; color: #f3f4f6; } .gr-button { font-weight: bold; }") as demo:
    gr.Markdown("# 🧠 تشخیص بیماری با هوش مصنوعی پارسی + توضیح Gemma")

    with gr.Row():
        inp = gr.Textbox(placeholder="مثلاً: سرفه خشک، تب، گلودرد", label="علائم شما")

    with gr.Row():
        out1 = gr.Textbox(label="نتیجه تشخیص اولیه")
    with gr.Row():
        out2 = gr.Textbox(label="توضیح بیماری و پیشنهادات اولیه")

    with gr.Row():
        status = gr.Markdown("⏳ لطفاً منتظر بمانید...", visible=False)

    btn = gr.Button("🔍 بررسی و تحلیل علائم")
    clr = gr.Button("🧹 پاک کردن همه موارد")

    # هنگام کلیک، وضعیت به visible و پس از اجرا به invisible
    def wrapper(symptoms):
        status.update(visible=True)
        return predict_and_explain(symptoms)

    btn.click(fn=predict_and_explain, inputs=inp, outputs=[out1, out2, status])
    clr.click(fn=lambda: ("", "", gr.update(visible=False)), inputs=[], outputs=[inp, out1, out2, status])

    gr.Markdown(f"""
    ---
    📞 برای مشاوره تخصصی با پزشک به [فرم تماس](https://diginoron.com/contact) مراجعه کنید.  
    🕒 آخرین به‌روزرسانی: {datetime.datetime.now().strftime('%Y/%m/%d')}
    """)

demo.launch()