from collections import deque
from datetime import datetime, timedelta

class ZigZagDev(QCAlgorithm):
    
    def Initialize(self):
        self.myResolution1 = Resolution.Hour
        self.myResolution2 = Resolution.Daily
        
        self.zzPriod1 = 200
        self.zzPriod2 = 200
        self.ticker ="SPY"
        
        self.SetStartDate(2019,8,1)
        #self.SetEndDate(2019,8,31)
        self.SetEndDate(datetime.now())
        self.SetWarmUp(TimeSpan.FromDays(self.zzPriod2))
        
        self.symbol = self.AddEquity(self.ticker, self.myResolution1).Symbol
        
        self.atr1 = AverageTrueRange(15)
        self.RegisterIndicator(self.symbol, self.atr1, self.myResolution1)
        self.ZZ1 = ZigZag(self, 'ZZ1', self.zzPriod1, 5, 1, 0.05, self.atr1)
        self.RegisterIndicator(self.symbol, self.ZZ1, self.myResolution1)
        
        self.atr2 = AverageTrueRange(15)
        self.RegisterIndicator(self.symbol, self.atr2, self.myResolution2)
        self.ZZ2 = ZigZag(self, 'ZZ2', self.zzPriod2, 5, 2, 5, self.atr2)
        self.RegisterIndicator(self.symbol, self.ZZ2, self.myResolution2)
        
        self.rsi1 = RelativeStrengthIndex(35)
        self.RegisterIndicator(self.ticker, self.rsi1, self.myResolution1)
        self.rsi1_zz = ZigZag(self, 'rsi1_zz', self.zzPriod1, 5, 1, 0.1, self.atr1)
        self.rsi1.Updated += self.rsi1_zz.IndicatorUpdate

    def OnData(self, data):
        if self.IsWarmingUp: return
        if not self.Portfolio.Invested:
            self.SetHoldings(self.ticker, 1)

    def OnEndOfAlgorithm(self):
        if False: self.ZZ1.ListZZ()
        if True: self.ZZ1.listZZPoints()
        if False: self.ZZ2.ListZZ()
        if False: self.ZZ2.listZZPoints()
        if False: self.rsi1_zz.ListZZ()
        if False: self.rsi1_zz.listZZPoints()
        
    def MyDebug(self, debugString):
        message = str(self.Time) + "  " + str(self.symbol) + ": "+ debugString
        self.Debug(message)
 
