dev_mode = False # for Shile's use, keep False
if dev_mode:
    from AlgorithmImports import *
resolution = Resolution.Hour

# ---indicies---
market = 'SPY'
silver = 'SLV'
gold = 'GLD'
utility = 'XLU'
industrial = 'XLI'
safe = 'FXF'  # safe currency 
risk = 'FXA'  # risk currency
debt_short = 'SHY'
debt_inflation = 'TIP'
metal = 'DBB'
inp = 'IGE' # input
cash = 'UUP'

# ---equities CHANGEABLE
equities = ['SPY', 'QQQ', 'AAPL']

# ---safeties CHANGEABLE
safeties = ['GLD', 'FXF']

# ---in and out parameters
# parameters found from file from you
bull = True # set False for bear
inOutLookbackBull = 252
inOutLookbackBear = 126

waitDaysConstant = 80 # WAITD_CONSTANT from your file
iniWaitDays = 15 # INI_WAIT_DAYS from your file
minWaitDays = 60 # 60 from the `min(60, self.WDadjvar)` from your file

# ---supertrend parameters CHANGEABLE
superTrendPeriod = 10
superTrendMultiple = 3

# ---squeeze parameters CHANGEABLE
squeezeTrendPeriod = 20
squeezeBBMultiple = 2 # BollingerBands
squeezeKeltMultiple = 1.5 # Kelter Channel

# ---portfolio parameters
# for the returns based portfolio allocation
max_drawdown = .1 # max drawdown allowed before liquidation is signaled
max_alloc = .4 # max allocation to any given stock
returns_lookback = 100 # lookback for returns/drawdown calculations 

#from configs import *
if dev_mode:
    from AlgorithmImports import *
    
from collections import deque
from datetime import datetime 

class AdaptableRedSnake(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 11)
        self.SetCash(100000) 
        
        self.supertrends = {}
        self.squeezes = {}

        for equity in equities:
            sym = self.AddEquity(equity, resolution).Symbol

            superTrend = MySuperTrend(self, superTrendPeriod, superTrendMultiple)
            self.RegisterIndicator(sym, superTrend, resolution)
            self.supertrends[sym] = superTrend

            squeeze = Squeeze(squeezeTrendPeriod, bollinger_multiple=squeezeBBMultiple,  kelt_multiple=squeezeKeltMultiple)
            self.RegisterIndicator(sym, squeeze, resolution)
            self.squeezes[sym] = squeeze

            self.SetWarmUp(20)
        
        self.safeties = [self.AddEquity(symbol, resolution) for symbol in safeties]

    def OnData(self, data):
        ''' OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''
        if self.IsWarmingUp:
            return

        for symbol in self.supertrends:
            if not data.Bars.ContainsKey(symbol):
                return
            if self.supertrends[symbol].Value < data[symbol].Close and self.squeezes[symbol]:
                self.SetHoldings(symbol, 1 / len(self.supertrends)) 
            else:
                self.Liquidate(symbol)
                for safety in self.safeties:
                    self.SetHoldings(symbol, 1 / len(self.supertrends) / len(self.safeties)) 
           
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

class Squeeze:
    '''
    .Value = 1 iff "squeezed" else .Value = 0
    
    Tells us if we are in or out of squeeze

    Is Squeeze: lower BB > lower Keltner and upper BB < upper Keltner
    '''

    def __init__(self, period, bollinger_multiple=2, kelt_multiple=1.5, movingAverageType=MovingAverageType.Simple):
        '''
        .Value = 1 iff "squeezed" else .Value = 0
        '''
        self.Name = "SuperTrend"
        self.Time = datetime.min
        self.Value = 0
        
        self.bb = BollingerBands(period, bollinger_multiple, movingAverageType)
        self.kelt = KeltnerChannels(period, kelt_multiple, movingAverageType)

    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.kelt.Update(input)
        self.bb.Update(input.EndTime, input.Close)

        isSqueeze = self.bb.LowerBand.Current.Value > self.kelt.LowerBand.Current.Value and self.bb.UpperBand.Current.Value <  self.kelt.UpperBand.Current.Value
        self.Value = int(isSqueeze)
        return self.IsReady 

    @property
    def IsReady(self):
        return self.kelt.IsReady and self.bb.IsReady