#region imports
from AlgorithmImports import *
#endregion
# https://quantpedia.com/Screener/Details/51


class MomentumShortTermReversalAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2010, 1, 1)  
        self.set_end_date(2018, 5, 10)    
        self.set_cash(100000)         

        self.universe_settings.resolution = Resolution.DAILY
        self.set_security_initializer(BrokerageModelSecurityInitializer(
            self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))

        spy = Symbol.create("SPY", SecurityType.EQUITY, Market.USA)
        date_rule = self.date_rules.month_start(spy)
        self.universe_settings.schedule.on(date_rule)
        self.add_universe(self._coarse_selection_function)
        self.schedule.on(date_rule, self.time_rules.midnight, self._rebalance)
        self._symbol_price = {}
        self._decrease_winner = None
        self._increase_loser = None
        self.set_warmup(timedelta(365))
        
    def _coarse_selection_function(self, coarse):
        for i in coarse:
            if i.symbol not in self._symbol_price:
                self._symbol_price[i.symbol] = SymbolData(self, i.symbol)
            self._symbol_price[i.symbol].window.add(float(i.adjusted_price))
            if self._symbol_price[i.symbol].window.is_ready:
                price = np.array([i for i in self._symbol_price[i.symbol].window])
                returns = (price[:-1]-price[1:])/price[1:]
                self._symbol_price[i.symbol].yearly_return = (price[0]-price[-1])/price[-1]
                GARR_12 = np.prod([(1+i)**(1/12) for i in returns])-1
                GARR_1 = (1+returns[0])**(1/12)-1
                self._symbol_price[i.symbol].garr_ratio = GARR_1 / GARR_12

        ready_symbol_price = {symbol: symbol_data for symbol, symbol_data in self._symbol_price.items() if symbol_data.window.is_ready}
        if ready_symbol_price and len(ready_symbol_price)>50:        
            sorted_by_return = sorted(ready_symbol_price, key=lambda x: ready_symbol_price[x].yearly_return, reverse=True)
            winner = sorted_by_return[:int(len(sorted_by_return)*0.3)]
            loser = sorted_by_return[-int(len(sorted_by_return)*0.3):]
            self._decrease_winner = sorted(winner, key=lambda x: ready_symbol_price[x].garr_ratio)[:50] 
            self._increase_loser = sorted(loser, key=lambda x: ready_symbol_price[x].garr_ratio, reverse=True)[:50]
            return self._decrease_winner+self._increase_loser
        else:
            return []

    def _get_tradable_assets(self, symbols):
        tradable_assets = []
        for symbol in symbols:
            security = self.securities[symbol]
            if security.price and security.is_tradable:
                tradable_assets.append(symbol)
        return tradable_assets[:15]

    def _rebalance(self):
        if self.is_warming_up:
            return

        self._decrease_winner = self._get_tradable_assets(self._decrease_winner)
        self._increase_loser = self._get_tradable_assets(self._increase_loser)
        if not (self._decrease_winner and self._increase_loser):
            return

        stocks_invested = [x.key for x in self.portfolio]
        for i in stocks_invested:
            if i not in self._decrease_winner+self._increase_loser:
                self.liquidate(i) 
        
        short_weight = 0.25/len(self._increase_loser)
        for j in self._increase_loser:
            self.set_holdings(j, -short_weight)
            
        long_weight = 0.25/len(self._decrease_winner)
        for i in self._decrease_winner:
            self.set_holdings(i, long_weight)

 
class SymbolData:
    def __init__(self, algorithm, symbol):
        self.symbol = symbol
        self.window = RollingWindow[float](13)
        self.garr_ratio = None
        self.yearly_return = None