import numpy as np
# ----------------------------------------------------------
STOCKS = ['QQQ', 'FDN']; BONDS = ['TLT', 'TLH']; 
VOLA = 126; BASE_RET = 83; RET = 252; EXCL = 21; LEV = 1.00; 
# ----------------------------------------------------------

class DualMomentumInOut(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2010, 2, 5)
        self.SetEndDate(2010, 3, 5)
        self.cap = 100000
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(
            self.SelectCoarse, self.SelectFine
        )
        
        self.BND1 = self.AddEquity('TLT', Resolution.Hour).Symbol
        self.BND2 = self.AddEquity('TLH', Resolution.Hour).Symbol

        self.SLV = self.AddEquity('SLV', Resolution.Daily).Symbol  
        self.GLD = self.AddEquity('GLD', Resolution.Daily).Symbol  
        self.XLI = self.AddEquity('XLI', Resolution.Daily).Symbol 
        self.XLU = self.AddEquity('XLU', Resolution.Daily).Symbol
        self.DBB = self.AddEquity('DBB', Resolution.Daily).Symbol  
        self.UUP = self.AddEquity('UUP', Resolution.Daily).Symbol  
        self.MKT = self.AddEquity('SPY', Resolution.Daily).Symbol          

        self.pairs = [self.SLV, self.GLD, self.XLI, self.XLU, self.DBB, self.UUP]
        
        self.bull = 1        
        self.count = 0 
        self.outday = 0        
        self.wt = {}
        self.real_wt = {}
        self.mkt = []
        self.SetWarmUp(timedelta(350))
        
        self.selected_bond = self.BND1

        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen('SPY', 100),
            self.calculate_signal)
           
            
        symbols = [self.MKT] + self.pairs
        for symbol in symbols:
            self.consolidator = TradeBarConsolidator(timedelta(days=1))
            self.consolidator.DataConsolidated += self.consolidation_handler
            self.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        
        self.history = self.History(symbols, VOLA + 1, Resolution.Daily)
        if self.history.empty or 'close' not in self.history.columns:
            return
        self.history = self.history['close'].unstack(level=0).dropna()

    def SelectCoarse(self, coarse):
        dollarvolume_array = np.array([ x.DollarVolume for x in coarse ])
        dollarvolume_percentile = np.percentile(dollarvolume_array, 75)
        # Select top 25% volume
        coarse = [x.Symbol for x in coarse if x.DollarVolume >= dollarvolume_percentile]
        return coarse

    def SelectFine(self, fine):
        # Select top 5% historical Returns over
        fine = [self.add_returns_to_security(s, RET, EXCL) for s in fine]
        returns = np.array([s.returns for s in fine])
        nan_mask = np.all(np.isnan(returns), axis=0)
        returns_p95 = np.percentile(returns[~nan_mask], 95)
        # Get top 10 momentum
        fine = sorted([s for s in fine], key=lambda s: s.returns, reverse=True)
        self.longs = [s.Symbol for s in fine if s.returns >= returns_p95][:10]
        return self.longs
        
    def OnSecuritiesChanged(self, changes):
        if self.IsWarmingUp:
            return
        # Subscribe to minute bars for added equities
        for instrument in changes.AddedSecurities:
            self.AddEquity(instrument.Symbol, Resolution.Minute)
            self.Securities[instrument.Symbol].FeeModel = ConstantFeeModel(0.0)
        # Remove equities
        for instrument in changes.RemovedSecurities:
            self.RemoveSecurity(instrument.Symbol)

    def consolidation_handler(self, sender, consolidated):
        self.history.loc[consolidated.EndTime, consolidated.Symbol] = consolidated.Close
        self.history = self.history.iloc[-(VOLA + 1):]  

    def add_returns_to_security(self, sec, period, excl):
        sec.returns = self.returns(sec.Symbol, period, excl)
        return sec
   
    def returns(self, symbol, period, excl):
        prices = self.History(symbol, TimeSpan.FromDays(period + excl), Resolution.Daily)
        if len(prices.index) < excl:
            return 0
        return prices['close'].iloc[-excl] / prices['close'].iloc[0]
        
        
    def calculate_signal(self):
        vola = self.history[[self.MKT]].pct_change().std() * np.sqrt(252)
        wait_days = int(vola * BASE_RET)
        period = int((1.0 - vola) * BASE_RET)        
        r = self.history.pct_change(period).iloc[-1]

        exit = ((r[self.SLV] < r[self.GLD]) and (r[self.XLI] < r[self.XLU]) and  (r[self.DBB] < r[self.UUP]))

        if exit:
            self.bull = False
            self.outday = self.count
        if self.count >= self.outday + wait_days:
            self.bull = True
        self.count += 1


        if self.returns(self.BND1, RET, EXCL) < self.returns(self.BND2, RET, EXCL):
            self.selected_bond = self.BND2
            
        elif self.returns(self.BND1, RET, EXCL) > self.returns(self.BND2, RET, EXCL):
            self.selected_bond = self.BND1
    
        if not self.bull:
            self.trade([self.selected_bond])

        elif self.bull:
            self.trade(self.longs)    

                    
    def trade(self, new_securities):
        existing = [sym for sym, sec in self.Securities.items() if sec.Invested]
        self.Debug("long: {}, rebalancing into {}".format(self.bull, new_securities))
        for security in set(existing+new_securities):
            if security not in new_securities:
                self.Liquidate(security)
            else:
                self.SetHoldings(security, LEV/len(new_securities))

    def OnEndOfDay(self):
        self.Plot('Security Count', 'Invested', len([sym for sym, sec in self.Securities.items() if sec.Invested]))
        self.Plot('Security Count', 'Universe', len(self.ActiveSecurities))
        self.Plot('Security Count', 'Longs', len(self.longs))
        account_leverage = self.Portfolio.TotalAbsoluteHoldingsCost / self.Portfolio.TotalPortfolioValue
        self.Plot('Leverage', 'Total Account', account_leverage)