from AlgorithmImports import *
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from ta.momentum import RSIIndicator
from ta.volatility import BollingerBands
from ta.trend import MACD

class TSLAPredictionStrategy(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2022, 1, 1)
        self.SetEndDate(2023, 12, 31)
        self.SetCash(100000)
        
        # Add TSLA equity
        self.tsla = self.AddEquity("TSLA", Resolution.Daily).Symbol
        
        # Initialize indicators
        self.sma = self.SMA(self.tsla, 10, Resolution.Daily)
        self.rsi = self.RSI(self.tsla, 14, MovingAverageType.Wilders, Resolution.Daily)
        self.momentum = self.MOM(self.tsla, 10, Resolution.Daily)
        self.volatility = self.STD(self.tsla, 10, Resolution.Daily)
        self.macd = self.MACD(self.tsla, 12, 26, 9, MovingAverageType.Wilders, Resolution.Daily)
        self.bollinger = self.BB(self.tsla, 20, 2, MovingAverageType.Simple, Resolution.Daily)
        
        # Schedule trades on a weekly basis
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Friday), self.TimeRules.AfterMarketOpen(self.tsla, 30), self.PredictAndTrade)
        
        # Initialize model
        self.model = RandomForestClassifier(n_estimators=100, random_state=42)
        self.isModelTrained = False
        
        # Historical data for model training
        self.history = self.History([self.tsla], 365 * 10, Resolution.Daily)
        self.df = self.history.loc[self.tsla]
        
    def OnData(self, slice: Slice):
        pass
    
    def PredictAndTrade(self):
        if not self.isModelTrained:
            self.TrainModel()
        
        # Get latest data for prediction
        features = self.GetFeatures(self.Time)
        
        if features is not None:
            prediction = self.model.predict([features])[0]
            self.Debug(f"Prediction for {self.Time}: {'Bullish' if prediction == 1 else 'Bearish'}")
            
            if prediction == 1:
                self.SetHoldings(self.tsla, 1)
                self.Debug(f"Going long on TSLA at {self.Time}")
            else:
                self.SetHoldings(self.tsla, -1)
                self.Debug(f"Going short on TSLA at {self.Time}")
    
    def TrainModel(self):
        df = self.df.copy()
        
        # Calculate technical indicators
        df['SMA10'] = df['close'].rolling(window=10).mean()
        df['Return'] = df['close'].pct_change()
        df['Volatility'] = df['close'].rolling(window=10).std()
        df['Momentum'] = df['close'].pct_change(10)
        df['RSI'] = RSIIndicator(df['close'], window=14).rsi()
        macd = MACD(df['close'], window_slow=26, window_fast=12, window_sign=9)
        df['MACD'] = macd.macd()
        df['MACD_Signal'] = macd.macd_signal()
        bb = BollingerBands(df['close'], window=20, window_dev=2)
        df['BB_High'] = bb.bollinger_hband()
        df['BB_Low'] = bb.bollinger_lband()
        
        # Shift the indicators to align with the prediction target
        df['SMA10'] = df['SMA10'].shift(-5)
        df['Return'] = df['Return'].shift(-5)
        df['Volatility'] = df['Volatility'].shift(-5)
        df['Momentum'] = df['Momentum'].shift(-5)
        df['RSI'] = df['RSI'].shift(-5)
        df['MACD'] = df['MACD'].shift(-5)
        df['MACD_Signal'] = df['MACD_Signal'].shift(-5)
        df['BB_High'] = df['BB_High'].shift(-5)
        df['BB_Low'] = df['BB_Low'].shift(-5)
        
        # Determine the actual next week's direction
        df['NextWeekChange'] = df['close'].pct_change(5).shift(-5)
        df['Actual'] = np.where(df['NextWeekChange'] > 0, 1, -1)
        
        # Drop NaN values
        df.dropna(inplace=True)
        
        # Prepare the features and target
        features = ['SMA10', 'Return', 'Volatility', 'Momentum', 'RSI', 'MACD', 'MACD_Signal', 'BB_High', 'BB_Low']
        X = df[features]
        y = df['Actual']
        
        # Train the Random Forest Classifier
        self.model.fit(X, y)
        self.isModelTrained = True
        self.Debug("Model trained successfully.")
    
    def GetFeatures(self, time):
        data = self.History([self.tsla], 20, Resolution.Daily).loc[self.tsla]
        
        if len(data) < 20:
            return None
        
        sma10 = data['close'].rolling(window=10).mean().iloc[-1]
        ret = data['close'].pct_change().iloc[-1]
        vol = data['close'].rolling(window=10).std().iloc[-1]
        mom = data['close'].pct_change(10).iloc[-1]
        rsi = RSIIndicator(data['close'], window=14).rsi().iloc[-1]
        macd = MACD(data['close'], window_slow=26, window_fast=12, window_sign=9)
        macd_val = macd.macd().iloc[-1]
        macd_signal = macd.macd_signal().iloc[-1]
        bb = BollingerBands(data['close'], window=20, window_dev=2)
        bb_high = bb.bollinger_hband().iloc[-1]
        bb_low = bb.bollinger_lband().iloc[-1]
        
        return [sma10, ret, vol, mom, rsi, macd_val, macd_signal, bb_high, bb_low]