#region imports
from AlgorithmImports import *
#endregion
# 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):

        # Setting starting cap to 100000, which will be used in the SPY benchmark chart
        self.cap = 100000

        self.SetCash(self.cap)
        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)
        
        # Plot SPY on Equity Graph
        self.BNC = "SPY"
        self.mkt = []
        
        self.syl_objs = []
        
        # we'll hold some computed data in these guys
        for symbol in list(self.GrowthSymbols):
            self.syl_objs.append(self.AddSecurity(SecurityType.Equity, symbol, Resolution.Minute).Symbol)
        for syl_obj in self.syl_objs:
            syl_obj.lookbackMovingAverage = self.SMA(syl_obj, 21*self.lookback, Resolution.Daily)
            
        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 syl_obj in self.syl_objs:
            price = self.Securities[syl_obj].Price
            sma = syl_obj.lookbackMovingAverage.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
        #
     
        N = 0
        for syl_obj in self.syl_objs:
            price = self.Securities[syl_obj].Price
            sma = syl_obj.lookbackMovingAverage.Current.Value
            syl_obj.MOM = (price / sma) - 1
            if syl_obj.MOM > 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([i.MOM for i in self.syl_objs], reverse=True)[n_eq - 1]
        
        if frac_eq > 0.0:
            for syl_obj in self.syl_objs:
                if syl_obj.MOM >= float(mom_threshold):
                    self.SetHoldings(syl_obj, w_eq)
                else:
                    if self.Portfolio[syl_obj].HoldStock:
                        self.Liquidate(syl_obj)
            self.SetHoldings(self.SafetySymbols[0], w_safe)
        else:
            for syl_obj in self.syl_objs:
                if self.Portfolio[syl_obj].HoldStock:
                    self.Liquidate(syl_obj)
            self.SetHoldings(self.SafetySymbols[0], 1.0)


    def OnEndOfDay(self):

        if not self.LiveMode:
            mkt_price = self.Securities[self.BNC].Close

        # the below fixes the divide by zero error in the MKT plot
        if mkt_price > 0 and mkt_price is not None:
            self.mkt.append(mkt_price)

        if len(self.mkt) >= 2 and not self.IsWarmingUp:
            mkt_perf = self.mkt[-1] / self.mkt[0] * self.cap
            self.Plot('Strategy Equity', self.BNC, mkt_perf)