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

# Function to calculate Hull Moving Average (HMA)
def hull_moving_average(data, window):
    half_length = int(window / 2)
    sqrt_length = int(np.sqrt(window))
    wma_half = data['Close'].rolling(half_length).apply(lambda x: np.dot(x, range(1, half_length+1)) / sum(range(1, half_length+1)), raw=True)
    wma_full = data['Close'].rolling(window).apply(lambda x: np.dot(x, range(1, window+1)) / sum(range(1, window+1)), raw=True)
    hma = 2 * wma_half - wma_full
    hma = hma.rolling(sqrt_length).apply(lambda x: np.dot(x, range(1, sqrt_length+1)) / sum(range(1, sqrt_length+1)), raw=True)
    return hma

# Function to calculate signals based on HMA crossover
def calculate_signals_hma(data, short_window, long_window):
    data['short_hma'] = hull_moving_average(data, short_window)
    data['long_hma'] = hull_moving_average(data, long_window)

    data['signal'] = np.where(data['short_hma'] > data['long_hma'], 1, -1)  # 1 for Buy, -1 for Sell
    data['positions'] = data['signal'].diff()  # Difference to detect signal changes
    return data

# Function to calculate accuracy of the strategy with adjustable day threshold
def calculate_accuracy(data, buy_threshold, sell_threshold):
    buy_signals = data[data['positions'] == 2]  # 2 indicates a Buy signal after a Sell
    sell_signals = data[data['positions'] == -2]  # -2 indicates a Sell signal after a Buy

    buy_accuracy = (data['Close'].shift(-buy_threshold)[buy_signals.index] > buy_signals['Close']).mean()
    sell_accuracy = (data['Close'].shift(-sell_threshold)[sell_signals.index] < sell_signals['Close']).mean()

    overall_accuracy = (buy_accuracy + sell_accuracy) / 2
    return overall_accuracy, buy_accuracy, sell_accuracy

# Function to optimize HMA parameters based on accuracy
def optimize_hma(data, short_windows, long_windows, buy_threshold, sell_threshold):
    results = []
    best_accuracy = 0
    best_params = None

    for short_window, long_window in product(short_windows, long_windows):
        if short_window >= long_window:
            continue
        temp_data = calculate_signals_hma(data.copy(), short_window, long_window)
        accuracy, buy_accuracy, sell_accuracy = calculate_accuracy(temp_data, buy_threshold, sell_threshold)

        results.append((short_window, long_window, accuracy))

        if accuracy > best_accuracy:
            best_accuracy = accuracy
            best_params = (short_window, long_window, buy_accuracy, sell_accuracy)

    results_df = pd.DataFrame(results, columns=['Short_HMA', 'Long_HMA', 'Accuracy'])
    return best_params, best_accuracy, results_df

# Plotting function with win rates in the legend next to buy/sell signals
def plot_results(data, best_short_window, best_long_window, horizon_name, best_accuracy, buy_accuracy, sell_accuracy):
    data = calculate_signals_hma(data.copy(), best_short_window, best_long_window)

    fig = go.Figure()

    # Add price and HMA lines
    fig.add_trace(go.Scatter(x=data.index, y=data['Close'], mode='lines', name='Price', hovertemplate='%{x|%Y-%m-%d}'))
    fig.add_trace(go.Scatter(x=data.index, y=data['short_hma'], mode='lines', name=f'Short HMA ({best_short_window})', hovertemplate='%{x|%Y-%m-%d}'))
    fig.add_trace(go.Scatter(x=data.index, y=data['long_hma'], mode='lines', name=f'Long HMA ({best_long_window})', hovertemplate='%{x|%Y-%m-%d}'))

    # Add Buy/Sell signals with increased marker size
    buy_signals = data[data['positions'] == 2]  # 2 indicates a Buy signal after a Sell
    sell_signals = data[data['positions'] == -2]  # -2 indicates a Sell signal after a Buy
    fig.add_trace(go.Scatter(x=buy_signals.index, y=buy_signals['Close'], mode='markers',
                             marker=dict(color='green', size=15, symbol='triangle-up'),
                             name=f'Buy Signal (Win Rate: {buy_accuracy:.2f})', hovertemplate='%{x|%Y-%m-%d}'))
    fig.add_trace(go.Scatter(x=sell_signals.index, y=sell_signals['Close'], mode='markers',
                             marker=dict(color='red', size=15, symbol='triangle-down'),
                             name=f'Sell Signal (Win Rate: {sell_accuracy:.2f})', hovertemplate='%{x|%Y-%m-%d}'))

    # Set title and layout, including more detailed date formatting for x-axis
    fig.update_layout(
        title=f'{horizon_name} Horizon: Price and HMA with Buy/Sell Signals (Best Accuracy: {best_accuracy:.2f})',
        xaxis_title='Date',
        yaxis_title='Price',
        xaxis=dict(
            tickformat="%b %Y",
            dtick="M1",
            tickangle=45,
        ),
        autosize=True
    )
    return fig

