import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import pandas as pd
import streamlit as st
from pandas.plotting import register_matplotlib_converters

register_matplotlib_converters()

def plot_stock_data_with_signals(data):
    # First, ensure the DataFrame's index is in datetime format.
    data.index = pd.to_datetime(data.index, errors='coerce')

    # Creating a figure and axes explicitly
    fig, ax = plt.subplots(figsize=(14, 7))

    # Plotting stock 'Close' prices
    ax.plot(data.index, data['Close'], label='Close Price', color='black', lw=2)

    # Check and plot EMAs if they exist
    if 'EMA_20' in data.columns and 'EMA_50' in data.columns:
        ax.plot(data.index, data['EMA_20'], label='EMA 20', color='blue', lw=1.5)
        ax.plot(data.index, data['EMA_50'], label='EMA 50', color='red', lw=1.5)

    # Check and plot Bollinger Bands if they exist
    if 'BB_Upper' in data.columns and 'BB_Lower' in data.columns:
        ax.fill_between(data.index, data['BB_Lower'], data['BB_Upper'], color='grey', alpha=0.1, label='Bollinger Bands')

    # Highlight buy/sell signals if they exist
    if 'Combined_Signal' in data.columns:
        buy_signals = data[data['Combined_Signal'] == 'buy']
        sell_signals = data[data['Combined_Signal'] == 'sell']
        ax.scatter(buy_signals.index, buy_signals['Close'], label='Buy Signal', marker='^', color='green', alpha=1)
        ax.scatter(sell_signals.index, sell_signals['Close'], label='Sell Signal', marker='v', color='red', alpha=1)

    ax.set_title('Stock Price with Buy/Sell Signals')
    ax.set_xlabel('Date')
    ax.set_ylabel('Price')
    ax.legend()

    # Improve the x-axis date format
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
    ax.xaxis.set_major_locator(mdates.MonthLocator())
    fig.autofmt_xdate()  # Rotation

    # Use the figure object in st.pyplot()
    st.pyplot(fig)