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"]
)

#pages_up = ["🏠 Home", "📊 Tabular Data", "📝 Clinical Text Notes", "🔀 Ensemble Prediction"]

styles = {
    "nav": {
        "background-color": "rgba(0, 0, 0, 0.5)",
        "display": "flex",
        "justify-content": "center",
        "align-items": "center",
        "flex-wrap": "nowrap",  # Prevents wrapping to a new line
        "padding": "0.5rem",
    },
    "div": {
        "max-width": "48rem",  # Increased width to prevent wrapping
        "white-space": "nowrap",  # Ensure text stays in one line
        "overflow": "hidden",
    },
    "span": {
        "border-radius": "0.26rem",
        "color": "rgb(255, 255, 255)",
        "margin": "0 0.1rem",  # Reduced margin
        "padding": "0.25rem 0.5rem",  # Reduced padding
        "white-space": "nowrap",  # Prevents text wrapping inside buttons
    },
    "active": {
        "background-color": "rgba(220, 66, 56, 0.95)",
    },
    "hover": {
        "background-color": "rgba(255, 255, 255, 0.95)",
    },
}


#page = st_navbar(pages, styles=styles)

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'>Innovation in healthcare isn’t just about technology, it’s about making a real difference where it matters most</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 works as an alarm system for predicting the risk of a patient being readmitted. "
    "Upload data, analyze clinical notes, and see predictions from our ensemble model.\n\n"
    "For more information or collaboration, contact us:\n"
    "[LinkedIn Feed](https://www.linkedin.com/feed/), "
    "[André Brito](https://www.linkedin.com/in/andrembrito), "
    "[Gonçalo Correia](https://www.linkedin.com/in/gonçalo-correia-b3774a336)")


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

    # Exemplo de nota clínica
    clinical_note_example = """Name: ___ Unit No: ___ Admission Date: ___ Discharge Date: ___ Date of Birth: ___ Sex: F Service: MEDICINE Allergies: No Known Allergies / Adverse Drug Reactions Attending: ___ Chief Complaint:Worsening ABD distension and pain Major Surgical or Invasive Procedure:Paracentesis History of Present Illness:___ HCV cirrhosis c/b ascites, hiv on ART, h/o IVDU, COPD, bioplar, PTSD, presented from OSH ED with worsening abd distension over past week. Pt reports self-discontinuing lasix and spirnolactone ___ weeks ago, because she feels like "they don't do anything" and that she "doesn't want to put more chemicals in her." She does not follow Na-restricted diets. In the past week, she notes that she has been having worsening abd distension and discomfort. She denies ___ edema, or SOB, or orthopnea. She denies f/c/n/v, d/c, dysuria. She had food poisoning a week ago from eating stale cake (n/v 20 min after food ingestion), which resolved the same day. She denies other recent illness or sick contacts. She notes that she has been noticing gum bleeding while brushing her teeth in recent weeks. she denies easy bruising, melena, BRBPR, hemetesis, hemoptysis, or hematuria. Because of her abd pain, she went to OSH ED and was transferred to ___ for further care. Per ED report, pt has brief period of confusion - she did not recall the ultrasound or bloodwork at osh. She denies recent drug use or alcohol use. She denies feeling confused, but reports that she is forgetful at times. In the ED, initial vitals were 98.4 70 106/63 16 97%RA Labs notable for ALT/AST/AP ___ ___: ___, Tbili1.6, WBC 5K, platelet 77, INR 1.6 Past Medical History:1. HCV Cirrhosis 2. No history of abnormal Pap smears. 3. She had calcification in her breast, which was removed previously and per patient not, it was benign. 4. For HIV disease, she is being followed by Dr. ___ Dr. ___. 5. COPD 6. Past history of smoking. 7. She also had a skin lesion, which was biopsied and showed skin cancer per patient report and is scheduled for a complete removal of the skin lesion in ___ of this year. 8. She also had another lesion in her forehead with purple discoloration. It was biopsied to exclude the possibility of ___'s sarcoma, the results is pending. 9. A 15 mm hypoechoic lesion on her ultrasound on ___ and is being monitored by an MRI. 10. History of dysplasia of anus in ___. 11. Bipolar affective disorder, currently manic, mild, and PTSD. 12. History of cocaine and heroin use. Social History:___Family History:She a total of five siblings, but she is not talking to most of them. She only has one brother that she is in touch with and lives in ___. She is not aware of any known GI or liver disease in her family. Her last alcohol consumption was one drink two months ago. No regular alcohol consumption. Last drug use ___ years ago. She quit smoking a couple of years ago. Physical Exam:VS: 98.1 107/61 78 18 97RA General: in NAD HEENT: CTAB, anicteric sclera, OP clear Neck: supple, no LAD CV: RRR,S1S2, no m/r/g Lungs: CTAb, prolonged expiratory phase, no w/r/r Abdomen: distended, mild diffuse tenderness, +flank dullness, cannot percuss liver/spleen edge ___ distension GU: no foley Ext: wwp, no c/e/e, + clubbing Neuro: AAO3, converse normally, able to recall 3 times after 5 minutes, CN II-XII intact Discharge:PHYSICAL EXAMINATION: VS: 98 105/70 95General: in NAD HEENT: anicteric sclera, OP clear Neck: supple, no LAD CV: RRR,S1S2, no m/r/g Lungs: CTAb, prolonged expiratory phase, no w/r/r Abdomen: distended but improved, TTP in RUQ, GU: no foley Ext: wwp, no c/e/e, + clubbing Neuro: AAO3, CN II-XII intact Pertinent Results:___ 10:25PM GLUCOSE-109* UREA N-25* CREAT-0.3* SODIUM-138 POTASSIUM-3.4 CHLORIDE-105 TOTAL CO2-27 ANION GAP-9___ 10:25PM estGFR-Using this___ 10:25PM ALT(SGPT)-100* AST(SGOT)-114* ALK PHOS-114* TOT BILI-1.6*___ 10:25PM LIPASE-77*___ 10:25PM ALBUMIN-3.3*___ 10:25PM WBC-5.0# RBC-4.29 HGB-14.3 HCT-42.6 MCV-99* MCH-33.3* MCHC-33.5 RDW-15.7*___ 10:25PM NEUTS-70.3* LYMPHS-16.5* MONOS-8.1 EOS-4.2* BASOS-0.8___ 10:25PM PLT COUNT-71*___ 10:25PM ___ PTT-30.9 ______ 10:25PM ___.CXR: No acute cardiopulmonary process. U/S: 1. Nodular appearance of the liver compatible with cirrhosis. Signs of portal hypertension including small amount of ascites and splenomegaly. 2. Cholelithiasis. 3. Patent portal veins with normal hepatopetal flow. Diagnostic para attempted in the ED, unsuccessful. On the floor, pt c/o abd distension and discomfort. Brief Hospital Course:___ HCV cirrhosis c/b ascites, hiv on ART, h/o IVDU, COPD, bioplar, PTSD, presented from OSH ED with worsening abd distension over past week and confusion. # Ascites - p/w worsening abd distension and discomfort for last week. likely ___ portal HTN given underlying liver disease, though no ascitic fluid available on night of admission. No signs of heart failure noted on exam. This was ___ to med non-compliance and lack of diet restriction. SBP negativediuretics: > Furosemide 40 mg PO DAILY > Spironolactone 50 mg PO DAILY, chosen over the usual 100mg dose d/t K+ of 4.5. CXR was wnl, UA negative, Urine culture blood culture negative. Pt was losing excess fluid appropriately with stable lytes on the above regimen. Pt was scheduled with current PCP for ___ check upon discharge. Pt was scheduled for new PCP with Dr. ___ at ___ and follow up in Liver clinic to schedule outpatient screening EGD and ___. Medications on Admission:The Preadmission Medication list is accurate and complete.1. Furosemide 20 mg PO DAILY 2. Spironolactone 50 mg PO DAILY 3. Albuterol Inhaler 2 PUFF IH Q4H:PRN wheezing, SOB 4. Raltegravir 400 mg PO BID 5. Emtricitabine-Tenofovir (Truvada) 1 TAB PO DAILY 6. Nicotine Patch 14 mg TD DAILY 7. Ipratropium Bromide Neb 1 NEB IH Q6H SOB Discharge Medications:1. Albuterol Inhaler 2 PUFF IH Q4H:PRN wheezing, SOB 2. Emtricitabine-Tenofovir (Truvada) 1 TAB PO DAILY 3. Furosemide 40 mg PO DAILY RX *furosemide 40 mg 1 tablet(s) by mouth Daily Disp #*30 Tablet Refills:*34. Ipratropium Bromide Neb 1 NEB IH Q6H SOB 5. Nicotine Patch 14 mg TD DAILY 6. Raltegravir 400 mg PO BID 7. Spironolactone 50 mg PO DAILY 8. Acetaminophen 500 mg PO Q6H:PRN pain Discharge Disposition:Home Discharge Diagnosis:Ascites from Portal HTN Discharge Condition:Mental Status: Clear and coherent.Level of Consciousness: Alert and interactive.Activity Status: Ambulatory - Independent. Discharge Instructions:Dear Ms. ___,It was a pleasure taking care of you! You came to us with stomach pain and worsening distension. While you were here we did a paracentesis to remove 1.5L of fluid from your belly. We also placed you on you 40 mg of Lasix and 50 mg of Aldactone to help you urinate the excess fluid still in your belly. As we discussed, everyone has a different dose of lasix required to make them urinate and it's likely that you weren't taking a high enough dose. Please take these medications daily to keep excess fluid off and eat a low salt diet. You will follow up with Dr. ___ in liver clinic and from there have your colonoscopy and EGD scheduled. Of course, we are always here if you need us. We wish you all the best!Your ___ Team. Followup Instructions:___
    """
    
    # Function to set the user's choice
    def set_note(choice):
        if choice == "Use Example":
            st.session_state["clinical_note"] = clinical_note_example
        else:
            st.session_state["clinical_note"] = ""
    
    # If it's the first time, show the pop-up
    if "show_popup" not in st.session_state:
        st.session_state["show_popup"] = True
        st.session_state["clinical_note"] = ""  # Initial empty note
    
    # Show the pop-up
    if st.session_state["show_popup"]:
        with st.popover("Choose how you want to fill in the clinical note"):
            st.write("Would you like to start with an example or create your own note?")
            if st.button("Use Example"):
                set_note("Use Example")
                st.session_state["show_popup"] = False
            if st.button("Create Note from Scratch"):
                set_note("Create from Scratch")
                st.session_state["show_popup"] = False
    
    # Text area for the clinical note
    clinical_note = st.text_area(
    "✍️ Enter Clinical Note", 
    value=st.session_state["clinical_note"], 
    height=400,
    placeholder="Type your clinical note here or use the example"
    )

    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.")