import streamlit as st
import yfinance as yf
from statsmodels.tsa.arima.model import ARIMA
import pandas as pd
import matplotlib.pyplot as plt

# Title and description
st.title("S&P 500 Price Prediction in USD/Share with ARIMA")
st.write("This app predicts the price of the S&P 500 index in USD/share using the ARIMA model.")

# Sidebar for user inputs
st.sidebar.title("Prediction Settings")
start_date = st.sidebar.date_input("Start Date", value=pd.to_datetime("2015-01-01"))
end_date = st.sidebar.date_input("End Date", value=pd.to_datetime("2023-12-31"))
n_days = st.sidebar.slider("Days to Predict", min_value=1, max_value=30, value=7)

# Fetch historical data
st.write("### Historical Data in USD/Share")
try:
    data = yf.download("^GSPC", start=start_date, end=end_date)
    data = data["Close"]
    st.line_chart(data)
except Exception as e:
    st.error(f"Error fetching data: {e}")

# Fit ARIMA model
if st.button("Predict"):
    st.write("### Predictions in USD/Share")

    # Check for sufficient data
    if len(data) < 30:
        st.error("Not enough data to train the model. Please select a longer date range.")
    else:
        try:
            # Fit ARIMA model
            model = ARIMA(data, order=(5, 1, 0))
            fitted_model = model.fit()

            # Make predictions
            future_index = pd.date_range(start=data.index[-1], periods=n_days + 1, freq="B")[1:]
            forecast = fitted_model.forecast(steps=n_days)
            forecast_df = pd.DataFrame({"Date": future_index, "Predicted Price (USD/Share)": forecast}).set_index("Date")

            # Plot results
            plt.figure(figsize=(10, 6))
            plt.plot(data, label="Historical Data (USD/Share)")
            plt.plot(forecast_df, label="Predicted Data (USD/Share)", linestyle="--")
            plt.legend()
            plt.title("S&P 500 Price Prediction in USD/Share")
            plt.xlabel("Date")
            plt.ylabel("Price (USD/Share)")
            st.pyplot(plt)

            # Show prediction table
            st.write("### Predicted Prices in USD/Share")
            st.write(forecast_df)
        except Exception as e:
            st.error(f"Error in prediction: {e}")