# Importing the necessary libraries
import streamlit as st
import pandas as pd 
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import altair as alt
import plotly.express as px
from sklearn.model_selection import train_test_split
import joblib
import shap
from xgboost import XGBClassifier
from sklearn.preprocessing import OneHotEncoder
from streamlit_shap import st_shap
from duckduckgo_search import DDGS

# Setting up the pace icon
st.set_page_config(page_icon="📊")

# Cache the data to avoid loading it multiple times
@st.cache_resource
def load_data():
    data = pd.read_csv('data_cleaned_new.csv')
    return data

# Cache the model and encoder to avoid loading them multiple times
@st.cache_resource
def load_model_and_encoder():
    xgb_model = joblib.load('xgb_model.joblib')
    ohe = joblib.load('ohe.joblib')
    return xgb_model, ohe

xgb_model, ohe = load_model_and_encoder()

# SHAP explainer
explainer = shap.TreeExplainer(xgb_model)

# Load the data
data = load_data()

# Making and naming the sidebars
st.sidebar.title("Explore Financial Insights and AI-Powered Tools")
option = st.sidebar.radio(
    "Select section:",
    ("Home", 'Description of Variables', "Regional-Based Analysis", "Income-Based Analysis", "Gender-Based Analysis", "Financial Advice", "Financial AI Helper", "Predict Financial Savings Behavior")
)            

# Addding a summary of FINDEX at the bottom of the sidebar with link

st.sidebar.markdown("""
    <div style='margin-bottom: 20px;'>
        <h3 style='margin-bottom: -15px;'>What is FINDEX?</h3>
        <p style='margin-bottom: 20px;'>
        The Global Findex database provides comprehensive data on how adults worldwide save, borrow, make payments, and manage risk. 
        Launched with support from the Bill & Melinda Gates Foundation, the database is updated every three years and is the world’s most 
        detailed dataset on how adults use formal and informal financial services. It offers insights into the financial behaviors and 
        access to financial systems globally.
        </p>
    </div>
""", unsafe_allow_html=True)

st.sidebar.image('Findex.png', use_column_width=True)
st.sidebar.markdown('For more information, visit:<br>[Global Findex Database](https://globalfindex.worldbank.org/)', unsafe_allow_html=True)

# Main section logic
if option == "Home":
    # First display the Plotly globe with the title
    # Create the globe visualization
    economy_data = data['Country_Economy'].value_counts(normalize=True) * 100
    economy_df = economy_data.reset_index()
    economy_df.columns = ['Country_Economy', 'percentage']

    # Round the percentage to 2 decimal places for display
    economy_df['percentage'] = economy_df['percentage'].round(2)

    # Create a choropleth map using Plotly with a green color scheme to make look like the earth
    fig = px.choropleth(
        economy_df,
        locations='Country_Economy',
        locationmode='country names',
        color='percentage',
        hover_name='Country_Economy',
        hover_data={'percentage': ':.2f'},  # Format hover data to 2 decimal places
        color_continuous_scale='Greens',
    )

    # Update hover text to add the percentage sign
    fig.update_traces(
        hovertemplate="<b>%{hovertext}</b><br>" +
                      "percentage=%{z:.2f}%<extra></extra>",
        hovertext=economy_df['Country_Economy']
    )

    # Add the title to the Plotly chart itself, which also functions as the headline for the homepage making
    fig.update_layout(
        title=dict(
            text="FINDEX 2021<br><span style='font-size:24px;'>Data Visualization and AI Driven Financial Recommendations</span>",  # Title with subtitle
            font=dict(size=49, color='black', family="Raleway, sans-serif"),  # Stylish font and bigger size
            x=0.5,  # Center the title
            xanchor='center',
            y=0.95,  # Adjust positioning
            yanchor='top',
            pad=dict(t=20),  # Add padding to reduce space
        ),
        geo=dict(
            showframe=True,  # Show a frame around the map
            framecolor="black",  # Frame color
            showcoastlines=True,  # Keep coastlines visible
            coastlinecolor="Black",  # Set coastlines color to black
            projection_type='orthographic',  # Change projection to orthographic for a globe effect
            projection_scale=0.85,  # Zoom out more by reducing the scale
            center=dict(lat=10, lon=0),  # Center the globe around the equator
            lataxis_range=[-85, 85],  # Strictly limit the vertical dragging
            lonaxis_range=[-180, 180],  # Strictly limit the horizontal dragging
            oceancolor='lightblue',  # Set the color of the oceans
            showocean=True,  # Ensure oceans are displayed
        ),
        coloraxis_colorbar=dict(
            title="Participation (%)",
            len=0.5,
            thickness=15,
            tickvals=[0.5, 1, 1.5, 2],
            ticks="outside",
        ),
        width=1000,
        height=800,
        margin={"r":50,"t":50,"l":0,"b":0}
    )

    # Display the Plotly chart first
    st.plotly_chart(fig, use_container_width=True, config={'displayModeBar': False})

    # Discription on home page
    st.markdown("""
    This application leverages the Global FINDEX 2021 dataset, with over 140,000 participants, to explore financial inclusion and behavior across various economies worldwide. 
    Instantly visualize the percentage of respondents from each region who participate in various financial services and gain insights into financial trends and behaviors.
    
    Features of this application include:
    - **Regional Analysis:** Explore financial trends and behaviors by country and region along with education level, identifying disparities in access to financial systems.

    - **Income-Based Analysis:** Analyze financial behaviors like savings, borrowing, and digital payments across different income levels.

    - **Gender-Based Analysis:** Compare financial inclusion patterns between genders, looking into variables such as account ownership, borrowing, and savings behavior.

    - **Financial Advice:** Receive tailored financial advice based on inputs related to the FINDEX dataset, offering insights into financial behaviors and decision-making.

    - **Financial AI Helper:** Receive personalized financial guidance and recommendations based on individual inputs or questions, leveraging AI to provide actionable advice.

    - **Predict Financial Savings Behavior:** Use a Supervised Machine Learning model to predict whether an individual has saved money based on socioeconomic factors, with AI-driven insights explaining the outcome.
""")

