#region imports
from AlgorithmImports import *
#endregion


class SeasonalitySignalAlgorithm(QCAlgorithm):
    '''
    A strategy that takes long and short positions based on historical same-calendar month returns
    Paper: https://www.nber.org/papers/w20815.pdf
    '''

    def initialize(self):
        self.set_start_date(2012, 1, 1)   # Set Start Date
        self.set_end_date(2019, 8, 1)     # Set End Date
        self.set_cash(100000)            # Set Strategy Cash

        self._num_coarse = 100           # Number of equities for coarse selection
        self._num_long = 5               # Number of equities to long
        self._num_short = 5              # Number of equities to short
        self._long_symbols = []           # Contain the equities we'd like to long
        self._short_symbols = []          # Contain the equities we'd like to short

        self.universe_settings.resolution = Resolution.DAILY     # Resolution of universe selection
        self._universe = self.add_universe(self._same_month_return_selection)         # Universe selection based on historical same-calendar month returns

        self._next_rebalance = self.time  # Next rebalance time

    def _same_month_return_selection(self, coarse):
        '''
        Universe selection based on historical same-calendar month returns
        '''
        # Before next rebalance time, just remain the current universe
        if self.time < self._next_rebalance:
            return Universe.UNCHANGED

        # Sort the equities with prices > 5 in DollarVolume decendingly
        selected = sorted([x for x in coarse if x.price > 5],
                          key=lambda x: x.dollar_volume, reverse=True)

        # Get equities after coarse selection
        symbols = [x.symbol for x in selected[:self._num_coarse]]

        # Get historical close data for coarse-selected symbols of the same calendar month
        start = self.time.replace(day = 1, year = self.time.year-1)
        end = Expiry.end_of_month(start) - timedelta(1)
        history = self.history(symbols, start, end, Resolution.DAILY).close.unstack(level=0)

        # Get the same calendar month returns for the symbols
        monthly_return = {ticker: prices.iloc[-1]/prices.iloc[0] for ticker, prices in history.items()}

        # Sorted the values of monthly return
        sorted_return = sorted(monthly_return.items(), key=lambda x: x[1], reverse=True)

        # Get the symbols to long / short
        self._long_symbols = [x[0] for x in sorted_return[:self._num_long]]
        self._short_symbols = [x[0] for x in sorted_return[-self._num_short:]]

        # Note that self._long_symbols/self._short_symbols contains strings instead of symbols
        return [x for x in symbols if str(x) in self._long_symbols + self._short_symbols]

    def _get_tradable_assets(self, symbols):
        return [s for s in symbols if self.securities[s].is_tradable]

    def on_data(self, data):
        '''
        Rebalance every month based on same-calendar month returns effect
        '''
        # Before next rebalance, do nothing
        if self.time < self._next_rebalance:
            return

        long_symbols = self._get_tradable_assets(self._long_symbols)
        short_symbols = self._get_tradable_assets(self._short_symbols)
        # Open long positions
        for symbol in long_symbols:
            self.set_holdings(symbol, 0.5/len(long_symbols))

        # Open short positions
        for symbol in short_symbols:
            self.set_holdings(symbol, -0.5/len(short_symbols))

        # Rebalance at the end of every month
        self._next_rebalance = Expiry.end_of_month(self.time) - timedelta(1)

    def on_securities_changed(self, changes):
        '''
        Liquidate the stocks that are not in the universe
        '''
        for security in changes.removed_securities:
            if security.invested:
                self.liquidate(security.symbol, 'Removed from Universe')