import numpy as np

class BetaAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 8, 20)   # Set Start Date
        self.SetEndDate(2020, 10, 26)     # Set End Date
        self.SetCash(10000)             # Set Strategy Cash
        
        self.ReductionCoeff = 0.98 # Coefficient to avoid leverage
        self.SetWarmup(22)
        self.PositionsNumber = 5 # Max. number of active positions for this strategy
        
        self.spy = self.AddEquity("QQQ", Resolution.Daily) # Add equity to use for trends

        # Dow 30 companies. 
        self.symbols = [self.AddEquity(ticker).Symbol
            
            for ticker in [
                'FVRR', # FIVERR
                'GH',   # Guardant Health
                'GWRE', # Guidewire Software
                'HUBS', # HubSpot
                'HUYA', # HUYA
                'IBKR', # Interactive Brokers Group
                'INCY', # Incyte Corporation
                'IONS', # Ionis Pharmaceuticals
                'IOVA', # Iovance Biotherapeutics
                'IRDM', # Iridium Communications
                'KTOS', # Kratos Defense & Security Solutions
                'LMND', # Lemonade
                'LSPD', # Lightspeed POS
                'MCRB', # Seres Therapeutics
                'MKL',  # Markel Corporation
                ] ]  

        # Benchmark - Security used as reference to beat. In this case QQQ (NASDAQ index)
        self.benchmark = Symbol.Create('QQQ', SecurityType.Equity, Market.USA)

        # Set number days to trace back
        self.lookback = 1

        # Schedule Event: trigger the event at the begining of each month.
        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.AfterMarketOpen(self.symbols[0]),
                         self.Rebalance)

        #QQQ trend filter: only trades if benchmark trends
        # Very fast moving average of benchmark
        self.spy_ma_vfast = self.EMA("QQQ", 5)
        # Fast moving average of benchmark
        self.spy_ma_fast = self.EMA("QQQ", 13)
        # Slow moving average of benchmark
        self.spy_ma_slow = self.EMA("QQQ", 21)


    def Rebalance(self):
        
        # Defines benchmark in a trend (fast ma over slow ma)
        self.trend_up = self.spy_ma_vfast > self.spy_ma_fast and self.spy_ma_fast > self.spy_ma_slow

        # Fetch the historical data to perform the linear regression

        history = self.History(

            self.symbols + [self.benchmark], 
            self.lookback,
            Resolution.Daily).close.unstack(level=0)
            
        symbols = self.SelectSymbols(history)

        # Liquidate positions that are not held by selected symbols and liquidates if benchmark is not in a trend
        for holdings in self.Portfolio.Values:
            try:
                symbol = holdings.Symbol
                if symbol not in symbols and holdings.Invested:
                    self.Liquidate(symbol)
            except:
                self.Debug("Error in: " + str(symbol))
                continue

        # Invest in the selected symbols is benchmark/reference is in a trend
        for symbol in symbols:
            try:
                if self.trend_up:
                    # Position size is calculated based on active strategies, max positions number and some coeffs to avoid leverage
                    self.SetHoldings(symbol, (1*self.ReductionCoeff)/self.PositionsNumber)
            
                if not self.trend_up:
                    self.Liquidate(symbol)
             
            except:
                self.Debug("Error in: " + str(symbol))
                continue

    def SelectSymbols(self, history):
        '''Select symbols with the highest intercept/alpha to the benchmark'''
        
        alphas = dict()

        # Get the benchmark returns
        benchmark = history[self.benchmark].pct_change().dropna()

        # Conducts linear regression for each symbol and save the intercept/alpha
        
        for symbol in self.symbols:
            try:
                # Get the security returns
                returns = history[symbol].pct_change().dropna()
                returns = np.vstack([returns, np.ones(len(returns))]).T
    
                # Simple linear regression function in Numpy
                result = np.linalg.lstsq(returns, benchmark)
                alphas[symbol] = result[0][1]
            except:
                self.Debug("Error in: " + str(symbol))
                continue

        # Select symbols and number of them with the highest intercept/alpha to the benchmark
        selected = sorted(alphas.items(), key=lambda x: x[1], reverse=True)[:self.PositionsNumber]
        return [x[0] for x in selected]