app.py

import streamlit as st
import numpy as np
import pandas as pd
import joblib
import torch
from transformers import AutoTokenizer, AutoModel
from xgboost import XGBClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.metrics import precision_recall_curve, roc_curve, confusion_matrix, classification_report
import matplotlib.pyplot as plt
import shap
import plotly.express as px
import streamlit as st
import pandas as pd
import datetime
import json
import requests
from streamlit_lottie import st_lottie
import streamlit.components.v1 as components
from streamlit_navigation_bar import st_navbar
from transformers import AutoTokenizer, AutoModel
import re
from tqdm import tqdm
import torch
import os
from hugchat.login import Login
from hugchat import hugchat
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch.nn as nn
import time


# Defining wide mode
st.set_page_config(layout="wide")

dark_theme = """
    <style>
        body, .stApp {
            background-color: #0e1117;
            color: white;
        }
        .stTextInput, .stButton>button {
            background-color: #222;
            color: white;
        }
        .stMarkdown, .stTextArea, .stSelectbox, .stCheckbox {
            color: white;
        }
    </style>
"""
st.markdown(dark_theme, unsafe_allow_html=True)

# Criar o menu na barra lateral
st.sidebar.title("📌 Menu")
page = st.sidebar.radio(
    "Selecione uma opção:",
    ["🏠 Home", "📊 Tabular Data", "📝 Clinical Text Notes", "🔀 Ensemble Prediction"]
)
if page=="🏠 Home":
   
    st.markdown("""
        <style>
            .title {
                text-align: center;
                font-size: 36px;
                font-weight: bold;
                color: #2C3E50;
            }
            .subtitle {
                text-align: center;
                font-size: 22px;
                color: #7F8C8D;
            }
            .box {
                background-color: #ECF0F1;
                padding: 15px;
                border-radius: 10px;
                text-align: center;
                margin-bottom: 10px;
                font-size: 18px;
            }
        </style>
    """, unsafe_allow_html=True)

    # Header
    st.markdown("<h1 class='title'>📊 AI Clinical Readmission Predictor</h1>", unsafe_allow_html=True)
    st.markdown("<h2 class='subtitle'>Using Machine Learning for Better Patient Outcomes</h2>", unsafe_allow_html=True)
    image_1 ='https://content.presspage.com/uploads/2110/4970f578-5f20-4675-acc2-3b2cda25fa96/1920_ai-machine-learning-cedars-sinai.jpg?10000'
    image_2 = 'https://med-tech.world/app/uploads/2024/10/AI-Hospitals.jpg.webp'


    st.image(image_2, width=1450)  # Hospital Icon

    st.write("This app helps predict patient readmission risk using machine learning models. "
            "Upload data, analyze clinical notes, and see predictions from our ensemble model.")

    # Navigation Buttons
    st.markdown("---")
    st.markdown("<h3 style='text-align: center;'>🚀 Explore the App</h3>", unsafe_allow_html=True)

