from datetime import datetime
from collections import *


class DualMomentumGem(QCAlgorithm):
    def Initialize(self):
        self.SetBenchmark("SPY")
        self.SetStartDate(2010, 1, 1)  # Set Start Date
        self.SetEndDate(2019, 1, 1)    # Set Start Date       
        self.SetCash(100000)           # Set Strategy Cash
        self.bonds = self.AddEquity("AGG", Resolution.Daily).Symbol
        self.lookBackPeriod = 100
        #create symbol data for US (SPY), EU (EFA) and emerging market (EEM) ETFs
        self.symbolObjects = [SymbolData(self,symbolString, self.lookBackPeriod) for symbolString in ["SPY","EFA","EEM"]]
        self.tBill = SymbolData(self,"SHV", self.lookBackPeriod)
        self.Schedule.On(self.DateRules.MonthStart(self.symbolObjects[0].Symbol),self.TimeRules.AfterMarketOpen(self.symbolObjects[0].Symbol), self.Rebalance)
        self.Portfolio.MarginCallModel = MarginCallModel.Null
        self.leverage = 2
        self.SetWarmUp(self.lookBackPeriod)
            
    def Rebalance(self):
        self.symbolObjects.sort(key=lambda symbolObject: symbolObject.Momentum.Current.Value, reverse=True)
        currentLong = self.symbolObjects[0].Symbol
        if self.symbolObjects[0].Momentum.Current.Value < self.tBill.Momentum.Current.Value:
            currentLong  = self.bonds
        self.SetHoldings(currentLong, self.leverage, True)
        
        
class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self,algorithm, symbolString,lookBackPeriod):
        self.Symbol = algorithm.AddEquity(symbolString, Resolution.Daily).Symbol
        self.Momentum = algorithm.MOM(self.Symbol, lookBackPeriod, Resolution.Daily)