#region imports
from AlgorithmImports import *

from collections import deque
#endregion
# https://quantpedia.com/Screener/Details/66


class CombiningMomentumEffectWithVolume(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2012, 1, 1)
        self.set_end_date(2018, 8, 1)
        self.set_cash(100_000)

        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._coarse_selection_function, self._fine_selection_function)
        self._data_dict = {}
        # 1/3 of the portfolio is rebalanced every month
        self._portfolios = deque(maxlen=3)
        self.schedule.on(
            self.date_rules.month_start(Symbol.create("SPY", SecurityType.EQUITY, Market.USA)),
            self.time_rules.at(0, 0), self._rebalance)
        # the lookback period for return calculation
        self._lookback = 252
        self._filtered_fine = None
        self._monthly_rebalance = False
        self.set_warmup(self._lookback + 1)

    def _coarse_selection_function(self, coarse):
        for i in coarse:
            if i.symbol not in self._data_dict:
                self._data_dict[i.symbol] = SymbolData(i.symbol, self._lookback)
            self._data_dict[i.symbol].roc.update(i.end_time, i.adjusted_price)
            self._data_dict[i.symbol].volume = i.volume
        
        if self._monthly_rebalance:
            # drop stocks which have no fundamental data 
            return [x.symbol for x in coarse if x.has_fundamental_data]
        return []

    def _fine_selection_function(self, fine):
        if self._monthly_rebalance:
            data_ready = {symbol: symbolData for (symbol, symbolData) in self._data_dict.items() if symbolData.roc.is_ready}
            if len(data_ready) < 100: 
                self._filtered_fine = []
            else:
                sorted_fine = [i for i in fine if i.earning_reports.basic_average_shares.three_months != 0 and i.symbol in data_ready]
                sorted_fine_symbols = [i.symbol for i in sorted_fine]
                filtered_data = {symbol: symbolData for (symbol, symbolData) in data_ready.items() if symbol in sorted_fine_symbols}
                for i in sorted_fine:
                    if i.symbol in filtered_data and filtered_data[i.symbol].volume != 0:
                        filtered_data[i.symbol].turnover = i.earning_reports.basic_average_shares.three_months / filtered_data[i.symbol].volume
                sorted_by_roc = sorted(filtered_data.values(), key = lambda x: x.roc.current.value, reverse = True)
                top_roc = sorted_by_roc[:int(len(sorted_by_roc)*0.2)]
                bottom_roc = sorted_by_roc[-int(len(sorted_by_roc)*0.2):]
                
                high_turnover_top_roc = sorted(top_roc, key = lambda x: x.turnover, reverse = True)
                high_turnover_bottom_roc = sorted(bottom_roc, key = lambda x: x.turnover, reverse = True)
                
                self._long =  [i.symbol for i in high_turnover_top_roc[:int(len(high_turnover_top_roc)*0.01)]]
                self._short = [i.symbol for i in high_turnover_bottom_roc[:int(len(high_turnover_bottom_roc)*0.01)]]
                self._filtered_fine = self._long + self._short
                self._portfolios.append(self._filtered_fine)
        else:
            self._filtered_fine = []
        
        if not self._filtered_fine:
            self._monthly_rebalance = False

        return self._filtered_fine
        
    def _rebalance(self):
        self._monthly_rebalance = True
           
    def on_data(self, data):
        if self._monthly_rebalance and self._filtered_fine:
            self._filtered_fine = None
            self._monthly_rebalance = False
            
            # 1/3 of the portfolio is rebalanced every month
            if len(self._portfolios) == self._portfolios.maxlen:
                for i in list(self._portfolios)[0]:
                    self.liquidate(i)
            
            # stocks are equally weighted and held for 3 months 
            short = [s for s in self._short if s in data.Bars]
            short_weight = 1/len(short)
            for i in short:
                self.set_holdings(i, -1/3*short_weight)
            
            long_ = [s for s in self._long if s in data.Bars]
            long_weight = 1/len(long_)
            for i in long_:
                self.set_holdings(i, 1/3*long_weight)            
 

class SymbolData:
    
    def __init__(self, symbol, lookback):
        self.symbol = symbol
        self.roc = RateOfChange(lookback)
        self.volume = None