import streamlit as st
import pickle 
import pandas as pd
import numpy as np
from prophet import Prophet 

MODEL_PATH = 'src/prophet_covid_model.pkl' 

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

# --- Streamlit Interface ---
st.set_page_config(page_title="COVID-19 Cases Forecast", layout="centered")
st.title("🦠 Global COVID-19 Case Forecast")
st.markdown("Predict the trend of daily COVID-19 cases for the upcoming period.")

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=90, value=30)
    
    if st.button(f"Generate Forecast for {n_periods} Days"):
        with st.spinner(f'Generating forecast...'):
            try:
                future = model.make_future_dataframe(periods=n_periods)
                
                forecast = model.predict(future)
                
                forecasted_data = forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].tail(n_periods)
                
                st.success("Forecast Successful!")

                st.subheader("Predicted Daily Cases")
                
                plot_data = forecasted_data[['ds', 'yhat']].rename(columns={'ds': 'Date', 'yhat': 'Predicted Cases'})
                plot_data = plot_index = plot_data.set_index('Date')
                
                st.line_chart(plot_data)

                st.subheader("Forecast Data (Next 5 Days)")
                st.dataframe(forecasted_data[['ds', 'yhat']].head(), hide_index=True)
                
            except Exception as e:
                st.error(f"Error during forecasting process. Error: {e}")