elif page== "📊 Tabular Data":
    
    # Function to load Lottie animation
    def load_lottie(url):
        response = requests.get(url)
        if response.status_code != 200:
            return None
        return response.json()
    
    # Load Lottie Animation
    lottie_hello = load_lottie("https://assets7.lottiefiles.com/packages/lf20_jcikwtux.json")
    if lottie_hello:
        st_lottie(lottie_hello, speed=1, loop=True, height=200)
    
    # Load dataset
    df = pd.read_csv('/home/user/app/ensemble_test.csv')
    
    # Streamlit App Header
    st.title('🏥 Hospital Readmission Prediction')
    st.markdown("""
    <h3 style='text-align: center; color: gray;'>Predict ICU hospital readmission using Artificial Intelligence</h3>
    """, unsafe_allow_html=True)
    st.markdown("---")
    
    # Helper Functions
    def get_age_group(age):
        """Classify age into predefined groups with correct column names."""
        if 36 <= age <= 50:
            return "age_group_36-50 (Middle-Aged Adults)"
        elif 51 <= age <= 65:
            return "age_group_51-65 (Older Middle-Aged Adults)"
        elif 66 <= age <= 80:
            return "age_group_66-80 (Senior Adults)"
        elif age >= 81:
            return "age_group_81+ (Elderly)"
        return "age_group_Below_36"
    
    
    def get_period(hour):
        """Determine admission/discharge period."""
        return "Morning" if 6 <= hour < 18 else "Night"
    
    # **User Inputs**
    st.subheader("📌 Select the admission's Characteristics")
    
    admission_type = st.selectbox("🛑 Type of Admission", df.columns[df.columns.str.startswith('admission_type_')])
    admission_location = st.selectbox("📍 Admission Location", df.columns[df.columns.str.startswith('admission_location_')])
    discharge_location = st.selectbox("🏥 Discharge Location", df.columns[df.columns.str.startswith('discharge_location_')])
    insurance = st.selectbox("💰 Insurance Type", df.columns[df.columns.str.startswith('insurance_')])

    st.sidebar.subheader("📊 Patient Information")
    language = st.sidebar.selectbox("🗣 Language", df.columns[df.columns.str.startswith('language_')])
    marital_status = st.sidebar.selectbox("💍 Marital Status", df.columns[df.columns.str.startswith('marital_status_')])
    race = st.sidebar.selectbox("🧑 Race", df.columns[df.columns.str.startswith('race_')])
    sex = st.sidebar.selectbox("⚧ Sex", ['gender_M', 'gender_F'])
    age = st.sidebar.slider("📅 Age", 18, 100, 50)
    
    admission_time = st.time_input("⏳ Admission Time", value=datetime.time(12, 0))
    discharge_time = st.time_input("⏳ Discharge Time", value=datetime.time(12, 0))
    
    # Laboratory & Clinical Values
    st.subheader("📈 Clinical Values")
    numerical_features = ['los_days', 'previous_stays', 'n_meds', 'drg_severity', 'drg_mortality', 'time_since_last_stay', 
                          'blood_cells', 'hemoglobin', 'glucose', 'creatine', 'plaquete']
    numeric_inputs = {}
    cols = st.columns(len(numerical_features))
    
    # General Numerical Values
    st.subheader("📊 General Hosptal Information")
    general_numerical_features = ['los_days', 'previous_stays', 'n_meds', 'drg_severity', 
                                  'drg_mortality', 'time_since_last_stay']
    
    general_inputs = {}
    cols = st.columns(3)  # Three columns for general values
    
    for i, feature in enumerate(general_numerical_features):
        col_index = i % 3  # Distribute across columns
        min_val, max_val = df[feature].min(), df[feature].max()
        
        with cols[col_index]:
            general_inputs[feature] = st.slider(
                f"📌 {feature.replace('_', ' ').title()}",
                float(min_val),
                float(max_val),
                float((min_val + max_val) / 2)
            )
    
    # Laboratory Values
    st.subheader("🧪 Laboratory Test Results")
    lab_numerical_features = ['blood_cells', 'hemoglobin', 'glucose', 
                              'creatine', 'plaquete']
    
    lab_inputs = {}
    lab_cols = st.columns(3)  # Three columns for lab values
    
    for i, feature in enumerate(lab_numerical_features):
        col_index = i % 3  # Distribute across columns
        min_val, max_val = df[feature].min(), df[feature].max()
        
        with lab_cols[col_index]:
            lab_inputs[feature] = st.slider(
                f"🩸 {feature.replace('_', ' ').title()}",
                float(min_val),
                float(max_val),
                float((min_val + max_val) / 2)
            )
    min_val, max_val = df["cci_score"].min(), df["cci_score"].max()
    lab_inputs["cci_score"] = st.sidebar.slider(
        f"📌 CCI Score",
        float(min_val),
        float(max_val),
        float((min_val + max_val) / 2)
)
    
    # Process Inputs into Features
    feature_vector = {col: 0 for col in df.columns}
    feature_vector.update({
        admission_type: 1,
        admission_location: 1,
        discharge_location: 1,
        insurance: 1,
        language: 1,
        marital_status: 1,
        race: 1,
        "gender_M": 1 if sex == "gender_M" else 0,
        f"admit_period_{get_period(admission_time.hour)}": 1,
        f"discharge_period_{get_period(discharge_time.hour)}": 1
    })
    age_group = get_age_group(age)  # This function now returns correct dataset column names
    
    # Use the exact column names from the dataset
    for group in [
        "age_group_36-50 (Middle-Aged Adults)", 
        "age_group_51-65 (Older Middle-Aged Adults)", 
        "age_group_66-80 (Senior Adults)", 
        "age_group_81+ (Elderly)"
    ]:
        feature_vector[group] = 1 if group == age_group else 0  # Set selected group to 1, others to 0
    
    feature_vector.update(numeric_inputs)
    # Display Processed Data
    st.markdown("---")
    
    # Load XGBoost model 
    tabular_model_path = "/home/user/app/final_xgboost_model.pkl"
    tabular_model = joblib.load(tabular_model_path)
    print("✅ XGBoost Tabular Model loaded successfully!")
    
    # Load dataset columns (use the same order as training)
    expected_columns = [
        col for col in df.columns if col not in ["Unnamed: 0", "subject_id", "hadm_id", "probs"]
    ]
    
    # Define correct dataset column names for age groups
    age_group_mapping = {
        "age_group_36-50": "age_group_36-50 (Middle-Aged Adults)",
        "age_group_51-65": "age_group_51-65 (Older Middle-Aged Adults)",
        "age_group_66-80": "age_group_66-80 (Senior Adults)",
        "age_group_81+": "age_group_81+ (Elderly)",
    }
    
    # Process Inputs into Features
    feature_vector = {col: 0 for col in df.columns}
    
    # Set selected categorical features to 1
    feature_vector.update({
        admission_type: 1,
        admission_location: 1,
        discharge_location: 1,
        insurance: 1,
        language: 1,
        marital_status: 1,
        race: 1,
        "gender_M": 1 if sex == "gender_M" else 0,
        f"admit_period_{get_period(admission_time.hour)}": 1,
        f"discharge_period_{get_period(discharge_time.hour)}": 1
    })
    
    # Set correct age group
    age_group = get_age_group(age)
    for group in [
        "age_group_36-50 (Middle-Aged Adults)", 
        "age_group_51-65 (Older Middle-Aged Adults)", 
        "age_group_66-80 (Senior Adults)", 
        "age_group_81+ (Elderly)"
    ]:
        feature_vector[group] = 1 if group == age_group else 0
    
    # Update with numerical inputs
    feature_vector.update(general_inputs)
    feature_vector.update(lab_inputs)
    
    # Ensure feature order matches expected model input
    fixed_feature_vector = {age_group_mapping.get(k, k): v for k, v in feature_vector.items()}
    feature_df = pd.DataFrame([fixed_feature_vector]).reindex(columns=expected_columns, fill_value=0)
    
    #st.write(feature_df)
    # Predict probability of readmission
    prediction_proba = tabular_model.predict_proba(feature_df)[:, 1]
    probability = float(prediction_proba[0])  # Convert NumPy array to scalar
    st.session_state["XGBoost probability"] = probability
    prediction = (prediction_proba >= 0.5).astype(int)
    
    import shap
    import matplotlib.pyplot as plt
    import streamlit.components.v1 as components  # Required for displaying SHAP force plot
    
    st.write(f"Raw Prediction Probability: {probability:.4f}")
    
    # Prediction Button
    if st.button("🚀 Predict Readmission"):
        with st.spinner("🔍 Processing Prediction..."):
            st.subheader("🎯 Prediction Results")
            col1, col2 = st.columns(2)
    
            with col1:
                st.metric(label="🧮 Readmission Probability", value=f"{probability:.2%}")
    
            with col2:
                if prediction == 1:
                    st.error("⚠️ High Risk of Readmission")
                else:
                    st.success("✅ Low Risk of Readmission")      

    # Feature Importance Button
    if st.button("🔍 Feature Importance for Prediction"):
        st.metric(label="🧮 Readmission Probability", value=f"{probability:.2%}")
        # ✅ Initialize SHAP Explainer for XGBoost
        explainer = shap.TreeExplainer(tabular_model)
        shap_values = explainer.shap_values(feature_df)  # SHAP values for all samples
    
        # ✅ Convert SHAP values into a DataFrame (Sorting First)
        shap_df = pd.DataFrame({
            "Feature": feature_df.columns,
            "SHAP Value": shap_values[0]  # SHAP values for the first instance
        })
    
        # ✅ Select **Top 10 Most Important Features** (Sorted by Absolute SHAP Value)
        shap_df["abs_SHAP"] = shap_df["SHAP Value"].abs()  # Add column with absolute values
        shap_df = shap_df.sort_values(by="abs_SHAP", ascending=False).head(10)  # Top 10
    
        # Get top features and their SHAP impact values (shap_df assumed to be available)
        top_features = sorted(zip(shap_df['Feature'], shap_df['SHAP Value']), key=lambda x: abs(x[1]), reverse=True)
    
        # Create a formatted string for `top_factors` to be shown in the UI
        top_factors = "\n".join([f"- {feat}: {round(value, 2)} impact" for feat, value in top_features])
    
        import time
        import random
        from textwrap import dedent

