#The investment universe consists of NYSE, AMEX and NASDAQ firms that have stock returns data in CRSP. Financial and utility firms with SIC codes from 6000 to 6999
#and from 4900 to 4949 are excluded.
#Firstly, the earnings acceleration is calculated as a difference of two fractions, where the first fraction is Earnings per share (EPS) of stock i at quarter t minus the EPS of
#stock i at quarter t-4 divided by the stock price price at the end of quarter t-1. The second fraction is a difference of EPS of stock i at quarter t-1 and EPS of stock i at
#quarter t-5 divided by the stock price at the end of quarter t-2.
#Long the highest earnings acceleration decile and short the lowest earnings acceleration decile. Holding period is one month
#Porfolio is value-weighted.

class January_Effect(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2014, 1, 1)
        self.SetEndDate(2019, 7, 1)
        self.SetCash(100000)
        
        self.quarters_count = 6
        self.num_coarse = 100
        self.num_fine = 50
        
        self.symbol_ea = {} # contain the earnings acceleration for every stock
        self.symbol_rw = {} # contain RollingWindow objects for every stock
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.AddEquity("SPY", Resolution.Daily)
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY"), self.Rebalance) # update quarterly_rebalance
        self.quarterly_rebalance = True
        
    def CoarseSelectionFunction(self, coarse):
        if self.quarterly_rebalance:
            selected = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 5)]
            filtered = sorted(selected, key=lambda x: x.DollarVolume, reverse=True) 
            self.filtered_coarse = [ x.Symbol for x in filtered[:self.num_coarse]]
            return self.filtered_coarse
        else: 
            return self.filtered_coarse
        
    def FineSelectionFunction(self, fine):
        if self.quarterly_rebalance:
            # pre-select
            selected = [x for x in fine if x.EarningReports.BasicEPS.ThreeMonths > 0]
            for stock in selected:
                if not stock.Symbol in self.symbol_rw.keys():
                    self.symbol_rw[stock.Symbol] = RollingWindow[float](self.quarters_count)
                # update rolling window for every stock
                self.symbol_rw[stock.Symbol].Add(stock.EarningReports.BasicEPS.ThreeMonths)
                
                if self.symbol_rw[stock.Symbol].IsReady:
                    rw = self.symbol_rw[stock.Symbol]
                    eps_fraction1 = (rw[0] - rw[4]) / rw[1]
                    eps_fraction2 = (rw[1] - rw[5]) / rw[2]
                    self.symbol_ea[stock.Symbol] = eps_fraction1 - eps_fraction2 # That's the earnings acceleration we want
            
            sorted_dict = sorted(self.symbol_ea.items(), key = lambda x: x[1], reverse = True)
            self.filtered_fine = [x[0] for x in sorted_dict[:self.num_fine]]
            self.quarterly_rebalance = False
            
            #Log for validate
            self.Log([x.Value for x in self.filtered_fine])
            
            return self.filtered_fine
        else:
            return self.filtered_fine

    def Rebalance(self):
        if self.Time.month % 3 == 0:
            self.quarterly_rebalance = True
        else:
            self.quarterly_rebalance = False
            
    def OnData(self, data):

        if self.quarterly_rebalance:
            
            if len(self.symbol_ea) == 0:
                return
            
            decile_len = int(len(self.symbol_ea) / 10)
            
            long_stocks = self.filtered_fine[:decile_len]
            short_stocks = self.filtered_fine[-decile_len:]
            
            # Close positions first
            stocks_invested = [x.Key for x in self.Portfolio if x.Value.Invested]
            for stock in stocks_invested:
                if (stock not in long_stocks) and (stock not in short_stocks):
                    self.Liquidate(stock)
                    
            for long_stock in long_stocks:
                if self.Portfolio[long_stock].IsShort:
                    self.Liquidate(long_stock)
                if not self.Portfolio[long_stock].IsLong:
                    self.SetHoldings(long_stock, 1/(2*decile_len))
                    
            for short_stock in short_stocks:
                if self.Portfolio[short_stock].IsLong:
                    self.Liquidate(short_stock)
                if not self.Portfolio[short_stock].IsShort:
                    self.SetHoldings(short_stock, -1/(2*decile_len))
                    
            self.StartTime = self.Time
        
        stocks_invested = [x.Key for x in self.Portfolio if x.Value.Invested]
        
        # holding period is 1 month
        if len(stocks_invested) > 0 and (self.Time - self.StartTime).days > 30:
            for stock in stocks_invested:
                self.Liquidate(stock)