# https://quantpedia.com/strategies/momentum-and-reversal-combined-with-volatility-effect-in-stocks/
#
# The investment universe consists of NYSE, AMEX, and NASDAQ stocks with prices higher than $5 per share. At the beginning of each month, 
# the sample is divided into equal halves, at the size median, and only larger stocks are used. Then each month, realized returns and realized 
# (annualized) volatilities are calculated for each stock for the past six months. One week (seven calendar days) prior to the beginning of 
# each month is skipped to avoid biases due to microstructures. Stocks are then sorted into quintiles based on their realized past returns 
# and past volatility. The investor goes long on stocks from the highest performing quintile from the highest volatility group and short on 
# stocks from the lowest-performing quintile from the highest volatility group. Stocks are equally weighted and held for six months 
# (therefore, 1/6 of the portfolio is rebalanced every month).
#
# QC implementation changes:
#   - Universe consists of top 3000 US stock by market cap from NYSE, AMEX and NASDAQ.
    
from numpy import sqrt
from collections import deque
import numpy as np

class MomentumReversalCombinedWithVolatilityEffectinStocks(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 1)
        self.SetCash(100000) 

        self.symbol = self.AddEquity('SPY', Resolution.Daily).Symbol
        
        # EW Tranching.
        self.holding_period = 6
        self.managed_queue = deque(maxlen = self.holding_period + 1)

        # Daily price data.
        self.data = {}
        self.period = 6 * 21
        
        self.coarse_count = 500
        self.selection_flag = True
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.Schedule.On(self.DateRules.MonthEnd(self.symbol), self.TimeRules.AfterMarketOpen(self.symbol), self.Selection)

    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            security.SetFeeModel(CustomFeeModel(self))
            security.SetLeverage(5)
            
    def CoarseSelectionFunction(self, coarse):
        # Update the rolling window every day.
        for stock in coarse:
            symbol = stock.Symbol

            # Store monthly price.
            if symbol in self.data:
                self.data[symbol].update(stock.AdjustedPrice)

        if not self.selection_flag:
            return Universe.Unchanged
        
        # selected = [x for x in coarse if x.HasFundamentalData and x.Market == 'usa' and x.Price > 5]
        
        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Market == 'usa' and x.Price > 5],    \
                key = lambda x: x.DollarVolume, reverse = True)[:self.coarse_count]
        
        # Warmup price rolling windows.
        for stock in selected:
            symbol = stock.Symbol
            if symbol in self.data:
                continue
            
            self.data[symbol] = SymbolData(symbol, self.period)
            history = self.History(symbol, self.period, Resolution.Daily)
            if history.empty:
                self.Log(f"Not enough data for {symbol} yet.")
                continue
            closes = history.loc[symbol].close
            for time, close in closes.iteritems():
                self.data[symbol].update(close)
                
        return [x.Symbol for x in selected if self.data[x.Symbol].is_ready()]

    def FineSelectionFunction(self, fine):
        fine = [x for x in fine if x.MarketCap != 0 and \
                    ((x.SecurityReference.ExchangeId == "NYS") or (x.SecurityReference.ExchangeId == "NAS") or (x.SecurityReference.ExchangeId == "ASE"))]
                    
        sorted_by_market_cap = sorted(fine, key = lambda x: x.MarketCap, reverse=True)
        half = int(len(sorted_by_market_cap) / 2)
        top_by_market_cap = [x.Symbol for x in sorted_by_market_cap][:half]
        
        # Performance and volatility tuple.
        perf_volatility = {}
        for symbol in top_by_market_cap:
            performance = self.data[symbol].performance()
            annualized_volatility = self.data[symbol].volatility()
            perf_volatility[symbol] = (performance, annualized_volatility)
        
        sorted_by_perf = sorted(perf_volatility.items(), key = lambda x: x[1][0], reverse = True)
        quintile = int(len(sorted_by_perf) / 5)
        top_by_perf = [x[0] for x in sorted_by_perf[:quintile]]
        low_by_perf = [x[0] for x in sorted_by_perf[-quintile:]]
        
        sorted_by_vol = sorted(perf_volatility.items(), key = lambda x: x[1][1], reverse = True)
        quintile = int(len(sorted_by_vol) / 5)
        top_by_vol = [x[0] for x in sorted_by_vol[:quintile]]
        low_by_vol = [x[0] for x in sorted_by_vol[-quintile:]]
        
        long = [x for x in top_by_perf if x in top_by_vol if not self.IsInvested(x)]
        short = [x for x in low_by_perf if x in top_by_vol if not self.IsInvested(x)]

        self.managed_queue.append(RebalanceQueueItem(long, short))
        
        return long + short
        
    def OnData(self, data):
        if not self.selection_flag:
            return
        self.selection_flag = False
       
        # Trade execution.
        if len(self.managed_queue) == 0: return
    
        # Liquidate first items if queue is full.
        if len(self.managed_queue) == self.managed_queue.maxlen:
            item_to_liquidate = self.managed_queue.popleft()
            for symbol in item_to_liquidate.long_symbols + item_to_liquidate.short_symbols:
                self.Liquidate(symbol)
            
        curr_stock_set = self.managed_queue[-1]
        if curr_stock_set.count == 0: return
    
        weight = 1 / self.holding_period
        
        # Open new trades.
        for symbol in curr_stock_set.long_symbols:
            self.SetHoldings(symbol, weight / len(curr_stock_set.long_symbols))
            
        for symbol in curr_stock_set.short_symbols:
            self.SetHoldings(symbol, -weight / len(curr_stock_set.short_symbols))
    
    def Selection(self):
        self.selection_flag = True
        
    def IsInvested(self, symbol):
        return self.Securities.ContainsKey(symbol) and self.Portfolio[symbol].Invested

class RebalanceQueueItem():
    def __init__(self, long_symbols, short_symbols):
        self.long_symbols = long_symbols
        self.short_symbols = short_symbols
        self.count = len(long_symbols + short_symbols)

class SymbolData():
    def __init__(self, symbol, period):
        self.Symbol = symbol
        self.Price = RollingWindow[float](period)
    
    def update(self, value):
        self.Price.Add(value)
    
    def is_ready(self):
        return self.Price.IsReady
    
    def update(self, close):
        self.Price.Add(close)
        
    def volatility(self):
        closes = np.array([x for x in self.Price][5:]) # Skip last week.
        daily_returns = closes[:-1] / closes[1:] - 1
        return np.std(daily_returns) * sqrt(252 / (len(closes)))
        
    def performance(self):
        closes = [x for x in self.Price][5:] # Skip last week.
        return (closes[0] / closes[-1] - 1)

# Custom fee model.
class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters):
        fee = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))