import streamlit as st
import numpy as np
import pandas as pd
import yfinance as yf
import plotly.graph_objects as go
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential
from keras.layers import LSTM, Dense, Dropout
import warnings
warnings.filterwarnings("ignore")
import datetime

# Define the StockPredictor class with the plot_predictions function exactly as in your code
class StockPredictor:
    def __init__(self, ticker, start_date, end_date, look_back=60):
        self.ticker = ticker
        self.start_date = start_date
        self.end_date = end_date
        self.look_back = look_back
        self.data = None
        self.scaler = MinMaxScaler(feature_range=(0, 1))
        self.model = None
        self.X, self.Y = None, None
        self.predictions = None
        self.predicted_dates = None
        self.features = None  # To store the list of features

    def load_data(self):
        self.data = yf.download(self.ticker, start=self.start_date, end=self.end_date, auto_adjust=False)
        if isinstance(self.data.columns, pd.MultiIndex):
            self.data.columns = self.data.columns.get_level_values(0)
        if self.data.empty:
            raise ValueError(f"No data retrieved for {self.ticker}")
        if len(self.data) < self.look_back + 1:  # Ensure enough data for look_back period
            raise ValueError(f"Insufficient data points for {self.ticker}. Need at least {self.look_back + 1} days.")
        print(f'Data loaded for {self.ticker} from {self.start_date} to {self.end_date}')

    def calculate_wma(self, prices, window):
        weights = np.arange(1, window + 1)
        return prices.rolling(window).apply(lambda x: np.dot(x, weights) / weights.sum(), raw=True)

    def calculate_hull_moving_average(self, close_prices, window=14):
        """Function to calculate the Hull Moving Average (HMA)."""
        half_length = int(window / 2)
        sqrt_length = int(np.sqrt(window))

        wma_half = self.calculate_wma(close_prices, half_length)
        wma_full = self.calculate_wma(close_prices, window)
        raw_hma = 2 * wma_half - wma_full
        hma = self.calculate_wma(raw_hma, sqrt_length)
        return hma

    def calculate_rolling_zscore(self, close_prices, window=20):
        """Function to calculate rolling Z-Score."""
        rolling_mean = close_prices.rolling(window=window).mean()
        rolling_std = close_prices.rolling(window=window).std()
        z_score = (close_prices - rolling_mean) / rolling_std
        return z_score

    def calculate_technical_indicators(self):
        """Calculate technical indicators for the historical data."""
        # Hull Moving Average
        self.data['HMA'] = self.calculate_hull_moving_average(self.data['Close'], window=14)

        # Bollinger Bands
        rolling_std = self.data['Close'].rolling(window=20).std()
        self.data['BB_upper'] = self.data['HMA'] + (rolling_std * 2)
        self.data['BB_lower'] = self.data['HMA'] - (rolling_std * 2)

        # Relative Strength Index (RSI)
        delta = self.data['Close'].diff(1)
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        rs = avg_gain / avg_loss
        self.data['RSI'] = 100 - (100 / (1 + rs))

        # Rolling Z-Score
        self.data['Z_Score'] = self.calculate_rolling_zscore(self.data['Close'], window=20)

        # Fill missing values
        self.data.fillna(method='bfill', inplace=True)

    def prepare_data(self):
        self.calculate_technical_indicators()

        # Select features
        features = ['Close', 'HMA', 'BB_upper', 'BB_lower', 'RSI', 'Z_Score']
        self.features = features
        data = self.data[features]
        data_scaled = self.scaler.fit_transform(data)

        X, Y = [], []
        for i in range(self.look_back, len(data_scaled)):
            X.append(data_scaled[i - self.look_back:i])
            Y.append(data_scaled[i, 0])  # Assuming 'Close' is the target
        self.X, self.Y = np.array(X), np.array(Y)
        print('Data prepared for training.')

    def build_model(self):
        self.model = Sequential()
        self.model.add(LSTM(units=300, return_sequences=True, input_shape=(self.X.shape[1], self.X.shape[2])))
        #self.model.add(Dropout(0.2))
        self.model.add(LSTM(units=300))
        #self.model.add(Dropout(0.2))
        self.model.add(Dense(1))
        self.model.compile(loss='mean_squared_error', optimizer='adam')
        print('Model built and compiled.')

    def train_model(self, epochs=10, batch_size=64):
        self.model.fit(self.X, self.Y, epochs=epochs, batch_size=batch_size, verbose=1)
        print('Model trained.')

    def make_predictions(self):
        self.predictions = self.model.predict(self.X)
        # Inverse transform predictions
        # We need to inverse transform the Close price
        # Since we have multiple features, we need to create a full array
        predictions_extended = np.zeros((self.predictions.shape[0], len(self.features)))
        predictions_extended[:, 0] = self.predictions[:, 0]
        predictions_inverse = self.scaler.inverse_transform(predictions_extended)
        self.predictions = predictions_inverse[:, 0]

        # Similarly for Y
        Y_extended = np.zeros((self.Y.shape[0], len(self.features)))
        Y_extended[:, 0] = self.Y
        Y_inverse = self.scaler.inverse_transform(Y_extended)
        self.Y = Y_inverse[:, 0]
        # Store corresponding dates
        self.predicted_dates = self.data.index[self.look_back:]
        print('Predictions made and stored.')

    def forecast_future(self, days=5, n_mc=100):
        # Start with the last 'look_back' periods from data
        last_sequence = self.data[self.features].iloc[-self.look_back:].copy()
        future_predictions = []
        future_std = []

        for _ in range(days):
            # Scale the input data
            last_sequence_scaled = self.scaler.transform(last_sequence)
            # Reshape to (1, look_back, num_features)
            input_data = last_sequence_scaled.reshape(1, self.look_back, len(self.features))

            # Perform N_mc stochastic forward passes
            predictions = []
            for _ in range(n_mc):
                prediction = self.model(input_data, training=True)
                predictions.append(prediction.numpy()[0, 0])

            predictions = np.array(predictions)
            mean_prediction = predictions.mean()
            std_prediction = predictions.std()
            future_std.append(std_prediction)

            # Create an array to inverse transform
            prediction_extended = np.zeros((1, len(self.features)))
            prediction_extended[0, 0] = mean_prediction
            predicted_price = self.scaler.inverse_transform(prediction_extended)[0, 0]
            future_predictions.append(predicted_price)

            # Prepare the next input
            # Create a new row with the predicted price and updated technical indicators
            new_row = {}
            new_row['Close'] = predicted_price

            # Append the new predicted price to the sequence
            # Use pd.concat instead of append
            new_row_df = pd.DataFrame([new_row], index=[last_sequence.index[-1] + pd.Timedelta(days=1)])
            last_sequence = pd.concat([last_sequence, new_row_df])

            # Recalculate technical indicators on the updated last_sequence
            last_close_prices = last_sequence['Close']
            # HMA
            hma_series = self.calculate_hull_moving_average(last_close_prices)
            last_sequence['HMA'] = hma_series
            # Bollinger Bands
            rolling_std = last_close_prices.rolling(window=20).std()
            last_sequence['BB_upper'] = last_sequence['HMA'] + (rolling_std * 2)
            last_sequence['BB_lower'] = last_sequence['HMA'] - (rolling_std * 2)
            # RSI
            delta = last_close_prices.diff(1)
            gain = delta.where(delta > 0, 0)
            loss = -delta.where(delta < 0, 0)
            avg_gain = gain.rolling(window=14).mean()
            avg_loss = loss.rolling(window=14).mean()
            rs = avg_gain / avg_loss
            last_sequence['RSI'] = 100 - (100 / (1 + rs))
            # Z_Score
            z_score_series = self.calculate_rolling_zscore(last_close_prices)
            last_sequence['Z_Score'] = z_score_series

            # Fill any missing values
            last_sequence.fillna(method='bfill', inplace=True)

            # Keep only the last 'look_back' periods
            last_sequence = last_sequence.iloc[-self.look_back:]

        # Generate future dates
        future_dates = pd.date_range(start=self.data.index[-1] + pd.Timedelta(days=1), periods=days, freq='B')

        # Calculate confidence intervals
        future_std = np.array(future_std).reshape(-1, 1)
        # Scale the standard deviations appropriately
        future_std_scaled = future_std * (self.scaler.data_max_[0] - self.scaler.data_min_[0])

        # Create a DataFrame with future dates and predicted prices
        future_df = pd.DataFrame({
            'Date': future_dates,
            'Predicted_Close': future_predictions,
            'Std_Dev': future_std_scaled.flatten()
        })

        # Calculate 68% and 95% confidence intervals
        future_df['CI_68_Lower'] = future_df['Predicted_Close'] - future_df['Std_Dev']
        future_df['CI_68_Upper'] = future_df['Predicted_Close'] + future_df['Std_Dev']
        future_df['CI_95_Lower'] = future_df['Predicted_Close'] - 2 * future_df['Std_Dev']
        future_df['CI_95_Upper'] = future_df['Predicted_Close'] + 2 * future_df['Std_Dev']

        print('Future predictions made and stored with confidence intervals.')
        return future_df

    def plot_predictions(self, future_predictions_df):
        # Create plotly figure
        fig = go.Figure()

        # Plot actual prices
        fig.add_trace(go.Scatter(
            x=self.data.index,
            y=self.data['Close'],
            mode='lines',
            name='Actual Price'
        ))

        # Plot predicted prices for historical data
        fig.add_trace(go.Scatter(
            x=self.predicted_dates,
            y=self.predictions,
            mode='lines',
            name='Predicted Price'
        ))

        # Plot future predictions
        fig.add_trace(go.Scatter(
            x=future_predictions_df['Date'],
            y=future_predictions_df['Predicted_Close'],
            mode='lines',
            name=f'Future {len(future_predictions_df)} Days Prediction',
            line=dict(dash='dash')
        ))

        # Plot 95% confidence interval as shaded area
        fig.add_trace(go.Scatter(
            x=future_predictions_df['Date'].tolist() + future_predictions_df['Date'][::-1].tolist(),
            y=future_predictions_df['CI_95_Upper'].tolist() + future_predictions_df['CI_95_Lower'][::-1].tolist(),
            fill='toself',
            fillcolor='rgba(255, 192, 0, 0.2)',
            line=dict(color='rgba(255,255,255,0)'),
            hoverinfo='skip',
            showlegend=True,
            name='95% Confidence Interval'
        ))

        # Plot 68% confidence interval as shaded area (on top of 95%)
        fig.add_trace(go.Scatter(
            x=future_predictions_df['Date'].tolist() + future_predictions_df['Date'][::-1].tolist(),
            y=future_predictions_df['CI_68_Upper'].tolist() + future_predictions_df['CI_68_Lower'][::-1].tolist(),
            fill='toself',
            fillcolor='rgba(0, 100, 80, 0.2)',
            line=dict(color='rgba(255,255,255,0)'),
            hoverinfo='skip',
            showlegend=True,
            name='68% Confidence Interval'
        ))

        # Determine if the final prediction is bullish or bearish
        last_actual_price = self.data['Close'].iloc[-1]
        final_predicted_price = future_predictions_df['Predicted_Close'].iloc[-1]
        bullish = final_predicted_price > last_actual_price

        # Add a triangle marker at the end of the actual price plot, colored based on bullish/bearish
        fig.add_trace(go.Scatter(
            x=[self.data.index[-1]],
            y=[last_actual_price],
            mode='markers',
            marker=dict(
                color='green' if bullish else 'red',
                size=12,
                symbol='triangle-up' if bullish else 'triangle-down'
            ),
            name='Bullish' if bullish else 'Bearish'
        ))

        # Add labels for the final confidence intervals and final price, placed next to the lines
        annotations = []

        # Set x offset for labels to avoid overlapping with the data points
        x_offset = 40  # Adjust as needed

        # Final predicted price annotation
        annotations.append(dict(
            x=future_predictions_df['Date'].iloc[-1],
            y=final_predicted_price,
            xref='x', yref='y',
            text=f"Final Predicted Price: ${final_predicted_price:.2f}",
            showarrow=False,
            xanchor='left',
            xshift=x_offset
        ))

        # 68% Confidence Interval Upper Bound Annotation
        annotations.append(dict(
            x=future_predictions_df['Date'].iloc[-1],
            y=future_predictions_df['CI_68_Upper'].iloc[-1],
            xref='x', yref='y',
            text=f"68% CI Upper: ${future_predictions_df['CI_68_Upper'].iloc[-1]:.2f}",
            showarrow=False,
            xanchor='left',
            xshift=x_offset
        ))

        # 68% Confidence Interval Lower Bound Annotation
        annotations.append(dict(
            x=future_predictions_df['Date'].iloc[-1],
            y=future_predictions_df['CI_68_Lower'].iloc[-1],
            xref='x', yref='y',
            text=f"68% CI Lower: ${future_predictions_df['CI_68_Lower'].iloc[-1]:.2f}",
            showarrow=False,
            xanchor='left',
            xshift=x_offset
        ))

        # 95% Confidence Interval Upper Bound Annotation
        annotations.append(dict(
            x=future_predictions_df['Date'].iloc[-1],
            y=future_predictions_df['CI_95_Upper'].iloc[-1],
            xref='x', yref='y',
            text=f"95% CI Upper: ${future_predictions_df['CI_95_Upper'].iloc[-1]:.2f}",
            showarrow=False,
            xanchor='left',
            xshift=x_offset
        ))

        # 95% Confidence Interval Lower Bound Annotation
        annotations.append(dict(
            x=future_predictions_df['Date'].iloc[-1],
            y=future_predictions_df['CI_95_Lower'].iloc[-1],
            xref='x', yref='y',
            text=f"95% CI Lower: ${future_predictions_df['CI_95_Lower'].iloc[-1]:.2f}",
            showarrow=False,
            xanchor='left',
            xshift=x_offset
        ))

        # Update layout with annotations
        fig.update_layout(
            title=f'Stock Price Prediction for {self.ticker}',
            xaxis_title='Date',
            yaxis_title='Price',
            hovermode='x unified',
            annotations=annotations
        )

        # Return the figure
        return fig