# Título estilizado
        st.markdown("""
            <h1 style='text-align: center; color: #2c3e50;'>🩺 AI-Powered Patient Readmission Analysis</h1>
            <hr style='border: 2px solid #3498db;'>
        """, unsafe_allow_html=True)
        
        # Função para animar uma aranha
        placeholder = st.empty()
        
        for i in range(40):  # 4 segundos (40 iterações de 0.1s)
            time.sleep(0.1)
        placeholder.empty()
        
        with st.spinner("🤖 Analyzing..."):
            time.sleep(2)  # Simular carregamento
            try:
                ai_output = """
                # 🩺 AI-Powered Patient Readmission Analysis
                
                ## 🤖 Understanding the Model's Prediction
                The feature impacts indicate how much each factor contributes to the model's decision to predict patient readmission. The magnitude and direction of the impact determine its importance.
                
                ---
                
                ## 🔍 Key Characteristics Influencing the Prediction:
                - **Previous Stays** (1.56 impact): This feature has the highest positive impact, suggesting that patients with a history of previous hospital stays are more likely to be readmitted. The large impact indicates that the model places significant weight on this factor, likely because repeated hospital visits can signify chronic conditions or complications that are not fully resolved.
                - **Number of Medications (n_meds)** (0.17 impact): Patients taking a higher number of medications are at a higher risk of readmission. This could be due to the complexity of their medical conditions, potential side effects, or interactions between medications that may lead to further health issues.
                - **Time Since Last Stay** (0.16 impact): The time elapsed since the patient's last hospital stay also positively influences the prediction of readmission. This might indicate that patients who have been discharged recently are at a higher risk of returning, possibly due to incomplete recovery or the nature of their condition requiring ongoing care.
                - **Real Age** (0.14 impact): The patient's age is another factor that increases the likelihood of readmission. Older patients may have more complex health issues, diminished physiological reserve, and a higher likelihood of comorbid conditions, all of which can contribute to the need for repeat hospitalizations.
                - **Length of Stay (los_days)** (0.07 impact): Although less influential than the top factors, a longer hospital stay during the current or previous admission(s) slightly increases the risk of readmission. This could be indicative of more severe illness, complications, or the need for prolonged recovery periods.
                
                ---
                
                ## 📉 Less Influential but Still Relevant Factors:
                - **CCI Score** (0.05 impact): The Charlson Comorbidity Index (CCI) score predicts the ten-year mortality for a patient with comorbid conditions such as heart disease, diabetes, or cancer. A higher score indicates a higher risk of mortality and, by extension, potentially a higher risk of readmission due to the complexity of the patient's health conditions.
                - **DRG Severity** (-0.04 impact): The Diagnosis-Related Group (DRG) severity categorizes hospital cases based on expected resource use. Its negative impact suggests that higher severity cases might actually have a slightly lower risk of readmission, possibly due to more intensive treatment and monitoring during their initial stay.
                - **Race** (race_WHITE: -0.04 impact, race_UNKNOWN: 0.03 impact): The impacts of race are relatively small and might reflect underlying socio-economic or healthcare access disparities rather than direct biological factors. However, interpreting these impacts requires caution due to the potential for confounding variables and the ethical considerations surrounding race in healthcare outcomes.
                - **Blood Cells** (0.03 impact): This factor, likely referring to some measure of blood cell count or health, has a minor positive impact, suggesting that abnormalities in blood cell counts could slightly increase the risk of readmission, potentially due to underlying conditions affecting the blood or bone marrow.
                
                ---
                
                ## ✅ Conclusion:
                The model's prediction of patient readmission is most strongly influenced by the patient's history of previous hospital stays, indicating a potential for chronic or recurring health issues. The number of medications, time since the last stay, and the patient's age are also significant factors, highlighting the complexity of the patient's health condition and the potential for ongoing care needs. 
                
                While other factors such as CCI score, DRG severity, race, and blood cell health play a role, their impacts are less pronounced. Understanding these factors can help healthcare providers identify high-risk patients and implement targeted interventions to reduce the likelihood of readmission.

                """
        
                # 🎭 **Show AI Response in a Stylish Chat Format**
                with st.chat_message("assistant"):
                    st.markdown(f"**💡 AI Explanation:**\n\n")
                    st.markdown(dedent(ai_output))
            except Exception as e:
                st.error(f"⚠️ Error retrieving response: {e}")
                st.stop()

        # Show Top 10 Features
        #st.write(shap_df[["Feature", "SHAP Value"]])  # Display only relevant columns
    
        # ✅ SHAP Bar Plot (Corrected for Top 10 Selection)
        fig, ax = plt.subplots(figsize=(8, 6))
        shap.bar_plot(shap_df["SHAP Value"].values, shap_df["Feature"].values)  # Correct Top 10
        st.pyplot(fig)
    
        # 🎯 SHAP Force Plot (How Features Affected the Prediction)
        st.subheader("🎯 SHAP Force Plot (How Features Affected the Prediction)")
    
        # ✅ Fix: Use explainer.expected_value (single scalar)
        force_plot = shap.force_plot(
            explainer.expected_value, shap_values[0], feature_df.iloc[0], matplotlib=False
        )
    
        # ✅ Convert SHAP force plot to HTML
        shap_html = f"<head>{shap.getjs()}</head><body>{force_plot.html()}</body>"
    
        # ✅ Render SHAP force plot in Streamlit
        components.html(shap_html, height=400)
                   
