#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.
from datetime import datetime

class January_Effect(QCAlgorithm):
    def Initialize(self):
        
        self.SetStartDate(1999, 1, 5) # Set Start Date
        self.SetEndDate(2019, 8, 6)   # Set End Date
        self.SetCash(100000)          # Set Strategy Cash
        
        self.quarters_count = 6       # Number of quarters to calculate the earning acceleration indicator
        self.num_coarse = 100         # Number of symbols selected at Coarse Selection
        self.num_fine = 50            # Number of symbols selected at Fine Selection
        
        self.epsBySymbol = {}         # Contains RollingWindow objects for every stock
        self.longSymbols = []         # Contains the stocks we'd like to long
        self.shortSymbols = []        # Contains the stocks we'd like to short
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.lastTradeTime = datetime.now() # Initialize last trade time
        
        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):
        '''Drop securities which have no fundamental data or have too low prices.
        Select those with highest by dollar volume'''
        
        if not self.quarterly_rebalance:
            return Universe.Unchanged

        self.quarterly_rebalance = False
        
        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 5],
            key=lambda x: x.DollarVolume, reverse=True)
            
        return [x.Symbol for x in selected[:self.num_coarse]]
        
    def FineSelectionFunction(self, fine):
        '''Select security with highest earnings acceleration'''

        eaBySymbol = dict()
        
        # pre-select
        selected = [x for x in fine if x.EarningReports.BasicEPS.ThreeMonths > 0]
        for stock in selected:
            if not stock.Symbol in self.epsBySymbol:
                self.epsBySymbol[stock.Symbol] = RollingWindow[float](self.quarters_count)
            # update rolling window for every stock
            self.epsBySymbol[stock.Symbol].Add(stock.EarningReports.BasicEPS.ThreeMonths)

            if self.epsBySymbol[stock.Symbol].IsReady:
                rw = self.epsBySymbol[stock.Symbol]
                eps_fraction1 = (rw[0] - rw[4]) / rw[1]
                eps_fraction2 = (rw[1] - rw[5]) / rw[2]
                
                eaBySymbol[stock.Symbol] = eps_fraction1 - eps_fraction2 # That's the earnings acceleration we want

        sorted_dict = sorted(eaBySymbol.items(), key = lambda x: x[1], reverse = True)

        symbols = [x[0] for x in sorted_dict[:self.num_fine]]

        decile_len = int(len(symbols) / 10)
        self.longSymbols = symbols[:decile_len]
        self.shortSymbols = symbols[-decile_len:]

        #Log for validate
        self.Log(','.join(sorted([x.Value for x in self.longSymbols + self.shortSymbols])))
        return self.longSymbols + self.shortSymbols


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

        # 31 days after we open the positions, we liquidate.
        if (self.Time - self.lastTradeTime).days > 31:
            for holding in self.Portfolio.Values:
                if holding.Invested:
                    self.Liquidate(holding.Symbol)

        if not self.quarterly_rebalance:
            return

        count = len(self.longSymbols) + len(self.shortSymbols)

        for symbol in self.longSymbols:
            self.SetHoldings(symbol, 1/count)

        for symbol in self.shortSymbols:
            self.SetHoldings(symbol, -1/count)

        # Last Trade time
        self.lastTradeTime = self.Time