# Second section for the description of variables
elif option == "Description of Variables":
    st.markdown("<h2 style='text-align: center;'>Description of Variables</h2>", unsafe_allow_html=True)
    st.markdown("""
- **Country_Economy**: The name of the country or economy.
- **Country_Code**: ISO 3-digit code representing each economy.
- **WorldBank_Region**: World Bank region classification (e.g., Sub-Saharan Africa, East Asia, etc.).
- **Adult_Population**: The population of adults (aged 15+) in the economy.
- **Respondent_ID**: A unique identifier for each respondent in the dataset.
- **Survey_Weight**: Survey weight for each respondent, used to make the sample representative of the population.
- **Gender**: Gender of the respondent (1 if female, 2 if male).
- **Respondent_Age**: Age of the respondent.
- **Education_Level**: Respondent’s education level from level 1 to 3.
- **Income_Quintile**: Income quintile of the respondent’s household.
- **Employment_Status**: Employment status of the respondent.
- **Account_At_Financial_Or_Mobile_Money_Provider**: Whether the respondent has an account at a financial institution or with a mobile money service provider.
- **Account_At_Formal_Financial_Institution**: Whether the respondent has an account at a formal financial institution.
- **Has_Debit_Card**: Has a debit card.
- **Used_Mobile_Money**: Whether the respondent used mobile money.
- **Paid_Bills_Online**: Made bill payments online using the Internet.
- **Sent_Money_To_Relative_Friend_Online**: Sent money to a relative or friend online using the Internet.
- **Bought_Something_Online**: Bought something online using the Internet.
- **Saved_For_Old_Age**: Saved for old age.
- **Saved_At_Formal_Financial_Institution**: Saved using an account at a financial institution.
- **Borrowed_For_Medical_Purposes**: Borrowed for medical purposes.
- **Borrowed_From_Formal_Financial_Institution**: Borrowed from a financial institution.
- **Borrowed_From_Family_Or_Friends**: Borrowed from family or friends.
- **Main_Source_Of_Emergency_Funds_30_Days**: Main source of emergency funds in 30 days.
- **Paid_Utility_Bill**: Paid a utility bill.
- **Received_Wage_Payments**: Received wage payments.
- **Received_Government_Transfer**: Received a government transfer.
- **Received_Government_Pension**: Received a government pension.
- **Financial_Worry_Old_Age**: Financially worried: old age.
- **Financial_Worry_Medical_Cost**: Financially worried: medical cost.
- **Financial_Worry_Bills**: Financially worried: bills.
- **Financial_Worry_Education**: Financially worried: education.
- **Saved_Money_Past_12_Months**: Saved money in the past 12 months.
- **Borrowed_Money_Past_12_Months**: Borrowed money in the past 12 months.
- **Received_Wage_Payment_And_Method**: Received a wage payment and method.
- **Received_Gov_Transfer_Or_Aid_And_Method**: Received government transfers or aid payments and method.
- **Received_Gov_Pension_Payments_And_Method**: Received government pension payments and method.
- **Paid_Utility_Bills_And_Method**: Paid utility bills and method.
- **Owns_Mobile_Phone**: Whether the respondent owns a mobile phone.
- **Has_Internet_Access**: Whether the respondent has access to the internet.
- **Made_Digital_Payment**: Whether the respondent made any digital payment.
- **Data_Collection_Year**: The year of the data collection.
""")

   # Third section for the regional-based analysis
