class DefensiveAssetAllocation(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2008, 1, 1)   # Set Start Date
        self.SetEndDate(2020,1,1)     # Set End Date
        self.SetCash(100000)            # Set Strategy Cash
        
        self.numberOfGrowth = 4

        self.AddEquity("SPY", Resolution.Minute)
        self.SetBenchmark("SPY")
        
        # Set Brokerage Model
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # These are the growth symbol(s) we choose from
        #self.growthSymbols = ["SPY","QQQ","IWN","VGK","EWJ","VWO","GSG","GLD","VNQ","HYG","TLT","LQD"]
        self.growthSymbols = ["SSO","QLD","UWM","VGK","EWJ","VWO","GSG","GLD","VNQ","HYG","UBT","LQD"]
        
        # These are the safety symbol(s) we choose from
        # self.safetySymbols = ["IEF","SHY"]
        self.safetySymbols = ["IEF","SHY"]
        
        # These are the canary universe symbols
        self.canarySymbols = ["VWO", "BND"]
        
        # Combine all tickers
        self.allSymbols = self.growthSymbols + self.safetySymbols + self.canarySymbols
        
        self.Schedule.On(self.DateRules.MonthStart("SPY"), \
                 self.TimeRules.AfterMarketOpen("SPY", 10), \
                 self.Rebalance)
                 
        # Storing data in the SymbolData object
        self.GrowthData = []
        for symbol in list(self.growthSymbols):
            self.AddSecurity(SecurityType.Equity, symbol, Resolution.Minute)
            self.oneMonthPerformance = self.MOMP(symbol, 21, Resolution.Daily)
            self.threeMonthPerformance = self.MOMP(symbol, 63, Resolution.Daily)
            self.sixMonthPerformance = self.MOMP(symbol, 126, Resolution.Daily)
            self.twelveMonthPerformance = self.MOMP(symbol, 252, Resolution.Daily)
            self.GrowthData.append([symbol, self.oneMonthPerformance, self.threeMonthPerformance, self.sixMonthPerformance, self.twelveMonthPerformance])
            
        self.SafetyData = []
        for symbol in list(self.safetySymbols):
            self.AddSecurity(SecurityType.Equity, symbol, Resolution.Minute)
            self.oneMonthPerformance = self.MOMP(symbol, 21, Resolution.Daily)
            self.threeMonthPerformance = self.MOMP(symbol, 63, Resolution.Daily)
            self.sixMonthPerformance = self.MOMP(symbol, 126, Resolution.Daily)
            self.twelveMonthPerformance = self.MOMP(symbol, 252, Resolution.Daily)
            self.SafetyData.append([symbol, self.oneMonthPerformance, self.threeMonthPerformance, self.sixMonthPerformance, self.twelveMonthPerformance])
        
        self.CanaryData = []
        for symbol in list(self.canarySymbols):
            self.AddSecurity(SecurityType.Equity, symbol, Resolution.Minute)
            self.oneMonthPerformance = self.MOMP(symbol, 21, Resolution.Daily)
            self.threeMonthPerformance = self.MOMP(symbol, 63, Resolution.Daily)
            self.sixMonthPerformance = self.MOMP(symbol, 126, Resolution.Daily)
            self.twelveMonthPerformance = self.MOMP(symbol, 252, Resolution.Daily)
            self.CanaryData.append([symbol, self.oneMonthPerformance, self.threeMonthPerformance, self.sixMonthPerformance, self.twelveMonthPerformance])
            
        
        self.SetWarmUp(252, Resolution.Daily)


    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
        '''
        pass

    def Rebalance(self):
        if self.IsWarmingUp:
            return
        
        # Rank all the Growth and Safety symbol groups by momentum score
        orderedGrowthScores = sorted(self.GrowthData, key=lambda x: Score(x[1].Current.Value,x[2].Current.Value,x[3].Current.Value,x[4].Current.Value).ObjectiveScore(), reverse=True)
        bestGrowth = orderedGrowthScores[0:self.numberOfGrowth]
        bestGrowth = [x[0] for x in bestGrowth]
        orderedSafetyScores = sorted(self.SafetyData, key=lambda x: Score(x[1].Current.Value,x[2].Current.Value,x[3].Current.Value,x[4].Current.Value).ObjectiveScore(), reverse=True)
        bestSafety = orderedSafetyScores[0][0]
        
        # Calculate the momentum score of the Canary symbols
        CanaryScores = [Score(x[1].Current.Value,x[2].Current.Value,x[3].Current.Value,x[4].Current.Value).ObjectiveScore() for x in self.CanaryData]
        
        # check if orderedCanaryScores are both negative
        if CanaryScores[0] <0 or CanaryScores[1] <0:
                
            # We are in safety mode. Liquidate any positions and invest in the best safety symbol
            self.Log(f'Safety mode - {CanaryScores}')
            # Check that we are not already invested in the correct Safety symbol
            # If we are invested, return
            if self.Securities[bestSafety].Invested:
                return
            self.Liquidate()
            self.SetHoldings(orderedSafetyScores[0][0],1)
            return
        
        # Canary Symbol gives positive signal
        # We are in Growth Mode
        self.Log(f'Growth mode - {CanaryScores}')

        # Check each symbol we own to see if it is in the bestGrowth list. If not, Liquidate it.
        invested = [x.Symbol.Value for x in self.Portfolio.Values if x.Invested]
        for x in invested:
            if x not in bestGrowth:
                self.Liquidate(x)
        # Check each symbol in bestGrowth list to see if we own it. If not, then buy it.
        for symbol in bestGrowth:
            if not self.Securities[symbol].Invested:
                self.SetHoldings(symbol, 1/self.numberOfGrowth)



class Score(object):
    
    def __init__(self,oneMonthPerformanceValue,threeMonthPerformanceValue,sixMonthPerformanceValue,twelveMonthPerformanceValue):
        self.oneMonthPerformance = oneMonthPerformanceValue
        self.threeMonthPerformance = threeMonthPerformanceValue
        self.sixMonthPerformance = sixMonthPerformanceValue
        self.twelveMonthPerformance = twelveMonthPerformanceValue
        
    def ObjectiveScore(self):
        weight1 = 12
        weight2 = 4
        weight3 = 2
        return (weight1 * self.oneMonthPerformance) + (weight2 * self.threeMonthPerformance) + (weight3 * self.sixMonthPerformance) + self.twelveMonthPerformance
    
class FRED(PythonQuandl):
    '''Custom quandl data type for setting customized value column name. Value column is used for the primary trading calculations and charting.'''
    def __init__(self):
        self.ValueColumnName = "Value"