import streamlit as st
import yfinance as yf
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from fastdtw import fastdtw
from scipy.spatial.distance import euclidean
import ta
from sklearn.decomposition import PCA
from sklearn.impute import SimpleImputer
from datetime import datetime, timedelta

# Function to normalize the time series 
def normalize(ts):
    return (ts - ts.min()) / (ts.max() - ts.min())

# Function to calculate DTW distance
def dtw_distance(ts1, ts2):
    ts1_normalized = normalize(ts1)
    ts2_normalized = normalize(ts2)
    distance, _ = fastdtw(ts1_normalized.reshape(-1, 1), ts2_normalized.reshape(-1, 1), dist=euclidean)
    return distance

# Function to calculate correlation
def correlation(ts1, ts2):
    return np.corrcoef(ts1, ts2)[0, 1]

# Function to find most similar patterns using DTW
def find_most_similar_pattern_dtw(price_data_pct_change, n_days, subsequent_days):
    current_window = price_data_pct_change[-n_days:].values

    min_distances = [(float('inf'), -1), (float('inf'), -1), (float('inf'), -1)]
    for start_index in range(len(price_data_pct_change) - 2 * n_days - subsequent_days):
        past_window = price_data_pct_change[start_index:start_index + n_days].values
        distance = dtw_distance(current_window, past_window)

        for i, (min_distance, _) in enumerate(min_distances):
            if distance < min_distance:
                min_distances[i] = (distance, start_index)
                break

    return min_distances

# Function to find most similar patterns using correlation
def find_most_similar_pattern_corr(price_data_pct_change, n_days, subsequent_days, pre_days):
    current_window = price_data_pct_change[-n_days:].values

    max_correlations = [(float('-inf'), -1), (float('-inf'), -1), (float('-inf'), -1)]
    for start_index in range(len(price_data_pct_change) - 2 * n_days - subsequent_days):
        past_window = price_data_pct_change[start_index:start_index + n_days].values
        corr = correlation(current_window, past_window)

        for i, (max_corr, _) in enumerate(max_correlations):
            if corr > max_corr:
                max_correlations[i] = (corr, start_index)
                break

    return max_correlations

# Add technical analysis features
def add_ta_features(data):
    data['trend_ichimoku_conv'] = ta.trend.ichimoku_a(data['High'], data['Low'])
    data['trend_ema_slow'] = ta.trend.ema_indicator(data['Close'], 50)
    data['momentum_kama'] = ta.momentum.kama(data['Close'])
    data['trend_psar_up'] = ta.trend.psar_up(data['High'], data['Low'], data['Close'])
    data['volume_vwap'] = ta.volume.VolumeWeightedAveragePrice(data['High'], data['Low'], data['Close'], data['Volume']).volume_weighted_average_price()
    data['trend_ichimoku_a'] = ta.trend.ichimoku_a(data['High'], data['Low'])
    data['volatility_kcl'] = ta.volatility.KeltnerChannel(data['High'], data['Low'], data['Close']).keltner_channel_lband()
    data['trend_ichimoku_b'] = ta.trend.ichimoku_b(data['High'], data['Low'])
    data['trend_ichimoku_base'] = ta.trend.ichimoku_base_line(data['High'], data['Low'])
    data['trend_sma_fast'] = ta.trend.sma_indicator(data['Close'], 20)
    data['volatility_dcm'] = ta.volatility.DonchianChannel(data['High'], data['Low'], data['Close']).donchian_channel_mband()
    data['volatility_bbl'] = ta.volatility.BollingerBands(data['Close']).bollinger_lband()
    data['volatility_bbm'] = ta.volatility.BollingerBands(data['Close']).bollinger_mavg()
    data['volatility_kcc'] = ta.volatility.KeltnerChannel(data['High'], data['Low'], data['Close']).keltner_channel_mband()
    data['volatility_kch'] = ta.volatility.KeltnerChannel(data['High'], data['Low'], data['Close']).keltner_channel_hband()
    data['trend_sma_slow'] = ta.trend.sma_indicator(data['Close'], 200)
    data['trend_ema_fast'] = ta.trend.ema_indicator(data['Close'], 20)
    data['volatility_dch'] = ta.volatility.DonchianChannel(data['High'], data['Low'], data['Close']).donchian_channel_hband()
    data['others_cr'] = ta.others.cumulative_return(data['Close'])
    data['Adj Close'] = data['Close']
    
    return data

