import decimal as d
from datetime import datetime, timedelta
from decimal import Decimal


class VigilantAssetAllocationAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetCash(25000)
        self.SetStartDate(2007,1,1)
        self.LastRotationTime = datetime.min
        self.RotationInterval = timedelta(days=30)
        self.first = True

        # these are the growth symbols we'll rotate through
        GrowthSymbols =["SPY", 
                        "EFA",
                        "EEM",
                        "AGG"]

        # these are the safety symbols we go to when things are looking bad for growth
        SafetySymbols = ["LQD", 
                         "IEF", 
                         "SHY"]
                         
        # I split the indicators into two different sets to make it easier for illustrative purposes below.
        # Storing all risky asset data into SymbolData object
        self.SymbolData = []
        for symbol in list(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.SymbolData.append([symbol, self.oneMonthPerformance, self.threeMonthPerformance, self.sixMonthPerformance, self.twelveMonthPerformance])
            
        # Storing all risk-free data into SafetyData object
        self.SafetyData = []
        for symbol in list(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])
        
    def OnData(self, data):
        if self.first:
            self.first = False
            self.LastRotationTime = self.Time
            return
        delta = self.Time - self.LastRotationTime
        if delta > self.RotationInterval:
            self.LastRotationTime = self.Time
            
            ##Using the Score class at the bottom, compute the score for each risky asset.
            ##This approach overweights the front month momentum value and progressively underweights older momentum values
            
            orderedObjScores = sorted(self.SymbolData, key=lambda x: Score(x[1].Current.Value,x[2].Current.Value,x[3].Current.Value,x[4].Current.Value).ObjectiveScore(), reverse=True)
            
            ##Using the Score class at the bottom, compute the score for each risk-free asset.
            orderedSafeScores = 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)
            
            ##Count the number of risky assets with negative momentum scores and store in N. If all four of the offensive assets exhibit positive momentum scores, 
            ##select the offensive asset with the highest score and allocate 100% of the portfolio to that asset at the close
            N = 0
            for x in orderedObjScores:
                self.Log(">>SCORE>>" + x[0] + ">>" + str(Score(x[1].Current.Value,x[2].Current.Value,x[3].Current.Value,x[4].Current.Value).ObjectiveScore()))
                if Score(x[1].Current.Value,x[2].Current.Value,x[3].Current.Value,x[4].Current.Value).ObjectiveScore() < 0:
                    N += 1
                   
            # pick which one is best from risky and risk-free symbols and store for use below
            bestGrowth = orderedObjScores[0]
            bestSafe = orderedSafeScores[0]
           
            ## If any of the four risky assets exhibit negative momentum scores, select the risk-free asset (LQD, IEF or SHY) with the highest score 
            ## (regardless of whether the score is > 0) and allocate 100% of the portfolio to that asset at the close. 
            if N > 0:
                self.Log("PREBUY>>LIQUIDATE>>")
                if not self.Portfolio[bestSafe[0]].Invested:
                    self.Liquidate()
                self.Log(">>BUY>>" + str(bestSafe[0]) + "@" + str(Decimal(100) * bestSafe[1].Current.Value))
                self.SetHoldings(bestSafe[0], 1) 
            ## If none of the risky assets come back with negative momentum scores, allocation 100% to the best scoring risky asset and hold until month end
            else:                
                self.Log("PREBUY>>LIQUIDATE>>")
                if not self.Portfolio[bestGrowth[0]].Invested:
                    self.Liquidate()
                self.Log(">>BUY>>" + str(bestGrowth[0]) + "@" + str(Decimal(100) * bestGrowth[1].Current.Value))
                self.SetHoldings(bestGrowth[0], 1) 

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