# https://quantpedia.com/strategies/12-month-cycle-in-cross-section-of-stocks-returns/
#
# The top 30% of firms based on their market cap from NYSE and AMEX are part of the investment universe. Every month, stocks are grouped 
# into ten portfolios (with an equal number of stocks in each portfolio) according to their performance in one month one year ago. Investors
# go long in stocks from the winner decile and shorts stocks from the loser decile. The portfolio is equally weighted and rebalanced every month.
#
# QC implementation changes:
#   - Universe consists of top 3000 US stock by market cap from NYSE, AMEX and NASDAQ.
#   - Portfolio is value weighted.

from AlgorithmImports import *
from typing import List, Dict, Tuple

class Month12CycleinCrossSectionofStocksReturns(QCAlgorithm):

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

        self.symbol:Symbol = self.AddEquity('SPY', Resolution.Daily).Symbol
        
        self.coarse_count:int = 500
        self.leverage:int = 5
        self.quantile:int = 10
        
        # Monthly close data.
        self.data:Dict[Symbol, SymbolData] = {}
        self.period:int = 13
        
        self.weight:Dict[Symbol, float] = {}

        self.selection_flag:bool = False
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.Schedule.On(self.DateRules.MonthEnd(self.symbol), self.TimeRules.BeforeMarketClose(self.symbol), self.Selection)

    def OnSecuritiesChanged(self, changes:SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            security.SetFeeModel(CustomFeeModel())
            security.SetLeverage(self.leverage)
            
    def CoarseSelectionFunction(self, coarse:List[CoarseFundamental]) -> List[Symbol]:
        if not self.selection_flag:
            return Universe.Unchanged

        # Update the rolling window every month.
        for stock in coarse:
            symbol:symbol = stock.Symbol

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

        # selected = [x.Symbol for x in coarse if x.HasFundamentalData and x.Market == 'usa']
        selected:List[Symbol] = [x.Symbol
            for x in sorted([x for x in coarse if x.HasFundamentalData and x.Market == 'usa'],
                key = lambda x: x.DollarVolume, reverse = True)[:self.coarse_count]]
        
        # Warmup price rolling windows.
        for symbol in selected:
            if symbol in self.data:
                continue
            
            self.data[symbol] = SymbolData(self.period)
            history:DataFrame = self.History(symbol, self.period*30, Resolution.Daily)
            if history.empty:
                self.Log(f"Not enough data for {symbol} yet.")
                continue
            closes:pd.Series = history.loc[symbol].close
            
            closes_len:int = len(closes.keys())
            # Find monthly closes.
            for index, time_close in enumerate(closes.iteritems()):
                # index out of bounds check.
                if index + 1 < closes_len:
                    date_month:int = time_close[0].date().month
                    next_date_month:int = closes.keys()[index + 1].month
                
                    # Found last day of month.
                    if date_month != next_date_month:
                        self.data[symbol].update(time_close[1])
            
        return [x for x in selected if self.data[x].is_ready()]    
        
    def FineSelectionFunction(self, fine:List[FineFundamental]) -> List[Symbol]:
        fine:List[Symbol] = [x for x in fine if x.MarketCap != 0 and x.CompanyReference.IsREIT != 1 and  \
                    ((x.SecurityReference.ExchangeId == "NYS") or (x.SecurityReference.ExchangeId == "NAS") or (x.SecurityReference.ExchangeId == "ASE"))]
                    
        if len(fine) > self.coarse_count:
            sorted_by_market_cap:List = sorted(fine, key = lambda x: x.MarketCap, reverse=True)
            top_by_market_cap:List[Symbol] = sorted_by_market_cap[:self.coarse_count]
        else:
            top_by_market_cap = fine

        # Performance sorting. One month performance, one year ago with market cap data.
        performance_market_cap:Dict[Symbol, Tuple[float, float]] = { x.Symbol : (self.data[x.Symbol].performance(), x.MarketCap) for x in top_by_market_cap if x.Symbol in self.data and self.data[x.Symbol].is_ready()}
        
        long:List[Tuple] = []
        short:List[Tuple] = []

        if len(performance_market_cap) >= self.quantile:
            sorted_by_perf:List[Tuple] = sorted(performance_market_cap.items(), key = lambda x:x[1][0], reverse = True)
            quantile:int = int(len(sorted_by_perf) / self.quantile)
            long = [x for x in sorted_by_perf[:quantile]]
            short = [x for x in sorted_by_perf[-quantile:]]
        
        total_market_cap_long:float = sum([x[1][1] for x in long])
        for symbol, perf_market_cap in long:
            self.weight[symbol] = perf_market_cap[1] / total_market_cap_long

        total_market_cap_short = sum([x[1][1] for x in short])
        for symbol, perf_market_cap in short:
            self.weight[symbol] = perf_market_cap[1] / total_market_cap_short
        
        return [x[0] for x in self.weight.items()]

    def OnData(self, data:Slice) -> None:
        if not self.selection_flag:
            return
        self.selection_flag = False
        
        # Trade execution.
        stocks_invested:List[Symbol] = [x.Key for x in self.Portfolio if x.Value.Invested]
        for symbol in stocks_invested:
            if symbol not in self.weight:
                self.Liquidate(symbol)
        
        for symbol, w in self.weight.items():
            if symbol in data and data[symbol]:
                self.SetHoldings(symbol, w)

        self.weight.clear()
    
    def Selection(self) -> None:
        self.selection_flag = True

class SymbolData():
    def __init__(self, period:int):
        self.Window:RollingWindow = RollingWindow[float](period)
    
    def update(self, value:float) -> None:
        self.Window.Add(value)
    
    def is_ready(self) -> bool:
        return self.Window.IsReady
        
    # One month performance, one year ago.
    def performance(self) -> float:
        values:float = [x for x in self.Window]
        return (values[-2] / values[-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"))