if option == "Regional-Based Analysis":
    st.markdown("<h2 style='text-align: center;'>Regional-Based Analysis</h2>", unsafe_allow_html=True)
    st.write("This section allows for exploration of financial trends and behaviors, including savings, borrowing, and digital payments, across various regions and education levels. It's possible to access how financial systems differs between regions and examine disparities in financial inclusion globally.")

# Creating a dictionary mapping to get rid of underscores in the variable names making them more readable
    variable_labels = {
        'Account_At_Financial_Or_Mobile_Money_Provider': 'Account at Financial or Mobile Money Provider',
        'Saved_Money_Past_12_Months': 'Saved Money Past 12 Months',
        'Borrowed_Money_Past_12_Months': 'Borrowed Money Past 12 Months',
        'Paid_Bills_Online': 'Paid Bills Online',
        'Owns_Mobile_Phone': 'Owns Mobile Phone',
        'Has_Internet_Access': 'Has Internet Access' 
    }

    # List of regions from the dataset
    regions = data['WorldBank_Region'].unique()

    # Multiselect for region selection
    selected_regions = st.multiselect("Select regions to compare", options=regions, default=regions[0])

    # Filter data based on selected regions
    regional_data = data[data['WorldBank_Region'].isin(selected_regions)]

    # Allow user to choose which variable they want to analyze
    variable_to_compare = st.selectbox(
        "Select variable to analyze:",
        options=list(variable_labels.keys()),
        format_func=lambda x: variable_labels[x]  # Replaces underscores with spaces in dropdown
    )

    # Summarize the data for the selected regions and variable, including education level (educ_label)
    summary = regional_data.groupby(['WorldBank_Region', 'educ_label'])[variable_to_compare].mean().reset_index()
    summary.columns = ['WorldBank_Region', 'Education_Level', f'Average {variable_to_compare}']

    # Multiply the average by 100 to display percentages
    summary[f'Average {variable_to_compare}'] = summary[f'Average {variable_to_compare}'].mul(100).round(2)

    # Create an interactive Plotly bar chart to compare the regions and education levels
    fig = px.bar(summary, x='WorldBank_Region', y=f'Average {variable_to_compare}', 
                 color='Education_Level',
                 color_continuous_scale='Teal',
                 title=f"Comparison of {variable_labels[variable_to_compare]} Across Selected Regions and Education Levels",
                 labels={'WorldBank_Region': 'Region', f'Average {variable_to_compare}': f'Average {variable_labels[variable_to_compare]} (%)'},
                 barmode='group')

    # Update layout for better aesthetics
    fig.update_layout(
        xaxis_title="Region",
        yaxis_title=f"Average {variable_labels[variable_to_compare]} (%)",
        showlegend=True,
        width=800,
        height=500,
        margin={"r":0,"t":50,"l":0,"b":50}   
    )

    # Show the chart in Streamlit
    st.plotly_chart(fig)

    # Summary of the analysis (formatting variable name)
    st.markdown(f"### Summary")
    st.write("Key takeaways:")
    for region in selected_regions:
        region_data = summary[summary['WorldBank_Region'] == region]
        for educ_level in region_data['Education_Level'].unique():
            avg_value = region_data[region_data['Education_Level'] == educ_level][f'Average {variable_to_compare}'].values[0]
            st.write(f"- **In {region}, individuals with {educ_level} have an average {variable_labels[variable_to_compare].lower()} of {avg_value:.0f}%.**")

