#region imports
from AlgorithmImports import *
#endregion
# Import packages
import numpy as np
import pandas as pd
import scipy as sc
import pickle
from scipy import stats


class InOut(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2012, 1, 1)  # Set Start Date
        self.cap = 10000
        self.SetCash(self.cap)  # Set Strategy Cash
        res = Resolution.Hour
        
        # Holdings
        ### 'Out' holdings and weights
        self.HLD_OUT = {self.AddEquity('TLT', res).Symbol: 1} 
        ### 'In' holdings and weights 
        self.HLD_IN = {self.AddEquity('QQQ', res).Symbol: 1} 

        # Market and list of signals based on ETFs
        self.MRKT = self.AddEquity('QQQ', res).Symbol  # market; QQQ
        self.PRDC = self.AddEquity('XLI', res).Symbol  # production (industrials)
        self.METL = self.AddEquity('DBB', res).Symbol  # input prices (metals)
        self.NRES = self.AddEquity('IGE', res).Symbol  # input prices (natural res)
        self.DEBT = self.AddEquity('SHY', res).Symbol  # cost of debt (bond yield)
        self.USDX = self.AddEquity('UUP', res).Symbol  # safe haven (USD)
        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.PRDC  
        
        ##
        self.QQQ = self.AddEquity('QQQ', Resolution.Daily).Symbol
        self.VEU = self.AddEquity('VEU', Resolution.Daily).Symbol
        self.TLT = self.AddEquity('TLT', Resolution.Daily).Symbol
        self.SPDN = self.AddEquity('SPDN', Resolution.Daily).Symbol
        self.IEF = self.AddEquity('IEF', Resolution.Daily).Symbol

        #self.QQQ = self.AddData(QQQ, "QQQ", Resolution.Daily).Symbol
        #self.VEU = self.AddData(VEU, "VEU", Resolution.Daily).Symbol
        #self.TLT = self.AddData(TLT, "TLT", Resolution.Daily).Symbol
        #self.SPDN = self.AddData(SPDN, "SPDN", Resolution.Daily).Symbol
        #self.IEF = self.AddData(IEF, "IEF", Resolution.Daily).Symbol
        
        self.indicator = self.AddData(MOMENTUM, "MOMENTUM", Resolution.Daily).Symbol

        self.SetWarmUp(timedelta(126))

        #self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)


        self.SIGNALS = [self.PRDC, self.METL, self.NRES, self.USDX, self.DEBT, self.MRKT]
        self.FORPAIRS = [self.GOLD, self.SLVA, self.UTIL, self.INDU]
        self.pairlist = ['G_S', 'U_I']
        
        self.basket_in = ['QQQ','VEU']  
        self.riskassetsmomentum = {}

        self.basket_out = ['TLT','IEF','SPDN']  
        self.safehavensmomentum = {}
        
        self.position = -1
        
        for ticker in self.basket_in: 
            self.AddEquity(ticker, res)
            self.Securities[ticker].SetFeeModel(CustomFeeModel())
            self.riskassetsmomentum[ticker] = CombinedMomentum(self, ticker)
            
        for ticker in self.basket_out: 
            self.AddEquity(ticker, res)
            self.Securities[ticker].SetFeeModel(CustomFeeModel())
            self.safehavensmomentum[ticker] = CombinedMomentum(self, ticker)
            
        # Initialize constants and variables
        self.INI_WAIT_DAYS, self.lookback, self.be_in, self.dcount, self.outday, self.portf_val = [5, 252*5, [1], 0, 0, [self.cap]] 
        
        # Symbols for charts
        self.SPY = self.AddEquity('SPY', res).Symbol
        self.QQQ = self.MRKT
        
        symbols = list(set(self.SIGNALS + [self.MRKT] + self.FORPAIRS + list(self.HLD_OUT.keys()) + list(self.HLD_IN.keys()) + [self.SPY] + [self.QQQ]))
        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.lookback, Resolution.Daily)
        if self.history.empty or 'close' not in self.history.columns:
            return
        self.history = self.history['close'].unstack(level=0).dropna()
        self.update_history_shift()
        
        # Benchmarks for charts
        self.benchmarks = [self.history[self.SPY].iloc[-2], self.history[self.QQQ].iloc[-2]]   
        
    def shiftAssets(self, target):
        if not (self.Portfolio[target].Invested):
            for symbol in self.Portfolio.Keys:
                self.Liquidate(symbol)
            if not self.Portfolio.Invested:
                self.SetHoldings(target, 1)
        
    def consolidation_handler(self, sender, consolidated):
        self.history.loc[consolidated.EndTime, consolidated.Symbol] = consolidated.Close
        self.history = self.history.iloc[-self.lookback:]
        self.update_history_shift()
        
    def update_history_shift(self):
        self.history_shift = self.history.rolling(11, center=True).mean().shift(60)

    def inout_check(self):
        if self.history.empty: return
    
        # Load saved dcount and outday (for live interruptions):
        if (self.dcount==0) and (self.outday==0) and (self.ObjectStore.ContainsKey('OS_counts')):
            OS_counts = self.ObjectStore.ReadBytes('OS_counts')
            OS_counts = pickle.loads(bytearray(OS_counts))
            self.dcount, self.outday = [OS_counts['dcount'], OS_counts['outday']]
    
        # Returns sample to detect extreme observations
        returns_sample = (self.history / self.history_shift - 1)
        # Reverse code USDX: sort largest changes to bottom
        returns_sample[self.USDX] = returns_sample[self.USDX] * (-1)
        # For pairs, take returns differential, reverse coded
        returns_sample['G_S'] = -(returns_sample[self.GOLD] - returns_sample[self.SLVA])
        returns_sample['U_I'] = -(returns_sample[self.UTIL] - returns_sample[self.INDU])

        # Extreme observations; statistical significance = 5%
        extreme_b = returns_sample.iloc[-1] < np.nanpercentile(returns_sample, 5, axis=0)
        
        # Re-assess/disambiguate double-edged signals
        abovemedian = returns_sample.iloc[-1] > np.nanmedian(returns_sample, axis=0)
        ### Interest rate expectations (cost of debt) may increase because the economic outlook improves (showing in rising input prices) = actually not a negative signal
        extreme_b.loc[self.DEBT] = np.where((extreme_b.loc[self.DEBT].any()) & (abovemedian[[self.METL, self.NRES]].any()), False, extreme_b.loc[self.DEBT])
        
        # Determine whether 'in' or 'out' of the market
        if (extreme_b[self.SIGNALS + self.pairlist]).any():
            self.be_in.append(0)
            self.outday = self.dcount
        if self.dcount >= self.outday + self.INI_WAIT_DAYS:
            self.be_in.append(1)

        current_portfolio = self.Portfolio.Keys
        # Swap to 'out' assets if applicable
        
        topriskassets = sorted(self.riskassetsmomentum.items(), key=lambda x: x[1].getValue(), reverse=True)
        topsafehavens = sorted(self.safehavensmomentum.items(), key=lambda x: x[1].getValue(), reverse=True)
        
        
        if not self.be_in[-1]:
            if self.position == -1 or self.position == 1:
                if not (self.Portfolio[topsafehavens[0][0]].Invested):
                    self.Liquidate()
                    self.SetHoldings(topsafehavens[0][0], 1)
                    self.position = 0
        elif self.be_in[-1]:
            if self.position == -1 or self.position == 0:
                if not (self.Portfolio[topriskassets[0][0]].Invested):
                    self.Liquidate()    
                    self.SetHoldings(topriskassets[0][0], 1)
                    self.position = 1

        self.charts(extreme_b)
        self.dcount += 1
        
        # Save data: day counts from live trading for interruptions (note: backtest saves data at the end so that it's available for live trading).      
        if self.LiveMode: self.SaveData_Counts()
        
    def OnData(self, data):
            if self.IsWarmingUp: return

            if (self.Time.year <= 2016) :        
                if data.ContainsKey(self.indicator):
                    ticker = data[self.indicator].GetProperty('Indicator')
                    if (ticker =="VEU"):
                            self.Securities["VEU"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.VEU)
                            self.be_in.append(1)
    
                    elif (ticker =="QQQ"):
                            self.Securities["QQQ"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.QQQ)
                            self.be_in.append(1)
    
                    elif (ticker =="TLT"):
                            self.Securities["TLT"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.TLT)
                            self.be_in.append(0)
    
                    elif (ticker =="IEF"):
                            self.Securities["IEF"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.IEF)
                            self.be_in.append(0)

            
            elif (self.Time.year <= 2021) :        
                if data.ContainsKey(self.indicator):
                    ticker = data[self.indicator].GetProperty('Indicator')
                    if (ticker =="VEU"):
                            self.Securities["VEU"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.VEU)
                            self.be_in.append(1)
    
                    elif (ticker =="QQQ"):
                            self.Securities["QQQ"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.QQQ)
                            self.be_in.append(1)
    
                    elif (ticker =="TLT"):
                            self.Securities["TLT"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.TLT)
                            self.be_in.append(0)
    
                    elif (ticker =="IEF"):
                            self.Securities["IEF"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.IEF)
                            self.be_in.append(0)

                    elif (ticker =="SPDN"):
                            self.Securities["SPDN"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.SPDN)
                            self.be_in.append(0)
                    
            
            elif (self.Time.year == 2022 and self.Time.month <= 3):        
                if data.ContainsKey(self.indicator):
                    ticker = data[self.indicator].GetProperty('Indicator')
                    if (ticker =="VEU"):
                            self.Securities["VEU"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.VEU)
                            self.be_in.append(1)
    
                    elif (ticker =="QQQ"):
                            self.Securities["QQQ"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.QQQ)
                            self.be_in.append(1)
    
                    elif (ticker =="TLT"):
                            self.Securities["TLT"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.TLT)
                            self.be_in.append(0)
    
                    elif (ticker =="IEF"):
                            self.Securities["IEF"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.IEF)
                            self.be_in.append(0)

                    elif (ticker =="SPDN"):
                            self.Securities["SPDN"].SetFeeModel(CustomFeeModel())
                            self.shiftAssets(self.SPDN)
                            self.be_in.append(0)
            
                            
            elif (self.Time.hour == 15):
                self.inout_check()

    def charts(self, extreme_b):
        # Market comparisons
        spy_perf = self.history[self.SPY].iloc[-1] / self.benchmarks[0] * self.cap
        qqq_perf = self.history[self.QQQ].iloc[-1] / self.benchmarks[1] * self.cap
        #self.Plot('Strategy Equity', 'SPY', spy_perf)
        #self.Plot('Strategy Equity', 'QQQ', qqq_perf)
        
        # Signals
        self.Plot("In Out", "in_market", self.be_in[-1])
        
        #self.Plot("Signals", "PRDC", int(extreme_b[self.SIGNALS + self.pairlist][0]))
        #self.Plot("Signals", "METL", int(extreme_b[self.SIGNALS + self.pairlist][1]))
        #self.Plot("Signals", "NRES", int(extreme_b[self.SIGNALS + self.pairlist][2]))
        #self.Plot("Signals", "USDX", int(extreme_b[self.SIGNALS + self.pairlist][3]))
        #self.Plot("Signals", "DEBT", int(extreme_b[self.SIGNALS + self.pairlist][4]))
        #self.Plot("Signals", "G_S", int(extreme_b[self.SIGNALS + self.pairlist][5]))
        #self.Plot("Signals", "U_I", int(extreme_b[self.SIGNALS + self.pairlist][6]))
        
        self.Plot("QQQ", "Held", self.Portfolio["QQQ"].Quantity)
        self.Plot("VEU", "Held", self.Portfolio["VEU"].Quantity)
        self.Plot("TLT", "Held", self.Portfolio["TLT"].Quantity)
        self.Plot("IEF", "Held", self.Portfolio["IEF"].Quantity)

        
            
    def trade(self, weight_by_sec):
        # sort: execute largest sells first, largest buys last
        hold_wt = {k: (self.Portfolio[k].Quantity*self.Portfolio[k].Price)/self.Portfolio.TotalPortfolioValue for k in self.Portfolio.Keys}
        order_wt = {k: weight_by_sec[k] - hold_wt.get(k, 0) for k in weight_by_sec}
        weight_by_sec = {k: weight_by_sec[k] for k in dict(sorted(order_wt.items(), key=lambda item: item[1]))}
        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
            # Only trade if holdings fundamentally change
            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 SaveData_Counts(self):
        counts = {"dcount": self.dcount, "outday": self.outday}
        self.ObjectStore.SaveBytes('OS_counts', pickle.dumps(counts))
        
            
