#region imports
from AlgorithmImports import *

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


class MomentumReversalCombinedWithVolatility(QCAlgorithm):

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

        self.set_security_initializer(BrokerageModelSecurityInitializer(
            self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._coarse_selection_function, self._fine_selection_function)
        self._data_dict = {}
        # 1/6 of the portfolio is rebalanced every month
        self._portfolios = deque(maxlen=6)
        spy = Symbol.create("SPY", SecurityType.EQUITY, Market.USA)
        self.schedule.on(self.date_rules.month_start(spy), self.time_rules.after_market_open(spy), self._rebalance)
        # the lookback period for volatility and return is six months
        self._lookback = 20*6
        self._filtered_fine = None
        self._monthly_rebalance = False
        self.set_warm_up(self._lookback)

    def _coarse_selection_function(self, coarse):
        # update the price of stocks in universe everyday
        for i in coarse:
            if i.symbol not in self._data_dict:
                self._data_dict[i.symbol] = SymbolData(self._lookback)
            self._data_dict[i.symbol].update(i.adjusted_price)
        
        if self._monthly_rebalance:
            # drop stocks which have no fundamental data or have too low prices
            return [x.symbol for x in coarse if (x.has_fundamental_data) and (float(x.price) > 5)]
        else:
            return []

    def _fine_selection_function(self, fine):
        if self._monthly_rebalance:  
            sorted_fine = sorted(fine, key=lambda x: x.earning_reports.basic_average_shares.value * self._data_dict[x.symbol].price, reverse=True)
            # select stocks with large size 
            top_fine = sorted_fine[:int(0.5*len(sorted_fine))]
            self._filtered_fine = [x.symbol for x in top_fine]
            return self._filtered_fine
        else:
            return []
        
    def _rebalance(self):
        self._monthly_rebalance = True
           
    def on_data(self, data):
        if self._monthly_rebalance and self._filtered_fine and not self.is_warming_up:
            filtered_data = {symbol: symbolData for (symbol, symbolData) in self._data_dict.items() if symbol in self._filtered_fine and symbolData.is_ready() and symbol in data.bars}
            
            self._filtered_fine = None
            self._monthly_rebalance = False
            
            # if the dictionary is empty, then return    
            if len(filtered_data) < 100: 
                return
            # sort the universe by volatility and select stocks in the top high volatility quintile
            sorted_by_vol = sorted(filtered_data.items(), key=lambda x: x[1].volatility(), reverse=True)[:int(0.2*len(filtered_data))]
            sorted_by_vol = dict(sorted_by_vol)
            # sort the stocks in top-quintile by realized return
            sorted_by_return = sorted(sorted_by_vol, key=lambda x: sorted_by_vol[x].return_(), reverse=True)
            long_ = sorted_by_return[:int(0.2*len(sorted_by_return))]
            short = sorted_by_return[-int(0.2*len(sorted_by_return)):]
            
            self._portfolios.append(short + long_)
            
            # 1/6 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 6 months 
            short_weight = 1/len(short)
            for i in short:
                self.set_holdings(i, -1/6*short_weight)
            
            long_weight = 1/len(long_)
            for i in long_:
                self.set_holdings(i, 1/6*long_weight)
                

class SymbolData:

    def __init__(self, lookback):
        self.price = None
        self._history = deque(maxlen=lookback)
        
    def update(self, value):
        # update yesterday's close price
        self.price = value
        # update the history price series 
        self._history.append(float(value))

    def is_ready(self):
        return len(self._history) == self._history.maxlen

    def volatility(self):
        # one week (5 trading days) prior to the beginning of each month is skipped 
        prices = np.array(self._history)[:-5]
        returns = (prices[1:]-prices[:-1])/prices[:-1]
        # calculate the annualized realized volatility
        return np.std(returns)*np.sqrt(250/len(returns))
    
    def return_(self):
        # one week (5 trading days) prior to the beginning of each month is skipped 
        prices = np.array(self._history)[:-5]
        # calculate the annualized realized return
        return (prices[-1]-prices[0])/prices[0]