#region imports
from AlgorithmImports import *
#endregion

class SectorMomentum(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 1)  
        self.SetCash(100000) 
        
        # daily ROC data
        self.data: Dict[str, RateOfChange] = {}
        
        self.roc_period: int = 6 * 21
        self.SetWarmUp(self.roc_period, Resolution.Daily)
        self.selected_symbol_count: int = 3 # long symbol count
        self.no_trading_fees: bool = True

        self.long_universe: List[str] = [
            "SPY",  # SPDR S&P 500 ETF Trust (large cap US stocks)
            "EEM",  # iShares MSCI Emerging Markets ETF (emerging market stocks)
            "EFA",  # iShares MSCI EAFE ETF (EAFE stocks)
            "VNQ",  # Vanguard Real Estate Index Fund ETF (real estate, REITs)
            "AGG",  # iShares Core U.S. Aggregate Bond ETF (fixed income ETF)
            "GLD",  # SPDR Gold Shares (GLD) (gold)
            "DBC",  # Invesco DB Commodity Index Tracking Fund (broad commodity index)
            "IWM",  # iShares Russell 2000 ETF (small cap US stocks)
        ]

        # NOTE long and short universe are the same for now, might be different in the future
        self.short_universe: List[str] = self.long_universe.copy()
        # self.short_symbols:List[str] = [
        #     "VNQ",  # Vanguard Real Estate Index Fund
        #     "XLK",  # Technology Select Sector SPDR Fund
        #     "XLE",  # Energy Select Sector SPDR Fund
        #     "XLV",  # Health Care Select Sector SPDR Fund
        #     "XLF",  # Financial Select Sector SPDR Fund
        #     "XLI",  # Industrials Select Sector SPDR Fund
        #     "XLB",  # Materials Select Sector SPDR Fund
        #     "XLY",  # Consumer Discretionary Select Sector SPDR Fund
        #     "XLP",  # Consumer Staples Select Sector SPDR Fund
        #     "XLU"   # Utilities Select Sector SPDR Fund
        # ]
        for ticker in set(self.long_universe + self.short_universe):
            data: Equity = self.AddEquity(ticker, Resolution.Daily)

            if ticker in self.long_universe:
                data.SetLeverage(5)

                if self.no_trading_fees:
                    data.SetFeeModel(CustomFeeModel())
                
                self.data[ticker] = self.ROC(ticker, self.roc_period, Resolution.Daily)
        
        self.data[self.long_universe[0]].Updated += self.OnROCUpdated
        self.recent_month: int = -1
        self.rebalance_flag: bool = False

    def OnROCUpdated(self, sender, updated) -> None:
        # set rebalance flag
        if self.recent_month != self.Time.month:
            self.recent_month = self.Time.month
            self.rebalance_flag = True
        
    def OnData(self, data: Slice) -> None:
        if self.IsWarmingUp: return

        # rebalance once a month
        if self.rebalance_flag:
            self.rebalance_flag = False
            
            # sort long universe by momentum
            sorted_by_momentum:List = sorted(
                [x for x in self.data.items() if x[1].IsReady and \
                x[0] in self.long_universe and \
                x[0] in data and data[x[0]]], \
                key = lambda x: x[1].Current.Value, reverse = True
            )

            if len(sorted_by_momentum) < self.selected_symbol_count:
                self.Liquidate()
                return

            long: List[str] = [x[0] for x in sorted_by_momentum[:self.selected_symbol_count]]
            
            # trade execution
            self.Liquidate()
            
            for ticker in long:
                quantity: int = self.Portfolio.TotalPortfolioValue // len(long) // data[ticker].Close
                self.MarketOrder(ticker, quantity)
            
            # short EW selected ETF
            short_universe: List[str] = [x for x in self.short_universe if x in data and data[x]]
            for ticker in short_universe:
                quantity: int = self.Portfolio.TotalPortfolioValue // len(short_universe) // data[ticker].Close
                self.MarketOrder(ticker, -quantity)
                
# Custom fee model
class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters):
        fee = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.0
        return OrderFee(CashAmount(fee, "USD"))