# fourth section for the income-based analysis
elif option == "Income-Based Analysis":
    st.markdown("<h2 style='text-align: center;'>Income-Based Analysis</h2>", unsafe_allow_html=True)
    st.write("This section allows for comparasions of financial behaviors such as having an financial acount aswell as savings and borrowing across different income levels.")

    # Create a dictionary mapping original column names to remove underscores and make them more readable
    variable_labels_income = {
        'Account_At_Financial_Or_Mobile_Money_Provider': 'Account at Financial or Mobile Money Provider',
        'Saved_Money_Past_12_Months': 'Saved Money Past 12 Months',
        'Borrowed_Money_Past_12_Months': 'Borrowed Money Past 12 Months'
    }

    # Select Income Quintile
    income_quintile = st.selectbox("Select Income Quintile:", data['Income_Quintile'].unique())

    # Filter the data based on the selected income quintile
    filtered_data_income = data[data['Income_Quintile'] == income_quintile]

    # Multi-select for financial indicators (displayed without underscores)
    selected_indicators_income = st.multiselect(
        "Select Financial Indicators to Analyze:",
        options=list(variable_labels_income.keys()),
        format_func=lambda x: variable_labels_income[x],  # Format options without underscores
        default='Account_At_Financial_Or_Mobile_Money_Provider'  # Default is financial account ownership
    )

    st.markdown(f"### Analysis for Income Quintile {income_quintile}")

    # Initialize a dictionary to store the summary for income analysis
    income_summary_dict = {}

    # Loop through selected indicators and create a chart for each
    for indicator in selected_indicators_income:
        # Normalize and calculate the percentage for the selected indicator
        income_indicator_chart = filtered_data_income[indicator].value_counts(normalize=True).mul(100).reset_index()
        income_indicator_chart.columns = [indicator, 'Percentage']

        # Get the percentage of people with the selected financial indicator
        has_indicator_income = income_indicator_chart[income_indicator_chart[indicator] == 1]['Percentage'].values[0] if 1 in income_indicator_chart[indicator].values else 0
        income_summary_dict[indicator] = has_indicator_income

        # Create a bar chart for each selected indicator (labels without underscores)
        fig_income = px.bar(
            income_indicator_chart,
            x=indicator,
            y='Percentage',
            title=f"{variable_labels_income[indicator]} for Income Quintile {income_quintile}",
            labels={indicator: variable_labels_income[indicator]},
            color=indicator,
            color_continuous_scale='Teal'
        )

        st.plotly_chart(fig_income)

    # Print out the summary text at the bottom for income analysis
    st.markdown("### Summary")
    for indicator, percentage in income_summary_dict.items():
        st.write(f"**{percentage:.1f}% of respondents in Income Quintile {income_quintile} have {variable_labels_income[indicator]}**.")

