import streamlit as st
import joblib
import pandas as pd

MODEL_PATH = 'src/prophet_birth_model.joblib' 

@st.cache_resource
def load_prophet_model():
    try:
        model = joblib.load(MODEL_PATH)
        return model
    except Exception as e:
        st.error(f"Error loading the Prophet model. Check if '{MODEL_PATH}' is uploaded and if the 'prophet' library is installed. Error: {e}")
        return None

# --- Streamlit Interface ---
st.set_page_config(page_title="Births Forecast App", layout="centered")
st.title("👶 Daily Births Forecasting (Prophet Model)")
st.markdown("Enter the number of future days you want to predict.")

model = load_prophet_model()

if model is not None:
    st.sidebar.header("Prediction Settings")

    n_periods = st.sidebar.slider("Future Days to Forecast:", min_value=1, max_value=365, value=30)
    
    if st.button(f"Generate Forecast for {n_periods} Days"):
        with st.spinner(f'Generating forecast...'):
            try:
                # 1. Generate future dates (using Prophet's built-in method)
                future = model.make_future_dataframe(periods=n_periods)
                
                # 2. Predict the values
                forecast = model.predict(future)
                
                # Filter to show only the new forecasted data (the last N rows)
                forecasted_data = forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].tail(n_periods)
                
                st.success("Forecast Successful!")

                # --- Visualization ---
                st.subheader("Predicted Daily Births")
                
                # Prepare data for Streamlit's simple line chart
                plot_data = forecasted_data[['ds', 'yhat']].rename(columns={'ds': 'Date', 'yhat': 'Predicted Births'})
                plot_data = plot_data.set_index('Date')
                
                st.line_chart(plot_data)

                # --- Data Display ---
                st.subheader("Raw Forecast Data (Next 5 Days)")
                st.dataframe(forecasted_data[['ds', 'yhat', 'yhat_upper']].head(), hide_index=True)
                
            except Exception as e:
                st.error(f"Error during forecasting process. Check the model input requirements. Error: {e}")