import numpy as np

class FamaFrenchFiveFactorsAlgorithm(QCAlgorithm):
    ''' Stocks Selecting Strategy based on Fama French 5 Factors Model
        Reference: https://tevgeniou.github.io/EquityRiskFactors/bibliography/FiveFactor.pdf
    '''
    def Initialize(self):
        self.SetStartDate(2010, 1, 1)  # Set Start Date
        self.SetEndDate(2019, 8, 1)    # Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.num_coarse = 200         # Number of symbols selected at Coarse Selection

        self.num_long = 5     # Number of stocks to long
        self.num_short = 5    # Number of stocks to short
        
        self.longSymbols = []           # Contains the stocks we'd like to long
        self.shortSymbols = []          # Contains the stocks we'd like to short
        
        self.nextLiquidate = self.Time     # Initialize last trade time
        self.rebalance_days = 30


    def CoarseSelectionFunction(self, coarse):
        '''Drop securities which have no fundamental data or have too low prices.
        Select those with highest by dollar volume'''

        if self.Time < self.nextLiquidate:
            return Universe.Unchanged
        
        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 securities with highest score on Fama French 5 factors'''
        
        # select stocks with these 5 factors
        
        # Operation profit margin: Quality
        # Book value per share: Value
        # ROE: Profitability
        # TotalEquity: Size
        # TotalAssetsGrowth: Investment Pattern
        filtered = [x for x in fine if x.OperationRatios.OperationMargin.Value
                                    and x.ValuationRatios.BookValuePerShare
                                    and x.OperationRatios.ROE
                                    and x.FinancialStatements.BalanceSheet.TotalEquity
                                    and x.OperationRatios.TotalAssetsGrowth]
                                    
        
        # sort by factors
        sortedByFactor1 = sorted(filtered, key=lambda x: x.OperationRatios.OperationMargin.Value, reverse=True)
        sortedByFactor2 = sorted(filtered, key=lambda x: x.ValuationRatios.BookValuePerShare, reverse=True)
        sortedByFactor3 = sorted(filtered, key=lambda x: x.OperationRatios.ROE.Value, reverse=True)
        sortedByFactor4 = sorted(filtered, key=lambda x: x.FinancialStatements.BalanceSheet.TotalEquity.Value, reverse=True)
        sortedByFactor5 = sorted(filtered, key=lambda x: x.OperationRatios.TotalAssetsGrowth.Value, reverse=False)
        
        stockBySymbol = {}
        
        # get the rank based on 5 factors for every stock
        for index, stock in enumerate(sortedByFactor1):
            rank1 = index
            rank2 = sortedByFactor2.index(stock)
            rank3 = sortedByFactor3.index(stock)
            rank4 = sortedByFactor4.index(stock)
            rank5 = sortedByFactor5.index(stock)
            avgRank = np.mean([rank1,rank2,rank3,rank4,rank5])
            stockBySymbol[stock.Symbol] = avgRank
            
        sorted_dict = sorted(stockBySymbol.items(), key = lambda x: x[1], reverse = True)
        symbols = [x[0] for x in sorted_dict]
        
        # pick the stocks with the highest scores to long
        self.longSymbols= symbols[:self.num_long]
        # pick the stocks with the lowest scores to short
        self.shortSymbols = symbols[-self.num_short:]
        
        return self.longSymbols + self.shortSymbols
    
    
    def OnData(self, data):
        '''Rebalance Every self.rebalance_days'''

        # Liquidate stocks in the end of every month
        if self.Time >= self.nextLiquidate:
            for holding in self.Portfolio.Values:
                # if the holding is in the long/short list for the next month, don't liquidate
                if holding.Symbol in self.longSymbols or holding.Symbol in self.shortSymbols:
                    continue
                # the holding not in the list
                if holding.Invested:
                    self.Liquidate(holding.Symbol)
                    
        count = len(self.longSymbols + self.shortSymbols)
        
        # It means the long & short lists for the month have been cleared
        if count == 0: 
            return
    
        # open long position at the start of every month
        for symbol in self.longSymbols:
            self.SetHoldings(symbol, 1/count)
        
        # open short position at the start of every month    
        for symbol in self.shortSymbols:
            self.SetHoldings(symbol, -1/count)
            
        # Set the Liquidate Date
        self.nextLiquidate = self.Time + timedelta(self.rebalance_days)
        
        # After opening positions, clear the long & short symbol list until next universe selection
        self.longSymbols.clear()
        self.shortSymbols.clear()