# Fifth section 
elif option == "Gender-Based Analysis":
    st.markdown("<h2 style='text-align: center;'>Gender-Based Analysis</h2>", unsafe_allow_html=True)
    st.write("Here it's possible to visualize financial behaviors such as savings and borrowing for selected gender and age groups.")

    # Create a dictionary mapping original column names to remove underscores and make them more readable
    variable_labels_gender = {
        'Account_At_Financial_Or_Mobile_Money_Provider': 'Account at Financial or Mobile Money Provider',
        'Saved_Money_Past_12_Months': 'Saved Money Past 12 Months',
        'Borrowed_Money_Past_12_Months': 'Borrowed Money Past 12 Months'
    }

    # Gender selection
    gender = st.radio("Select Gender:", ("Female", "Male"))

    # Age group selection
    age_group = st.selectbox("Select Age Group:", data['age_group'].unique())

    # Convert gender to appropriate coding
    gender_code = 1 if gender == "Female" else 2

    # Filter the data based on gender and age group
    filtered_data = data[(data['Gender'] == gender_code) & (data['age_group'] == age_group)]

    # Multi-select for financial indicators (displayed without underscores)
    selected_indicators = st.multiselect(
        "Select Financial Indicators to Analyze:",
        options=list(variable_labels_gender.keys()),
        format_func=lambda x: variable_labels_gender[x],  # Format options without underscores
        default=['Account_At_Financial_Or_Mobile_Money_Provider']  # Default is financial account ownership
    )

    st.markdown(f"### Analysis for {gender}s in {age_group} Age Group")

    # Initialize a dictionary to store the summary
    summary_dict = {}

    # Loop through selected indicators and create a chart for each
    for indicator in selected_indicators:
        # Normalize and calculate the percentage for the selected indicator
        indicator_chart = filtered_data[indicator].value_counts(normalize=True).mul(100).reset_index()
        indicator_chart.columns = [indicator, 'Percentage']

        # Get the percentage of people with the selected financial indicator
        has_indicator = indicator_chart[indicator_chart[indicator] == 1]['Percentage'].values[0] if 1 in indicator_chart[indicator].values else 0
        summary_dict[indicator] = has_indicator

        # Create a bar chart for each selected indicator (without underscores in labels)
        fig = px.bar(
            indicator_chart,
            x=indicator,
            y='Percentage',
            title=f"{variable_labels_gender[indicator]} for {gender}s in {age_group} Age Group",
            labels={indicator: variable_labels_gender[indicator]},
            color=indicator,
            color_continuous_scale='Teal'
        )

        st.plotly_chart(fig)

    # Print out the summary text at the bottom
    st.markdown("### Summary")
    for indicator, percentage in summary_dict.items():
        st.write(f"**{percentage:.1f}% of {gender}s in the {age_group} age group have {variable_labels_gender[indicator]}**.")

# Sixth section for the financial advice    
elif option == "Financial Advice":
    st.markdown("<h2 style='text-align: center;'>Financial Advice</h2>", unsafe_allow_html=True)

    st.write("""Based on the information provided, this section offers financial advice to help with financial decisions, derived from the FINDEX dataset. The advice is generated from general trends in financial behavior.

For more personalized financial advice tailored to individual circumstances, the AI Financial Helper provides deeper, AI-driven recommendations.""")
    
    # Define the enhanced_recommender function inside the elif block
    def enhanced_recommender(age, income_quintile, has_debit_card, uses_mobile_money, financial_goal, savings_habit, investment_interest):
        recommendations = []
        
        # Financial product recommendations
        if has_debit_card == "No":
            recommendations.append("Consider getting a debit card. Debit cards offer secure, convenient access to your money and can help you manage day-to-day transactions.")
        
        if uses_mobile_money == "No":
            recommendations.append("Mobile money services are a great way to manage transactions remotely and even save small amounts. Consider trying them for increased financial flexibility.")
        
        # Recommendations based on age
        if age < 30:
            recommendations.append("Starting early is key to long-term financial success! Focus on building a savings habit and avoid unnecessary debt.")
        elif 30 <= age <= 50:
            recommendations.append("This is the perfect time to focus on increasing savings and planning for long-term goals like buying a home or preparing for children's education.")
        elif age > 50:
            recommendations.append("As you near retirement, focus on low-risk investments and savings. Consider discussing retirement plans with a financial advisor.")
        
        # Recommendations based on income quintile
        if income_quintile < 3:
            recommendations.append("You might be eligible for government support programs or financial assistance tailored to lower-income groups. Explore these options to improve your financial stability.")
        else:
            recommendations.append("With a higher income level, consider diversifying your investments, including retirement savings and possibly high-return investments like stocks or real estate.")
        
        # Financial goals recommendations
        if financial_goal == "Retirement":
            recommendations.append("It's important to have a solid retirement plan. Focus on long-term, stable investments like pension funds or bonds.")
        elif financial_goal == "Home Ownership":
            recommendations.append("Buying a home is a big goal. Consider saving aggressively or exploring mortgage options to make your goal achievable.")
        elif financial_goal == "Education":
            recommendations.append("Education savings can benefit from high-interest savings accounts or tax-advantaged education accounts.")
        
        # Savings habit recommendations
        if savings_habit == "No":
            recommendations.append("Starting a savings habit, even if it's a small amount each month, can build your financial security over time.")
        else:
            recommendations.append("Good job on saving! Consider increasing the amount or exploring higher-yield savings accounts or investments.")
        
        # Investment interest recommendations
        if investment_interest == "Yes":
            recommendations.append("Since you're interested in investing, explore stocks, mutual funds, or real estate. A financial advisor can help you find the right options.")
        else:
            recommendations.append("If you're unsure about investing, start small with safer options like government bonds or index funds.")
        
        return recommendations

    # Collect user inputs
    age = st.slider("Your Age", 18, 100, 30)
    income_quintile = st.slider("Income Quintile (1 = Lowest, 5 = Highest)", 1, 5, 3)
    has_debit_card = st.radio("Do you have a debit card?", ("Yes", "No"))
    uses_mobile_money = st.radio("Have you used mobile money?", ("Yes", "No"))
    financial_goal = st.radio("What is your main financial goal?", ("Retirement", "Home Ownership", "Education"))
    savings_habit = st.radio("Do you have a savings habit?", ("Yes", "No"))
    investment_interest = st.radio("Are you interested in investing?", ("Yes", "No"))

        # Initialize recommendations as an empty list
    recommendations = []

    if st.button("Get Advice"):
        # Call the enhanced_recommender function and generate recommendations
        recommendations = enhanced_recommender(age, income_quintile, has_debit_card, uses_mobile_money, financial_goal, savings_habit, investment_interest)

    # Check if the recommendations were generated successfully and display them
    if recommendations:
        st.markdown("### Advice:")
        for rec in recommendations:
            st.write(f"- {rec}")
    else:
        st.write("Click the 'Get Advice' button to receive financial advice based on the FINDEX dataset.")