# Streamlit App
st.set_page_config(layout="wide")
st.title("Deep Learning Asset Price Forecasting")

# Include a short description in the main body of the app
st.markdown("""
This tool forecasts future stock prices and cryptocurrency pairs using a deep neural network with a large number of proprietary historical external variables. 
The model provides both 68% and 95% confidence intervals to represent uncertainty. Zoom into the forecast area of the plot for a clearer view.
""")

with st.expander("Additional Trading Tools and Analysis", expanded=False):
    st.markdown("""
    To implement the forecasts into trading, users are advised to further combine the predictions with the following tools to assess the latest asset price patterns:
    - [Expected Stock Price Movement Using Volatility Multipliers](https://entreprenerdly.com/expected-stock-price-movement-using-volatility-multipliers/)
    - [Future Stock Price Movements with Monte Carlo Simulations](https://entreprenerdly.com/future-stock-price-movements-with-monte-carlo-simulations/)
    - [Technical Analysis with Trading Indicators](https://entreprenerdly.com/technical-analysis/)

    Additionally, comparing the results with Entreprenerdly's algorithmic indicators can provide further insights for making the final decision.
    """)

# Sidebar
st.sidebar.title("Input Parameters")

with st.sidebar.expander("How to Use", expanded=False):
    st.write("""
    **Instructions:**
    - Enter the stock ticker symbol or crypto pair you want to forecast (e.g., AAPL or BTC-USD).
    - Select the number of days ahead you want to forecast using the slider. We recommend choosing a small number.
    - Click on "Run Model" to generate the predictions and view the results.
    """)

