#region imports
from AlgorithmImports import *
#endregion
#This is a Template of dynamic stock selection.
#You can try your own fundamental factor and ranking method by editing the _select method

class BasicTemplateAlgorithm(QCAlgorithm):
    _long = []
    _short = []

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100000)
        self.settings.seed_initial_prices = True

        date_rules = self.date_rules.month_start("SPY")
        self.universe_settings.schedule.on(date_rules)
        self.add_universe(self._select)

        # Schedule the rebalance function to execute at the begining of each month
        self.schedule.on(date_rules, self.time_rules.after_market_open("SPY", 5), self._rebalance)

    def _select(self, fundamental):
        top_volume = 250
        top_rank = 10

        # Drop stocks which don't have the information we need or have too low prices
        # You can try replacing those factor with your own factors here
        fundamental = sorted(
            [x for x in fundamental if x.price > 5
                and not any([np.isnan(x.valuation_ratios.price_change_1m),
                np.isnan(x.valuation_ratios.book_value_per_share),
                np.isnan(x.operation_ratios.operation_margin.value)]) ],
            key=lambda x: x.dollar_volume, reverse=True)

        self.plot('Universe', '-', len(fundamental))

        fundamental = fundamental[:top_volume]
        
        # Rank stocks by three factor.
        sorted_by_factor1 = sorted(fundamental, key=lambda x: x.valuation_ratios.price_change_1m, reverse=True)
        sorted_by_factor2 = sorted(fundamental, key=lambda x: x.valuation_ratios.book_value_per_share, reverse=True)
        sorted_by_factor3 = sorted(fundamental, key=lambda x: x.operation_ratios.operation_margin.value, reverse=True)

        # Assign a score to each stock, you can also change the rule of scoring here.
        stock_dict = {}
        for rank1, ele in enumerate(sorted_by_factor1):
            rank2 = sorted_by_factor2.index(ele)
            rank3 = sorted_by_factor3.index(ele)
            stock_dict[ele] = .2*rank1 + .4*rank2 + .4*rank3

        # Store the top stocks into the long_list and the bottom ones into the short_list
        sorted_stock = sorted(stock_dict.items(), key=lambda d:d[1],reverse=False)
        self._long = [x.symbol for x,_ in sorted_stock[:top_rank]]
        self._short = [x.symbol for x,_ in sorted_stock[-top_rank:]]

        return self._long + self._short

    def _rebalance(self):
        # 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()

        # In this example, we will have a long bias
        targets  = [PortfolioTarget(s,  1.0/len(self._long)) for s in self._long]
        targets += [PortfolioTarget(s, -0.5/len(self._short)) for s in self._short]
        self.set_holdings(targets, liquidate_existing_holdings=True)