import streamlit as st
import yfinance as yf
import numpy as np
import pandas as pd
import plotly.graph_objs as go
from itertools import product
from datetime import datetime, timedelta

# Set the app layout to wide
st.set_page_config(layout="wide", page_title="Volatility Squeeze Detection")

# Main page title and description
st.title("Volatility Squeeze Detection")
st.markdown("""
This tool automatically identifies volatility squeezes in real-time using a combination of Bollinger Bands, Keltner Channels, and various technical indicators.
The methodology involves detecting periods of low volatility (squeeze) where the Bollinger Bands move inside the Keltner Channels. These squeezes often precede significant price movements.
You can analyze both stock tickers and cryptocurrency pairs by entering the appropriate asset symbol, setting the parameters, and running the analysis.
""")

# Cache the stock/crypto data using the new st.cache_data with yfinance adjustments
@st.cache_data
def get_stock_data(ticker, start_date, end_date):
    data = yf.download(ticker, start=start_date, end=end_date, auto_adjust=False)
    if isinstance(data.columns, pd.MultiIndex):
        data.columns = data.columns.get_level_values(0)
    if data.empty:
        raise ValueError(f"No data fetched for {ticker} from {start_date} to {end_date}.")
    return data

# Function Definitions
def bollinger_bands(data, window=20, no_of_std=2):
    rolling_mean = data['Close'].rolling(window).mean()
    rolling_std = data['Close'].rolling(window).std()
    data['Bollinger_High'] = rolling_mean + (rolling_std * no_of_std)
    data['Bollinger_Low'] = rolling_mean - (rolling_std * no_of_std)
    return data

def keltner_channel(data, window=20, atr_mult=1.5):
    data['Previous_Close'] = data['Close'].shift(1)
    data['TR'] = data[['High', 'Low', 'Previous_Close']].apply(lambda x: max(x[0] - x[1], abs(x[0] - x[2]), abs(x[1] - x[2])), axis=1)
    data['ATR'] = data['TR'].rolling(window).mean()
    rolling_mean = data['Close'].rolling(window).mean()
    data['Keltner_High'] = rolling_mean + (data['ATR'] * atr_mult)
    data['Keltner_Low'] = rolling_mean - (data['ATR'] * atr_mult)
    return data

def moving_average_filter(data, window=50):
    data['MA'] = data['Close'].rolling(window=window).mean()
    return data

def add_confirmation_indicators(data, rsi_window=14, macd_short=12, macd_long=26, macd_signal=9):
    data['RSI'] = 100 - (100 / (1 + data['Close'].diff(1).clip(lower=0).rolling(window=rsi_window).mean() /
                                  -data['Close'].diff(1).clip(upper=0).rolling(window=rsi_window).mean()))
    data['MACD'] = data['Close'].ewm(span=macd_short, adjust=False).mean() - data['Close'].ewm(span=macd_long, adjust=False).mean()
    data['MACD_Signal'] = data['MACD'].ewm(span=macd_signal, adjust=False).mean()
    return data

def identify_volatility_squeeze(data, base_threshold=0.045, rsi_min=30, rsi_max=70, volatility_window=20):
    recent_volatility = data['ATR'].rolling(volatility_window).mean()
    dynamic_squeeze_threshold = base_threshold * recent_volatility / recent_volatility.mean()
    squeeze_on = (data['Bollinger_High'] < data['Keltner_High'] * (1 + dynamic_squeeze_threshold)) & \
                 (data['Bollinger_Low'] > data['Keltner_Low'] * (1 - dynamic_squeeze_threshold)) & \
                 (data['RSI'] > rsi_min) & (data['RSI'] < rsi_max) & \
                 (data['MACD'] > data['MACD_Signal']) & \
                 (abs(data['Close'] - data['MA']) / data['MA'] < 0.015)
    data['Squeeze'] = np.where(squeeze_on, 1, 0)
    return data

