import gradio as gr
from transformers import pipeline
import pandas as pd
import os

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

# Load ICD9 codes data
icd9_data = pd.read_csv('D_ICD_DIAGNOSES.csv')
icd9_data.columns = ['ROW_ID', 'ICD9_CODE', 'SHORT_TITLE', 'LONG_TITLE']

def preprocessing(test_df):
    test_df.loc[
        test_df['ICD9_CODE'].str.startswith("V"), 'ICD9_CODE'] = test_df.ICD9_CODE.apply(
        lambda x: x[:4])
    test_df.loc[
        test_df['ICD9_CODE'].str.startswith("E"), 'ICD9_CODE'] = test_df.ICD9_CODE.apply(
        lambda x: x[:4])
    test_df.loc[(~test_df.ICD9_CODE.str.startswith("E")) & (
        ~test_df.ICD9_CODE.str.startswith("V")), 'ICD9_CODE'] = test_df.ICD9_CODE.apply(
        lambda x: x[:3])
    return test_df

icd9_data = preprocessing(icd9_data)

def classify_symptoms(text):
    try:
        results = classifier(text, top_k=5)
        formatted_results = []
        for result in results:
            code = result['label']
            code_info = icd9_data[icd9_data['ICD9_CODE'] == code]
            formatted_results.append({
                "ICD9 Code": code,
                "Short Title": code_info['SHORT_TITLE'].iloc[0] if not code_info.empty else "N/A",
                "Long Title": code_info['LONG_TITLE'].iloc[0] if not code_info.empty else "N/A",
                "Confidence": f"{result['score']:.2%}"
            })
        return formatted_results
    except Exception as e:
        return f"Error processing classification: {str(e)}"

# Enhanced CSS with fullscreen layout and violet theme
custom_css = """
.gradio-container {
    width: 100% !important;
    max-width: 100% !important;
    margin: 0 !important;
    padding: 0 !important;
    min-height: 100vh !important;
    display: flex !important;
    flex-direction: column !important;
    background-color: #f5f3f7 !important;
}
.main-container {
    text-align: center;
    padding: 2rem;
    margin: 0;
    background: #ffffff;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
    width: 100%;
}
.content-wrapper {
    max-width: 1400px;
    margin: 0 auto;
    padding: 0 2rem;
    width: 100%;
    box-sizing: border-box;
}
h1 {
    color: #4a148c !important;
    font-size: 3rem !important;
    margin-bottom: 0.5rem !important;
}
h3 {
    color: #6a1b9a !important;
    font-size: 1.4rem !important;
    font-weight: normal !important;
}
.input-container {
    background: white !important;
    padding: 2rem !important;
    border-radius: 10px !important;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1) !important;
    margin: 2rem 0 !important;
    width: 100% !important;
}
textarea {
    background: white !important;
    color: #000000 !important;
    border: 2px solid #7b1fa2 !important;
    border-radius: 8px !important;
    padding: 1rem !important;
    font-size: 1.2rem !important;
    min-height: 150px !important;
    width: 100% !important;
}
.submit-btn {
    background-color: #6a1b9a !important;
    color: white !important;
    padding: 1rem 3rem !important;
    border-radius: 8px !important;
    font-size: 1.2rem !important;
    margin-top: 1.5rem !important;
    transition: background-color 0.3s ease !important;
    width: auto !important;
}
.submit-btn:hover {
    background-color: #4a148c !important;
}
.output-container {
    background: white !important;
    padding: 2rem !important;
    border-radius: 10px !important;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1) !important;
    margin: 2rem 0 !important;
    width: 100% !important;
}
.output-container pre {
    background: #f8f9fa !important;
    color: #000000 !important;
    border-radius: 8px !important;
    padding: 1.5rem !important;
    font-size: 1.1rem !important;
}
.examples-container {
    background: white !important;
    padding: 2rem !important;
    border-radius: 10px !important;
    margin: 2rem 0 !important;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1) !important;
    width: 100% !important;
}
.footer {
    text-align: center;
    padding: 2rem;
    background: white;
    margin-top: auto;
    width: 100%;
    box-shadow: 0 -2px 8px rgba(0, 0, 0, 0.1);
    color: #4a148c;
}
"""

with gr.Blocks(css=custom_css) as demo:
    with gr.Row(elem_classes=["main-container"]):
        with gr.Column(elem_classes=["content-wrapper"]):
            gr.Markdown(
                """
                # 🏥 Clinical Symptom ICD9 Classifier
                ### AI-Powered Medical Diagnosis Code Suggestion Tool
                """
            )
    
            with gr.Row():
                with gr.Column(elem_classes=["input-container"]):
                    input_text = gr.Textbox(
                        label="Clinical Symptom Description",
                        placeholder="Enter detailed patient symptoms and clinical observations...",
                        lines=5
                    )
                    submit_btn = gr.Button("Analyze Symptoms", elem_classes=["submit-btn"])
            
            with gr.Row(elem_classes=["output-container"]):
                output = gr.JSON(
                    label="Suggested ICD9 Diagnostic Codes with Descriptions"
                )
            
            with gr.Row(elem_classes=["examples-container"]):
                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,
                    label="Example Clinical Cases"
                )
    
    submit_btn.click(fn=classify_symptoms, inputs=input_text, outputs=output)
    input_text.submit(fn=classify_symptoms, inputs=input_text, outputs=output)
    
    with gr.Row(elem_classes=["footer"]):
        gr.Markdown(
            """
            ⚕️ <strong>Medical Disclaimer:</strong> This AI tool is designed to assist medical professionals in ICD9 code classification.
            Always verify suggestions with clinical judgment and consult appropriate medical resources.
            """
        )

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