elif page == "📝 Clinical Text Notes":
    # Set Streamlit Page Title
    st.subheader("📝 Clinical Text Note")

    # Utility Functions

    def clean_text(text):
        """Cleans input text by removing non-ASCII characters, extra spaces, and unwanted symbols."""
        text = re.sub(r"[^\x20-\x7E]", " ", text)
        text = re.sub(r"_{2,}", "", text)
        text = re.sub(r"\s+", " ", text)
        text = re.sub(r"[^\w\s.,:;*%()\[\]-]", "", text)
        return text.lower().strip()

    import re

    def extract_fields(text):
        """Extracts key fields from clinical notes using regex patterns."""
        patterns = {
            "Discharge Medications": r"Discharge Medications[:\-]?\s*(.+?)\s+(?:Discharge Disposition|Discharge Condition|Discharge Instructions|Followup Instructions|$)",
            "Discharge Diagnosis": r"Discharge Diagnosis[:\-]?\s*(.+?)\s+(?:Discharge Condition|Discharge Medications|Discharge Instructions|Followup Instructions|$)",
            "Discharge Instructions": r"Discharge Instructions[:\-]?\s*(.*?)\s+(?:Followup Instructions|Discharge Disposition|Discharge Condition|$)",
            "History of Present Illness": r"History of Present Illness[:\-]?\s*(.+?)\s+(?:Past Medical History|Social History|Family History|Physical Exam|$)",
            "Past Medical History": r"Past Medical History[:\-]?\s*(.+?)\s+(?:Social History|Family History|Physical Exam|$)"
        }

        extracted_data = {}

        for field, pattern in patterns.items():
            match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
            if match:
                extracted_data[field] = match.group(1).strip()

        return extracted_data

    def extract_features(texts, model, tokenizer, device, batch_size=8):
        """Extracts CLS token embeddings from the Clinical-Longformer model."""
        all_features = []
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i+batch_size]
            inputs = tokenizer(batch_texts, return_tensors="pt", truncation=True, padding=True, max_length=4096).to(device)
            global_attention_mask = torch.zeros_like(inputs["input_ids"]).to(device)
            global_attention_mask[:, 0] = 1  # Set global attention for CLS token
            
            with torch.no_grad():
                outputs = model(**inputs, global_attention_mask=global_attention_mask)
            
            all_features.append(outputs.last_hidden_state[:, 0, :])
        
        return torch.cat(all_features, dim=0)


    def extract_entities(text, pipe, entity_group):
        """Extracts specific entities from the clinical note using a NER pipeline."""
        entities = pipe(text)
        return [ent['word'] for ent in entities if ent['entity_group'] == entity_group] or ["No relevant entities found"]

    # Load Model and Tokenizer
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    @st.cache_resource()
    def load_models():
        """Loads transformer models for text processing and NER."""
        longformer_tokenizer = AutoTokenizer.from_pretrained("yikuan8/Clinical-Longformer")
        longformer_model = AutoModel.from_pretrained("yikuan8/Clinical-Longformer").to(device).eval()
        
        ner_tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
        ner_model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
        ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple")
        
        return longformer_tokenizer, longformer_model, ner_pipe

    longformer_tokenizer, longformer_model, ner_pipe = load_models()

    # Text Input
    clinical_note = st.text_area("✍️ Enter Clinical Note", placeholder="Write the clinical note here...")

    if clinical_note:
        cleaned_note = clean_text(clinical_note)
        #st.write("### 📝 Cleaned Clinical Note:")
        #st.write(cleaned_note)
        
        # Extract Fields
        extracted_data = extract_fields(cleaned_note)
        st.write("### Extracted Fields")
        st.write(extracted_data)
        
        # Extract Embeddings
        with st.spinner("🔍 Extracting embeddings..."):
            embeddings = extract_features([cleaned_note], longformer_model, longformer_tokenizer, device)
        #st.write("### Extracted Embeddings")
        #st.write(embeddings)
        # Definir a classe RobustMLPClassifier
        class RobustMLPClassifier(nn.Module):
            def __init__(self, input_dim, hidden_dims=[256, 128, 64], dropout=0.3, activation=nn.ReLU()):
                super(RobustMLPClassifier, self).__init__()
                layers = []
                current_dim = input_dim
                
                for h in hidden_dims:
                    layers.append(nn.Linear(current_dim, h))
                    layers.append(nn.BatchNorm1d(h))
                    layers.append(activation)
                    layers.append(nn.Dropout(dropout))
                    current_dim = h
                
                layers.append(nn.Linear(current_dim, 1))
                self.net = nn.Sequential(*layers)
                
            def forward(self, x):
                return self.net(x)
            
        # --- Load MLP Model and PCA ---
        mlp_model_path = "/home/user/app/best_mlp_model_full.pth"
        pca_path = "/home/user/app/best_pca_model.pkl"

        best_mlp_model = torch.load(mlp_model_path, weights_only=False)
        best_mlp_model.to(device)
        best_mlp_model.eval()

        pca = joblib.load(pca_path)

        def predict_readmission(texts):
            """Predicts hospital readmission probability using Clinical-Longformer embeddings and MLP."""
            embeddings = extract_features(texts, longformer_model, longformer_tokenizer, device)
            embeddings_pca = pca.transform(embeddings.cpu().numpy())  # Apply PCA

            inputs = torch.FloatTensor(embeddings_pca).to(device)
            
            with torch.no_grad():
                logits = best_mlp_model(inputs)
                probabilities = torch.sigmoid(logits).cpu().numpy()
            
            return probabilities

        # Extract Medical Entities
        with st.spinner("🔍 Identifying medical entities..."):
            extracted_data["Extracted Medications"] = extract_entities(
                extracted_data.get("Discharge Medications", ""), ner_pipe, "Medication"
            )

            extracted_data["Extracted Diseases"] = extract_entities(
                extracted_data.get("Discharge Diagnosis", ""), ner_pipe, "Disease_disorder"
            )

            extracted_data["Extracted Diseases (Past Medical History)"] = extract_entities(
                extracted_data.get("Past Medical History", ""), ner_pipe, "Disease_disorder"
            )

            extracted_data["Extracted Diseases (History of Present Illness)"] = extract_entities(
                extracted_data.get("History of Present Illness", ""), ner_pipe, "Disease_disorder"
            )

            # Extração de sintomas agora inclui "History of Present Illness"
            extracted_data["Extracted Symptoms"] = extract_entities(
                extracted_data.get("Review of Systems", "") + " " + extracted_data.get("History of Present Illness", ""), 
                ner_pipe, "Sign_symptom"
            )


        def clean_entities(entities):
            """Reconstruct fragmented tokens and remove duplicates."""
            cleaned = []
            temp = ""

            for entity in entities:
                if entity.startswith("##"):  # Fragmented token
                    temp += entity.replace("##", "")  
                else:
                    if temp:
                        cleaned.append(temp)  # Add the reconstructed token
                    temp = entity
            if temp:
                cleaned.append(temp)  # Add the last processed token

            # Filter out irrelevant short words and special characters
            cleaned = [word for word in cleaned if len(word) > 2 and not re.match(r"^[\W_]+$", word)]
            
            return sorted(set(cleaned))  # Remove duplicates and sort

        # Clean extracted diseases and symptoms
        diseases_cleaned = clean_entities(
            extracted_data.get("Extracted Diseases", []) +
            extracted_data.get("Extracted Diseases (Past Medical History)", []) +
            extracted_data.get("Extracted Diseases (History of Present Illness)", [])
        )
        # Clean and reconstruct medication names
        medications_cleaned = clean_entities(extracted_data.get("Extracted Medications", []))

        # Store cleaned data in the main dictionary
        extracted_data["Extracted Medications Cleaned"] = medications_cleaned

        symptoms_cleaned = clean_entities(extracted_data.get("Extracted Symptoms", []))

        # Display extracted entities
        def display_list(title, items, icon="📌"):
            """Display extracted medical entities in an expandable list."""
            with st.expander(f"**{title} ({len(items)})**"):
                if items:
                    for item in items:
                        st.markdown(f"- {icon} **{item}**")
                else:
                    st.markdown("_No information available._")


        # Layout Header
        st.markdown("## 🏥 **Patient Medical Analysis**")
        st.markdown("---")

        # Creating columns for metrics
        col1, col2, col3 = st.columns(3)

        # Medications Metrics
        num_medications = len(medications_cleaned )
        col1.metric(label="💊 Total Medications", value=num_medications)

        # Diseases Metrics
        num_diseases = len(diseases_cleaned)
        col2.metric(label="🦠 Total Diseases", value=num_diseases)

        # Symptoms Metrics
        num_symptoms = len(symptoms_cleaned)
        col3.metric(label="🤒 Total Symptoms", value=num_symptoms)

        st.markdown("---")

        # Organizing lists in two columns
        col1, col2 = st.columns(2)

        # Display Medications List
        with col1:
            st.markdown("### 💊 **Medications**")
            display_list("Medication List", medications_cleaned , icon="💊")

        # Display Diseases List
        with col2:
            st.markdown("### 🦠 **Diseases**")
            display_list("Disease List", diseases_cleaned, icon="🦠")

        # Symptoms Section
        st.markdown("### 🤒 **Symptoms**")
        display_list("Symptoms List", symptoms_cleaned, icon="🤒")

        st.markdown("---")

        # Load tokenizer
        tokenizer = AutoTokenizer.from_pretrained("yikuan8/Clinical-Longformer")

        # Functions for token count and truncation
        def count_tokens(text):
            tokens = tokenizer.tokenize(text)
            return len(tokens)

        def trunced_text(nr):
            return 1 if nr > 4096 else 0

        # Dictionary of diseases with synonyms (matching capitalization in the image)
        disease_synonyms = {
            "Pneumonia": ["pneumonia", "pneumonitis"],
            "Diabetes": ["diabetes", "diabetes mellitus", "dm"],
            "CHF": ["CHF", "congestive heart failure", "heart failure"],
            "Septicemia": ["septicemia", "sepsis", "blood infection"],
            "Cirrhosis": ["cirrhosis", "liver cirrhosis", "hepatic cirrhosis"],
            "COPD": ["COPD", "chronic obstructive pulmonary disease"],
            "Renal_Failure": ["renal failure", "kidney failure", "chronic kidney disease", "CKD"]
        }

        # Extract relevant fields (assuming extract_fields is defined elsewhere)
        extracted_data = extract_fields(clinical_note)

        # Compute token counts
        number_of_tokens = count_tokens(clinical_note)
        number_of_tokens_med = count_tokens(extracted_data.get("Discharge Medications", ""))
        number_of_tokens_dis = count_tokens(extracted_data.get("Discharge Diagnosis", ""))
        trunced = trunced_text(number_of_tokens)

        # Convert diagnosis text to lowercase for case-insensitive matching
        full_diagnosis_text = extracted_data.get("Discharge Diagnosis", "").lower()

        # Function to check for any synonym in the diagnosis text
        def check_disease_presence(disease_list, text):
            return int(any(re.search(rf"\b{synonym}\b", text, re.IGNORECASE) for synonym in disease_list))

        # Create binary columns for each disease based on synonyms
        disease_flags = {disease: check_disease_presence(synonyms, full_diagnosis_text) 
                        for disease, synonyms in disease_synonyms.items()}

        # Count total diseases found
        disease_flags["total_conditions"] = sum(disease_flags.values())

        # Create DataFrame with a single row (matching column names from the image)
        df = pd.DataFrame([{
            'number_of_tokens_dis': number_of_tokens_dis,
            'number_of_tokens': number_of_tokens,
            'number_of_tokens_med': number_of_tokens_med,
            'diagnostic_count': num_diseases,  # Ensuring column name matches
            'total_conditions': disease_flags["total_conditions"],  # Matching name
            'trunced': trunced,
            **{disease: disease_flags[disease] for disease in disease_synonyms.keys()}  # Disease presence flags
        }])

        # Display DataFrame
        #st.write(df)

        #load lighGBoost model
        light_path = '/home/user/app/best_lgbm_model.pkl'
        light_model = joblib.load(light_path)
        #st.write("LightGBoost Model loaded sucessfully!")

        # Ensure df is already created from previous steps
        # Select only the columns that match the model input
        model_features = light_model.feature_name_

        # Check if all required features are in df
        missing_features = [feat for feat in model_features if feat not in df.columns]
        if missing_features:
            st.write(f"⚠️ Warning: Missing features in df: {missing_features}")

        # Fill missing columns with 0 (if needed, assuming binary features)
        for feat in missing_features:
            df[feat] = 0  # Default to 0 for missing binary disease indicators

        # Reorder df to match model features exactly
        df = df[model_features]

        # Convert df to NumPy array for LightGBM prediction
        X = df.to_numpy()

        # Make prediction
        # Get probability of readmission
        light_probability = light_model.predict_proba(X)[:, 1]  # Get probability for class 1 (readmission)
        # Armazenar no session_state
        st.session_state["lightgbm probability"] = light_probability

        # Output results
        #st.write(f"🔹 Readmission Prediction: {probability}")

        # Prediction Button
        if st.button("🚀 Predict Readmission"):
            with st.spinner("🔍 Extracting embeddings and predicting readmission..."):
                readmission_prob = predict_readmission([cleaned_note])[0][0]  # Compute only once
                st.session_state["MLP probability"] = readmission_prob
                prediction = 1 if readmission_prob > 0.5 else 0  # Define prediction value

            # Display Results
            st.subheader("🎯 Prediction Results")
            col1, col2 = st.columns(2)

            with col1:
                st.metric(label="🧮 Readmission Probability", value=f"{readmission_prob:.2%}")

            with col2:
                if prediction == 1:
                    st.error("⚠️ High Risk of Readmission")
                else:
                    st.success("✅ Low Risk of Readmission")

            # Display Readmission Probability with Centered Styling
            st.markdown(f"""
                <div style="text-align:center; padding: 20px; background-color: #f8f9fa; border-radius: 10px;">
                    <h3>📊 Readmission Probability</h3>
                    <h2 style="color: {'red' if readmission_prob > 0.5 else 'green'};">{readmission_prob:.2%}</h2>
                </div>
            """, unsafe_allow_html=True)

