#region imports
from AlgorithmImports import *
#endregion

class BasicTemplateAlgorithm(QCAlgorithm):
    
    def __init__(self):
        # set the flag for rebalance
        self._reb = 1
        # Number of stocks to pass CoarseSelection process
        self._num_coarse = 200
        # Number of stocks to long/short
        self._num_fine = 10
        self._symbols = None

    def initialize(self):
        self.set_cash(100000)
        self.set_start_date(2021,1,1)
        self.set_end_date(2023,1,1)
        
        self.set_security_initializer(
            BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices))
        )
        self._spy = self.add_equity("SPY", Resolution.DAILY).symbol
        
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._coarse_selection_function,self._fine_selection_function)
        
        # Schedule the rebalance function to execute at the begining of each month
        self.schedule.on(
            self.date_rules.month_start(self._spy), 
            self.time_rules.after_market_open(self._spy,5), 
            self._rebalance
        )
    
    def _coarse_selection_function(self, coarse):
        # if the rebalance flag is not 1, return null list to save time.
        if self._reb != 1:
            return self._long + self._short
            
        # make universe selection once a month
        # drop stocks which have no fundamental data or have too low prices
        selected = [x for x in coarse if (x.has_fundamental_data) 
                    and (float(x.price) > 5)]
        
        sorted_by_dollar_volume = sorted(selected, key=lambda x: x.dollar_volume, reverse=True) 
        top = sorted_by_dollar_volume[:self._num_coarse]
        return [i.symbol for i in top]

    def _fine_selection_function(self, fine):
        # return null list if it's not time to rebalance
        if self._reb != 1:
            return self._long + self._short
            
        self._reb = 0
            
        # drop stocks which don't have the information we need.
        # you can try replacing those factor with your own factors here
        filtered_fine = [x for x in fine if x.operation_ratios.operation_margin.value
                                        and x.valuation_ratios.price_change_1m 
                                        and x.valuation_ratios.book_value_per_share]
                                        
        self.log('remained to select %d'%(len(filtered_fine)))
        
        # rank stocks by three factor.
        sorted_by_factor1 = sorted(filtered_fine, key=lambda x: x.operation_ratios.operation_margin.value, reverse=True)
        sorted_by_factor2 = sorted(filtered_fine, key=lambda x: x.valuation_ratios.price_change_1m, reverse=True)
        sorted_by_factor3 = sorted(filtered_fine, key=lambda x: x.valuation_ratios.book_value_per_share, reverse=True)
        
        stock_dict = {}
        
        # assign a score to each stock, you can also change the rule of scoring here.
        for i,ele in enumerate(sorted_by_factor1):
            rank1 = i
            rank2 = sorted_by_factor2.index(ele)
            rank3 = sorted_by_factor3.index(ele)
            score = sum([rank1*0.25,rank2*0.5,rank3*0.25])
            stock_dict[ele] = score
        
        # sort the stocks by their scores
        sorted_stock = sorted(stock_dict.items(), key=lambda d:d[1],reverse=False)
        sorted_symbol = [x[0] for x in sorted_stock]

        # sotre the top stocks into the long_list and the bottom ones into the short_list
        self._long = [x.symbol for x in sorted_symbol[:self._num_fine]]
        self._short = [x.symbol for x in sorted_symbol[-self._num_fine:]]
    
        return self._long + self._short
    
    def _rebalance(self):
        # if this month the stock are not going to be long/short, liquidate it.
        long_short_list = self._long + self._short
        for symbol, security_holding in self.portfolio.items():
            if (security_holding.invested) and (symbol not in long_short_list):
                self.liquidate(symbol)
                
        # Alternatively, you can liquidate all the stocks at the end of each month.
        # Which method to choose depends on your investment philosiphy 
        # if you prefer to realized the gain/loss each month, you can choose this method.
    
        #self.liquidate()
        
        # Assign each stock equally. Alternatively you can design your own portfolio construction method
        for i in self._long:
            self.set_holdings(i, 0.9/self._num_fine)
        
        for i in self._short:
            self.set_holdings(i, -0.9/self._num_fine)

        self._reb = 1