import streamlit as st
import pandas as pd
import numpy as np
import joblib

st.set_page_config(
    page_title="Obesity Risk Prediction",
    page_icon="🏥",
    layout="centered")

@st.cache_resource
def load_assets():
    try:
        return joblib.load('src/model.pkl')
    except FileNotFoundError:
        st.error("Model file not found. Please upload it to the Files section.")
        return None

assets = load_assets()

if assets:
    model = assets['model']
    encoders = assets['encoders']
    target_encoder = assets['target_encoder']

    st.title("🏥 Obesity Risk Level Prediction")
    st.markdown("""
    This application predicts the obesity risk level based on physical and behavioral habits.
    """)

    st.divider()

    col1, col2 = st.columns(2)

    with col1:
        st.subheader("Personal Information")
        gender = st.selectbox("Gender", ["Male", "Female"])
        age = st.number_input("Age", min_value=10, max_value=80, value=25)
        height = st.number_input("Height (meters)", min_value=1.00, max_value=2.30, value=1.70, step=0.01)
        weight = st.number_input("Weight (kg)", min_value=30.0, max_value=200.0, value=70.0, step=0.1)
        family_history = st.selectbox("Family History with Overweight", ["yes", "no"])

        st.subheader("Eating Habits")
        favc = st.selectbox("Frequent consumption of high caloric food (FAVC)", ["yes", "no"])
        fcvc = st.slider("Frequency of consumption of vegetables (FCVC)", 1.0, 3.0, 2.0, help="1: Never, 2: Sometimes, 3: Always")
        ncp = st.slider("Number of main meals (NCP)", 1.0, 4.0, 3.0)
        caec = st.selectbox("Consumption of food between meals (CAEC)", ["no", "Sometimes", "Frequently", "Always"])

    with col2:
        st.subheader("Lifestyle & Activity")
        smoke = st.selectbox("Do you smoke? (SMOKE)", ["yes", "no"])
        ch2o = st.slider("Consumption of water daily (CH2O) (Liters)", 1.0, 3.0, 2.0)
        scc = st.selectbox("Calories consumption monitoring (SCC)", ["yes", "no"])
        faf = st.slider("Physical activity frequency (FAF)", 0.0, 3.0, 1.0, help="0: None, 1: 1-2 days, 2: 2-4 days, 3: 4+ days")
        tue = st.slider("Time using technology devices (TUE)", 0.0, 2.0, 1.0, help="0: 0-2 hours, 1: 3-5 hours, 2: >5 hours")
        calc = st.selectbox("Consumption of alcohol (CALC)", ["no", "Sometimes", "Frequently", "Always"])
        mtrans = st.selectbox("Transportation used (MTRANS)", ["Public_Transportation", "Walking", "Automobile", "Motorbike", "Bike"])

    if st.button("Predict Obesity Risk", type="primary"):
        
        input_data = pd.DataFrame({
            'Gender': [gender], 
            'Age': [age], 
            'Height': [height], 
            'Weight': [weight],
            'family_history_with_overweight': [family_history], 
            'FAVC': [favc],
            'FCVC': [fcvc], 
            'NCP': [ncp], 
            'CAEC': [caec], 
            'SMOKE': [smoke],
            'CH2O': [ch2o], 
            'SCC': [scc], 
            'FAF': [faf], 
            'TUE': [tue],
            'CALC': [calc], 
            'MTRANS': [mtrans]
        })

        input_data['BMI'] = input_data['Weight'] / (input_data['Height'] ** 2)
        input_data['BMI_2'] = input_data['BMI'] ** 2
        
        input_data['Age'] = input_data['Age'].round().astype(int)
        input_data['IsYoung'] = input_data['Age'].apply(lambda x: 1 if x < 25 else 0)
        input_data['IsAging'] = input_data['Age'].apply(lambda x: 1 if 25 <= x < 40 else 0)
        input_data['IsOld'] = input_data['Age'].apply(lambda x: 1 if 40 <= x <= 61 else 0)

        bmi_bool_str = (input_data['BMI'] > 25).astype(str)
        age_bool_str = (input_data['Age'] > 30).astype(str)
        
        input_data['Family_BMI_Interaction'] = input_data['family_history_with_overweight'].astype(str) + "_" + bmi_bool_str
        input_data['Gender_Age_Interaction'] = input_data['Gender'].astype(str) + "_" + age_bool_str

        integer_cols = ['FCVC', 'NCP', 'CH2O', 'FAF', 'TUE']
        for col in integer_cols:
            input_data[col] = input_data[col].round().astype(int)

        categorical_cols = ['Gender', 'family_history_with_overweight', 'FAVC', 'CAEC', 'SMOKE', 'SCC', 'CALC', 'MTRANS',
                            'Family_BMI_Interaction', 'Gender_Age_Interaction']
        
        try:
            for col in categorical_cols:
                le = encoders[col]
                input_data[col] = le.transform(input_data[col])
                
            prediction_idx = model.predict(input_data)[0]
            prediction_label = target_encoder.inverse_transform([prediction_idx])[0]
            
            formatted_label = prediction_label.replace("_", " ")

            st.success(f"Prediction: **{formatted_label}**")
            
            bmi_val = input_data['BMI'].values[0]
            st.info(f"Calculated BMI: {bmi_val:.2f}")

            if "Obesity" in formatted_label:
                st.warning("⚠️ High risk detected. Consider consulting a specialist.")
            elif "Overweight" in formatted_label:
                st.warning("⚠️ Moderate risk detected.")
            else:
                st.success("✅ Weight is within or near normal range.")
        except Exception as e:
            st.error(f"An error occurred during processing: {e}")