from clr import AddReference
AddReference("QuantConnect.Algorithm.Framework")
AddReference("QuantConnect.Indicators")
AddReference("QuantConnect.Common")

from QuantConnect import *
from QuantConnect.Indicators import *
from QuantConnect.Algorithm.Framework.Alphas import *
from datetime import timedelta

class MomentumBasedTacticalAllocation(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2014, 10, 1) 
        #self.SetEndDate(2010, 8, 1) 
        self.SetCash(100000)   
        self.SetBrokerageModel(BrokerageName.AlphaStreams)

        self.spy = self.AddEquity("SPY", Resolution.Daily)
        self.bnd = self.AddEquity("BND", Resolution.Daily)
        
        self.spyMomentum = self.MOMP("SPY", 50, Resolution.Daily)
        self.bndMomentum = self.MOMP("BND", 50, Resolution.Daily)
        
        self.reference = "SPY"
        self.SetBenchmark("SPY")  
        self.SetWarmUp(50) 

    def Update(self, algorithm, data):

        insights = []

        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbolData.CanEmit:

                direction = InsightDirection.Flat
                magnitude = symbolData.Return
                if magnitude > 0: direction = InsightDirection.Up
                if magnitude < 0: direction = InsightDirection.Down

                insights.append(Insight.Price(symbol, self.predictionInterval, direction, magnitude, None))
                
        self.EmitInsights(insights)
        return insights

    def OnData(self, data):
        
        # Don't place trades until our indicators are warmed up:
        if self.IsWarmingUp:
            return
        
        #1. If SPY has more upward momentum than BND, then we liquidate our holdings in BND and allocate 100% of our equity to SPY
        if self.spyMomentum.Current.Value > self.bndMomentum.Current.Value:
            self.Liquidate("BND")
            self.SetHoldings("SPY",1)

        #2. Otherwise we liquidate our holdings in SPY and allocate 100% to BND
        else:
            self.Liquidate("SPY")
            self.SetHoldings("BND",1)