import streamlit as st
import yfinance as yf
import numpy as np
import pandas as pd
import plotly.graph_objects as go
from bayes_opt import BayesianOptimization
import warnings
from datetime import datetime, timedelta

warnings.filterwarnings('ignore')

# Set the app layout to wide
st.set_page_config(layout="wide", page_title="ATR-based Dynamic Stop Loss Estimation")

# Initialize session state for data persistence
if 'data' not in st.session_state:
    st.session_state.data = None
if 'optimized_params' not in st.session_state:
    st.session_state.optimized_params = None
if 'atr_multiplier' not in st.session_state:
    st.session_state.atr_multiplier = None
if 'atr_window' not in st.session_state:
    st.session_state.atr_window = None

# Sidebar for user input and navigation
st.sidebar.title("Input Parameters")

with st.sidebar.expander("Stop-Loss Strategy", expanded=True):
    page = st.radio("Select Mode", ["Run Strategy", "Optimize Parameters"])

# Common user inputs in the sidebar with collapsible sections
with st.sidebar:
    with st.expander("How to use", expanded=False):
        st.header("How to Use")
        st.markdown("""
        1. Select **Run Strategy** to calculate stop-loss thresholds using specified ATR parameters.
        2. Select **Optimize Parameters** to find the optimal ATR settings using Bayesian Optimization.
        3. Set the symbol, date range, and strategy parameters in the **Parameters** section.
        4. Click **Run Strategy** or **Run Optimization** to execute the selected operation.
        """)
    
    # Symbol and Date Section
    with st.expander("Symbol and Date Range", expanded=True):
        symbol = st.text_input("Asset Symbol", value="ASML", help="Enter the stock or cryptocurrency symbol, e.g., AAPL, BTC-USD")
        start_date = st.date_input("Start Date", value=pd.to_datetime("2022-01-01"), help="Select the start date for the data")
        end_date = st.date_input("End Date", value=datetime.now() + timedelta(days=1), help="Select the end date for the data")
    
    # ATR Parameters Section
    with st.expander("ATR Parameters", expanded=True):
        atr_window = st.slider("ATR Window", min_value=5, max_value=60, value=14, step=1, help="Adjust the ATR window length")
        atr_multiplier = st.slider("ATR Multiplier", min_value=1.0, max_value=5.0, value=2.0, step=0.1, help="Set the ATR multiplier for stop-loss calculation")
    
    # Strategy Parameters Section
    with st.expander("Strategy Parameters", expanded=True):
        min_holding_period = st.slider("Minimum Holding Period (days)", min_value=1, max_value=10, value=5, step=1, help="Set the minimum holding period in days. This would find the best stop loss given how many days you want to hold the asset")

    if page == "Run Strategy":
        run_strategy_button = st.button("Run Strategy")
    elif page == "Optimize Parameters":
        run_optimization_button = st.button("Run Optimization")

# Main page description
st.header("ATR-based Dynamic Stop Loss Estimation")
st.markdown("""
This application esimates dynamic stop-loss levels using the Average True Range (ATR) indicator. You can either run a strategy with specific ATR parameters or optimize those parameters to find the best stop-loss settings. The app works with both stocks and cryptocurrency pairs.
""")

# Function to fetch data with yfinance adjustments
@st.cache_data(show_spinner=False)
def fetch_data(symbol, start, end):
    data = yf.download(symbol, start=start, end=end, 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 {symbol} from {start} to {end}.")
    return data

# Function to compute Relative Strength Index (RSI)
def compute_RSI(series, period=14):
    delta = series.diff(1)
    gain = (delta.where(delta > 0, 0)).rolling(window=period).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=period).mean()
    RS = gain / loss
    RSI = 100 - (100 / (1 + RS))
    return RSI