# Plotting function for strategy performance over time
def plot_strategy_over_time(data, best_short_window, best_long_window):
    data = calculate_signals_hma(data.copy(), best_short_window, best_long_window)
    
    # Rolling accuracy calculation
    window_size = 252  # Using a 1-year window for rolling accuracy
    data['rolling_accuracy'] = data['signal'].rolling(window=window_size).apply(lambda x: (x.shift(-1) * x > 0).mean(), raw=False)
    
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=data.index, y=data['rolling_accuracy'], mode='lines', name='Rolling Accuracy', hovertemplate='%{x|%Y-%m-%d}'))

    fig.update_layout(
        title='Strategy Accuracy Over Time',
        xaxis_title='Date',
        yaxis_title='Rolling Accuracy',
        autosize=True
    )
    return fig

# Streamlit app layout
st.set_page_config(layout="wide")

# Sidebar configuration
with st.sidebar:
    st.header("Input Parameters")

    with st.expander("How to Use", expanded=False):
        st.write("""
        - Select the stock ticker.
        - Set the start and end dates.
        - Click 'Run' to execute the strategy.
        """)

    with st.expander("Ticker Parameters", expanded=True):
        ticker = st.text_input("Stock Ticker", value="AAPL", help="Enter the stock ticker symbol (e.g., AAPL, TSLA)")
        start_date = st.date_input("Start Date", value=datetime(2019, 1, 1), 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")

    with st.expander("Select Horizon", expanded=True):
        st.radio("Horizon", ["Short-Term", "Medium-Term", "Long-Term"], key='horizon_page')

    # Load appropriate horizon settings based on the selected page
    horizons = {
        'Short-Term': {'short_windows': range(5, 20, 2), 'long_windows': range(20, 50, 3), 'buy_threshold': 1, 'sell_threshold': 1},
        'Medium-Term': {'short_windows': range(20, 50, 2), 'long_windows': range(50, 100, 5), 'buy_threshold': 5, 'sell_threshold': 5},
        'Long-Term': {'short_windows': range(50, 100, 5), 'long_windows': range(100, 200, 10), 'buy_threshold': 10, 'sell_threshold': 10},
    }
    selected_horizon = horizons[st.session_state.horizon_page]

    # Run button at the bottom of the sidebar
    run_button = st.button("Run Strategy")

# Title based on the selected page
st.title(f"Hull Moving Average Cross-Over Strategy Optimizer - {st.session_state.horizon_page}")

# Explanation with LaTeX formulas
st.write("""
This application optimizes a trading strategy based on the Hull Moving Average. The strategy uses a cross-over method to generate buy and sell signals by finding the best MA parameters in a given horizon. 
""")

with st.expander("Hull Moving Average Methodology", expanded=False):
    st.latex(r"""
    \text{HMA} = \text{WMA}(2 \times \text{WMA}(n/2) - \text{WMA}(n), \sqrt{n})
    """)
    
    st.write("""
    The cross-over signals are generated based on the following rule:
    """)
    st.latex(r"""
    \text{Signal} = 
    \begin{cases} 
    \text{Buy} & \text{if } \text{Short HMA} > \text{Long HMA} \\
    \text{Sell} & \text{if } \text{Short HMA} < \text{Long HMA}
    \end{cases}
    """)
    st.write("""
    To read more about moving averages methodologies, visit [this link](https://entreprenerdly.com/top-36-moving-averages-methods-for-stock-prices-in-python/).
    """)

# Main application logic
if run_button:
    try:
        if 'data' not in st.session_state or st.session_state.get('ticker') != ticker or st.session_state.get('start_date') != start_date or st.session_state.get('end_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 retrieved for {ticker}")
            if len(data) < max(selected_horizon['short_windows']) + max(selected_horizon['long_windows']):
                raise ValueError(f"Insufficient data points for {ticker}. Need at least {max(selected_horizon['short_windows']) + max(selected_horizon['long_windows'])} days.")
            st.session_state['data'] = data
            st.session_state['ticker'] = ticker
            st.session_state['start_date'] = start_date
            st.session_state['end_date'] = end_date

        data = st.session_state['data']

        # Cache optimization results for each horizon
        if f'{st.session_state.horizon_page}_results' not in st.session_state:
            st.session_state[f'{st.session_state.horizon_page}_results'] = optimize_hma(data, selected_horizon['short_windows'], selected_horizon['long_windows'], selected_horizon['buy_threshold'], selected_horizon['sell_threshold'])

        # Unpack the results from the session state
        best_params, best_accuracy, results_df = st.session_state[f'{st.session_state.horizon_page}_results']
        best_short_window, best_long_window, buy_accuracy, sell_accuracy = best_params

        # Display results
        st.write(f"**{st.session_state.horizon_page} Horizon - Best Short HMA**: {best_short_window}, **Best Long HMA**: {best_long_window}, **Best Accuracy**: {best_accuracy:.2f}")
        st.write(f"**Buy Win Rate**: {buy_accuracy:.2f}, **Sell Win Rate**: {sell_accuracy:.2f}")

        # Plot results within a container to limit the height
        with st.container():
            fig = plot_results(data, best_short_window, best_long_window, st.session_state.horizon_page, best_accuracy, buy_accuracy, sell_accuracy)
            st.plotly_chart(fig, use_container_width=True, height=600)

        # Plot strategy performance over time within a container to limit the height
        st.write("Strategy Performance Over Time")
        with st.container():
            strategy_fig = plot_strategy_over_time(data, best_short_window, best_long_window)
            st.plotly_chart(strategy_fig, use_container_width=True, height=400)

        # Display heatmap of accuracy with annotations
        st.write(f"{st.session_state.horizon_page} Horizon: Accuracy Heatmap of HMA Combinations")
        heatmap_df = results_df.pivot(index='Short_HMA', columns='Long_HMA', values='Accuracy')

        # Create the heatmap with annotations
        heatmap_fig = go.Figure(data=go.Heatmap(
            z=heatmap_df.values,
            x=heatmap_df.columns,
            y=heatmap_df.index,
            colorscale='YlGnBu',
            text=heatmap_df.values,
            texttemplate="%{text:.2f}",
            hovertemplate="Short HMA: %{y}<br>Long HMA: %{x}<br>Accuracy: %{text:.2f}<extra></extra>",
            showscale=True
        ))

        heatmap_fig.update_layout(
            title=f'{st.session_state.horizon_page} Horizon: Accuracy Heatmap of HMA Combinations',
            xaxis_title='Long HMA',
            yaxis_title='Short HMA',
            autosize=True
        )

        with st.container():
            st.plotly_chart(heatmap_fig, use_container_width=True, height=600)

    except Exception as e:
        st.error(f"An error occurred while running the analysis: {e}")

# Re-display the results if they exist and user switches pages without re-running
else:
    if f'{st.session_state.horizon_page}_results' in st.session_state:
        # Unpack the results from the session state
        best_params, best_accuracy, results_df = st.session_state[f'{st.session_state.horizon_page}_results']
        best_short_window, best_long_window, buy_accuracy, sell_accuracy = best_params

        # Display results
        st.write(f"**{st.session_state.horizon_page} Horizon - Best Short HMA**: {best_short_window}, **Best Long HMA**: {best_long_window}, **Best Accuracy**: {best_accuracy:.2f}")
        st.write(f"**Buy Win Rate**: {buy_accuracy:.2f}, **Sell Win Rate**: {sell_accuracy:.2f}")

        # Plot results within a container to limit the height
        with st.container():
            fig = plot_results(st.session_state['data'], best_short_window, best_long_window, st.session_state.horizon_page, best_accuracy, buy_accuracy, sell_accuracy)
            st.plotly_chart(fig, use_container_width=True, height=600)

        # Plot strategy performance over time within a container to limit the height
        st.write("Strategy Performance Over Time")
        with st.container():
            strategy_fig = plot_strategy_over_time(st.session_state['data'], best_short_window, best_long_window)
            st.plotly_chart(strategy_fig, use_container_width=True, height=400)

        # Display heatmap of accuracy with annotations
        st.write(f"{st.session_state.horizon_page} Horizon: Accuracy Heatmap of HMA Combinations")
        heatmap_df = results_df.pivot(index='Short_HMA', columns='Long_HMA', values='Accuracy')

        # Create the heatmap with annotations
        heatmap_fig = go.Figure(data=go.Heatmap(
            z=heatmap_df.values,
            x=heatmap_df.columns,
            y=heatmap_df.index,
            colorscale='YlGnBu',
            text=heatmap_df.values,
            texttemplate="%{text:.2f}",
            hovertemplate="Short HMA: %{y}<br>Long HMA: %{x}<br>Accuracy: %{text:.2f}<extra></extra>",
            showscale=True
        ))

        heatmap_fig.update_layout(
            title=f'{st.session_state.horizon_page} Horizon: Accuracy Heatmap of HMA Combinations',
            xaxis_title='Long HMA',
            yaxis_title='Short HMA',
            autosize=True
        )

        with st.container():
            st.plotly_chart(heatmap_fig, use_container_width=True, height=600)

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