import streamlit as st
import yfinance as yf
import pandas as pd
import plotly.graph_objs as go


def fetch_data(ticker, start_date, end_date):
    data = yf.download(ticker, start=start_date, end=end_date)
    return data

def fetch_weekly_data(ticker, start_date, end_date):
    data = yf.download(ticker, start=start_date, end=end_date, interval='1wk')
    return data

def calculate_indicators(data):
    # Bollinger Bands
    data['Middle Band'] = data['Close'].rolling(window=20).mean()
    data['Upper Band'] = data['Middle Band'] + 1.96 * data['Close'].rolling(window=20).std()
    data['Lower Band'] = data['Middle Band'] - 1.96 * data['Close'].rolling(window=20).std()
    
    # Moving Averages
    data['MA5'] = data['Close'].rolling(window=5).mean()
    data['MA10'] = data['Close'].rolling(window=10).mean()
    
    return data

def identify_signals(data, weekly_data):
    # Calculate Buy and Sell signals on daily data
    data['Buy Signal'] = ((data['Close'] < data['Lower Band']) & (data['Close'].shift(1) > data['Lower Band'])) | \
                         ((data['Close'] > data['MA5']) & (data['Close'].shift(1) < data['MA5']))
    data['Sell Signal'] = ((data['Close'] > data['Upper Band']) & (data['Close'].shift(1) < data['Upper Band'])) | \
                          ((data['Close'] < data['MA5']) & (data['Close'].shift(1) > data['MA5']))

    # Filter signals by volume
    avg_volume = data['Volume'].rolling(window=20).mean()
    data['Buy Signal'] = data['Buy Signal'] & (data['Volume'] > avg_volume)
    data['Sell Signal'] = data['Sell Signal'] & (data['Volume'] > avg_volume)

    # Confirm signals with weekly data
    weekly_data = calculate_indicators(weekly_data)
    data['Weekly Buy Signal'] = weekly_data['Buy Signal'].reindex(data.index, method='ffill')
    data['Weekly Sell Signal'] = weekly_data['Sell Signal'].reindex(data.index, method='ffill')
    data['Buy Signal'] = data['Buy Signal'] & data['Weekly Buy Signal']
    data['Sell Signal'] = data['Sell Signal'] & data['Weekly Sell Signal']

    return data

def plot_data(data):
    fig = go.Figure()

    # Adding Close price trace
    fig.add_trace(go.Scatter(x=data.index, y=data['Close'], name='Close Price', line=dict(color='blue', width=2)))

    # Adding Bollinger Bands traces
    fig.add_trace(go.Scatter(x=data.index, y=data['Upper Band'], name='Upper Bollinger Band', line=dict(color='red', dash='dash')))
    fig.add_trace(go.Scatter(x=data.index, y=data['Middle Band'], name='Middle Bollinger Band', line=dict(color='white', dash='dash')))
    fig.add_trace(go.Scatter(x=data.index, y=data['Lower Band'], name='Lower Bollinger Band', line=dict(color='red', dash='dash')))

    # Adding Moving Averages traces
    fig.add_trace(go.Scatter(x=data.index, y=data['MA5'], name='5-Day MA', line=dict(color='green', dash='dot')))
    fig.add_trace(go.Scatter(x=data.index, y=data['MA10'], name='10-Day MA', line=dict(color='orange', dash='dot')))

    # Adding Buy and Sell signals
    buys = data[data['Buy Signal']]
    sells = data[data['Sell Signal']]
    fig.add_trace(go.Scatter(x=buys.index, y=buys['Close'], mode='markers', name='Buy Signal', marker=dict(symbol='triangle-up', size=10, color='green')))
    fig.add_trace(go.Scatter(x=sells.index, y=sells['Close'], mode='markers', name='Sell Signal', marker=dict(symbol='triangle-down', size=10, color='red')))

    # Layout updates
    fig.update_layout(title='Stock Price and Trading Signals', xaxis_title='Date', yaxis_title='Price', template='plotly_dark')
    fig.update_xaxes(rangeslider_visible=True)
    
    return fig

def main():
    st.title("OMA Ally BBMA Trading Strategy Visualization")
    ticker = st.text_input("Enter the ticker symbol, e.g., 'AAPL'")
    start_date = st.date_input("Select the start date")
    end_date = st.date_input("Select the end date")
    
    if st.button("Analyze"):
        data = fetch_data(ticker, start_date, end_date)
        weekly_data = fetch_weekly_data(ticker, start_date, end_date)
        data = calculate_indicators(data)
        data = identify_signals(data, weekly_data)
        fig = plot_data(data)
        st.plotly_chart(fig, use_container_width=True)

if __name__ == "__main__":
    main()