import pandas as pd

def get_supertrend(high:pd.Series, low:pd.Series, close:pd.Series, lookback:int, multiplier:int):
    
    # ATR
    
    tr1 = pd.DataFrame(high - low)
    tr2 = pd.DataFrame(abs(high - close.shift(1)))
    tr3 = pd.DataFrame(abs(low - close.shift(1)))
    frames = [tr1, tr2, tr3]
    tr = pd.concat(frames, axis = 1, join = 'inner').max(axis = 1)
    atr = tr.ewm(lookback).mean()
    
    # H/L AVG AND BASIC UPPER & LOWER BAND
    
    hl_avg = (high + low) / 2
    upper_band = (hl_avg + multiplier * atr).dropna()
    lower_band = (hl_avg - multiplier * atr).dropna()
    
    # FINAL UPPER BAND
    
    final_bands = pd.DataFrame(columns = ['upper', 'lower'])
    final_bands.iloc[:,0] = [x for x in upper_band - upper_band]
    final_bands.iloc[:,1] = final_bands.iloc[:,0]
    
    for i in range(len(final_bands)):
        if i == 0:
            final_bands.iloc[i,0] = 0
        else:
            if (upper_band[i] < final_bands.iloc[i-1,0]) | (close[i-1] > final_bands.iloc[i-1,0]):
                final_bands.iloc[i,0] = upper_band[i]
            else:
                final_bands.iloc[i,0] = final_bands.iloc[i-1,0]
    
    # FINAL LOWER BAND
    
    for i in range(len(final_bands)):
        if i == 0:
            final_bands.iloc[i, 1] = 0
        else:
            if (lower_band[i] > final_bands.iloc[i-1,1]) | (close[i-1] < final_bands.iloc[i-1,1]):
                final_bands.iloc[i,1] = lower_band[i]
            else:
                final_bands.iloc[i,1] = final_bands.iloc[i-1,1]
    
    # SUPERTREND
    
    supertrend = pd.DataFrame(columns = [f'supertrend_{lookback}'])
    supertrend.iloc[:,0] = [x for x in final_bands['upper'] - final_bands['upper']]
    
    for i in range(len(supertrend)):
        if i == 0:
            supertrend.iloc[i, 0] = 0
        elif supertrend.iloc[i-1, 0] == final_bands.iloc[i-1, 0] and close[i] < final_bands.iloc[i, 0]:
            supertrend.iloc[i, 0] = final_bands.iloc[i, 0]
        elif supertrend.iloc[i-1, 0] == final_bands.iloc[i-1, 0] and close[i] > final_bands.iloc[i, 0]:
            supertrend.iloc[i, 0] = final_bands.iloc[i, 1]
        elif supertrend.iloc[i-1, 0] == final_bands.iloc[i-1, 1] and close[i] > final_bands.iloc[i, 1]:
            supertrend.iloc[i, 0] = final_bands.iloc[i, 1]
        elif supertrend.iloc[i-1, 0] == final_bands.iloc[i-1, 1] and close[i] < final_bands.iloc[i, 1]:
            supertrend.iloc[i, 0] = final_bands.iloc[i, 0]
    
    supertrend = supertrend.set_index(upper_band.index)
    supertrend = supertrend.dropna()[1:]
    
    # ST UPTREND/DOWNTREND
    
    upt = []
    dt = []
    close = close.iloc[len(close) - len(supertrend):]

    for i in range(len(supertrend)):
        if close[i] > supertrend.iloc[i, 0]:
            upt.append(supertrend.iloc[i, 0])
            dt.append(np.nan)
        elif close[i] < supertrend.iloc[i, 0]:
            upt.append(np.nan)
            dt.append(supertrend.iloc[i, 0])
        else:
            upt.append(np.nan)
            dt.append(np.nan)
            
    st, upt, dt = pd.Series(supertrend.iloc[:, 0]), pd.Series(upt), pd.Series(dt)
    upt.index, dt.index = supertrend.index, supertrend.index
    
    return st, upt, dt


# from AlgorithmImports import *
from collections import deque
from datetime import datetime 
import math

class BasicNerdEquityAlgorithm(QCAlgorithm):

    def Initialize(self):
        superTrendPeriod = 10
        superTrendMultiple = 3

        self.SetCash(100000)
        self.SetStartDate(2018, 4, 1)

        self.sym = self.AddEquity("SPY", Resolution.Minute).Symbol

        # define custom indicator
        self.supertrend = MySuperTrend(self, superTrendPeriod, superTrendMultiple)
        self.RegisterIndicator(self.sym, self.supertrend, Resolution.Minute)
        self.SetWarmUp(superTrendPeriod)

    def OnData(self, data):
        if not self.supertrend.IsReady: return
        
        self.Plot('Custom', 'spy', data[self.sym].Close)
        self.Plot('Custom', 'Supertrend', self.supertrend.Value)
        if self.supertrend.Value < data[self.sym].Close:
            self.SetHoldings(self.sym, 1)
        else:
            self.SetHoldings(self.sym, 0)
        
class MySuperTrend:
    def __init__(self, algorithm, period, multiple, movingAverageType=MovingAverageType.Simple):
        
        self.Name = "Custom Indicator"
        self.Time = datetime.min
        self.Value = 0
        
        self.multiplier = multiple
        self.atr = AverageTrueRange(period, movingAverageType)
        self.values = deque(maxlen=period)

        self.previousTrailingLowerBand = 0
        self.previousTrailingUpperBand = 0
        self.previousClose = 0
        self.previousTrend = 0

        
    def __repr__(self):
        return "{0} -> IsReady: {1}. Time: {2}. Value: {3}".format(self.Name, self.IsReady, self.Time, self.Value)

    def Update(self, input:TradeBar):
        self.Time = input.EndTime
        self.atr.Update(input)
        superTrend = 0

        currentClose = input.Close
        currentBasicLowerBand = (input.Low + input.High) / 2 - self.multiplier * self.atr.Current.Value
        currentBasicUpperBand = (input.Low + input.High) / 2 + self.multiplier * self.atr.Current.Value

        if self.previousClose > self.previousTrailingLowerBand:
            currentTrailingLowerBand = max(currentBasicLowerBand, self.previousTrailingLowerBand)
        else:
            currentTrailingLowerBand = currentBasicLowerBand

        if self.previousClose < self.previousTrailingUpperBand:
            currentTrailingUpperBand = min(currentBasicUpperBand, self.previousTrailingUpperBand)
        else:
            currentTrailingUpperBand = currentBasicUpperBand

        if currentClose > currentTrailingUpperBand:
            currentTrend = 1
        elif currentClose < currentTrailingLowerBand:
            currentTrend = -1
        else:
            currentTrend = self.previousTrend

        if currentTrend == 1:
            superTrend = currentTrailingLowerBand
        elif currentTrend == -1:
            superTrend = currentTrailingUpperBand
        
        self.previousTrailingLowerBand = currentTrailingLowerBand
        self.previousTrailingUpperBand = currentTrailingUpperBand
        self.previousClose = currentClose
        self.previousTrend = currentTrend

        if not self.atr.IsReady:
            return 0
        
        self.Value = superTrend
        return self.IsReady
    
    @property
    def IsReady(self):
        return self.atr.IsReady and self.Value != 0