import gradio as gr
import torch
import numpy as np
from PIL import Image
from models import load_vision_model, get_image_transform, load_segmentation_model, DEVICE
from diagnosis_module import diagnose_symptoms

# تحميل الموديلات
vision_model = load_vision_model()
seg_model = load_segmentation_model()
transform = get_image_transform()

# تحليل الأشعة
def analyze_image(image):
    if image is None:
        return "يرجى رفع صورة الأشعة"
    img = np.array(image.convert("RGB"))
    tensor = transform(img).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        output = vision_model(tensor)
        probs = torch.softmax(output, dim=1).cpu().numpy()[0]
    top_classes = np.argsort(probs)[::-1][:5]
    result = "\n".join([f"فئة {i}: احتمال {probs[i]:.3f}" for i in top_classes])
    return result

# تحليل الأعراض
def analyze_symptoms(symptom_text):
    if not symptom_text.strip():
        return "يرجى كتابة الأعراض"
    results = diagnose_symptoms(symptom_text)
    text = "🔍 الأمراض المحتملة:\n\n"
    for r in results:
        text += f"- {r['disease']} (احتمال: {r['score']:.2f})\n📘 المصدر: {r['source']}\n\n"
    return text

# واجهة Gradio
with gr.Blocks(title="MedAI Assistant") as app:
    gr.Markdown("# 🧠 MedAI — مساعد التحليل الطبي الذكي\n### ⚠️ لأغراض بحثية فقط، ليست بديلاً عن الطبيب")

    with gr.Tab("تحليل الأشعة"):
        image_input = gr.Image(type="pil", label="📤 ارفع صورة الأشعة")
        image_output = gr.Textbox(label="🔍 نتائج التحليل")
        analyze_button = gr.Button("تشغيل التحليل")
        analyze_button.click(analyze_image, inputs=image_input, outputs=image_output)

    with gr.Tab("تحليل الأعراض"):
        symptom_input = gr.Textbox(label="✏️ اكتب الأعراض (مثال: حمى، سعال، ضيق تنفس)")
        symptom_output = gr.Textbox(label="🩺 التشخيص المبدئي")
        analyze_symptoms_btn = gr.Button("تشغيل التحليل")
        analyze_symptoms_btn.click(analyze_symptoms, inputs=symptom_input, outputs=symptom_output)

app.launch()