# Function to calculate various technical indicators
def calculate_indicators(data, atr_window):
    data['TR'] = np.maximum((data['High'] - data['Low']),
                            np.maximum(abs(data['High'] - data['Close'].shift(1)),
                                       abs(data['Low'] - data['Close'].shift(1))))
    data['ATR'] = data['TR'].rolling(window=int(atr_window)).mean()
    data['EMA12'] = data['Close'].ewm(span=12, adjust=False).mean()
    data['EMA26'] = data['Close'].ewm(span=26, adjust=False).mean()
    data['MACD'] = data['EMA12'] - data['EMA26']
    data['Signal'] = data['MACD'].ewm(span=9, adjust=False).mean()
    data['RSI'] = compute_RSI(data['Close'])
    return data

# Function to calculate stop loss levels
def calculate_stop_loss(data, atr_multiplier):
    data['Buy_Stop_Loss'] = data['Close'] - (data['ATR'] * atr_multiplier)
    data['Sell_Stop_Loss'] = data['Close'] + (data['ATR'] * atr_multiplier)
    return data

# Function to enforce minimum holding period
def enforce_min_holding_period(signals, min_holding_period):
    hold_signals = signals.copy()
    for i in range(1, len(signals)):
        if signals.iloc[i]:
            hold_signals.iloc[i + 1:i + min_holding_period] = False
    return hold_signals

# Function to backtest the strategy
def backtest_strategy(atr_multiplier, atr_window, data, min_holding_period):
    data = calculate_indicators(data.copy(), atr_window)
    data = calculate_stop_loss(data, atr_multiplier)
    buy_signals = ((data['Close'] > data['Buy_Stop_Loss']) & (data['MACD'] > data['Signal']) & (data['RSI'] < 70)).shift(1).fillna(False)
    sell_signals = ((data['Close'] < data['Sell_Stop_Loss']) & (data['MACD'] < data['Signal']) & (data['RSI'] > 30)).shift(1).fillna(False)

    buy_signals = enforce_min_holding_period(buy_signals, min_holding_period)
    sell_signals = enforce_min_holding_period(sell_signals, min_holding_period)

    returns = data['Close'].pct_change().fillna(0)
    strategy_returns = returns * buy_signals.astype(int) - returns * sell_signals.astype(int)
    cumulative_returns = (1 + strategy_returns).cumprod()

    data['Cumulative_Strategy_Returns'] = cumulative_returns
    return data

# Function to plot results
def plot_results(data, symbol, show_strategy_returns=False):
    fig = go.Figure()

    fig.add_trace(go.Scatter(
        x=data.index, 
        y=data['Close'], 
        mode='lines', 
        name='Close Price', 
        line=dict(color='blue'),
        hovertemplate="Date: %{x}<br>Close: %{y:.2f}<extra></extra>"
    ))
    fig.add_trace(go.Scatter(
        x=data.index, 
        y=data['Buy_Stop_Loss'], 
        mode='lines', 
        name='Buy Stop Loss', 
        line=dict(color='red', dash='dash'),
        hovertemplate="Date: %{x}<br>Buy Stop Loss: %{y:.2f}<extra></extra>"
    ))
    fig.add_trace(go.Scatter(
        x=data.index, 
        y=data['Sell_Stop_Loss'], 
        mode='lines', 
        name='Sell Stop Loss', 
        line=dict(color='green', dash='dash'),
        hovertemplate="Date: %{x}<br>Sell Stop Loss: %{y:.2f}<extra></extra>"
    ))

    fig.update_layout(
        title=f'{symbol} Stop Loss Levels based on ATR', 
        xaxis_title='Date', 
        yaxis_title='Price', 
        width=1200,
        height=600
    )

    st.plotly_chart(fig, use_container_width=True)

    if show_strategy_returns:
        fig2 = go.Figure()
        fig2.add_trace(go.Scatter(
            x=data.index, 
            y=data['Cumulative_Strategy_Returns'], 
            mode='lines', 
            name='Strategy Returns', 
            line=dict(color='purple'),
            hovertemplate="Date: %{x}<br>Cumulative Return: %{y:.2f}<extra></extra>"
        ))

        fig2.update_layout(
            title=f'{symbol} Cumulative Strategy Returns', 
            xaxis_title='Date', 
            yaxis_title='Cumulative Returns', 
            width=1200,
            height=600
        )

        st.plotly_chart(fig2, use_container_width=True)

