import datetime

class SuperTrendTester(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 10)  # Start Date
        self.SetEndDate(2020, 1, 18)  # End Date
        self.symbolDataBySymbol = {}
        self.SetCash(100000)  # Set Strategy Cash
        self.ticker = self.AddEquity("SPY", Resolution.Minute, Market.USA).Symbol
        self.superTrend = SuperTrend(self, self.ticker, 2, 14, datetime.timedelta(minutes=15))
        self.Debug(" superTrend.Value is this: " + str(self.superTrend.Value))
        self.Debug("self.superTrend.atrUp is this: " + str(self.superTrend))
        self.trade = True

        
    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)
    
    
    
    def buySignals(self):
        self.Debug(" Post Calc superTrend.Value is this: " + str(self.superTrend.Value))
        if self.trade == False:
            return
        for self.superTrend, symbolData in self.symbolDataBySymbol.items():
            self.Debug("Er ror here: " + str(self.ticker))
        if self.superTrend.Value is upNow: return    
        self.SetHoldings(symbol, .10, False, "Buy Signal")
                
            
    def sellSignals(self):
        for self.superTrend, symbolData in self.symbolDataBySymbol.items():
            if self.superTrend.Value is downNow: return
            self.Liquidate(self.ticker, "Sell Signal")
        

        # if not self.Portfolio.Invested:
        #    self.SetHoldings("SPY", 1)

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