#region imports
from AlgorithmImports import *
#endregion
import numpy as np
import pandas as pd
import xgboost as xgb

class SMACrossover(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2005,1,1)
        self.SetEndDate(2022,1,1)
        self.SetCash(100000)

        # Add SPY and QQQ as securities
        self.spy = self.AddEquity("SPY", Resolution.Daily)
        self.qqq = self.AddEquity("QQQ", Resolution.Daily)

        # Set up XGBoost model
        self.model = xgb.XGBClassifier()

        # Set up 20 and 50-day moving averages
        self.spy_sma20 = self.SMA("SPY", 20, Resolution.Daily)
        self.spy_sma50 = self.SMA("SPY", 50, Resolution.Daily)

        self.qqq_sma20 = self.SMA("QQQ", 20, Resolution.Daily)
        self.qqq_sma50 = self.SMA("QQQ", 50, Resolution.Daily)

        # Set up scheduled function to run every day
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.Every(timedelta(hours=2)),
                 self.Trade)
        #self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 30), self.Trade)

    def Trade(self):
        # Get historical prices for SPY and QQQ
        history = self.History([self.spy.Symbol, self.qqq.Symbol], 50, Resolution.Daily)

        # Compute 20 and 50-day moving averages
        spy_sma20 = history.loc[self.spy.Symbol]['close'].rolling(window=20).mean()
        spy_sma50 = history.loc[self.spy.Symbol]['close'].rolling(window=50).mean()

        qqq_sma20 = history.loc[self.qqq.Symbol]['close'].rolling(window=20).mean()
        qqq_sma50 = history.loc[self.qqq.Symbol]['close'].rolling(window=50).mean()

        # Merge dataframes
        data = pd.concat([spy_sma20, spy_sma50, qqq_sma20, qqq_sma50], axis=1).dropna()

        # Prepare data for XGBoost model
        X = data.values
        y = np.where(data.iloc[:,0] > data.iloc[:,1], 1, 0)

        # Undersample the majority class
        from imblearn.under_sampling import RandomUnderSampler
        undersampler = RandomUnderSampler(sampling_strategy='majority')
        X_resampled, y_resampled = undersampler.fit_resample(X, y)

        # Train XGBoost model
        self.model.fit(X_resampled, y_resampled)

        # Predict using XGBoost model
        last_data = X[-1,:].reshape(1,-1)
        pred = self.model.predict(last_data)[0]

        # Check if SPY and QQQ are already invested
        spy_holdings = self.Portfolio[self.spy.Symbol].Quantity
        qqq_holdings = self.Portfolio[self.qqq.Symbol].Quantity

        # Buy or sell SPY and QQQ based on XGBoost prediction
        if pred == 1:
            if spy_holdings <= 0:
                self.SetHoldings(self.spy.Symbol, 0.5)
            if qqq_holdings <= 0:
                self.SetHoldings(self.qqq.Symbol, 0.5)
        elif pred == 0:
            if spy_holdings > 0:
                self.Liquidate(self.spy.Symbol)
            if qqq_holdings > 0:
                self.Liquidate(self.qqq.Symbol)