#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
import pandas as pd

class ProtectiveAssetAllocationAlgo(QCAlgorithm):

    def Initialize(self):
        self.SetCash(25000)
        self.SetStartDate(2008,1,1)
        
        ##Parameters for algorithm 
        #self.lookback = 4  ##Lookback in months
        self.lookback = 6  ##Lookback in months
        self.protection = 1 ##Protection factor = 0(low), 1, 2 (high)
        #self.topM = 6 ##topM is the max number of equities
        self.topM = 4 ##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.SafetySymbols = ["TIP", "IEF", "BIL", "TLT", "SPDN"] ##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", "RSP", "TLT",
                             "VGK", "EWJ", "EEM", 
                             "DBC"]
                        
        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)
        #self.AddSecurity(SecurityType.Equity, "BIL", Resolution.Minute)
        #self.AddSecurity(SecurityType.Equity, "SPDN", Resolution.Minute)
        
        self.syl_objs = []
        self.syl_objs_safe = []
        
        # 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 symbol in list(self.SafetySymbols):
            self.syl_objs_safe.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)
            syl_obj.mom21 = self.MOMP(syl_obj, 21, Resolution.Daily)
            syl_obj.mom64 = self.MOMP(syl_obj, 64, Resolution.Daily)
            syl_obj.mom128 = self.MOMP(syl_obj, 128, Resolution.Daily)
            
        #for syl_obj_safe in self.syl_objs_safe:
        for i, item in enumerate(self.syl_objs_safe):
            self.syl_objs_safe[i].Ticker = item.Value
            self.syl_objs_safe[i].lookbackMovingAverage = self.SMA(item, 21*self.lookback, Resolution.Daily)
            self.syl_objs_safe[i].mom21 = self.MOMP(item, 21, Resolution.Daily)
            self.syl_objs_safe[i].mom64 = self.MOMP(item, 64, Resolution.Daily)
            self.syl_objs_safe[i].mom128 = self.MOMP(item, 128, Resolution.Daily)

        #self.SetWarmup(21*self.lookback+1)
        self.SetWarmup(130+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
            #if combmom > 0: 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
            combmom = (syl_obj.mom64.Current.Value)
            syl_obj.MOM = combmom
            #syl_obj.MOM = (price / sma) - 1
            #if syl_obj.MOM > 0.0: N+=1
            if combmom > 0.0: N+=1
        
        for i, item in enumerate(self.syl_objs_safe):
            combmom = (item.mom64.Current.Value)
            self.syl_objs_safe[i].MOM = combmom

        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)
            sortedSafe = sorted([(item.Ticker, item.MOM) for item in self.syl_objs_safe], reverse=True)[-1]
            #self.SetHoldings(self.SafetySymbols[0], w_safe)
            self.SetHoldings(sortedSafe[0], w_safe)
        else:
            for syl_obj in self.syl_objs:
                if self.Portfolio[syl_obj].HoldStock:
                    self.Liquidate(syl_obj)

            df_safe = pd.DataFrame(self.syl_objs_safe)
            sortedSafe = sorted([(item.Ticker, item.MOM) for item in self.syl_objs_safe], reverse=True)[-1]
            self.SetHoldings(sortedSafe[0], 1)
            #self.SetHoldings(self.SafetySymbols[-1], .5)