def add_buy_sell_signals(data, buy_proximity_range, sell_proximity_range, n_days=5):
    best_buy_accuracy = -np.inf
    best_sell_accuracy = -np.inf
    best_buy_proximity = None
    best_sell_proximity = None

    for buy_proximity in buy_proximity_range:
        for sell_proximity in sell_proximity_range:
            data['Buy_Signal'] = np.where(
                (data['Squeeze'] == 1) & (data['Close'] >= data['Bollinger_High'].shift(1) * buy_proximity), 1, 0)
            data['Sell_Signal'] = np.where(
                (data['Squeeze'] == 1) & (data['Close'] <= data['Bollinger_Low'].shift(1) * sell_proximity), 1, 0)

            correct_buy_signals = sum(data['Buy_Signal'].iloc[i] == 1 and data['Close'].iloc[i + n_days] > data['Close'].iloc[i] for i in range(len(data) - n_days))
            correct_sell_signals = sum(data['Sell_Signal'].iloc[i] == 1 and data['Close'].iloc[i + n_days] < data['Close'].iloc[i] for i in range(len(data) - n_days))

            total_buy_signals = data['Buy_Signal'].sum()
            total_sell_signals = data['Sell_Signal'].sum()

            buy_accuracy = correct_buy_signals / total_buy_signals if total_buy_signals > 0 else 0
            sell_accuracy = correct_sell_signals / total_sell_signals if total_sell_signals > 0 else 0

            if buy_accuracy > best_buy_accuracy:
                best_buy_accuracy = buy_accuracy
                best_buy_proximity = buy_proximity

            if sell_accuracy > best_sell_accuracy:
                best_sell_accuracy = sell_accuracy
                best_sell_proximity = sell_proximity

    data['Buy_Signal'] = np.where(
        (data['Squeeze'] == 1) & (data['Close'] >= data['Bollinger_High'].shift(1) * best_buy_proximity), 1, 0)
    data['Sell_Signal'] = np.where(
        (data['Squeeze'] == 1) & (data['Close'] <= data['Bollinger_Low'].shift(1) * best_sell_proximity), 1, 0)

    return data, best_buy_accuracy, best_sell_accuracy, best_buy_proximity, best_sell_proximity

def plot_volatility_squeeze_plotly(data, ticker, show_signals=False, best_buy_proximity=None, best_buy_accuracy=None, best_sell_proximity=None, best_sell_accuracy=None):
    fig = go.Figure()

    # Add volume to the chart on the secondary y-axis
    fig.add_trace(go.Bar(x=data.index, y=data['Volume'], name='Volume', marker_color='lightgray', opacity=0.3, yaxis='y2'))

    fig.add_trace(go.Scatter(x=data.index, y=data['Close'], mode='lines', name='Close Price'))
    fig.add_trace(go.Scatter(x=data.index, y=data['Bollinger_High'], mode='lines', name='Bollinger High', line=dict(dash='dash')))
    fig.add_trace(go.Scatter(x=data.index, y=data['Bollinger_Low'], mode='lines', name='Bollinger Low', line=dict(dash='dash')))
    fig.add_trace(go.Scatter(x=data.index, y=data['Keltner_High'], mode='lines', name='Keltner High', line=dict(dash='dash')))
    fig.add_trace(go.Scatter(x=data.index, y=data['Keltner_Low'], mode='lines', name='Keltner Low', line=dict(dash='dash')))
    fig.add_trace(go.Scatter(x=data.index, y=data['MA'], mode='lines', name='Moving Average', line=dict(dash='dot')))

    squeeze_zones = data[data['Squeeze'] == 1]

    # Shade the squeeze areas by filling the area between the top and bottom Bollinger Bands during the squeeze
    for start, end in zip(squeeze_zones.index[:-1], squeeze_zones.index[1:]):
        if end - start == pd.Timedelta(days=1):  # Ensure contiguous regions
            fig.add_vrect(x0=start, x1=end, fillcolor="LightSalmon", opacity=0.5, layer="below", line_width=0)

    # Add squeeze to the legend
    fig.add_trace(go.Scatter(x=[None], y=[None], mode='markers',
                             marker=dict(size=10, color='LightSalmon'),
                             legendgroup="Squeeze",
                             showlegend=True,
                             name="Squeeze"))

    if show_signals:
        buy_signals = data[data['Buy_Signal'] == 1]
        sell_signals = data[data['Sell_Signal'] == 1]
        fig.add_trace(go.Scatter(x=buy_signals.index, y=buy_signals['Close'], mode='markers', name=f'Buy Signal (Proximity: {best_buy_proximity}, Accuracy: {best_buy_accuracy:.2%})', marker=dict(symbol='triangle-up', color='green', size=10)))
        fig.add_trace(go.Scatter(x=sell_signals.index, y=sell_signals['Close'], mode='markers', name=f'Sell Signal (Proximity: {best_sell_proximity}, Accuracy: {best_sell_accuracy:.2%})', marker=dict(symbol='triangle-down', color='red', size=10)))

    fig.update_layout(
        title=f"Volatility Squeeze for {ticker} with Buy/Sell Signals",
        xaxis_title="Date",
        yaxis_title="Price",
        yaxis2=dict(
            title="Volume",
            overlaying='y',
            side='right',
            showgrid=False,
        ),
        width=1600,
        height=600,
    )
    
    # Display the plot within a container
    with st.container():
        st.plotly_chart(fig, use_container_width=True)

