# Import libraries
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("Algorithmic Trading Bot with ARIMA")
st.write(
    "This app simulates an algorithmic trading bot using the ARIMA model for price prediction. "
    "The bot predicts future stock prices and simulates trades based on the predicted trends."
)

# Sidebar for user inputs
st.sidebar.title("Settings")
ticker = st.sidebar.text_input("Stock Ticker (e.g., AAPL, TSLA, ^GSPC):", value="^GSPC")
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("Prediction Horizon (days)", min_value=1, max_value=30, value=7)
initial_balance = st.sidebar.number_input("Initial Balance (USD):", value=10000.0)
arima_order = st.sidebar.text_input("ARIMA Order (p, d, q):", value="5,1,0")

# Parse ARIMA order
try:
    p, d, q = map(int, arima_order.split(","))
except ValueError:
    st.error("Invalid ARIMA order. Please enter in the format 'p,d,q'.")

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

# Trading Bot Simulation
if st.button("Run Trading Bot"):
    st.write("### Trading Bot Simulation")
    
    if len(data) < 30:
        st.error("Not enough data to train the model. Please select a longer date range.")
    else:
        try:
            # Train ARIMA model
            model = ARIMA(data, order=(p, d, q))
            fitted_model = model.fit()

            # Predict future prices
            future_index = pd.date_range(start=data.index[-1], periods=n_days + 1, freq="B")[1:]
            forecast = fitted_model.forecast(steps=n_days)
            st.write("### Debug: Forecasted Prices")
            st.write(pd.DataFrame({"Date": future_index, "Predicted Price": forecast}))

            # Simulate trading
            balance = initial_balance
            position = 0  # Number of shares held
            trades = []

            for i in range(1, len(forecast)):
                if forecast[i] > forecast[i - 1]:  # Buy signal
                    if position == 0:
                        position = balance / forecast[i]
                        balance = 0
                        trades.append((future_index[i], "BUY", forecast[i]))
                elif forecast[i] < forecast[i - 1]:  # Sell signal
                    if position > 0:
                        balance = position * forecast[i]
                        position = 0
                        trades.append((future_index[i], "SELL", forecast[i]))

            # Final balance
            final_balance = balance + (position * forecast[-1] if position > 0 else 0)
            profit = final_balance - initial_balance

            # Show results
            st.write(f"### Final Balance: ${final_balance:,.2f}")
            st.write(f"### Total Profit: ${profit:,.2f}")
            trades_df = pd.DataFrame(trades, columns=["Date", "Action", "Price"])
            st.write("### Trade History")
            st.write(trades_df)

            # Plot results
            plt.figure(figsize=(10, 6))
            plt.plot(data, label="Historical Data (USD/share)")
            plt.plot(future_index, forecast, label="Predicted Data (USD/share)", linestyle="--")
            for trade in trades:
                plt.scatter(trade[0], trade[2], label=trade[1], color="green" if trade[1] == "BUY" else "red")
            plt.legend()
            plt.title("Algorithmic Trading Bot Simulation")
            plt.xlabel("Date")
            plt.ylabel("Price (USD/share)")
            st.pyplot(plt)

        except Exception as e:
            st.error(f"Error in trading bot simulation: {e}")