#region imports
import pytz
from AlgorithmImports import *
from datetime import time
#endregion
class SimpleNQExample(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 1, 1)  # Set Start Date
        self.SetEndDate(2023, 4, 18)  
        self.SetCash(1000000)  # Set Strategy Cash
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.SetTimeZone("America/New_York")

        # Set up the future we're looking at
        future = Futures.Indices.NASDAQ100EMini
        self._continuousContract = self.AddFuture(future,
                                                  dataNormalizationMode = DataNormalizationMode.BackwardsRatio,
                                                  dataMappingMode = DataMappingMode.OpenInterest,
                                                  contractDepthOffset = 0,
                                                  extendedMarketHours = True)
        self.commod = self._continuousContract.Symbol
        
        self._continuousContract.SetFilter(timedelta(days=1), timedelta(days=60))


        # Set up a short-period EMA on 1-minute bar
        qqqequityconsolidator = TradeBarConsolidator(timedelta(minutes = 1))
        qqqequityconsolidator.DataConsolidated += self.qqqequityHandler        
        self.qqqema = ExponentialMovingAverage(9)
        self.RegisterIndicator(self.commod, self.qqqema, qqqequityconsolidator)
        self.qqqema_past = IndicatorExtensions.Of(Delay(1), self.qqqema)
        self.qqqema_slope = IndicatorExtensions.Minus(self.qqqema, self.qqqema_past)

        # Set up a long-period EMA on 1-minute bar      
        self.qqqemalong = ExponentialMovingAverage(150)
        self.RegisterIndicator(self.commod, self.qqqemalong, qqqequityconsolidator)
        self.qqqemalong_past = IndicatorExtensions.Of(Delay(1), self.qqqemalong)
        self.qqqemalong_slope = IndicatorExtensions.Minus(self.qqqemalong, self.qqqemalong_past)

        # Set up a 14-period RSI on 1-minute bar
        self.qqqrsi = RelativeStrengthIndex(14)
        self.RegisterIndicator(self.commod, self.qqqrsi, qqqequityconsolidator)
        self.sma_qqqrsi = IndicatorExtensions.SMA(self.qqqrsi, 5) 
        self.sma_qqqrsi_past = IndicatorExtensions.Of(Delay(1), self.sma_qqqrsi)
        self.qqqrsi_slope = IndicatorExtensions.Minus(self.sma_qqqrsi, self.sma_qqqrsi_past)
        
        # Set up a Momentum indicator on 1-minute bar
        self.qqqmomentum = Momentum(9)
        self.RegisterIndicator(self.commod, self.qqqmomentum, qqqequityconsolidator)
        self.ema_qqqmom = IndicatorExtensions.EMA(self.qqqmomentum, 3) 
        self.ema_qqqmom_past = IndicatorExtensions.Of(Delay(1), self.ema_qqqmom)
        self.qqqmom_slope = IndicatorExtensions.Minus(self.ema_qqqmom, self.ema_qqqmom_past)

        # Set up an Ultimate Oscillator indicator on 1-minute bar
        self.qqqUA = UltimateOscillator(7, 14, 28)
        self.RegisterIndicator(self.commod, self.qqqUA, qqqequityconsolidator)
        self.ema_qqqUA = IndicatorExtensions.EMA(self.qqqUA, 3) 
        self.ema_qqqUA_past = IndicatorExtensions.Of(Delay(1), self.ema_qqqUA)
        self.qqqUA_slope = IndicatorExtensions.Minus(self.ema_qqqUA, self.ema_qqqUA_past)

        self.qqqequityslopeshort = dict()
        self.qqqequityslopeshort[0] = 0

        self.emaslopelong = dict()
        self.emaslopelong[0] = 0

        self.qqqslope_rsi = dict()
        self.qqqslope_rsi[0] = 0

        self.qqqslope_mom = dict()
        self.qqqslope_mom[0] = 0
       
        self.qqqslope_UA = dict()
        self.qqqslope_UA[0] = 0

        self.tradingnumber = dict()
        self.tradingnumber[0] = 0 

    def emaUpdated(self, sender, updated):
        '''Adds updated values to rolling window'''
        self.emaWin.Add(updated)    

    def qqqequityHandler(self, sender, consolidated):
        if self.IsWarmingUp or not self.qqqema_slope.IsReady: return 
        
        self.qqqequityslopeshort[0] = self.qqqema_slope.Current.Value
        self.emaslopelong[0] = self.qqqemalong_slope.Current.Value
        self.qqqslope_rsi[0] = self.qqqrsi_slope.Current.Value
        self.qqqslope_mom[0] = self.qqqmom_slope.Current.Value
        self.qqqslope_UA[0] = self.qqqUA_slope.Current.Value
        
    def OnData(self, data: Slice): 
        if not data.ContainsKey(self.commod):
            return        
        
        futures_invested = [holding.Symbol for holding in self.Portfolio.Values if holding.Type == SecurityType.Future and holding.Invested]
        futures_invested_short = [(symbol, holding.Quantity) for symbol, holding in self.Portfolio.items() if holding.Type == SecurityType.Future and holding.Symbol.SecurityType == SecurityType.Future and holding.IsShort]
        futures_invested_long = [(symbol, holding.Quantity) for symbol, holding in self.Portfolio.items() if holding.Type == SecurityType.Future and holding.Symbol.SecurityType == SecurityType.Future and holding.IsLong]
        futures_invested_long1 = [symbol for symbol, holding in self.Portfolio.items() if holding.Type == SecurityType.Future and holding.Symbol.SecurityType == SecurityType.Future and holding.IsLong]
        futures_invested_short1 = [symbol for symbol, holding in self.Portfolio.items() if holding.Type == SecurityType.Future and holding.Symbol.SecurityType == SecurityType.Future and holding.IsShort]

        slopelong = self.emaslopelong[0]
        slopeshort = self.qqqequityslopeshort[0]
        slope_rsi = self.qqqslope_rsi[0] 
        momslope = self.qqqslope_mom[0]
        UAslope = self.qqqslope_UA[0]
        RSI = self.sma_qqqrsi.Current.Value
        

        current_time = self.Time.time()

        trading_start = time(8, 0)
        trading_end = time(19, 30)
        #trading_allowed = trading_start <= current_time <= trading_end
        trading_not_allowed = current_time > trading_end
        trading_allowed = trading_start <= current_time <= trading_end and self.Time.weekday() != 6
        
        if self.tradingnumber[0] == 0 and RSI < 30:
            self.tradingnumber[0] = 1

        if self.tradingnumber[0] == 0 and RSI > 70:
            self.tradingnumber[0] = 3


        if trading_allowed and self.tradingnumber[0] == 1:
            if slopeshort > 0 and slope_rsi > 0 and not futures_invested_short1:
                    self.MarketOrder(self._continuousContract.Mapped, -1)
                    self.Log("Going long")
                    self.tradingnumber[0] = 2

        if self.tradingnumber[0] == 2 and slopeshort < 0 and slope_rsi < 0:
            for symbol, quantity in futures_invested_short:
                self.MarketOrder(symbol, -quantity)
                self.tradingnumber[0] = 0


        if trading_allowed and self.tradingnumber[0] == 3:
            if slopeshort < 0 and slope_rsi < 0 and not futures_invested_long1:
                self.MarketOrder(self._continuousContract.Mapped, 1)
                self.Log("Going short")
                self.tradingnumber[0] = 4
        
        if self.tradingnumber[0] == 4 and slopeshort > 0 and slope_rsi > 0:
            for symbol, quantity in futures_invested_long:
                self.MarketOrder(symbol, -quantity)  # Liquidate the short position
                self.tradingnumber[0] = 0

        
        friday_close_time = time(16, 58)
        if self.Time.weekday() == 4 and current_time == friday_close_time:
            for symbol in futures_invested:
                self.MarketOrder(symbol, -self.Portfolio[symbol].Quantity)  # Liquidate the position
                self.tradingnumber[0] = 0

        #self.Log("10 EMA slope: {0}, 12RSI5SMA Slope: {1}, Histo: {2}".format(round(slope, 2), round(slope_rsi, 2), round(slope_macd_histogram, 2)))