import streamlit as st
import yfinance as yf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def fetch_data(ticker, start_date, end_date):
    data = yf.download(ticker, start=start_date, end=end_date)
    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):
    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']))
    return data

def plot_data(data):
    plt.figure(figsize=(10, 5))
    plt.plot(data['Close'], label='Close Price')
    plt.plot(data['Upper Band'], label='Upper Bollinger Band', linestyle='--')
    plt.plot(data['Middle Band'], label='Middle Bollinger Band', linestyle='--')
    plt.plot(data['Lower Band'], label='Lower Bollinger Band', linestyle='--')
    plt.plot(data['MA5'], label='5-Day MA', color='green', linestyle='-.')
    plt.plot(data['MA10'], label='10-Day MA', color='red', linestyle='-.')
    
    buy_signals = data[data['Buy Signal']]
    sell_signals = data[data['Sell Signal']]
    plt.scatter(buy_signals.index, buy_signals['Close'], marker='^', color='green', s=100, label='Buy Signal')
    plt.scatter(sell_signals.index, sell_signals['Close'], marker='v', color='red', s=100, label='Sell Signal')
    
    plt.title('Stock Price and Trading Signals')
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()
    plt.grid(True)
    plt.show()

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)
        data = calculate_indicators(data)
        data = identify_signals(data)
        plot_data(data)
        st.pyplot(plt)

if __name__ == "__main__":
    main()