# Seventh section for the financial AI helper
elif option ==  "Financial AI Helper":
    st.markdown("<h2 style='text-align: center;'>Financial AI Helper</h2>", unsafe_allow_html=True)

    st.write("""
    ### Personalized Financial Guidance
    Welcome to the **Financial AI Helper!** This smart **AI assistant** is designed to provide you with **personalized financial guidance** based on the information you provide. 
    It could either be that you're seeking a further explanation of the advice received in the **Financial Advice section**, 
    which is based on the **FINDEX** data variables, or that you're looking for more personalized financial advice from the **AI helper.**

    Simply enter your information below, and let the **AI** assist you with clear and actionable financial solutions!""")

    # Container for user input and chat button for the AI assistant
    with st.container():
        user_input = st.text_area("Financial AI Helper", height=200)
        if st.button("Get Financial Advice"):
            results = DDGS().chat(user_input, model='gpt-4o-mini')
            st.write(results)

# Eighth section for the financial savings behavior prediction
elif option == "Predict Financial Savings Behavior":
    st.markdown("<h2 style='text-align: center;'>Predict Financial Savings Behavior</h2>", unsafe_allow_html=True)


    # Introduction to the Predict Financial Savings Behavior page
    st.write("""
    ### Supervised Learning Model with AI-Powered Insights
    This page utilizes a Supervised Machine Learning model to predict whether an individual has saved money in the past year, using data from the FINDEX dataset. 
    Saving is crucial for financial stability, helping individuals manage unexpected expenses, avoid debt, and achieve long-term goals like home ownership or retirement. 
    Understanding saving behavior offers insights into financial habits and supports better financial planning. Based on the provided information, the model analyzes the input and makes the following prediction:
    
    - **Saving habit is likely:** If the model predicts that the individual is likely to have saved money.
    - **Saving habit is unlikely:** If the model predicts that the individual has not likely saved money.

    In addition to the prediction, an AI-driven analysis will provide insights and explain the key factors contributing to the outcome. 
    This will help clarify how variables such as age, income, employment status, and other socioeconomic factors impact the likelihood of saving. 
    Age group, income bracket, and other factors will be used to predict the saving behavior. 
    
    The age ranges are: Teen 13-18, Young Adult 18-25, Adult 25-35, Middle Age 35-45, Older Adult 45-55, Senior 55-65, Elder 65+""")
    
    # Collect user inputs
    female = st.radio('Gender', ['Female', 'Male'])
    place_of_living = st.selectbox('Place of Living', ['Urban Area', 'Rural Area', 'Unknown'])
    education_level = st.selectbox('Education Level', ['Primary', 'Secondary', 'Tertiary'])
    age_group = st.selectbox('Age Group', ['Adult', 'Middle Age', 'Older Adult', 'Senior', 'Teen', 'Young Adult'])
    income_bracket = st.selectbox('Income Bracket', [1, 2, 3, 4, 5])
    high_income_region = st.checkbox('High Income Region?')
    employed = st.checkbox('Employed?')
    is_mobileowner = st.checkbox('Is Mobile Owner?')
    has_internet_access = st.checkbox('Has Internet Access?')
    

    # Prepare categorical and numcerial/binary features
    cat_features = pd.DataFrame({
        'Place of living': [place_of_living],
        'Education Level': [education_level],
        'Age Group': [age_group]
    })

    num_features = pd.DataFrame({
        'Female': [1 if female == 'Female' else 0],
        'Is Mobileowner': [1 if is_mobileowner else 0],
        'Has Internet Access': [1 if has_internet_access else 0],
        'Employed': [1 if employed else 0],
        'High Income Region': [1 if high_income_region else 0],
        'Income Bracket': [income_bracket]  # Directly use Income Bracket as numerical
    })

    # One-hot encode categorical features
    cat_encoded = pd.DataFrame(ohe.transform(cat_features), 
                               columns=ohe.get_feature_names_out(['Place of living', 'Education Level', 'Age Group']))

    # Combine categorical and numerical features
    features = pd.concat([num_features, cat_encoded], axis=1)

    # Prediction button
    if st.button('Predict Saving Behavior'):
        # Make the prediction
        predicted_saved = xgb_model.predict(features)[0]

        # Display saving habit likelihood instead of probability
        if predicted_saved >= 0.5:
            st.success("Saving habit is likely.")
        else:
            st.error("Saving habit is unlikely.")    

        # SHAP explanation
        st.subheader('Feature Contributions 🤖')
        st.write("""- Blue bars push the probability lower
        while the Red bars push the probability higher""")
        st.write("- An analysis of the results will be visible below the plot after a short loading time")
        
        # Ensure SHAP values are handled properly (flatten if multi-dimensional)
        shap_values = explainer.shap_values(features)
        if isinstance(shap_values, list):
            shap_values = shap_values[0]  # For binary classification, select the first set of SHAP values

        st_shap(shap.force_plot(explainer.expected_value, shap_values, features), height=175, width=1750)

        # Extract important features based on SHAP values
        shap_df = pd.DataFrame({
            'Feature': features.columns,
            'SHAP Value': shap_values.flatten()  # Ensure SHAP values are flattened
        })

        # Sort by absolute SHAP value to get the most important features
        shap_df['Absolute SHAP Value'] = shap_df['SHAP Value'].abs()
        top_important_features = shap_df.sort_values(by='Absolute SHAP Value', ascending=False).head(4)

        # Generate summary of the most important features
        important_features_summary = "\n".join(
            [f"- **{row['Feature']}**: {'Positive' if row['SHAP Value'] > 0 else 'Negative'} contribution"
             for _, row in top_important_features.iterrows()]
        )

        # Pass the summary to the AI assistant for commentary
        with st.expander("AI's analysis of the results"):
            # Construct the AI input with explicit instructions
            user_input = (
                f"You are an AI financial expert, comment on why the saving habit was deemed likely or unlikely. "
                f"Interpret the features based on these rules: "
                f"1. Blue bars represent negative contributions to saving likelihood, and red bars represent positive contributions. "
                f"2. If 'Employment' is 0, the person is unemployed which is not good for saving money. "
                f"3. If 'Has Internet Access' is 0, the person does not have internet access which indicates that maybe the person does not save money as they do not have money for internet. "
                f"4. If 'Is Mobile Owner' is 0, the person does not have mobile access wich indicates that the person does not save money as the person can not afford a phone. "
                f"5. If 'High Income Region' is 0, the person is from a non-high-income region which is not good for the chances of saving. If the person is from a high income region they a larger change of saving up money "
                f"Here are the key factors and their contributions: {important_features_summary}"
            )
            # Generate AI response
            ai_response = DDGS().chat(user_input, model='gpt-4o-mini')
            st.markdown(ai_response)