# Sidebar layout
st.sidebar.title("Input Parameters")
with st.sidebar.expander("How to use", expanded=False):
    st.write("### How to Use")
    st.write("""
    1. Select the symbol (stock ticker or cryptocurrency pair), date range, and parameters.
    2. Click 'Run Analysis' to generate the results.
    3. Use the toggle inside the 'Trading Signals' section to display buy/sell signals.
    """)

# Input widgets
with st.sidebar.expander("Symbol Selection", expanded=True):  # Open by default
    ticker = st.text_input("Asset Symbol", value="BTC-USD", help="Enter the stock or cryptocurrency symbol, e.g., 'AAPL' for Apple, 'BTC-USD' for Bitcoin.")
    start_date = st.date_input("Start Date", value=pd.to_datetime("2020-01-01"), help="Select the start date for the analysis.")
    end_date = st.date_input("End Date", value=datetime.now() + timedelta(days=1), help="Select the end date for the analysis.")

with st.sidebar.expander("Analysis Parameters", expanded=True):
    window = st.slider("Rolling Window Size", min_value=10, max_value=100, value=20, step=1, help="Window size for the rolling calculations, e.g., Bollinger Bands.")
    atr_mult = st.slider("ATR Multiplier", min_value=1.0, max_value=3.0, value=1.5, step=0.1, help="Multiplier for the Average True Range in the Keltner Channel calculation.")
    base_threshold = st.slider("Squeeze Threshold", min_value=0.005, max_value=0.05, value=0.045, step=0.005, help="Threshold for determining a squeeze condition.")

with st.sidebar.expander("Confirmation Indicators", expanded=False):
    rsi_window = st.slider("RSI Window", min_value=5, max_value=30, value=14, step=1, help="Window size for the Relative Strength Index (RSI) calculation.")
    macd_short = st.slider("MACD Short Period", min_value=5, max_value=20, value=12, step=1, help="Short period for the MACD calculation.")
    macd_long = st.slider("MACD Long Period", min_value=20, max_value=50, value=26, step=1, help="Long period for the MACD calculation.")
    macd_signal = st.slider("MACD Signal Period", min_value=5, max_value=20, value=9, step=1, help="Signal period for the MACD calculation.")