def dtw_distance_with_ta(ts1, ts2, ts1_ta, ts2_ta, weight=0.8):
    ts1_normalized = normalize(ts1)
    ts2_normalized = normalize(ts2)
    distance_pct_change, _ = fastdtw(ts1_normalized.reshape(-1, 1), ts2_normalized.reshape(-1, 1), dist=euclidean)

    distance_ta, _ = fastdtw(ts1_ta, ts2_ta, dist=euclidean)
    distance = weight * distance_pct_change + (1 - weight) * distance_ta
    return distance

def extract_and_reduce_features(data, n_components=3):
    ta_features = data.drop(columns=['Open', 'High', 'Low', 'Close', 'Volume', 'Adj Close'])
    
    imputer = SimpleImputer(strategy='mean')
    imputed_ta_features = imputer.fit_transform(ta_features)
    
    pca = PCA(n_components=n_components)
    reduced_features = pca.fit_transform(imputed_ta_features)
    return reduced_features

# Streamlit app
st.set_page_config(page_title="Pattern Recognition", layout="wide")
st.title('Pattern Recognition in Asset Prices')

# Sidebar for method selection
st.sidebar.title("Input Parameters")

with st.sidebar.expander("How to Use", expanded=False):
    st.markdown("""
    1. Select the pattern recognition method you want to use.
    2. Set the stock ticker or crypto pair, date range, and other parameters.
    3. Click 'Run' to perform the analysis.
    """)

# Grouping and expanding sidebar options with tooltips
with st.sidebar.expander("Pattern Recognition Methods", expanded=True):
    selected = st.radio("Choose a method:",
                        ["DTW Pattern Recognition", "Correlation Pattern Recognition", "TA-Enhanced DTW Pattern Recognition"],
                        help="Select the pattern recognition method you want to use.")

with st.sidebar.expander("Input Parameters", expanded=True):
    ticker = st.text_input('Enter Stock Ticker or Crypto Pair', 'ASML.AS', help="Input the stock ticker (e.g., AAPL, TSLA) or Crypto Pair(e.g., BTC-USD).")
    start_date = st.date_input('Start Date', pd.to_datetime('2000-01-01'), help="Select the start date for the analysis.")
    end_date = st.date_input('End Date', datetime.now() + timedelta(days=1), help="Select the end date for the analysis.")
    subsequent_days = st.slider('Subsequent Days to Forecast', min_value=5, max_value=60, value=20, step=5, help="Number of days to forecast after identifying similar patterns.")
    n_days_options = st.multiselect('Days to Compare', options=[15, 20, 30, 40, 50], default=[15, 20, 40], help="Select the days to compare against the current pattern.")
    pre_days = st.slider('Days Prior to Similar Series', min_value=10, max_value=100, value=60, step=10, help="Number of days prior to the similar series for additional context.")

# Initialize session state variables
if 'data' not in st.session_state:
    st.session_state.data = None
if 'price_data_pct_change' not in st.session_state:
    st.session_state.price_data_pct_change = None
if 'data_with_ta' not in st.session_state:
    st.session_state.data_with_ta = None
if 'reduced_features' not in st.session_state:
    st.session_state.reduced_features = None
if 'results_dtw' not in st.session_state:
    st.session_state.results_dtw = None
if 'results_corr' not in st.session_state:
    st.session_state.results_corr = None
if 'results_ta_dtw' not in st.session_state:
    st.session_state.results_ta_dtw = None

