# Import packages
import numpy as np
import pandas as pd
import scipy as sc
from QuantConnect.Python import PythonQuandl


# ------------------------------------------------------------------
STK = ['QQQ']; BND = ['TLT']; VOLA = 126; BASE_RET = 85; LEV = 1.00; 
PAIRS = ['SLV', 'GLD', 'XLI', 'XLU', 'DBB', 'UUP']
# ------------------------------------------------------------------

class InOut(QCAlgorithm):

    def Initialize(self):
    
        self.quandlCode = "OECD/KEI_LOLITOAA_OECDE_ST_M"
        ## Optional argument - personal token necessary for restricted dataset
        #Quandl.SetAuthCode("PrzwuZR28Wqegvv1sdJ7")    
        
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2021, 1, 1)
        self.SetCash(25000)  # Set Strategy Cash
        self.cap = 12500

        self.UniverseSettings.Resolution = Resolution.Daily
        res = Resolution.Minute
        
        # Holdings
        ### 'Out' holdings and weights
        self.BND1 = self.AddEquity('TLT', res).Symbol #TLT; TMF for 3xlev
        self.BND2 = self.AddEquity('IEF', res).Symbol #IEF; TYD for 3xlev
        self.HLD_OUT = {self.BND1: .5, self.BND2: .5}
        ### 'In' holdings and weights (static stock selection strategy)
        self.STKS = self.AddEquity('QQQ', res).Symbol #SPY or QQQ; TQQQ for 3xlev
        self.HLD_IN = {self.STKS: 1}
        
        # Market and list of signals based on ETFs
        self.MRKT = self.AddEquity('SPY', res).Symbol  # market
        self.GOLD = self.AddEquity('GLD', res).Symbol  # gold
        self.SLVA = self.AddEquity('SLV', res).Symbol  # vs silver
        self.UTIL = self.AddEquity('XLU', res).Symbol  # utilities
        self.INDU = self.AddEquity('XLI', res).Symbol  # vs industrials
        self.METL = self.AddEquity('DBB', res).Symbol  # input prices (metals)
        self.USDX = self.AddEquity('UUP', res).Symbol  # safe haven (USD)
        
        self.FORPAIRS = [self.GOLD, self.SLVA, self.UTIL, self.INDU, self.METL, self.USDX]
        
        # set a warm-up period to initialize the indicators
        self.SetWarmUp(timedelta(350))
        
        # Specific variables
        self.DISTILLED_BEAR = 999
        self.BE_IN = 999
        self.VOLA_LOOKBACK = 126
        self.WAITD_CONSTANT = 85
        self.DCOUNT = 1 # count of total days since start
        self.OUTDAY = 1 # dcount when self.be_in=0
        
        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.AfterMarketOpen('SPY', 120),
            self.BearRebalance
        )        
        
        # Setup daily consolidation
        symbols = [self.MRKT] + self.FORPAIRS
        for symbol in symbols:
            self.consolidator = TradeBarConsolidator(timedelta(days=1))
            self.consolidator.DataConsolidated += self.consolidation_handler
            self.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        
        # Warm up history
        self.history = self.History(symbols, self.VOLA_LOOKBACK+1, Resolution.Daily)
        if self.history.empty or 'close' not in self.history.columns:
            return
        self.history = self.history['close'].unstack(level=0).dropna()
        self.derive_vola_waitdays()
        
        ### kei 
        #self.SetWarmup(100)
        self.SetBenchmark("SPY")
        self.init = True
        self.kei = self.AddData(QuandlCustomColumns, self.quandlCode, Resolution.Daily, TimeZones.NewYork).Symbol
        self.sma = self.SMA(self.kei, 1)
        self.mom = self.MOMP(self.kei, 2)
        #self.SPY = self.AddEquity('SPY', Resolution.Daily).Symbol
        self.stock = self.AddEquity('QQQ', Resolution.Hour).Symbol
        self.bond = self.AddEquity('TLT', Resolution.Hour).Symbol
        
        
        self.STK = self.AddEquity('QQQ', Resolution.Minute).Symbol
        self.BND = self.AddEquity('TLT', Resolution.Minute).Symbol

        self.ASSETS = [self.STK, self.BND]

        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.symbols = ['FAS', 'ERX', 'UYM', 'DUSL', 'WANT', 'UGE', 'UTSL', 'TECL', 'CURE', 'TENG', 'XLRE']
        
        #Leverged 3x
