import datetime

class SuperTrendTester(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Start Date
        self.SetEndDate(2020, 1, 12)  # End Date
        self.symbolDataBySymbol = {}
        self.SetCash(10000)  # Set Strategy Cash
        ticker = "SPY"
        
        # Profit vs Stop Loss
        self.tpPercent = 20/10000
        self.slPercent = 10/10000
        
        self.ticker = self.AddEquity(ticker, Resolution.Hour, Market.USA).Symbol
        self.superTrend = SuperTrend(self, self.ticker, 2, 14, datetime.timedelta(hours=4))
        
        # Rolling SMA window
        self.sma_high = self.SMA(ticker, 8, Resolution.Daily, Field.High)
        self.sma_low = self.SMA(ticker, 8, Resolution.Daily, Field.Low)
        
        self.window = RollingWindow[TradeBar](2)
        self.sma_high.Updated += self.SmaUpdated
        self.smaWin = RollingWindow[float](2)
        self.sma_low.Updated += self.SmaUpdated_low
        self.smaWinn = RollingWindow[float](2)
        #
        
        # create a 13 period exponential moving average
        self.fast = self.EMA(ticker, 13, Resolution.Hour);
        
        # create a 20 period simple moving average
        self.slow = self.SMA(ticker, 20, Resolution.Hour);
        
        
         # create the 4 hour data consolidator
        #fourHours = QuoteBarConsolidator(timedelta(hours=4))
        #fourHours.DataConsolidated += self.fourHoursConsolidator
        
        #self.duck1[symbol] = SimpleMovingAverage(60)
        #self.RegisterIndicator(symbol, self.duck1[symbol], fourHours)
        
        
        self.PlotIndicator("EMA", self.fast);
        self.PlotIndicator("SMA", self.slow);
        
        self.previous = None
        #
        
        
        self.trade = True

       
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.Every(timedelta(hours=1)),
                 self.buySignals)
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.Every(timedelta(hours=1)),
                 self.sellSignals)
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen("SPY"),
                 self.buySignals)
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen("SPY"),
                 self.sellSignals)
#

        # if not self.Portfolio.Invested:
        #    self.SetHoldings("SPY", 1)
        
    
        
    def SmaUpdated(self, sender, updated):
        self.smaWin.Add(self.sma_high.Current.Value)
        
    def SmaUpdated_low(self, sender, updated):
        self.smaWin.Add(self.sma_low.Current.Value)    
    
    def OnData(self, data):
        self.Plot("Chart", "SPY", data[self.ticker].Close)
        if self.superTrend.Value is None: return
        self.Plot("Chart", "SuperTrend", self.superTrend.Value)
        #self.Debug(" superTrend.Value is this: " + str(self.superTrend.Value))
        
    # Rolling Window code:
        if self.IsWarmingUp: return
        self.window.Add(data["SPY"])
        
        if not (self.window.IsReady and self.smaWin.IsReady): return
    #
    
    def tradeStart(self):
        self.trade = True
        
    def tradeEnd(self):
        self.trade = False
    
    def buySignals(self):
        #self.Debug(" Post Calc superTrend.Value is this: " + str(self.superTrend.Value))
        if self.trade == False or not self.superTrend.atrUp.IsReady:
            return
        self.Quit('buySignal')
        #for self.superTrend, symbolData in self.symbolDataBySymbol.items():
        #    self.Debug("Er ror here: " + str(self.ticker))
        if self.superTrend.Value == self.superTrend.atrUp[0] and self.fast.Current.Value > self.slow.Current.Value: return    
        self.SetHoldings(self.superTrend.symbol, 1, True, "Buy Signal") 
                
            
    def sellSignals(self):
        
        #self.Debug(" superTrend.Value is this: " + str(self.superTrend.Value))
        for self.superTrend, symbolData in self.symbolDataBySymbol.items():
            if not self.superTrend.atrDown.IsReady or self.superTrend.Value == self.superTrend.atrDown[0]: return
            #self.Liquidate(self.ticker, "Sell Signal")
            
            ## you need to assign to self.order or else variable will remain local
            self.order = self.MarketOrder(self.symbol, 1)
                
            ## you need to define fxClose or else the orders won't be properly placed
            fxClose = data[self.symbol].Close
            
            ## you need to replace sl with self.sl, as you want these to be global variables
            ## otherwise, they will not be accessible outside of OnData
            self.sl = self.StopMarketOrder(self.symbol, -1, fxClose - self.slPercent)  
            self.tp = self.LimitOrder(self.symbol, -1, fxClose + self.tpPercent)    
            


class SuperTrend:

    def __init__(self, algorithm:QCAlgorithm, symbol, multiplier:float, period:int, resolution:datetime.timedelta):
        self.algorithm = algorithm
        self.symbol = symbol
        self.multiplier = multiplier
        self.period = period
        self.resolution = resolution

        self.IsReady = False
        self.Value = None


        self.atr = None
        self.registerATR()


        self.Bars = RollingWindow[IBaseDataBar](2)
        self.atrWindow = RollingWindow[float](2)
        self.atrDown = RollingWindow[float](2)
        self.atrUp = RollingWindow[float](2)
        algorithm.Consolidate(self.symbol, self.resolution, self.OnStart)

        
    
    def OnStart(self, bar):
        self.Bars.Add(bar)

        if not self.atr.IsReady:
            self.IsReady = False
            return
        else: 
            self.atrWindow.Add(self.atr.Current.Value)

        if not self.Bars.IsReady or not self.atrWindow.IsReady:
            self.IsReady = False
            return
        else: 
            self.calculateAtrUpAtrDown()

        if not self.atrUp.IsReady or not self.atrDown.IsReady:
            self.IsReady = False
            return
        else:
            self.calculateSuperTrend()
        
    
        
        self.IsReady = True


    def registerATR(self):
        self.atr = AverageTrueRange(self.period) #self.algorithm.ATR(self.symbol, self.period)
        timeDeltaConsolidator = BaseDataConsolidator(self.resolution)
        self.algorithm.SubscriptionManager.AddConsolidator(self.symbol, timeDeltaConsolidator)
        self.algorithm.RegisterIndicator(self.symbol, self.atr, timeDeltaConsolidator)

    def calculateAtrUpAtrDown(self):
        self.algorithm.Quit('calculateAtrUpAtrDown')
        hltwo = (self.Bars[0].High + self.Bars[0].Low) / 2
        hltwoPrev = (self.Bars[1].High + self.Bars[1].Low) / 2

        downNow = hltwo - self.multiplier * self.atrWindow[0]
        downPrev = hltwoPrev - self.multiplier * self.atrWindow[1]
        atrDown = max(downPrev, downNow) if self.Bars[1].Close > downPrev else downNow
        self.atrDown.Add(atrDown)

        upNow = hltwo + self.multiplier * self.atrWindow[0]
        upPrev = hltwoPrev + self.multiplier * self.atrWindow[1]
        atrUp = min(upNow, upPrev) if self.Bars[1].Close < upPrev else upNow
        self.atrUp.Add(atrUp)
        
    def calculateSuperTrend(self):
        #self.Debug(" superTrend.Value is this: " + str(self.superTrend.Value))
        if self.Bars[0].Close > self.atrUp[1]:
            self.Value = self.atrDown[0]
        elif self.Bars[0].Close < self.atrDown[1]:
            self.Value = self.atrUp[0]
        #else:
        #    self.Debug(" superTrend.Value is this: " + str(self.superTrend.Value))
        #    pass