from clr import AddReference
AddReference("System.Core")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.UniverseSelection import *
import decimal as d
from datetime import datetime, timedelta
from decimal import Decimal
import numpy as np

class ProtectiveAssetAllocationAlgo(QCAlgorithm):

    def Initialize(self):
        self.SetCash(25000)
        self.SetStartDate(2008,1,1)
        
        ##Parameters for algorithm 
        self.lookback = 4  ##Lookback in months
        self.protection = 2 ##Protection factor = 0(low), 1, 2 (high)
        self.topM = 6 ##topM is the max number of equities
        self.n_levels = 2 ##number of discrete levels for bond_fraction (>=2)
        self.SafetySymbols = ["IEF"] ##risk free asset to move into for protection 
        self.N_safe = int(len(self.SafetySymbols))

        # these are the growth symbols we"ll rotate through
        self.GrowthSymbols =["SPY", "QQQ", "IWM",
                             "VGK", "EWJ", "EEM",
                             "VNQ", "DBC", "GLD",
                             "HYG", "LQD", "TLT"]
                        
        self.N_eq = len(self.GrowthSymbols)
        # these are the safety symbols we go to when things are looking bad for growth

        self.AddSecurity(SecurityType.Equity, "IEF", Resolution.Minute)
        # we'll hold some computed data in these guys
        self.SymbolData = []
        for symbol in list(self.GrowthSymbols):
            self.AddSecurity(SecurityType.Equity, symbol, Resolution.Minute)
            self.lookbackMovingAverage = self.SMA(symbol, 21*self.lookback, Resolution.Daily)
            self.SymbolData.append([symbol, self.lookbackMovingAverage])
            
        self.SetWarmup(21*self.lookback+1)
        
        self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(9,45),
                Action(self.Rebalance))

    def OnData(self, data):
        pass 
    
    def Rebalance(self):
        
        #
        # poll the Growth Symbols set to determine the number of assets with positive momentum
        #
    
        n = 0
        for i in range(self.N_eq):
            price = self.Securities[self.SymbolData[i][0]].Price
            sma = self.SymbolData[i][1].Current.Value
            if price > sma: n+=1
            
        # Calculate the bond fraction based on N_eq, prot, and n
        # This is the portion to be invested in safe harbor
        # Calculate equity fraction and weight per equity (frac_eq, w_eq) 
        # Limit bond_fraction to a discrete number of levels (n_levels >=2)
        
        n1 = int(int(self.protection) * int(self.N_eq) / 4.0)
        bond_fraction = float(min(1.0, float(float(self.N_eq) - float(n)) / float(float(self.N_eq) - float(n1))))
        #n_steps = float(self.n_levels) - 1.0
        #bond_fraction = float(bond_fraction*n_steps)/n_steps
    
        w_safe = float(bond_fraction)
        self.Log("Safe Weight "+str(w_safe))
        
        #
        # calculate the MOM for each equity
        # determine the number of equities to be purchases
        #

        MOM = {}
        N = 0
        for i in range(self.N_eq):
            price = self.Securities[self.SymbolData[i][0]].Price
            sma = self.SymbolData[i][1].Current.Value
            MOM[self.SymbolData[i][0]] = (price / sma) - 1
            if MOM[self.SymbolData[i][0]] > 0.0: N+=1
        
        frac_eq = float(1.0 - w_safe)
        n_eq = min(N, self.topM)
        w_eq = 0.0
        if N > 0: w_eq = float(float(frac_eq) / float(n_eq))
        mom_threshold = sorted(MOM.values(), reverse=True)[n_eq - 1]
        
        if frac_eq > 0.0:
            for i in range(self.N_eq):
                symbol = self.SymbolData[i][0]
                if MOM[symbol] >= float(mom_threshold):
                    self.SetHoldings(symbol, w_eq)
                else:
                    if self.Portfolio[symbol].HoldStock:
                        self.Liquidate(symbol)
            self.SetHoldings(self.SafetySymbols[0], w_safe)
        else:
            for i in range(self.N_eq):
                symbol = self.SymbolData[i][0]
                if self.Portfolio[symbol].HoldStock:
                    self.Liquidate(symbol)
            self.SetHoldings(self.SafetySymbols[0], 1.0)