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):
    try:
        results = classifier(text, top_k=5)
        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 without theme dependencies
custom_css = """
.gradio-container {
    background-color: #f8f9fa;
}

.contain {
    background-color: white !important;
    border-radius: 8px !important;
    padding: 20px !important;
    box-shadow: 0 2px 6px rgba(0, 0, 0, 0.1) !important;
    margin-bottom: 20px !important;
}

.textbox {
    border: 2px solid #4a90e2 !important;
    border-radius: 6px !important;
}

.label {
    color: #2c3e50 !important;
    font-weight: bold !important;
    font-size: 1.1em !important;
}

footer {
    text-align: center;
    padding: 20px;
    color: #666;
    font-size: 0.9em;
}
"""

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown(
        """
        # MedAI: Clinical Symptom ICD9 Classifier
        ### AI-Powered Diagnostic Code Suggestion Tool
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_text = 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
            )
            
            examples = gr.Examples(
                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"]
                ],
                inputs=input_text
            )
    
    output = gr.JSON(label="Diagnostic Suggestions")
    
    input_text.submit(fn=classify_symptoms, inputs=input_text, outputs=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()