import gradio as gr
from transformers import BlipProcessor, BlipForQuestionAnswering, MarianMTModel, MarianTokenizer
from PIL import Image
import torch

# تحميل نموذج BLIP المدرب مسبقًا
blip_model = BlipForQuestionAnswering.from_pretrained("mshsahmed/blip-vqa-finetuned-kvasir-v58849")
processor = BlipProcessor.from_pretrained("mshsahmed/blip-vqa-finetuned-kvasir-v58849")

# تحميل نماذج الترجمة
ar_en_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
ar_en_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-ar-en")

en_ar_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
en_ar_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-ar")

# دوال الترجمة
def translate_ar_to_en(text):
    inputs = ar_en_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = ar_en_model.generate(**inputs)
    return ar_en_tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

def translate_en_to_ar(text):
    inputs = en_ar_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = en_ar_model.generate(**inputs)
    return en_ar_tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

# قاموس ترجمة طبي يدوي
medical_terms = {
    "chest x-ray": "أشعة سينية للصدر",
    "x-ray": "أشعة سينية",
    "ct scan": "تصوير مقطعي محوسب",
    "mri": "تصوير بالرنين المغناطيسي",
    "ultrasound": "تصوير بالموجات فوق الصوتية",
    "normal": "طبيعي",
    "abnormal": "غير طبيعي",
    "brain": "الدماغ",
    "fracture": "كسر",
    "no abnormality detected": "لا توجد شذوذات",
    "left lung": "الرئة اليسرى",
    "right lung": "الرئة اليمنى"
}

# دالة الترجمة الذكية للإجابات
def translate_answer_medical(answer_en):
    key = answer_en.lower().strip()
    if key in medical_terms:
        return medical_terms[key]
    else:
        return translate_en_to_ar(answer_en)

# الدالة الرئيسية
def vqa_multilingual(image, question):
    if not image or not question.strip():
        return "يرجى رفع صورة وكتابة سؤال.", "", "", ""

    # هل السؤال بالعربية؟
    is_arabic = any('؀' <= c <= 'ۿ' for c in question)
    if is_arabic:
        question_ar = question.strip()
        question_en = translate_ar_to_en(question_ar)
    else:
        question_en = question.strip()
        question_ar = translate_en_to_ar(question_en)

    # تمرير السؤال الإنجليزي للنموذج
    inputs = processor(image, question_en, return_tensors="pt")
    with torch.no_grad():
        output = blip_model.generate(**inputs)
    answer_en = processor.decode(output[0], skip_special_tokens=True).strip()

    # ترجمة الإجابة باستخدام القاموس أو آليًا
    answer_ar = translate_answer_medical(answer_en)

    return (
        f"السؤال بالعربية:
{question_ar}",
        f"السؤال بالإنجليزية:
{question_en}",
        f"الإجابة بالعربية:
{answer_ar}",
        f"الإجابة بالإنجليزية:
{answer_en}"
    )

# واجهة Gradio
gr.Interface(
    fn=vqa_multilingual,
    inputs=[
        gr.Image(type="pil", label="🔍 ارفع صورة الأشعة"),
        gr.Textbox(label="💬 أدخل سؤالك (بالعربية أو الإنجليزية)")
    ],
    outputs=[
        gr.Textbox(label="🟠 السؤال بالعربية"),
        gr.Textbox(label="🟢 السؤال بالإنجليزية"),
        gr.Textbox(label="🟠 الإجابة بالعربية"),
        gr.Textbox(label="🟢 الإجابة بالإنجليزية")
    ],
    title="نموذج ثنائي اللغة (عربي - إنجليزي)",
    description="ارفع صورة طبية واسأل بالعربية أو الإنجليزية، وستحصل على الإجابة باللغتين."
).launch(share=True)