def run_dtw():
    min_distances = []
    figs = []
    for n_days in n_days_options:
        min_distances.append(find_most_similar_pattern_dtw(st.session_state.price_data_pct_change, n_days, subsequent_days))

        fig1 = go.Figure()
        # Plot the entire stock price data
        fig1.add_trace(go.Scatter(x=st.session_state.price_data.index, y=st.session_state.price_data, mode='lines', name='stock price', line=dict(color='blue')))
        colors = ['red', 'green', 'orange']
        for i, (_, start_index) in enumerate(min_distances[-1]):
            # Plot the pattern period
            past_window_start_date = st.session_state.price_data.index[start_index]
            past_window_end_date = st.session_state.price_data.index[start_index + n_days + subsequent_days]
            fig1.add_trace(go.Scatter(x=st.session_state.price_data[past_window_start_date:past_window_end_date].index,
                                      y=st.session_state.price_data[past_window_start_date:past_window_end_date],
                                      mode='lines', name=f"Pattern {i + 1}", line=dict(color=colors[i % len(colors)])))

        # Add labels and legend
        fig1.update_layout(title=f'{ticker} Stock Price Data',
                           xaxis_title='Date',
                           yaxis_title='Stock Price',
                           legend_title='Legend')

        reindexed_current_window = (st.session_state.price_data_pct_change[-n_days:] + 1).cumprod() * 100
        fig2 = go.Figure()
        for i, (_, start_index) in enumerate(min_distances[-1]):
            past_window = st.session_state.price_data_pct_change[start_index:start_index + n_days + subsequent_days]
            reindexed_past_window = (past_window + 1).cumprod() * 100
            fig2.add_trace(go.Scatter(x=list(range(n_days + subsequent_days)), y=reindexed_past_window,
                                      mode='lines', name=f"Past window {i + 1}",
                                      line=dict(color=colors[i % len(colors)], width=3 if i == 0 else 1)))

        fig2.add_trace(go.Scatter(x=list(range(n_days)), y=reindexed_current_window, mode='lines',
                                  name="Current window", line=dict(color='white', width=3)))

        fig2.update_layout(title=f"Most similar {n_days}-day patterns in {ticker} stock price history (aligned, reindexed)",
                           xaxis_title="Days",
                           yaxis_title="Reindexed Price",
                           legend_title="Legend")

        figs.append((fig1, fig2))

    st.session_state.results_dtw = figs

def run_corr():
    max_correlations = []
    figs = []
    for n_days in n_days_options:
        max_correlations.append(find_most_similar_pattern_corr(st.session_state.price_data_pct_change, n_days, subsequent_days, pre_days))

        fig1 = go.Figure()
        # Plot the entire stock price data
        fig1.add_trace(go.Scatter(x=st.session_state.price_data.index, y=st.session_state.price_data, mode='lines', name='stock price', line=dict(color='blue')))
        colors = ['red', 'green', 'orange']
        for i, (_, start_index) in enumerate(max_correlations[-1]):
            # Plot the previous period
            past_window_start_date = st.session_state.price_data.index[start_index - pre_days]
            past_window_end_date = st.session_state.price_data.index[start_index + n_days + subsequent_days]
            fig1.add_trace(go.Scatter(x=st.session_state.price_data[past_window_start_date:past_window_end_date].index,
                                      y=st.session_state.price_data[past_window_start_date:past_window_end_date],
                                      mode='lines', name=f"Pattern {i + 1}", line=dict(color=colors[i % len(colors)])))

        # Add labels and legend
        fig1.update_layout(title=f'{ticker} Stock Price Data',
                           xaxis_title='Date',
                           yaxis_title='Stock Price',
                           legend_title='Legend')

        reindexed_current_window = (st.session_state.price_data_pct_change[-n_days:] + 1).cumprod() * 100
        fig2 = go.Figure()
        for i, (_, start_index) in enumerate(max_correlations[-1]):
            past_window = st.session_state.price_data_pct_change[start_index:start_index + n_days + subsequent_days]
            reindexed_past_window = (past_window + 1).cumprod() * 100
            fig2.add_trace(go.Scatter(x=list(range(pre_days, pre_days + n_days + subsequent_days)), y=reindexed_past_window,
                                      mode='lines', name=f"Past window {i + 1}",
                                      line=dict(color=colors[i % len(colors)], width=3 if i == 0 else 1)))

        fig2.add_trace(go.Scatter(x=list(range(pre_days, pre_days + n_days)), y=reindexed_current_window, mode='lines',
                                  name="Current window", line=dict(color='white', width=3)))

        fig2.update_layout(title=f"Most similar {n_days}-day patterns in {ticker} stock price history (aligned, reindexed with correlation)",
                           xaxis_title="Days",
                           yaxis_title="Reindexed Price",
                           legend_title="Legend")

        figs.append((fig1, fig2))

    st.session_state.results_corr = figs