with st.sidebar.expander("Input Parameters", expanded=True):
    ticker = st.text_input(
        "Stock Ticker",
        value="AAPL",
        help="Enter the stock ticker symbol or Cryptocurrency pair you want to forecast (e.g., AAPL, BTC-USD)."
    )
    forecast_days = st.slider(
        "Forecast Horizon (Days)",
        min_value=1,
        max_value=10,
        value=5,
        help="Select the number of future days to forecast. We recommend choosing a smaller number, as longer horizons generally result in higher prediction errors."
    )

# Set end date to today's date plus one in the backend
end_date = (datetime.date.today() + datetime.timedelta(days=1)).strftime("%Y-%m-%d")
start_date = '2015-01-01'  # Fixed start date

# Run button
run_button = st.sidebar.button("Run Model")

if run_button:
    try:
        # Progress bar
        progress_bar = st.progress(0)
        status_text = st.empty()

        # Instantiate and run the predictor
        predictor = StockPredictor(
            ticker,
            start_date,
            end_date,
            look_back=60
        )
        status_text.text("Loading data...")
        predictor.load_data()
        progress_bar.progress(10)

        status_text.text("Preparing data...")
        predictor.prepare_data()
        progress_bar.progress(30)

        status_text.text("Building model...")
        predictor.build_model()
        progress_bar.progress(50)

        status_text.text("Training model...This may take a few minutes")
        predictor.train_model(epochs=10, batch_size=16)
        progress_bar.progress(80)

        status_text.text("Making predictions...")
        predictor.make_predictions()
        progress_bar.progress(90)

        status_text.text("Forecasting future prices...")
        future_predictions_df = predictor.forecast_future(days=forecast_days, n_mc=100)
        progress_bar.progress(100)

        status_text.text("Generating plot...")
        fig = predictor.plot_predictions(future_predictions_df)

        # Clear the progress bar and status text
        progress_bar.empty()
        status_text.empty()

        # Display the plot
        st.plotly_chart(fig)

        # Include a short interpretation of the results
        last_actual_price = predictor.data['Close'].iloc[-1]
        final_predicted_price = future_predictions_df['Predicted_Close'].iloc[-1]
        bullish = final_predicted_price > last_actual_price

        # Interpretation of results
        ci_68_lower = future_predictions_df['CI_68_Lower'].iloc[-1]
        ci_68_upper = future_predictions_df['CI_68_Upper'].iloc[-1]
        ci_95_lower = future_predictions_df['CI_95_Lower'].iloc[-1]
        ci_95_upper = future_predictions_df['CI_95_Upper'].iloc[-1]
        mean_prediction = future_predictions_df['Predicted_Close'].mean()
        
        st.write(f"### Final Predicted Price: ${final_predicted_price:.2f}")
        
        if bullish:
            st.success("The model predicts a **bullish** trend over the forecast horizon.")
        else:
            st.error("The model predicts a **bearish** trend over the forecast horizon.")
        
        # Build the interpretation text step by step
        interpretation_text = "**Interpretation of Results:**\n\n"
        interpretation_text += f"The model forecasts that the stock price is expected to be around **${final_predicted_price:.2f}** by the end of the forecast horizon.\n\n"
        interpretation_text += f"The **68% confidence interval** suggests that the price is likely to fall between **${ci_68_lower:.2f}** and **${ci_68_upper:.2f}**, "
        interpretation_text += f"while the broader **95% confidence interval** ranges from **${ci_95_lower:.2f}** to **${ci_95_upper:.2f}**.\n\n"
        interpretation_text += f"On average, the forecasted prices over the horizon have a mean value of **${mean_prediction:.2f}**.\n\n"
        interpretation_text += f"The predicted trend indicates a **{'bullish' if bullish else 'bearish'}** signal, suggesting the price is expected to **{'rise' if bullish else 'decline'}** "
        interpretation_text += f"compared to the last recorded price of **${last_actual_price:.2f}**.\n\n"
        interpretation_text += "The confidence intervals represent the uncertainty in the model's predictions, with wider intervals indicating higher uncertainty. "
        interpretation_text += "The shaded areas on the plot provide a visual representation of this uncertainty."
        
        st.markdown(interpretation_text)
    except Exception as e:
        st.error(f"An error occurred while running the analysis: {e}")
        progress_bar.empty()
        status_text.empty()

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