# https://quantpedia.com/strategies/momentum-factor-effect-in-stocks/
#
# The investment universe consists of NYSE, AMEX, and NASDAQ stocks. We define momentum as the past 12-month return, skipping the most 
# recent month’s return (to avoid microstructure and liquidity biases). To capture “momentum”, UMD portfolio goes long stocks that have 
# high relative past one-year returns and short stocks that have low relative past one-year returns.
#
# QC implementation changes:
#   - Instead of all listed stock, we select top 500 stocks by market cap from QC stock universe.

# region imports
from AlgorithmImports import *
from typing import List, Dict
from pandas.core.frame import DataFrame
# endregion

class MomentumFactorEffectinStocks(QCAlgorithm):

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

        self.AddEquity('SPY', Resolution.Daily).Symbol
        
        self.weight: Dict[Symbol, float] = {}
        self.data: Dict[Symbol, RollingWindow] = {}
        self.period: int = 12 * 21
        self.quantile: int = 5
        self.leverage: int = 5
        self.exchange_codes: List[str] = ['NYS', 'NAS', 'ASE']
        
        self.fundamental_count: int = 500
        self.fundamental_sorting_key = lambda x: x.DollarVolume

        self.selection_flag: bool = False
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.FundamentalSelectionFunction)
        
        self.settings.daily_precise_end_time = False
        self.settings.minimum_order_margin_portfolio_percentage = 0.

        market: Symbol = self.AddEquity('SPY', Resolution.Daily).Symbol
        self.schedule.on(self.date_rules.month_start(market),
                        self.time_rules.after_market_open(market),
                        self.selection)

    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            security.SetFeeModel(CustomFeeModel())
            security.SetLeverage(self.leverage)

    def FundamentalSelectionFunction(self, fundamental: List[Fundamental]) -> List[Symbol]:
        # update the rolling window every day
        [self.data[stock.Symbol].Add(stock.Price) for stock in fundamental if stock.Symbol in self.data]

        if not self.selection_flag:
            return Universe.Unchanged

        selected: List[Fundamental] = [
            x for x in fundamental if x.HasFundamentalData and x.MarketCap != 0 and x.Market == 'usa' and \
            x.SecurityReference.ExchangeId in self.exchange_codes
        ]
        if len(selected) > self.fundamental_count:
            selected = [x for x in sorted(selected, key=self.fundamental_sorting_key, reverse=True)[:self.fundamental_count]]

        # warmup price rolling windows
        for stock in selected:
            symbol: Symbol = stock.Symbol

            if symbol not in self.data:
                self.data[symbol] = RollingWindow[float](self.period)
                history: DataFrame = self.History(symbol, self.period, Resolution.Daily)

                if history.empty:
                    self.Log(f"Not enough data for {symbol} yet")
                    continue
                closes: Series = history.loc[symbol].close
                for time, close in closes.items():
                    self.data[symbol].Add(close)
            
        perf: Dict[Symbol, float] = {
            stock.Symbol : self.data[stock.Symbol][0] / self.data[stock.Symbol][self.period - 1] - 1 for stock in selected if self.data[stock.Symbol].IsReady
        }

        if len(perf) >= self.quantile:
            sorted_by_perf: List = sorted(perf, key=perf.get)
            quantile: int = int(len(sorted_by_perf) / self.quantile)
            long: List[Symbol] = sorted_by_perf[-quantile:]
            short: List[Symbol] = sorted_by_perf[:quantile]

            # calculate weights
            for i, portfolio in enumerate([long, short]):
                for symbol in portfolio:
                    self.weight[symbol] = ((-1) ** i) / len(portfolio)
        
        return list(self.weight.keys())
        
    def OnData(self, slice: Slice) -> None:
        if not self.selection_flag:
            return
        self.selection_flag = False

        # trade execution
        portfolio: List[PortfolioTarget] = [
            PortfolioTarget(symbol, w) for symbol, w in self.weight.items() if symbol in slice and slice[symbol]
        ]
        self.SetHoldings(portfolio, True)

        self.weight.clear()

    def selection(self) -> None:
        self.selection_flag = True

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