import gradio as gr
from transformers import pipeline

# Load the model
classifier = pipeline(
    "text-classification", 
    model="ashishkgpian/biobert_icd9_classifier_ehr"
)

def classify_symptoms(text):
    """
    Classify medical symptoms and return top ICD9 codes
    
    Args:
        text (str): Input medical symptom description
    
    Returns:
        dict: Top classification results with ICD9 codes and probabilities
    """
    try:
        # Get classification results
        results = classifier(text, top_k=5)
        
        # Format results for more readable output
        formatted_results = []
        for result in results:
            formatted_results.append({
                "ICD9 Code": result['label'], 
                "Confidence": f"{result['score']:.2%}"
            })
        
        return formatted_results
    
    except Exception as e:
        return f"Error processing classification: {str(e)}"

# Custom CSS for medical-themed UI
custom_css = """
/* Medical-inspired color palette and design */
.gradio-container {
    background-color: #f4f7f6;
    font-family: 'Arial', 'Helvetica Neue', sans-serif;
}

/* Styled input area */
.input-container {
    background-color: #ffffff;
    border: 2px solid #3498db;
    border-radius: 10px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    padding: 15px;
}

/* Result card styling */
.result-card {
    background-color: #ffffff;
    border-left: 5px solid #2ecc71;
    margin-bottom: 10px;
    padding: 10px;
    border-radius: 5px;
    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
}

/* Medical icon-like styling for headers */
.gradio-header {
    color: #2c3e50;
    text-align: center;
    font-weight: bold;
}

/* Soften the medical green theme */
.submit-button {
    background-color: #2ecc71 !important;
    color: white !important;
    border: none !important;
    transition: all 0.3s ease;
}

.submit-button:hover {
    background-color: #27ae60 !important;
    transform: scale(1.05);
}
"""

# Create Gradio interface with advanced medical-themed design
demo = gr.Interface(
    fn=classify_symptoms,
    inputs=gr.Textbox(
        label="Patient Symptom Description", 
        placeholder="Enter detailed patient symptoms (e.g., 'Patient reports persistent chest pain radiating to left arm, accompanied by shortness of breath')",
        lines=4,
        container=False,
        elem_classes=["input-container"]
    ),
    outputs=gr.JSON(
        label="Diagnostic Suggestions", 
        elem_classes=["result-card"]
    ),
    title="Symptom-to-ICD9 Classifier",
    description="An advanced AI-powered diagnostic assistance tool for converting clinical symptom descriptions into potential ICD9 diagnostic codes.",
    theme="huggingface",
    css=custom_css,
    examples=[
        ["45-year-old male experiencing severe chest pain, radiating to left arm, with shortness of breath and excessive sweating"],
        ["Persistent headache for 2 weeks, accompanied by dizziness and occasional blurred vision"],
        ["Diabetic patient reporting frequent urination, increased thirst, and unexplained weight loss"],
        ["Elderly patient with chronic knee pain, reduced mobility, and signs of inflammation"]
    ]
)

# Add some footer-like information
demo.layout = gr.Column(
    [
        demo.title,
        demo.description,
        demo.input,
        demo.output,
        gr.Markdown("**Disclaimer:** This is an AI-assisted diagnostic tool. Always consult with a healthcare professional for accurate diagnosis and treatment.")
    ]
)

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