from AlgorithmImports import *
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from sklearn.model_selection import train_test_split
import math
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error


class MomentumVolumeTrendStrategy(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2023, 1, 1)
        self.SetCash(1000000)

        # Add equities with daily resolution
        self.symbols = []
        self.symbols.append(self.AddEquity("AMD", Resolution.Daily).Symbol)
        self.symbols.append(self.AddEquity("JNJ", Resolution.Daily).Symbol)

        # Dictionary to store indicators for each stock
        self.indicators = {}
        for symbol in self.symbols:
            self.indicators[symbol] = {
                "sma50": self.SMA(symbol, 50, Resolution.Daily),
                "sma200": self.SMA(symbol, 200, Resolution.Daily),
                "rsi": self.RSI(symbol, 14, MovingAverageType.Simple, Resolution.Daily),
                "macd": self.MACD(symbol, 12, 26, 9, MovingAverageType.Exponential, Resolution.Daily),
                "bb": self.BB(symbol, 20, 2.0, Resolution.Daily),
                "obv": self.OBV(symbol, Resolution.Daily),
                "volume_ma": self.SMA(symbol, 20, Resolution.Daily, Field.Volume) 
            }

        # Warm up indicators
        self.SetWarmUp(200)

        # Initialize LSTM models for each symbol
        self.models = {}
        for symbol in self.symbols:
            self.models[symbol] = self.build_lstm_model()

    def OnData(self, data):
        if self.IsWarmingUp:
            return

        for symbol in self.symbols:
            ind = self.indicators[symbol]

            # Ensure indicators are ready
            if not all([ind["sma50"].IsReady, ind["sma200"].IsReady, ind["rsi"].IsReady, 
                        ind["macd"].IsReady, ind["bb"].IsReady, ind["obv"].IsReady, ind["volume_ma"].IsReady]):
                continue

            # Get values of indicators
            price = self.Securities[symbol].Price
            sma50 = ind["sma50"].Current.Value
            sma200 = ind["sma200"].Current.Value
            rsi = ind["rsi"].Current.Value
            macd = ind["macd"].Current.Value
            bb = ind["bb"]
            upperBB = bb.UpperBand.Current.Value
            middleBB = bb.MiddleBand.Current.Value
            lowerBB = bb.LowerBand.Current.Value
            obv = ind["obv"].Current.Value
            avg_volume = ind["volume_ma"].Current.Value
            current_volume = self.Securities[symbol].Volume

            # Define base position (trend-following component)
            baseWeight = 0.75 if sma50 > sma200 else -0.75

            # Volume-based confirmation
            volume_confirmation = current_volume > 1.5 * avg_volume  # Significant volume spike

            # LSTM prediction
            # Use the last 30 days of data to predict the next logreturn and risk
            historical_data = self.History(symbol, 30, Resolution.Daily)
            if len(historical_data) < 30:
                return  # Ensure we have enough data points for LSTM (skip if not enough data)

            df = pd.DataFrame(historical_data)
            features = ['logclose', 'close', 'volume']
            df['logclose'] = df['close'].apply(lambda x: math.log(x))
            df['logclose'] = df['logclose'] - df['logclose'].iloc[0]
            scaled_features = MinMaxScaler(feature_range=(0, 1)).fit_transform(np.array(df[features]))

            # Ensure there are exactly 30 rows (time steps) and 3 features (close, volume, logclose)
            if scaled_features.shape[0] != 30 or scaled_features.shape[1] != 3:
                return  # Skip if data shape is incorrect

            X = scaled_features.reshape(1, 30, 3)  # Reshape for LSTM input (1 batch, 30 time steps, 3 features)
            prediction = self.models[symbol].predict(X)
            predicted_return = prediction[0][0]  # The predicted logreturn
            predicted_risk = prediction[0][1]  # New risk prediction output

            # Bullish and Bearish signals
            bullishSignal = (price > upperBB and rsi > 55 and predicted_return > 0)
            bearishSignal = (price < lowerBB and rsi < 45 and predicted_return < 0)

            # Volume spike and MACD confirmation
            volume_signal = volume_confirmation and predicted_return > 0
            macd_signal = macd > 0

            # Define position adjustment based on signals
            moveFactor = 1  # Default "small move"

            # Check for medium or big moves
            if bullishSignal and macd_signal:
                moveFactor = 1.5  # Medium move for BB, RSI, LSTM & MACD
            if bullishSignal and macd_signal and volume_signal:
                moveFactor = 2  # Big move for BB, RSI, LSTM, MACD & Volume

            if bearishSignal and macd_signal:
                moveFactor = -1.5  # Medium bearish move
            if bearishSignal and macd_signal and volume_signal:
                moveFactor = -2  # Big bearish move

            # Dynamic Position Scaling based on predicted risk
            # If risk is high, reduce the position size
            risk_factor = max(0.1, 1 - predicted_risk)  # Ensure the risk factor doesn't go below 0.1

            if baseWeight > 0:
                targetWeight = min(baseWeight * risk_factor * moveFactor, 1) 
            if baseWeight < 0:
                targetWeight = max(baseWeight * risk_factor * moveFactor, -1)

            # Set holdings based on final calculated weight
            self.SetHoldings(symbol, targetWeight)

            # Debugging output
            self.Debug(f"{self.Time} {symbol.Value}: Price={price:.2f}, SMA50={sma50:.2f}, "
                       f"SMA200={sma200:.2f}, RSI={rsi:.2f}, MACD={macd:.2f}, "
                       f"UpperBB={upperBB:.2f}, MiddleBB={middleBB:.2f}, LowerBB={lowerBB:.2f}, "
                       f"OBV={obv:.2f}, Volume={current_volume}, Avg Volume={avg_volume}, "
                       f"Predicted Return={predicted_return:.2f}, Predicted Risk={predicted_risk:.2f}, "
                       f"Target Weight={targetWeight:.2f}")

    def build_lstm_model(self):
        # Build and return the LSTM model (modified for dual output: return and risk prediction)
        model = Sequential()
        model.add(LSTM(40, return_sequences=True, input_shape=(10, 3)))
        model.add(Dropout(0.2))
        model.add(LSTM(40, return_sequences=False))
        model.add(Dropout(0.2))
        model.add(Dense(2))  # 2 outputs: return and risk

        model.compile(optimizer='adam', loss='mean_squared_error')
        return model

    def OnEndOfDay(self):
        # You can use this function to save models or do post-analysis
        for symbol in self.symbols:
            self.models[symbol].save(f'{symbol.Value}_lstm_model.h5')