# Four ways to get the normalized ATR 

import numpy as np
import pandas as pd

# -----------------------------
STOCK = "SPY"; ATR_PERIOD = 14;
# -----------------------------

class MuscularGreenAlbatross(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 9, 1)  
        self.SetEndDate(2021, 10, 7) 
        self.stock = self.AddEquity(STOCK, Resolution.Daily).Symbol
        self.SetWarmUp(ATR_PERIOD + 1, Resolution.Daily)
        self.atr = self.ATR(self.stock, ATR_PERIOD, MovingAverageType.Simple, Resolution.Daily)
        self.close = self.Identity(self.stock, Resolution.Daily, Field.Close)
        self.yest_close =  IndicatorExtensions.Of(Delay(1), self.close)
        self.atr_norm_2 = IndicatorExtensions.Over(self.atr, self.yest_close)
        self.atr_norm_3 = symbolData(self, self.stock, ATR_PERIOD).atrOverPrice
        
        
    def OnData(self, data):
        if self.IsWarmingUp or not self.atr.IsReady or not self.atr_norm_2.IsReady: return

        atr_norm_1 = float(self.atr.Current.Value / self.yest_close.Current.Value)
        atr_norm_2 = float(self.atr_norm_2.Current.Value)
        atr_norm_3 = float(self.atr_norm_3.Current.Value)
        atr_norm_4 = float(self.atr_norm(self.stock, ATR_PERIOD))

        self.Plot("ATR", "atr_norm_1", atr_norm_1)
        self.Plot("ATR", "atr_norm_2", atr_norm_2)
        self.Plot("ATR", "atr_norm_3", atr_norm_3)
        self.Plot("ATR", "atr_norm_4", atr_norm_4)
        
        
    def atr_norm(self, symbol, atr_period):
        hist = self.History(symbol, atr_period + 1)
        range = hist.high - hist.low
        h_minus_p_close = np.abs(hist.high - hist.close.shift(1))
        l_minus_p_close = np.abs(hist.low - hist.close.shift(1))
        tr = [max(x,y,z) for x,y,z in zip(range, h_minus_p_close, l_minus_p_close)][1:]
        atr_norm = pd.Series(tr).rolling(atr_period).mean() / hist.close[-2]
    
        return atr_norm.values[-1]
        
        
class symbolData:
    
    def __init__(self, algo, symbol, atr_period):
        self.atr = algo.ATR(symbol, atr_period, MovingAverageType.Simple, Resolution.Daily)
        self.id = algo.Identity(symbol, Resolution.Daily, Field.Close)
        self.atrOverPrice = IndicatorExtensions.Over(self.atr, self.id)

        history = algo.History(symbol, atr_period + 1, Resolution.Daily)
        for bar in history.itertuples():
            tradeBar = TradeBar(bar.Index[1], bar.Index[0], bar.open, bar.high, 
                                bar.low, bar.close, bar.volume)
            self.atr.Update(tradeBar)

        history = algo.History(symbol, 2, Resolution.Daily).iloc[-2]
        self.id.Update(pd.to_datetime(history.name[1]), history.close)