import gradio as gr
from transformers import pipeline
import numpy as np

# Use a simpler approach that doesn't require sentence-transformers
MODEL_NAME = "samrawal/bert-base-uncased_clinical-ner"
CLASSES = [
    "Cardiology", "Neurology", "Oncology", "Pediatrics", 
    "Orthopedics", "Dermatology", "Gastroenterology", 
    "Endocrinology", "Psychiatry", "Pulmonology"
]

# Initialize model
classifier = pipeline(
    "text-classification", 
    model="bhadresh-savani/bert-base-uncased-emotion",
    tokenizer="bhadresh-savani/bert-base-uncased-emotion"
)

def predict_specialty(symptoms):
    """
    Predict the most relevant medical specialty based on symptoms.
    Simplified version without sentence-transformers.
    """
    # Get classification prediction
    pred = classifier(symptoms)
    predicted_class = pred[0]['label']
    
    # Simple mapping - in a real app you'd want more sophisticated logic
    specialty_map = {
        'sadness': 'Psychiatry',
        'joy': 'Pediatrics',  # Just example mapping
        'love': 'Cardiology',
        'anger': 'Neurology',
        'fear': 'Psychiatry',
        'surprise': 'Emergency Medicine'
    }
    
    primary = specialty_map.get(predicted_class, "General Practice")
    confidence = f"{pred[0]['score']*100:.1f}%"
    
    # Simple alternative suggestions
    alternatives = []
    if primary == "Psychiatry":
        alternatives = ["Neurology", "Endocrinology"]
    elif primary == "Cardiology":
        alternatives = ["Pulmonology", "Gastroenterology"]
    else:
        alternatives = ["General Practice", "Internal Medicine"]
    
    result = {
        "Primary Specialty": primary,
        "Confidence": confidence,
        "Alternative Suggestions": alternatives
    }
    
    return result

# Create Gradio interface
demo = gr.Interface(
    fn=predict_specialty,
    inputs=gr.Textbox(label="Describe your symptoms", placeholder="e.g., chest pain and shortness of breath..."),
    outputs=[
        gr.Label(label="Primary Specialty"),
        gr.Textbox(label="Confidence"),
        gr.JSON(label="Alternative Suggestions")
    ],
    examples=[
        ["chest pain and dizziness"],
        ["persistent headaches with nausea"],
        ["unexplained weight loss and fatigue"],
        ["skin rash and itching"]
    ],
    title="Medical Specialty Classifier",
    description="Enter your symptoms to find the most relevant medical specialty. Note: This is for educational purposes only and not a substitute for professional medical advice."
)

demo.launch()