class ZigZag():
    def __init__(self, algo, name, period, lookback=10, thresholdType=1, threshold=0.05, atr=None):
        self.algo = algo
        self.name = name
        self.period = period
        self.lookback = lookback
        self.thresholdType = thresholdType
        self.threshold = threshold
        self.currentThreshold = 0
        self.atr = None
        if atr != None: self.atr = atr  
        self.Time = datetime.min
        self.Value = 0
        self.IsReady = False

        self.lastLow = None
        self.lastHigh = None
        self.shortTrendCount = 0
        self.longTrendCount = 0
        self.trendDir = None
        self.trendChange = False
        self.bars = deque(maxlen=period)
        self.zzLow = deque(maxlen=period)
        self.shortTrendCount = deque(maxlen=period)
        self.zzHigh = deque(maxlen=period)
        self.longTrendCount = deque(maxlen=period)
        self.zigzag = deque(maxlen=period)
        self.barCount = 0
        self.count = 0
        
        self.zzPoints = []
        
    def __str__(self):
        return self.name
    
    def __repr__(self):
        return "Zig Zag Name:{}, IsReady:{}, Time:{}, Value:{}".format(self.Name, self.IsReady, self.Time, self.Value)

    # Update method is mandatory
    def IndicatorUpdate(self, caller, updated):
        bar = TradeBar()
        value = updated.Value
        bar.Open = value
        bar.High = value
        bar.Low = value
        bar.Close = value
        bar.EndTime = updated.EndTime
        self.Update(bar)
        
    # Update method is mandatory
    def Update(self, input):
        bar = input
        self.count +=1
        self.bars.appendleft(bar)
        self.barCount = len(self.bars)
        self.Time = bar.EndTime
        self.Value = self.count
        self.IsReady = self.barCount == self.bars.maxlen 
        
        #update currentThreshold value
        if self.thresholdType == 1:
            self.currentThreshold = self.bars[0].Close * self.threshold
        elif self.thresholdType == 2:
            self.currentThreshold = self.atr.Current.Value * self.threshold
        
        if self.barCount < self.lookback:
            #Not enough bars yet to calculte  
            self.zzLow.appendleft(0)
            self.shortTrendCount.appendleft(0)
            self.zzHigh.appendleft(0)
            self.longTrendCount.appendleft(0)
            self.zigzag.appendleft(0)
            return self.IsReady
        
        low, thisisLow , retraceLow = self.localLow()
        high, thisisHigh, retraceHigh = self.localHigh()
        self.lastLow = low
        self.lastHigh = high
        #previousTrendDir = self.trendDir
        self.trendChange = False
        
        if thisisLow or retraceLow:
            if self.trendDir == 1:
                self.trendChange = True
            self.trendDir = -1
        if thisisHigh or retraceHigh:
            if self.trendDir == -1:
                self.trendChange = True
            self.trendDir = 1
        
        #Update ZZ Low
        if thisisLow or retraceLow:
            self.zzLow.appendleft(low)
        else: self.zzLow.appendleft(0)
        #Update Short TrendCount
        shortTrendCount = self.shortTrendCount[0]
        if self.trendChange and self.trendDir == -1:
            shortTrendCount-=1
        self.shortTrendCount.appendleft(shortTrendCount)

        #Update ZZ High
        if thisisHigh or retraceHigh:
            self.zzHigh.appendleft(high)
        else: self.zzHigh.appendleft(0)
        #Update Long TrendCount
        longTrendCount = self.longTrendCount[0]
        if self.trendChange and self.trendDir == 1:
            longTrendCount+=1
        self.longTrendCount.appendleft(longTrendCount)
        
        #Update ZZ
        if self.zzLow[0]!=0:
            self.zigzag.appendleft(self.zzLow[0])
        else:
            self.zigzag.appendleft(self.zzHigh[0])
        self.Value = self.zigzag[0]
        
        #Erase precious values if any in this Short Trend
        i=1
        while i < len(self.bars)-1 and self.zzLow[0] !=0 and self.shortTrendCount[i] == shortTrendCount:
            if self.zzLow[i] == self.zigzag[i]:
                self.zigzag[i] = 0 
            self.zzLow[i] = 0
            i+=1        
        
        #Erase precious values if any in this Long Trend
        i=1
        while i < len(self.bars)-1 and self.zzHigh[0] != 0 and self.longTrendCount[i] == longTrendCount:
            if self.zzHigh[i] == self.zigzag[i]:
                self.zigzag[i] = 0            
            self.zzHigh[i] = 0
            i+=1              
        
        if False and self.IsReady :self.algo.MyDebug("Low:{}/{}/{}/{}/{}, High:{}/{}/{}/{}/{}, ZZ:{}/{}".format( \
                str(low),str(thisisLow),str(retraceLow),str(self.zzLow[0]),str(self.shortTrendCount[0]), \
                str(high),str(thisisHigh),str(retraceHigh),str(self.zzHigh[0]),str(self.longTrendCount[0]), \
                str(self.zigzag[0]),str(self.trendDir)))

        #Update zzPoints List
        del self.zzPoints[:]
        for i in range(0, len(self.zigzag)-1):
            if self.zigzag[i]!=0:
                self.zzPoints.append(ZZPoint(self,i))
        return self.IsReady
        
    #Local Low or Retracement
    def localLow(self):
        thisisLow = False
        retrace = False
        low = self.bars[0].Low
        for i in range(1, self.lookback):
            if self.bars[i].Low < low: low = self.bars[i].Low
        thisisLow = self.bars[0].Low == low
        
        #trendDir is not updated yet so it is the previous value
        if self.lastHigh != None and self.trendDir == 1 and (self.lastHigh - self.bars[0].High)>self.currentThreshold:
            retrace = True
            low = self.bars[0].Low 
        return low, thisisLow, retrace
    
    #Local High or Retracement
    def localHigh(self):
        thisisHigh = False
        retrace = False
        high = self.bars[0].High
        for i in range(1, self.lookback):
            if self.bars[i].High > high: high = self.bars[i].High
        thisisHigh = self.bars[0].High == high
        
        #trendDir is not updated yet so it is the previous value
        if self.lastLow != None and self.trendDir == -1 and (self.bars[0].Low - self.lastLow)>self.currentThreshold:
            retrace = True
            high = self.bars[0].High
        return high, thisisHigh, retrace 
        
    def ListZZ(self):
        printItems = min(70, len(self.bars)-1)
        i=printItems
        while i >= 0:
            if self.zzLow[i]==self.zigzag[i] and self.zigzag[i]!=0:
                trendCount = self.shortTrendCount[i]
            elif self.zzHigh[i]==self.zigzag[i] and self.zigzag[i]!=0:
                trendCount = self.longTrendCount[i]
            else:
                trendCount = ""
            self.algo.MyDebug("Date:{}, Close:{}, zzLow:{}/{}, zzHigh:{}/{}, ZZ:{}/{}".format( \
                str(self.bars[i].EndTime), str(self.bars[i].Close), \
                str(self.zzLow[i]),str(self.shortTrendCount[i]), \
                str(self.zzHigh[i]),str(self.longTrendCount[i]), \
                str(self.zigzag[i]),str(trendCount)))
            i-=1
        return
    
    def listZZPoints(self):
        i=1
        for level in self.zzPoints:
            self.algo.MyDebug("{}/{}. Value:{}, EndTime:{}, Index:{}, TrendCount:{}".format( \
                str(level.zz), str(i), \
                str(level.value), str(level.endTime), str(level.index), str(level.trendCount),))
            i+=1
        return    
    
class ZZPoint():
    def __init__(self, zz, index):
        self.zz = zz
        self.index = index
        self.endTime = zz.bars[index].EndTime
        self.value = zz.zigzag[index]
        if self.value == zz.zzLow[index]:
            self.trendCount = zz.shortTrendCount[index]
        else:
            self.trendCount = zz.longTrendCount[index]
        return