import gradio as gr
from gliner import GLiNER
import torch

"""
Arabic Mental Health Named Entity Recognition Interface
Using GLiNER model: AhmadDarif/Arabic-Mental-NER
"""

# Initialize the model
print("Loading GLiNER model...")
try:
    model = GLiNER.from_pretrained("AhmadDarif/Arabic-Mental-NER")
    print("Model loaded successfully!")
except:
    # Fallback to base model if fine-tuned model not available
    print("Fine-tuned model not found, using base model...")
    model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")

# Entity labels
LABELS = ["MEDICATION", "DOSAGE", "DURATION"]

def extract_entities(text):
    """
    Extract named entities from Arabic text using GLiNER
    """
    if not text.strip():
        return "يرجى إدخال نص للتحليل / Please enter text to analyze", ""
    
    try:
        # Use GLiNER to predict entities
        entities = model.predict_entities(text, LABELS)
        
        if not entities:
            return "لم يتم العثور على أي كيانات / No entities found", ""
        
        # Format the results
        entities_text = "الكيانات المستخرجة / Extracted Entities:\n\n"
        entities_list = []
        
        for entity in entities:
            entity_text = entity["text"]
            entity_label = entity["label"]
            confidence = entity.get("score", 0)
            
            entities_list.append(f"{entity_text} => {entity_label}")
            entities_text += f"{entity_text} => {entity_label} (confidence: {confidence:.3f})\n"
        
        # Simple format for copy-paste (matching your exact format)
        simple_output = "\n".join(entities_list)
        
        return entities_text, simple_output
        
    except Exception as e:
        error_msg = f"خطأ في التحليل / Analysis Error: {str(e)}"
        return error_msg, ""

def analyze_text(text):
    """
    Main function to analyze Arabic text and return formatted results
    """
    detailed_output, simple_output = extract_entities(text)
    return detailed_output, simple_output

# Create the Gradio interface
with gr.Blocks(
    title="Arabic Mental Health NER",
    theme=gr.themes.Soft(),
    css="""
    .arabic-text {
        direction: rtl;
        text-align: right;
        font-family: 'Arial', 'Tahoma', sans-serif;
    }
    """
) as demo:
    
    gr.Markdown("""
    # 🧠 Arabic Mental Health Named Entity Recognition
    ## استخراج الكيانات المسماة للصحة النفسية باللغة العربية
    
    This tool extracts mental health-related entities from Arabic text using GLiNER, including:
    - **MEDICATION** (الأدوية): Names of medications
    - **DOSAGE** (الجرعة): Dosage information  
    - **DURATION** (المدة): Duration of treatment
    
    أدخل النص العربي أدناه لاستخراج الكيانات المتعلقة بالصحة النفسية
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                label="النص العربي / Arabic Text",
                placeholder="أدخل النص العربي هنا... / Enter Arabic text here...",
                lines=5,
                elem_classes=["arabic-text"]
            )
            
            analyze_btn = gr.Button(
                "تحليل النص / Analyze Text", 
                variant="primary",
                size="lg"
            )
            
            # Add some example texts
            gr.Examples(
                examples=[
                    ["تناولت سبرالكس بجرعة خمسة مليجرام لمدة ثلاثة أسابيع"],
                    ["وصف الطبيب دواء بروزاك بجرعة عشرة مليجرام يومياً لمدة شهرين"],
                    ["أخذت نصف حبة من الدواء لمدة أسبوعين"],
                    ["الـ سبرالكس بجرعة بسيطة خمسة مليجرام لمدة عشرة أيام"]
                ],
                inputs=input_text,
                label="أمثلة / Examples"
            )
        
        with gr.Column(scale=1):
            detailed_output = gr.Textbox(
                label="النتائج التفصيلية / Detailed Results",
                lines=10,
                elem_classes=["arabic-text"],
                interactive=False
            )
            
            simple_output = gr.Textbox(
                label="تنسيق بسيط للنسخ / Simple Format for Copy",
                lines=5,
                elem_classes=["arabic-text"],
                interactive=False,
                info="النتائج بالتنسيق: نص => تصنيف"
            )
    
    # Event handlers
    analyze_btn.click(
        fn=analyze_text,
        inputs=[input_text],
        outputs=[detailed_output, simple_output]
    )
    
    input_text.submit(
        fn=analyze_text,
        inputs=[input_text],
        outputs=[detailed_output, simple_output]
    )
    
    gr.Markdown("""
    ---
    **ملاحظة / Note**: هذا النموذج مدرب خصيصاً على البيانات العربية للصحة النفسية باستخدام GLiNER
    
    This model is specifically fine-tuned on Arabic mental health data using GLiNER.
    
    **Entity Types:**
    - MEDICATION: أسماء الأدوية والعقاقير
    - DOSAGE: معلومات الجرعة والكمية  
    - DURATION: مدة العلاج والاستخدام
    """)

if __name__ == "__main__":
    demo.launch()