elif page == "🔀 Ensemble Prediction":

    # Load the ensemble model
    ensemble_model = joblib.load("/home/user/app/best_ensemble_model.pkl")
    #st.write("✅ Ensemble Model loaded successfully!")

    # Define models
    models = ["XGBoost", "lightgbm", "MLP"]

    # Retrieve stored probabilities from session state and ensure they are numeric
    probabilities = []
    for model in models:
        key = f"{model} probability"
        if key in st.session_state:
            try:
                prob = float(st.session_state[key])
                probabilities.append(prob)
            except ValueError:
                st.error(f"⚠️ Invalid probability value for {model}: {st.session_state[key]}")
                probabilities.append(None)
        else:
            probabilities.append(None)

    # Ensure all probabilities are valid before proceeding
    if None not in probabilities:
        st.write("### 🗳️ Voting Process in Progress...")

        progress_bar = st.progress(0)  # Progress bar
        voting_display = st.empty()  # Placeholder for voting animation

        votes = []
        for i, (model, prob) in enumerate(zip(models, probabilities)):
            time.sleep(1)  # Simulate suspense
            
            # Simulated blinking effect
            for _ in range(3):
                voting_display.markdown(f"⏳ {model} is deciding...")
                time.sleep(0.5)
                voting_display.markdown("")
                time.sleep(0.5)
            
            # Convert probability to label
            if prob < 0.33:
                vote = "🟢 Low"
            elif prob < 0.46:
                vote = "🟡 Medium"
            else:
                vote = "🔴 High"
            
            votes.append(vote)
            voting_display.markdown(f"✅ **{model} voted: {vote}**")
            progress_bar.progress((i + 1) / len(models))

        time.sleep(1)
        progress_bar.empty()

        # Create a DataFrame with numeric probabilities
        final_df = pd.DataFrame([probabilities], columns=['probs', 'probs_lgb', 'probs_mlp'])
        final_df = final_df.astype(float)  # Ensure all values are float

        # Fazer a predição final com o ensemble
        final_probability = ensemble_model.predict_proba(final_df)[:, 1][0]  # Probabilidade de classe 1
        final_prediction = 1 if final_probability >= 0.25 else 0  # Aplicando threshold de 0.25

        # Estilização do resultado final
        st.markdown("---")
        if final_prediction == 1:
            st.markdown(f"""
            <div style="text-align: center; background-color: #ffdddd; padding: 15px; border-radius: 10px;">
                <h2>🚨 <b>Final Prediction: 1</b> (Readmission Likely) </h2>
                <h3>🔍 Probability: {final_probability:.2f} (Threshold: 0.25)</h3>
            </div>
            """, unsafe_allow_html=True)
        else:
            st.markdown(f"""
            <div style="text-align: center; background-color: #ddffdd; padding: 15px; border-radius: 10px;">
                <h2>✅ <b>Final Prediction: 0</b> (No Readmission Risk) </h2>
                <h3>🔍 Probability: {final_probability:.2f} (Threshold: 0.25)</h3>
            </div>
        """, unsafe_allow_html=True)

        # 🎨 **Weight Visualization: How Much Each Model Contributed**
        st.write("### ⚖️ Model Contribution to Final Decision")
        fig, ax = plt.subplots()
        ax.bar(models, probabilities, color=["blue", "green", "red"])
        ax.set_ylabel("Probability")
        ax.set_title("Model Prediction Probabilities")
        st.pyplot(fig)

        # Show detailed voting breakdown
        st.write("### 📊 Voting Breakdown:")
        for model, vote in zip(models, votes):
            st.write(f"🔹 {model}: **{vote}** (Prob: {probabilities[models.index(model)]:.2f})")

    else:
        st.warning("⚠️ Some model predictions are missing. Please run all models before voting.")