#region imports
from AlgorithmImports import *
#endregion


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.set_start_date(2010, 1, 1)    # Set Start Date
        self.set_end_date(2019, 8, 1)      # Set End Date
        self.set_cash(100000)             # Set Strategy Cash

        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._coarse_selection_function, self._fine_selection_function)

        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._long_symbols = []            # Contains the stocks we'd like to long
        self._short_symbols = []           # Contains the stocks we'd like to short

        self._next_liquidate = self.time   # Initialize last trade time
        self._rebalance_days = 30

        # Set the weights of each factor
        self._beta_m = 1
        self._beta_s = 1
        self._beta_h = 1
        self._beta_r = 1
        self._beta_c = 1

    def _coarse_selection_function(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._next_liquidate:
            return Universe.UNCHANGED

        selected = sorted([x for x in coarse if x.has_fundamental_data and x.price > 5],
                          key=lambda x: x.dollar_volume, reverse=True)

        return [x.symbol for x in selected[:self._num_coarse]]

    def _fine_selection_function(self, fine):
        '''Select securities with highest score on Fama French 5 factors'''

        # Select stocks with these 5 factors:
        # MKT -- Book value per share: Value
        # SMB -- TotalEquity: Size
        # HML -- Operation profit margin: Quality
        # RMW -- ROE: Profitability
        # CMA -- TotalAssetsGrowth: Investment Pattern
        filtered = [x for x in fine if x.valuation_ratios.book_value_per_share
                                    and x.financial_statements.balance_sheet.total_equity
                                    and x.operation_ratios.operation_margin.value
                                    and x.operation_ratios.ROE
                                    and x.operation_ratios.total_assets_growth]

        # Sort by factors
        sorted_by_mkt = sorted(filtered, key=lambda x: x.valuation_ratios.book_value_per_share, reverse=True)
        sorted_by_smb = sorted(filtered, key=lambda x: x.financial_statements.balance_sheet.total_equity.value, reverse=True)
        sorted_by_hml = sorted(filtered, key=lambda x: x.operation_ratios.operation_margin.value, reverse=True)
        sorted_by_rmw = sorted(filtered, key=lambda x: x.operation_ratios.ROE.value, reverse=True)
        sorted_by_cma = sorted(filtered, key=lambda x: x.operation_ratios.total_assets_growth.value, reverse=False)

        stock_by_symbol = {}

        # Get the rank based on 5 factors for every stock
        for index, stock in enumerate(sorted_by_mkt):
            mkt_rank = self._beta_m * index
            smb_rank = self._beta_s * sorted_by_smb.index(stock)
            hml_rank = self._beta_h * sorted_by_hml.index(stock)
            rmw_rank = self._beta_r * sorted_by_rmw.index(stock)
            cma_rank = self._beta_c * sorted_by_cma.index(stock)
            avg_rank = np.mean([mkt_rank,smb_rank,hml_rank,rmw_rank,cma_rank])
            stock_by_symbol[stock.symbol] = avg_rank

        sorted_dict = sorted(stock_by_symbol.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._long_symbols= symbols[:self._num_long]
        # Pick the stocks with the lowest scores to short
        self._short_symbols = symbols[-self._num_short:]

        return self._long_symbols + self._short_symbols

    def on_data(self, data):
        '''Rebalance Every self._rebalance_days'''

        # Liquidate stocks in the end of every month
        if self.time >= self._next_liquidate:
            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._long_symbols or holding.symbol in self._short_symbols:
                    continue
                # If the holding is not in the list, liquidate
                if holding.invested:
                    self.liquidate(holding.symbol)

        count = len(self._long_symbols + self._short_symbols)

        # 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._long_symbols:
            self.set_holdings(symbol, 1/count)

        # Open short position at the start of every month    
        for symbol in self._short_symbols:
            self.set_holdings(symbol, -1/count)

        # Set the Liquidate Date
        self._next_liquidate = self.time + timedelta(self._rebalance_days)

        # After opening positions, clear the long & short symbol lists until next universe selection
        self._long_symbols.clear()
        self._short_symbols.clear()