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

# Page Configuration
st.set_page_config(
    page_title="Calorie Burn Predictor",
    page_icon="🔥",
    layout="centered"
)

# Load Model and Columns
@st.cache_resource
def load_assets():
    try:
        model = joblib.load('src/xgboost_model.pkl')
        model_columns = joblib.load('src/model_columns.pkl')
        return model, model_columns
    except FileNotFoundError:
        return None, None

model, model_columns = load_assets()

# App Header
st.title("🔥 Calorie Burn Predictor")
st.write("Enter your exercise details below to estimate the calories burned.")

if model is None:
    st.error("Error: Model files not found. Please upload 'xgboost_model.pkl' and 'model_columns.pkl' to the directory.")
    st.stop()

# User Inputs 
col1, col2 = st.columns(2)

with col1:
    gender = st.selectbox("Gender", ["Male", "Female"])
    age = st.number_input("Age", min_value=10, max_value=100, value=25)
    height = st.number_input("Height (cm)", min_value=100, max_value=250, value=175)
    weight = st.number_input("Weight (kg)", min_value=30, max_value=200, value=70)

with col2:
    duration = st.slider("Duration (minutes)", min_value=1, max_value=60, value=30)
    heart_rate = st.slider("Heart Rate (bpm)", min_value=60, max_value=200, value=100)
    body_temp = st.slider("Body Temperature (°C)", min_value=35.0, max_value=42.0, value=38.0, step=0.1)

# Prediction 
if st.button("Calculate Calories", type="primary"):
    
    input_data = pd.DataFrame({
        'Sex': [gender], 
        'Age': [age],
        'Height': [height],
        'Weight': [weight],
        'Duration': [duration],
        'Heart_Rate': [heart_rate],
        'Body_Temp': [body_temp]})

    # Feature Engineering 
    
    # BMI
    input_data['BMI'] = input_data['Weight'] / ((input_data['Height'] / 100) ** 2)

    # BMR 
    input_data['BMR'] = (10 * input_data['Weight']) + (6.25 * input_data['Height']) - (5 * input_data['Age'])

    # Effort 
    input_data['Effort'] = input_data['Duration'] * input_data['Heart_Rate']

    # Temperature Difference
    input_data['Temp_Diff'] = input_data['Body_Temp'] - 37.0

    # Intensity
    input_data['Intensity'] = input_data['Heart_Rate'] * input_data['Body_Temp']

    # Temp per Minute
    input_data['Temp_per_Minute'] = input_data['Body_Temp'] / input_data['Duration']

    # Encoding
    input_data['Sex'] = input_data['Sex'].map({'Male': 0, 'Female': 1})

    input_data = input_data.reindex(columns=model_columns, fill_value=0)

    # Prediction
    try:
        prediction = model.predict(input_data)[0]
        
        st.success(f"Estimated Calories Burned: **{prediction:.1f} kcal**")
        
        if prediction < 100:
            st.info("Nice warm-up! Keep moving! 🚶")
        elif prediction < 300:
            st.info("Great workout! 💪")
        else:
            st.balloons()
            st.info("Incredible effort! You are on fire! 🔥")
            
        with st.expander("See calculated features"):
            st.write(input_data)
            
    except Exception as e:
        st.error(f"An error occurred during prediction: {str(e)}")