import datetime

class SuperTrendTester(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 10)  # Start Date
        self.SetEndDate(2020, 1, 10)  # End Date
        
        self.SetCash(100000)  # Set Strategy Cash
        self.spy = self.AddEquity("SPY", Resolution.Minute, Market.USA).Symbol
        self.superTrend = SuperTrend(self, self.spy, 2, 14, datetime.timedelta(minutes=15))


    def OnData(self, data):
        self.Plot("Chart", "SPY", data[self.spy].Close)
        if self.superTrend.Value is None: return
        self.Plot("Chart", "SuperTrend", self.superTrend.Value)

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 and not self.atrWindow.IsReady:
            self.IsReady = False
            return
        else: 
            self.calculateAtrUpAtrDown()

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


    def registerATR(self):
        self.atr = 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):
        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):
        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:
            pass