#        self.XLF = self.AddEquity('FAS', Resolution.Hour).Symbol
#        self.XLE = self.AddEquity('ERX', Resolution.Hour).Symbol
#        self.XLB = self.AddEquity('UYM', Resolution.Hour).Symbol
#        self.XLI = self.AddEquity('DUSL', Resolution.Hour).Symbol
#        self.XLY = self.AddEquity('WANT', Resolution.Hour).Symbol
#        self.XLP = self.AddEquity('UGE', Resolution.Hour).Symbol
#        self.XLU = self.AddEquity('UTSL', Resolution.Hour).Symbol
#        self.XLK = self.AddEquity('TECL', Resolution.Hour).Symbol
#        self.XLV = self.AddEquity('CURE', Resolution.Hour).Symbol
#        self.XLC = self.AddEquity('TENG', Resolution.Hour).Symbol
#        self.XLRE = self.AddEquity('XLRE', Resolution.Hour).Symbol
        
        self.XLF = self.AddEquity('XLF', Resolution.Hour).Symbol
        self.XLE = self.AddEquity('XLE', Resolution.Hour).Symbol
        self.XLB = self.AddEquity('XLB', Resolution.Hour).Symbol
        self.XLI = self.AddEquity('XLI', Resolution.Hour).Symbol
        self.XLY = self.AddEquity('XLY', Resolution.Hour).Symbol
        self.XLP = self.AddEquity('XLP', Resolution.Hour).Symbol
        self.XLU = self.AddEquity('XLU', Resolution.Hour).Symbol
        self.XLK = self.AddEquity('XLK', Resolution.Hour).Symbol
        self.XLV = self.AddEquity('XLV', Resolution.Hour).Symbol
        self.XLC = self.AddEquity('XLC', Resolution.Hour).Symbol
        self.XLRE = self.AddEquity('XLRE', Resolution.Hour).Symbol




        #self.Schedule.On(self.DateRules.WeekStart(self.stock), self.TimeRules.AfterMarketOpen(self.stock, 31), 
        #    self.Rebalance)
        
        self.bull = 1        
        self.count = 0 
        self.outday = 50        
        self.wt = {}
        self.real_wt = {}
        self.mkt = []
        self.SetWarmUp(timedelta(350))
        
        self.Schedule.On(self.DateRules.EveryDay(self.stock), self.TimeRules.AfterMarketOpen(self.stock, 1), 
            self.Rebalance)

        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen('SPY', 100),
            self.daily_check)
            
        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()
        
        #self.AddRiskManagement(TrailingStopRiskManagementModel(0.05))
        ###
    
    
    def consolidation_handler(self, sender, consolidated):
        
        self.history.loc[consolidated.EndTime, consolidated.Symbol] = consolidated.Close
        self.history = self.history.iloc[-(VOLA + 1):] 
        

    def daily_check(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 not self.bull:
            for sec in self.ASSETS:    
                self.wt[sec] = LEV if sec is self.BND else 0 
            self.trade() 

        elif self.bull:
            for sec in self.ASSETS:
                self.wt[sec] = LEV if sec is self.STK else 0 
            self.trade()  

                    
    def trade(self):

        for sec, weight in self.wt.items():
            if weight == 0 and self.Portfolio[sec].IsLong:
                self.Liquidate(sec)
                
            cond1 = weight == 0 and self.Portfolio[sec].IsLong
            cond2 = weight > 0 and not self.Portfolio[sec].Invested
            if cond1 or cond2:
                self.SetHoldings(sec, weight)
            
                    
    def OnEndOfDay(self): 
        
        mkt_price = self.Securities[self.MKT].Close
        self.mkt.append(mkt_price)
        mkt_perf = self.mkt[-1] / self.mkt[0] * (self.cap * .5)
        self.Plot('Strategy Equity', 'SPY', mkt_perf)
        
        account_leverage = self.Portfolio.TotalHoldingsValue / self.Portfolio.TotalPortfolioValue
        self.Plot('Holdings', 'leverage', round(account_leverage, 1))
        for sec, weight in self.wt.items(): 
            self.real_wt[sec] = round(self.ActiveSecurities[sec].Holdings.Quantity * self.Securities[sec].Price / self.Portfolio.TotalPortfolioValue,4)
            self.Plot('Holdings', self.Securities[sec].Symbol, round(self.real_wt[sec], 3))
            
    def Rebalance(self):
        if self.IsWarmingUp or not self.mom.IsReady or not self.sma.IsReady: return
        initial_asset = self.stock if self.mom.Current.Value > 0 else self.bond
        
        if self.init:
            self.SetHoldings(initial_asset, .01)
            self.init = False
            

        keihist = self.History([self.kei], 1400)
        #keihist = keihist['Value'].unstack(level=0).dropna()
        keihistlowt = np.nanpercentile(keihist, 15)
        keihistmidt = np.nanpercentile(keihist, 50)
        keihisthight = np.nanpercentile(keihist, 90)
        kei = self.sma.Current.Value
        keimom = self.mom.Current.Value
            
        if (keimom < 0 and kei < keihistmidt and  kei > keihistlowt) and not (self.Securities[self.XLP].Invested):
            # DECLINE
            self.Liquidate()
            self.SetHoldings(self.XLP, .5)
            self.SetHoldings(self.XLV, .5)
            #self.SetHoldings(self.bond, 1)
            self.Debug("STAPLES {0} >> {1}".format(self.XLP, self.Time))
        
        elif (keimom > 0 and kei < keihistlowt) and not (self.Securities[self.XLB].Invested):
            # RECOVERY
            self.Liquidate()
            self.SetHoldings(self.XLB, .5)
            self.SetHoldings(self.XLY, .5)
            self.Debug("MATERIALS {0} >> {1}".format(self.XLB, self.Time))
            
        elif (keimom > 0 and kei > keihistlowt and kei < keihistmidt) and not (self.Securities[self.XLE].Invested):
            # EARLY
            self.Liquidate()
            self.SetHoldings(self.XLE, .33)
            self.SetHoldings(self.XLF, .33)
            self.SetHoldings(self.XLI, .33)
            self.Debug("ENERGY {0} >> {1}".format(self.XLE, self.Time))
            
        elif (keimom > 0 and kei > keihistmidt and kei < keihisthight) and not (self.Securities[self.XLU].Invested):
            # REBOUND
            self.Liquidate()
            self.SetHoldings(self.XLK, .5)
            self.SetHoldings(self.XLU, .5)
            self.Debug("UTILITIES {0} >> {1}".format(self.XLU, self.Time))
        
        elif (keimom < 0 and kei < keihisthight and kei > keihistmidt) and not (self.Securities[self.XLK].Invested):
            # LATE
            self.Liquidate()
            self.SetHoldings(self.XLK, .5)
            self.SetHoldings(self.XLC, .5)
            
            self.Debug("INFO TECH {0} >> {1}".format(self.XLK, self.Time))
        
        elif (keimom < 0 and kei < 100 and not self.Securities[self.bond].Invested):
            self.Liquidate()
            self.SetHoldings(self.bond, 1)

        self.Plot("LeadInd", "SMA(LeadInd)", self.sma.Current.Value)
        self.Plot("LeadInd", "THRESHOLD", 100)
        self.Plot("MOMP", "MOMP(LeadInd)", self.mom.Current.Value)
        self.Plot("MOMP", "THRESHOLD", 0)
    
        
    ## BEAR    
        
    def consolidation_handler(self, sender, consolidated):
        self.history.loc[consolidated.EndTime, consolidated.Symbol] = consolidated.Close
        self.history = self.history.iloc[-(self.VOLA_LOOKBACK+1):]
        self.derive_vola_waitdays()

    def derive_vola_waitdays(self):
        volatility = np.log1p(self.history[[self.MRKT]].pct_change()).std() * np.sqrt(252)
        wait_days = int(volatility * self.WAITD_CONSTANT)
        returns_lookback = int((1.0 - volatility) * self.WAITD_CONSTANT)
        return wait_days, returns_lookback
        
    def BearRebalance(self):
        wait_days, returns_lookback = self.derive_vola_waitdays()
        
        ## Check for Bear
        returns = self.history.pct_change(returns_lookback).iloc[-1]
    
        silver_returns = returns[self.SLVA]
        gold_returns = returns[self.GOLD]
        industrials_returns = returns[self.INDU]
        utilities_returns = returns[self.UTIL]
        metals_returns = returns[self.METL]
        dollar_returns = returns[self.USDX]
        
        self.DISTILLED_BEAR = (((gold_returns > silver_returns) and
                       (utilities_returns > industrials_returns)) and 
                       (metals_returns < dollar_returns)
                       )
        
        # Determine whether 'in' or 'out' of the market
        if self.DISTILLED_BEAR:
            self.BE_IN = False
            self.OUTDAY = self.DCOUNT
        if self.DCOUNT >= self.OUTDAY + wait_days:
            self.BE_IN = True
        self.DCOUNT += 1
        
        # Determine holdings
        if not self.BE_IN:
            # Only trade when changing from in to out
            self.trade({**dict.fromkeys(self.HLD_IN, 0), **self.HLD_OUT})
        elif self.BE_IN:
            # Only trade when changing from out to in
            self.trade({**self.HLD_IN, **dict.fromkeys(self.HLD_OUT, 0)})
            
    def trade(self, weight_by_sec):
        buys = []
        for sec, weight in weight_by_sec.items():
            # Check that we have data in the algorithm to process a trade
            if not self.CurrentSlice.ContainsKey(sec) or self.CurrentSlice[sec] is None:
                continue
            cond1 = weight == 0 and self.Portfolio[sec].IsLong
            cond2 = weight > 0 and not self.Portfolio[sec].Invested
            if cond1 or cond2:
                quantity = self.CalculateOrderQuantity(sec, weight)
                if quantity > 0:
                    buys.append((sec, quantity))
                elif quantity < 0:
                    self.Order(sec, quantity)
        for sec, quantity in buys:
            self.Order(sec, quantity)
            
    def charting(self, weight_inout_vs_dbear, weighted_be_in): 
        
        if self.dcount==1: self.benchmarks = [self.history[self.MRKT].iloc[-2], self.Portfolio.TotalPortfolioValue, self.history[self.QQQ].iloc[-2]]
        # reset portfolio value and qqq benchmark annually
        if self.Time.year!=self.year: self.benchmarks = [self.benchmarks[0], self.Portfolio.TotalPortfolioValue, self.history[self.QQQ].iloc[-2]]
        self.year = self.Time.year
        
        # SPY benchmark for main chart
        spy_perf = self.history[self.MRKT].iloc[-1] / self.benchmarks[0] * self.cap
        self.Plot('Strategy Equity', 'SPY', spy_perf)
        
        # Leverage gauge: cash level
        self.Plot('Cash level', 'cash', round(self.Portfolio.Cash+self.Portfolio.UnsettledCash, 0))
        
        # Annual saw tooth return comparison: Portfolio VS QQQ
        saw_portfolio_return = self.Portfolio.TotalPortfolioValue / self.benchmarks[1] - 1
        saw_qqq_return = self.history[self.QQQ].iloc[-1] / self.benchmarks[2] - 1
        self.Plot('Annual Saw Tooth Returns: Portfolio VS QQQ', 'Annual portfolio return', round(saw_portfolio_return, 4))
        self.Plot('Annual Saw Tooth Returns: Portfolio VS QQQ', 'Annual QQQ return', round(float(saw_qqq_return), 4))
        
        ### IN/Out indicator and wait days
        self.Plot("In Out", "inout", int(self.be_in_inout))
        self.Plot("In Out", "dbear", int(self.be_in_dbear))
        self.Plot("In Out", "rel_w_inout", float(weight_inout_vs_dbear))
        self.Plot("In Out", "pct_in_market", float(weighted_be_in))
        self.Plot("Wait Days", "waitdays", self.waitdays_inout)
        
    def consolidation_handler(self, sender, consolidated):
        
        self.history.loc[consolidated.EndTime, consolidated.Symbol] = consolidated.Close
        self.history = self.history.iloc[-(VOLA + 1):] 
        

    def daily_check(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 not self.bull:
            for sec in self.ASSETS:    
                self.wt[sec] = LEV if sec is self.BND else 0 
            self.trade() 

        elif self.bull:
            for sec in self.ASSETS:
                self.wt[sec] = LEV if sec is self.STK else 0 
            self.trade()  

                    
    def trade(self):

        for sec, weight in self.wt.items():
            if weight == 0 and self.Portfolio[sec].IsLong:
                self.Liquidate(sec)
                
            cond1 = weight == 0 and self.Portfolio[sec].IsLong
            cond2 = weight > 0 and not self.Portfolio[sec].Invested
            if cond1 or cond2:
                self.SetHoldings(sec, weight)
            
                    
    def OnEndOfDay(self): 
        
        mkt_price = self.Securities[self.MKT].Close
        self.mkt.append(mkt_price)
        mkt_perf = self.mkt[-1] / self.mkt[0] * (self.cap * .5)
        self.Plot('Strategy Equity', 'SPY', mkt_perf)
        
        account_leverage = self.Portfolio.TotalHoldingsValue / self.Portfolio.TotalPortfolioValue
        self.Plot('Holdings', 'leverage', round(account_leverage, 1))
        for sec, weight in self.wt.items(): 
            self.real_wt[sec] = round(self.ActiveSecurities[sec].Holdings.Quantity * self.Securities[sec].Price / self.Portfolio.TotalPortfolioValue,4)
            self.Plot('Holdings', self.Securities[sec].Symbol, round(self.real_wt[sec], 3))
            
    def Rebalance(self):
        if self.IsWarmingUp or not self.mom.IsReady or not self.sma.IsReady: return
        initial_asset = self.stock if self.mom.Current.Value > 0 else self.bond
        
        if self.init:
            self.SetHoldings(initial_asset, .01)
            self.init = False
            

        keihist = self.History([self.kei], 1400)
        #keihist = keihist['Value'].unstack(level=0).dropna()
        keihistlowt = np.nanpercentile(keihist, 15)
        keihistmidt = np.nanpercentile(keihist, 50)
        keihisthight = np.nanpercentile(keihist, 90)
        kei = self.sma.Current.Value
        keimom = self.mom.Current.Value
            
        if (keimom < 0 and kei < keihistmidt and  kei > keihistlowt) and not (self.Securities[self.XLP].Invested):
            # DECLINE
            self.Liquidate()
            self.SetHoldings(self.XLP, .5)
            self.SetHoldings(self.XLV, .5)
            #self.SetHoldings(self.bond, 1)
            self.Debug("STAPLES {0} >> {1}".format(self.XLP, self.Time))
        
        elif (keimom > 0 and kei < keihistlowt) and not (self.Securities[self.XLB].Invested):
            # RECOVERY
            self.Liquidate()
            self.SetHoldings(self.XLB, .5)
            self.SetHoldings(self.XLY, .5)
            self.Debug("MATERIALS {0} >> {1}".format(self.XLB, self.Time))
            
        elif (keimom > 0 and kei > keihistlowt and kei < keihistmidt) and not (self.Securities[self.XLE].Invested):
            # EARLY
            self.Liquidate()
            self.SetHoldings(self.XLE, .33)
            self.SetHoldings(self.XLF, .33)
            self.SetHoldings(self.XLI, .33)
            self.Debug("ENERGY {0} >> {1}".format(self.XLE, self.Time))
            
        elif (keimom > 0 and kei > keihistmidt and kei < keihisthight) and not (self.Securities[self.XLU].Invested):
            # REBOUND
            self.Liquidate()
            self.SetHoldings(self.XLK, .5)
            self.SetHoldings(self.XLU, .5)
            self.Debug("UTILITIES {0} >> {1}".format(self.XLU, self.Time))
        
        elif (keimom < 0 and kei < keihisthight and kei > keihistmidt) and not (self.Securities[self.XLK].Invested):
            # LATE
            self.Liquidate()
            self.SetHoldings(self.XLK, .5)
            self.SetHoldings(self.XLC, .5)
            
            self.Debug("INFO TECH {0} >> {1}".format(self.XLK, self.Time))
        
        elif (keimom < 0 and kei < 100 and not self.Securities[self.bond].Invested):
            self.Liquidate()
            self.SetHoldings(self.bond, 1)

        self.Plot("LeadInd", "SMA(LeadInd)", self.sma.Current.Value)
        self.Plot("LeadInd", "THRESHOLD", 100)
        self.Plot("MOMP", "MOMP(LeadInd)", self.mom.Current.Value)
        self.Plot("MOMP", "THRESHOLD", 0)
        
        
class QuandlCustomColumns(PythonQuandl):
    def __init__(self):
        # Define ValueColumnName: cannot be None, Empty or non-existant column name
        self.ValueColumnName = "Value"