class CombinedMomentum():
    def __init__(self, algo, symbol):
        self.one = algo.MOMP(symbol,  21, Resolution.Daily)
        self.three = algo.MOMP(symbol,  63, Resolution.Daily)
        self.six = algo.MOMP(symbol,  126, Resolution.Daily)
        
    def getValue(self):
        value = (self.one.Current.Value + self.three.Current.Value + self.six.Current.Value) / 3
        return value

        
class MOMENTUM(PythonData):
    def GetSource(self, config, date, isLiveMode):
        return SubscriptionDataSource("https://www.dropbox.com/s/c7dh6xdc55i9mmq/Indicator_SHY_XLI_with%20VEU.csv?dl=1", SubscriptionTransportMedium.RemoteFile)
    def Reader(self, config, line, date, isLive):
        
        if not (line.strip() and line[0].isdigit()): 
            return None
        
        index = MOMENTUM()
        index.Symbol = config.Symbol
        data = line.split(',')
        index.Time = datetime.strptime(data[0], "%Y-%m-%d")
        index.EndTime = index.Time + timedelta(days=1)
        index.SetProperty("Indicator", str(data[1]))
        
        return index
        
class CustomFeeModel:
    def GetOrderFee(self, parameters):
        self.trading_fee = 0.005    #Set fee per trade
        fee = parameters.Order.AbsoluteQuantity* self.trading_fee
        return OrderFee(CashAmount(fee, 'USD'))