# Trading signals expander with toggle and proximity settings
with st.sidebar.expander("Trading Signals", expanded=False):
    show_trading_signals = st.checkbox("Show Trading Signals", value=False, help="Toggle to display indicative buy and sell signals based on a trend-following strategy using volatility squeezes and proximity to Bollinger Bands. Buy signals occur when a squeeze is detected, and the price breaks above the Bollinger High band, suggesting an upward trend. Sell signals occur when the price drops below the Bollinger Low band, indicating a downward trend. Both signals allow adjustable proximity to fine-tune sensitivity.")

    buy_proximity_range = st.multiselect(
        "Buy Proximity Range",
        options=[0.91, 0.93, 0.95, 0.97, 0.98, 0.99],
        default=[0.93, 0.95, 0.97],
        help="Select proximity values for buy signals relative to the Bollinger High band."
    )
    sell_proximity_range = st.multiselect(
        "Sell Proximity Range",
        options=[1.01, 1.02, 1.03, 1.05, 1.07, 1.09],
        default=[1.03, 1.05, 1.07],
        help="Select proximity values for sell signals relative to the Bollinger Low band."
    )
    n_days = st.slider("Signal Validation Days", min_value=1, max_value=20, value=5, help="Number of days to validate the buy/sell signals after they're generated.")

# Run Button
if st.sidebar.button('Run Analysis'):
    # Only run analysis if the input has changed
    if (
        ticker != st.session_state.get('ticker_symbol') or 
        start_date != st.session_state.get('start_date') or 
        end_date != st.session_state.get('end_date') or
        window != st.session_state.get('window') or
        atr_mult != st.session_state.get('atr_mult') or
        base_threshold != st.session_state.get('base_threshold') or
        rsi_window != st.session_state.get('rsi_window') or
        macd_short != st.session_state.get('macd_short') or
        macd_long != st.session_state.get('macd_long') or
        macd_signal != st.session_state.get('macd_signal') or
        buy_proximity_range != st.session_state.get('buy_proximity_range') or
        sell_proximity_range != st.session_state.get('sell_proximity_range') or
        n_days != st.session_state.get('n_days')
    ):
        data = get_stock_data(ticker, start_date, end_date)
        data = bollinger_bands(data, window=window)
        data = keltner_channel(data, window=window, atr_mult=atr_mult)
        data = moving_average_filter(data, window=window)
        data = add_confirmation_indicators(data, rsi_window=rsi_window, macd_short=macd_short, macd_long=macd_long, macd_signal=macd_signal)
        data = identify_volatility_squeeze(data, base_threshold=base_threshold)
        data, best_buy_accuracy, best_sell_accuracy, best_buy_proximity, best_sell_proximity = add_buy_sell_signals(
            data,
            buy_proximity_range,
            sell_proximity_range,
            n_days
        )

        # Save results to session state
        st.session_state['data'] = data
        st.session_state['ticker_symbol'] = ticker
        st.session_state['start_date'] = start_date
        st.session_state['end_date'] = end_date
        st.session_state['window'] = window
        st.session_state['atr_mult'] = atr_mult
        st.session_state['base_threshold'] = base_threshold
        st.session_state['rsi_window'] = rsi_window
        st.session_state['macd_short'] = macd_short
        st.session_state['macd_long'] = macd_long
        st.session_state['macd_signal'] = macd_signal
        st.session_state['buy_proximity_range'] = buy_proximity_range
        st.session_state['sell_proximity_range'] = sell_proximity_range
        st.session_state['n_days'] = n_days
        st.session_state['best_buy_proximity'] = best_buy_proximity
        st.session_state['best_buy_accuracy'] = best_buy_accuracy
        st.session_state['best_sell_proximity'] = best_sell_proximity
        st.session_state['best_sell_accuracy'] = best_sell_accuracy

# Plot the results
if 'data' in st.session_state:
    data = st.session_state['data']

    plot_volatility_squeeze_plotly(
        data, 
        ticker, 
        show_signals=show_trading_signals, 
        best_buy_proximity=st.session_state.get('best_buy_proximity'),
        best_buy_accuracy=st.session_state.get('best_buy_accuracy'),
        best_sell_proximity=st.session_state.get('best_sell_proximity'),
        best_sell_accuracy=st.session_state.get('best_sell_accuracy')
    )