#https://www.youtube.com/watch?v=_BHi5Y2Aow4

from collections import deque

class SmoothBrownPanda(QCAlgorithm):

    def Initialize(self):
        self.SetTimeZone("America/Toronto")
        self.SetStartDate(2021, 9, 1)
        self.SetEndDate(2021, 12, 13)
        self.SetCash(100000)
        self.spy = self.AddEquity("MSFT", Resolution.Daily)
        self.spy.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        self.efi_raw = EFI()
        self.RegisterIndicator(self.spy.Symbol, self.efi_raw, Resolution.Daily)

        self.efi = IndicatorExtensions.EMA(self.efi_raw, 12)

        self.atr_ema = deque(maxlen=2)
        self.efi.Updated += self.EfiUpdated
        
        self.sig = IndicatorExtensions.EMA(self.efi_raw, 21)

    def EfiUpdated(self, sender, updated):
        if len(self.atr_ema) == 0:
            self.atr_ema.appendleft(abs(updated.Value))
        else:
            self.atr_ema.appendleft(abs(updated.Value - self.atr_ema[0]))
        

    def OnData(self, data):
        if not self.efi.IsReady or not self.sig.IsReady:
            return

        atr_out = sum(self.atr_ema) / len(self.atr_ema)
        
        atr_high1 = self.sig.Current.Value + atr_out * 1
        atr_low1 = self.sig.Current.Value - atr_out * 1
        
        atr_high2= self.sig.Current.Value + atr_out * 2
        atr_low2 = self.sig.Current.Value - atr_out * 2
        
        atr_high3 = self.sig.Current.Value + atr_out * 3
        atr_low3 = self.sig.Current.Value - atr_out * 3

        self.Plot("EFI", "EFI", self.efi.Current.Value)
        self.Plot("EFI", "Signal", self.efi.Current.Value)
        self.Plot("EFI", "ATR", self.efi.Current.Value)
        
        self.Plot("EFI", "ATR High 1", atr_high1)
        self.Plot("EFI", "ATR Low 1", atr_high1)
        
        self.Plot("EFI", "ATR High 2", atr_high2)
        self.Plot("EFI", "ATR Low 2", atr_low2)
        
        self.Plot("EFI", "ATR High 3", atr_high3)
        self.Plot("EFI", "ATR Low 3", atr_low3)

class EFI(PythonIndicator):
    
    def __init__(self):
        self.Time = datetime.min
        self.Value = 0
        self.deque = deque(maxlen=2)

    def Update(self, input):
        self.Time = input.EndTime

        count = len(self.deque)
        if count >= 1:
            self.Value = (input.Close - self.deque[0]) * input.Volume
        
        self.deque.appendleft(input.Close)
        return (count == self.deque.maxlen)