def run_ta_dtw():
    min_distance_indices = []
    figs = []
    for n_days in n_days_options:
        current_window = st.session_state.price_data_pct_change[-n_days:].values
        current_ta_window = st.session_state.reduced_features[-n_days:]

        distances = [dtw_distance_with_ta(current_window, st.session_state.price_data_pct_change[start_index:start_index + n_days].values, 
                                         current_ta_window, st.session_state.reduced_features[start_index:start_index + n_days]) 
                     for start_index in range(len(st.session_state.price_data_pct_change) - 2 * n_days - subsequent_days)]

        min_distance_indices.append(np.argsort(distances)[:3])  # find indices of 3 smallest distances

        fig1 = go.Figure()
        # Plot the entire stock price data
        fig1.add_trace(go.Scatter(x=st.session_state.data.index, y=st.session_state.data['Close'], mode='lines', name='stock price', line=dict(color='blue')))
        colors = ['red', 'green', 'orange']
        for i, start_index in enumerate(min_distance_indices[-1]):
            # Plot the pattern period
            past_window_start_date = st.session_state.data.index[start_index]
            past_window_end_date = st.session_state.data.index[start_index + n_days + subsequent_days]
            fig1.add_trace(go.Scatter(x=st.session_state.data['Close'][past_window_start_date:past_window_end_date].index,
                                      y=st.session_state.data['Close'][past_window_start_date:past_window_end_date],
                                      mode='lines', name=f"Pattern {i + 1}", line=dict(color=colors[i % len(colors)])))

        # Add labels and legend
        fig1.update_layout(title=f'{ticker} Stock Price Data',
                           xaxis_title='Date',
                           yaxis_title='Stock Price',
                           legend_title='Legend')

        reindexed_current_window = (st.session_state.price_data_pct_change[-n_days:] + 1).cumprod() * 100
        fig2 = go.Figure()
        for i, start_index in enumerate(min_distance_indices[-1]):
            past_window = st.session_state.price_data_pct_change[start_index:start_index + n_days + subsequent_days]
            reindexed_past_window = (past_window + 1).cumprod() * 100
            fig2.add_trace(go.Scatter(x=list(range(n_days + subsequent_days)), y=reindexed_past_window,
                                      mode='lines', name=f"Past window {i + 1}",
                                      line=dict(color=colors[i % len(colors)], width=3 if i == 0 else 1)))

        fig2.add_trace(go.Scatter(x=list(range(n_days)), y=reindexed_current_window, mode='lines',
                                  name="Current window", line=dict(color='white', width=3)))

        fig2.update_layout(title=f"Most similar {n_days}-day patterns in {ticker} stock price history (aligned, reindexed)",
                           xaxis_title="Days",
                           yaxis_title="Reindexed Price",
                           legend_title="Legend")

        figs.append((fig1, fig2))

    st.session_state.results_ta_dtw = figs

if st.sidebar.button('Run Analysis'):
    # Fetch data with auto_adjust=False and flatten columns if multi-indexed
    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 not data.empty:
        st.session_state.data = data
        st.session_state.price_data = st.session_state.data['Close']
        st.session_state.price_data_pct_change = st.session_state.price_data.pct_change().dropna()
        st.session_state.data_with_ta = add_ta_features(st.session_state.data.copy())
        st.session_state.reduced_features = extract_and_reduce_features(st.session_state.data_with_ta, n_components=2)

        if selected == "DTW Pattern Recognition":
            run_dtw()
        elif selected == "Correlation Pattern Recognition":
            run_corr()
        elif selected == "TA-Enhanced DTW Pattern Recognition":
            run_ta_dtw()
    else:
        st.error(f"No data returned for {ticker} from {start_date} to {end_date}")

