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

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

# 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)}"


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: #000000 !important;
    color: #ffffff !important;
}
.main-container {
    text-align: center;
    padding: 2rem;
    margin: 0;
    background: #000000;
    width: 100%;
    color: #ffffff;
}
.content-wrapper {
    max-width: 1400px;
    margin: 0 auto;
    padding: 0 2rem;
    width: 100%;
    box-sizing: border-box;
    background: #000000;
    color: #ffffff;
}
h1 {
    color: #b388ff !important;
    font-size: 2.5rem !important;  /* Reduced from 3rem */
    margin-bottom: 0.5rem !important;
    font-weight: 700 !important;
}
h3 {
    color: #9575cd !important;
    font-size: 1.2rem !important;  /* Reduced from 1.4rem */
    font-weight: 500 !important;
    margin-bottom: 2rem !important;
}
.input-container {
    background: #121212 !important;
    padding: 2rem !important;
    border-radius: 12px !important;
    box-shadow: 0 4px 6px rgba(255, 255, 255, 0.05) !important;
    margin: 2rem 0 !important;
    width: 100% !important;
    border: 1px solid #333333 !important;
}
.input-container label {
    color: #ffffff !important;
    font-weight: 600 !important;
    font-size: 1rem !important;  /* Reduced from 1.1rem */
    margin-bottom: 0.5rem !important;
    background: transparent !important;
}
textarea {
    background: #1e1e1e !important;
    color: #ffffff !important;
    border: 2px solid #673ab7 !important;
    border-radius: 8px !important;
    padding: 1rem !important;
    font-size: 1.1rem !important;  /* Reduced from 1.2rem */
    min-height: 150px !important;
    width: 100% !important;
}
.submit-btn {
    background-color: #673ab7 !important;
    color: white !important;
    padding: 1rem 3rem !important;
    border-radius: 8px !important;
    font-size: 1.1rem !important;  /* Reduced from 1.2rem */
    margin-top: 1.5rem !important;
    transition: all 0.3s ease !important;
    width: auto !important;
    font-weight: 600 !important;
    border: none !important;
}
.submit-btn:hover {
    background-color: #5e35b1 !important;
}
.output-container {
    background: #121212 !important;
    padding: 2rem !important;
    border-radius: 12px !important;
    box-shadow: 0 4px 6px rgba(255, 255, 255, 0.05) !important;
    margin: 2rem 0 !important;
    width: 100% !important;
    border: 1px solid #333333 !important;
    color: #ffffff !important;
}
.output-container label {
    color: #ffffff !important;
    font-weight: 600 !important;
    font-size: 1rem !important;  /* Reduced from 1.1rem */
    margin-bottom: 1rem !important;
    background: transparent !important;
}
.examples-container {
    background: #121212 !important;
    padding: 2rem !important;
    border-radius: 12px !important;
    margin: 2rem 0 !important;
    box-shadow: 0 4px 6px rgba(255, 255, 255, 0.05) !important;
    width: 100% !important;
    border: 1px solid #333333 !important;
    color: #ffffff !important;
}
.examples-container label {
    color: #ffffff !important;
    font-weight: 600 !important;
    font-size: 1rem !important;  /* Reduced from 1.1rem */
    background: transparent !important;
}
.footer {
    text-align: center;
    padding: 2rem;
    background: #000000;
    margin-top: auto;
    width: 100%;
    border-top: 1px solid #333333;
    color: #ffffff;
}
"""

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(elem_classes=["input-output-row"]):
                with gr.Column(elem_classes=["input-container"]):
                    gr.Markdown("Clinical Symptom Description")
                    with gr.Column(elem_classes=["inner-input-container"]):
                        input_text = gr.Textbox(
                            show_label=False,
                            placeholder="Enter detailed patient symptoms and clinical observations...",
                            lines=5
                        )
                    submit_btn = gr.Button("Analyze Symptoms", elem_classes=["submit-btn"])
                
                with gr.Column(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)