#region imports
from AlgorithmImports import *
#endregion
# https://quantpedia.com/Screener/Details/7
# The investment universe consists of global large cap stocks (or US large cap stocks). 
# At the end of the each month, the investor constructs equally weighted decile portfolios 
# by ranking the stocks on the past one year volatility of daily price. The investor 
# goes long stocks with the lowest volatility.


class ShortTermReversalAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2016, 12, 31) # Set Start Date
        self.set_end_date(2018, 6, 1)     # Set Start Date       
        self.set_cash(100000)             # Set Strategy Cash
        self._lookback = 252
    
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._coarse_selection_function, self._fine_selection_function)
        self._symbol_data_dict = {}
        self.add_equity("SPY", Resolution.DAILY)
        self.schedule.on(self.date_rules.month_start("SPY"),self.time_rules.after_market_open("SPY"), self._rebalance)

    def _coarse_selection_function(self, coarse):
        # 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)]
        # rank the stocks by dollar volume 
        filtered = sorted(selected, key=lambda x: x.dollar_volume, reverse=True) 
        return [x.symbol for x in filtered[:100]]

    def _fine_selection_function(self, fine):
        # filter stocks with the top market cap
        fine = [
            x for x in fine 
            if all([
                not np.isnan(factor) 
                for factor in [
                    x.earning_reports.basic_average_shares.three_months, 
                    x.earning_reports.basic_eps.twelve_months, 
                    x.valuation_ratios.pe_ratio
                ]
            ])
        ]
        top = sorted(fine, key=lambda x: x.earning_reports.basic_average_shares.three_months * (x.earning_reports.basic_eps.twelve_months*x.valuation_ratios.pe_ratio), reverse=True)
        return [x.symbol for x in top[:50]]
        
    def _rebalance(self):
        sorted_symbol_data = sorted(self._symbol_data_dict, key=lambda x: self._symbol_data_dict[x].volatility())
        # pick 5 stocks with the lowest volatility
        long_stocks = sorted_symbol_data[:5]
        stocks_invested = [x.key for x in self.portfolio if x.value.invested]
        # liquidate stocks not in the list 
        for i in stocks_invested:
            if i not in long_stocks:
                self.liquidate(i)
        # long stocks with the lowest volatility by equal weighting
        for i in long_stocks:
            self.set_holdings(i, 1/5)

    def on_data(self, data):
        for symbol, symbol_data in self._symbol_data_dict.items():
            if data.bars.contains_key(symbol):
                symbol_data.roc.update(self.time, self.securities[symbol].close)
        
    def on_securities_changed(self, changes):
        # clean up data for removed securities
        for removed in changes.removed_securities:
            symbol_data = self._symbol_data_dict.pop(removed.symbol, None)
            if symbol_data:
                symbol_data.dispose()

        # warm up the indicator with history price for newly added securities
        added_symbols = [x.symbol for x in changes.added_securities if x.symbol.value != "SPY"]
        history = self.history(added_symbols, self._lookback+1, Resolution.DAILY)

        for symbol in added_symbols:
            if symbol not in self._symbol_data_dict.keys():
                symbol_data = SymbolData(symbol, self._lookback)
                self._symbol_data_dict[symbol] = symbol_data
                if str(symbol) in history.index:             
                    symbol_data.warm_up_indicator(history.loc[str(symbol)])


class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    
    def __init__(self, symbol, lookback):
        self.symbol = symbol
        self.roc = RateOfChange(1)
        self.roc.updated += self._on_update
        self._roc_window = RollingWindow[IndicatorDataPoint](lookback)

    def _on_update(self, sender, updated):
        self._roc_window.add(updated)

    def warm_up_indicator(self, history):
        # warm up the RateOfChange indicator with the history request
        for t, row in history.iterrows():
            self.roc.update(t, row.close)
            
    def volatility(self):
        data = [float(x.value) for x in self._roc_window]
        return np.std(data)

    def dispose(self):
        self.roc.updated -= self._on_update
        self._roc_window.reset()