# Function to optimize parameters
def optimize_parameters(data, min_atr_multiplier, max_atr_multiplier, min_holding_period, progress_bar):
    def objective(atr_multiplier, atr_window):
        data_with_indicators = calculate_indicators(data.copy(), int(atr_window))
        data_with_returns = backtest_strategy(atr_multiplier, int(atr_window), data_with_indicators, min_holding_period)
        return data_with_returns['Cumulative_Strategy_Returns'].iloc[-1]

    optimizer = BayesianOptimization(
        f=objective,
        pbounds={"atr_multiplier": (min_atr_multiplier, max_atr_multiplier), "atr_window": (5, 60)},
        random_state=42
    )

    for i in range(1, 101):
        optimizer.maximize(init_points=1, n_iter=1)
        progress_bar.progress(i / 100)

    return optimizer.max['params']

# Page 1: Run Strategy with User-Specified Parameters
if page == "Run Strategy":
    st.subheader("Run Strategy with Specified Parameters")

    if 'run_strategy_button' in locals() and run_strategy_button:
        with st.spinner("Fetching data..."):
            data = fetch_data(symbol, start_date, end_date)
            st.session_state.data = data  # Store data in session state
        
        with st.spinner("Calculating indicators..."):
            data = calculate_indicators(st.session_state.data, atr_window)
            st.session_state.data = data
        
        with st.spinner("Running Strategy..."):
            data = backtest_strategy(atr_multiplier, atr_window, st.session_state.data, min_holding_period)
            st.session_state.data = data
        
        plot_results(st.session_state.data, symbol, show_strategy_returns=False)

    # If the data has already been processed, display it
    elif st.session_state.data is not None:
        plot_results(st.session_state.data, symbol, show_strategy_returns=False)

# Page 2: Optimize Parameters
elif page == "Optimize Parameters":
    st.subheader("Optimize Strategy Parameters")
    min_atr_multiplier = st.sidebar.slider("Minimum ATR Multiplier", min_value=1.0, max_value=3.0, value=1.5, step=0.1)
    max_atr_multiplier = st.sidebar.slider("Maximum ATR Multiplier", min_value=3.0, max_value=5.0, value=3.5, step=0.1)

    if 'run_optimization_button' in locals() and run_optimization_button:
        with st.spinner("Fetching data..."):
            data = fetch_data(symbol, start_date, end_date)
            st.session_state.data = data
        
        with st.spinner("Calculating indicators..."):
            data = calculate_indicators(st.session_state.data, atr_window=14)  # Use a default window to calculate indicators
            st.session_state.data = data
        
        with st.spinner("Optimizing Parameters..."):
            progress_bar = st.progress(0)
            optimized_params = optimize_parameters(st.session_state.data, min_atr_multiplier, max_atr_multiplier, min_holding_period, progress_bar)
            st.session_state.optimized_params = optimized_params
            st.session_state.atr_multiplier = optimized_params['atr_multiplier']
            st.session_state.atr_window = optimized_params['atr_window']
            st.success(f"Optimization complete! Best ATR Multiplier: {st.session_state.atr_multiplier}, Best ATR Window: {st.session_state.atr_window}")
        
        with st.spinner("Running Strategy with Optimized Parameters..."):
            data = backtest_strategy(st.session_state.atr_multiplier, int(st.session_state.atr_window), st.session_state.data, min_holding_period)
            st.session_state.data = data
        
        plot_results(st.session_state.data, symbol, show_strategy_returns=True)

    # If the optimization has already been run, display the results
    elif st.session_state.optimized_params is not None:
        plot_results(st.session_state.data, symbol, show_strategy_returns=True)
        st.success(f"Optimization complete! Best ATR Multiplier: {st.session_state.atr_multiplier}, Best ATR Window: {st.session_state.atr_window}")

# Hide Streamlit style
hide_streamlit_style = """
<style>
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
</style>
"""
st.markdown(hide_streamlit_style, unsafe_allow_html=True)