# Display results and descriptions based on the selected method
if selected == "DTW Pattern Recognition":
    st.markdown("""
    ### DTW Pattern Recognition
    Dynamic Time Warping (DTW) is a method that measures the similarity between two time series that may vary in time or speed. """)
    
    with st.expander("Click here to read more about the methodology"):
        st.markdown("""DTW aligns the time series by warping the time axis to minimize the distance between them. This method could potentially identify historical periods that have similar patterns to the current stock price pattern by comparing their shapes, regardless of possible distortions in the time axis.
        The DTW distance \( D \) between two time series \( X \) and \( Y \) is calculated as:
        """)
        st.latex(r'''
        D(i, j) = \text{dist}(X_i, Y_j) + \min \begin{cases} 
        D(i-1, j) \\
        D(i, j-1) \\
        D(i-1, j-1) 
        \end{cases}
        ''')
        st.markdown("""
        where {dist}(Xi, Yj) is the distance between points (Xi) and (Yj), and D(i, j) is the accumulated cost.
        """)
    
    st.markdown("""**Results:**
    The left chart shows the entire stock price data with the identified patterns highlighted. The right chart shows the reindexed price patterns for comparison.
    """)
    if st.session_state.results_dtw:
        for fig1, fig2 in st.session_state.results_dtw:
            col1, col2 = st.columns(2)
            with col1:
                st.plotly_chart(fig1)
            with col2:
                st.plotly_chart(fig2)

elif selected == "Correlation Pattern Recognition":
    st.markdown("""
    ### Correlation Pattern Recognition
    This method calculates the correlation between the current stock price pattern and historical patterns. Correlation measures the degree to which two time series move together. A higher correlation value indicates that the patterns are more similar. This method helps identify historical periods where stock prices behaved similarly to the current pattern by looking at how closely the price movements align.
    The correlation coefficient \( r \) is calculated as:
    """)
    st.latex(r'''
    r = \frac{ \sum (X_i - \bar{X})(Y_i - \bar{Y}) }{ \sqrt{ \sum (X_i - \bar{X})^2 } \sqrt{ \sum (Y_i - \bar{Y})^2 } }
    ''')
    st.markdown("""
    where \( X \) and \( Y \) are the two time series being compared.

    **Results:**
    The left chart shows the entire stock price data with the identified patterns highlighted. The right chart shows the reindexed price patterns for comparison.
    """)
    if st.session_state.results_corr:
        for fig1, fig2 in st.session_state.results_corr:
            col1, col2 = st.columns(2)
            with col1:
                st.plotly_chart(fig1)
            with col2:
                st.plotly_chart(fig2)

elif selected == "TA-Enhanced DTW Pattern Recognition":
    st.markdown("""
    ### TA-Enhanced DTW Pattern Recognition
    This method combines technical analysis features with Dynamic Time Warping to enhance pattern recognition. It integrates various indicators such as moving averages, momentum indicators, and volatility measures into the time series data. These indicators provide additional context to the price movements.
    Additionally, Principal Component Analysis (PCA) is used to reduce the dimensionality of the features. PCA transforms the indicators into a smaller set of uncorrelated components, capturing the most significant information with fewer variables. 
    The method uses a weighted mechanism to combine the DTW distance of the price changes and the TA features. The formula is:
    """)
    st.latex(r'''
    \text{Total Distance} = w \cdot \text{DTW distance of price changes} + (1 - w) \cdot \text{DTW distance of TA features}
    ''')
    st.markdown("""
    where \( w \) is a weight factor that balances the contribution of price changes and TA features.

    **Results:**
    The left chart shows the entire stock price data with the identified patterns highlighted. The right chart shows the reindexed price patterns for comparison.
    """)
    if st.session_state.results_ta_dtw:
        for fig1, fig2 in st.session_state.results_ta_dtw:
            col1, col2 = st.columns(2)
            with col1:
                st.plotly_chart(fig1)
            with col2:
                st.plotly_chart(fig2)

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