# region imports
from AlgorithmImports import *
# endregion


class FactorMomentumDistressAlgorithm(QCAlgorithm):
    """
    Long-only equity strategy using 10 momentum/distress factors.
    Universe: SPY ETF constituents. Monthly rebalance. $1B capital.

    Factors (6 momentum, 4 distress):
      1. 252-21 momentum  (12m return minus 1m return)  - momentum
      2. 63-day momentum  (3-month return)              - momentum
      3. 21-day momentum  (1-month return)              - momentum
      4. RSI 14 (oversold = distress)                   - distress
      5. Rate-of-change 126-day (6-month)               - momentum
      6. Price / SMA200 - 1                             - momentum
      7. STD 21-day (short-term vol distress)           - distress
      8. STD 63-day (volatility distress)               - distress
      9. MACD (line - signal)                           - momentum
     10. Bollinger %B (position within bands)           - distress
    """
    # Composite score: momentum factors positive, distress negative
    # Momentum (higher = better):  1(252-21), 2(63d), 3(21d), 5(ROC126), 6(Price/SMA200), 9(MACD)
    # Distress (higher = worse):    7(STD21), 8(STD63)
    # Distress (higher = better):   4(RSI), 10(BB%B)
    _signs = [1, 1, 1, 1, 1, 1, -1, -1, 1, 1]
    _NUM_LONG = 50  # The number of assets to hold.

    def initialize(self) -> None:
        self.set_start_date(2023, 1, 1)
        self.set_cash(1_000_000_000)
        self.settings.seed_initial_prices = True
        self.settings.minimum_order_margin_portfolio_percentage = 0
        # Add a universe that selects a subset of the SPY ETF Constituents to hold.
        self._selection_data_by_symbol = {}
        self.universe_settings.resolution = Resolution.DAILY
        date_rule = self.date_rules.month_start("SPY")
        self.universe_settings.schedule.on(date_rule)
        self._universe = self.add_universe(self.universe.etf("SPY"), self._select_assets)
        # Add a Scheduled Event to rebalance the portfolio at the start of each month.
        self._weight_by_symbol = {}
        self.schedule.on(date_rule, self.time_rules.at(8, 0), self._rebalance)
        # Add warm-up to prime the indicators.
        self.set_warm_up(timedelta(400))

    def _select_assets(self, fundamentals: list[Fundamental]) -> list[Symbol]:
        # Update the indicator of all stocks in the universe dataset and
        # get the subset of stocks that have their indicator ready.
        ready_stocks = [
            f for f in fundamentals
            if self._selection_data_by_symbol.setdefault(f.symbol, SelectionData(self, f)).update(f)
        ]
        # As assests leave the Fundamental dataset, delete their SelectionData object.
        for symbol in self._selection_data_by_symbol.keys() - {f.symbol for f in fundamentals}:
            del self._selection_data_by_symbol[symbol]
        # During warm-up, keep the universe empty.
        if self.is_warming_up:
            return []
        # Compute 10 factors for each symbol with enough history
        factor_data = {
            f.symbol: self._selection_data_by_symbol[f.symbol].compute_factors()
            for f in ready_stocks
        }
        # Standardize the raw factor values.
        scored = list(factor_data.items())
        X = np.array([fv for _, fv in scored])
        mean, std = X.mean(0), X.std(0)
        std[std == 0] = 1
        z_matrix = ((X - mean) / std).T
        # Aggregate the factor values so each stock has one composite score.
        composites = []
        for i, (sym, _) in enumerate(scored):
            score = sum(self._signs[j] * z_matrix[j][i] for j in range(10))
            composites.append((sym, score))
        # Select the n stocks with the largest scores.
        composites.sort(key=lambda x: x[1], reverse=True)
        selected = [s for s, _ in composites[:self._NUM_LONG]]
        # Calculate the target weight of each stock.
        weight = 1 / len(selected)
        self._weight_by_symbol = {s: weight for s in selected}
        # Select the assets that we want to hold.
        return list(self._weight_by_symbol.keys())

    def _rebalance(self) -> None:
        if not self._weight_by_symbol:
            return
        targets = [PortfolioTarget(s, w) for s, w in self._weight_by_symbol.items()]
        self.set_holdings(targets, liquidate_existing_holdings=True)


class SelectionData:

    def __init__(self, algorithm, f):
        self._algorithm = algorithm
        self._price_scale_factor = f.price_scale_factor
        # Define some indicators to help calculate the factor values.
        self._price = Identity()
        self._roc_by_period = {period: RateOfChange(period) for period in [1, 21, 63, 126, 252]}
        self._sma = SimpleMovingAverage(200)
        self._rsi = RelativeStrengthIndex(14, MovingAverageType.WILDERS)
        self._macd = MovingAverageConvergenceDivergence(12, 26, 9)
        self._bb = BollingerBands(20, 2.0)
        self._indicators = [self._price, self._sma, self._rsi, self._macd, self._bb] + list(self._roc_by_period.values())
        # Define the factors.
        self._factors = [
            # 1. 252-21 momentum (12m return minus 1m return)
            IndicatorExtensions.minus(self._roc_by_period[252], self._roc_by_period[21]),
            # 2. 63-day momentum
            self._roc_by_period[63],
            # 3. 21-day momentum
            self._roc_by_period[21],
            # 4. RSI 14 (Wilder's smoothing)
            self._rsi,
            # 5. ROC 126-day
            self._roc_by_period[126],
            # 6. Price / SMA200 - 1
            IndicatorExtensions.minus(IndicatorExtensions.over(self._price, self._sma), 1), 
            # 7. STD 21-day (daily return volatility)
            IndicatorExtensions.of(StandardDeviation(21), self._roc_by_period[1]), 
            # 8. STD 63-day (daily return volatility)
            IndicatorExtensions.of(StandardDeviation(63), self._roc_by_period[1]), 
            # 9. MACD histogram (EMA12 - EMA26, signal = EMA9 of MACD line)
            self._macd.histogram, 
            # 10. Bollinger %B = (price - lower) / (upper - lower)
            self._bb.percent_b
        ]

    def update(self, f):
        # If there hasn't been a split or dividend since the last trading
        # day, just update the indicator like normal.
        if f.price_scale_factor == self._price_scale_factor:
            return all([indicator.update(f.end_time, f.price) for indicator in self._indicators])
        # Otherwise, reset the indicator and warm it up with the new 
        # adjusted history.
        self._price_scale_factor = f.price_scale_factor
        for indicator in self._indicators:
            indicator.reset()
        history = self._algorithm.history[TradeBar](
            f.symbol, 
            253, 
            Resolution.DAILY, 
            data_normalization_mode=DataNormalizationMode.SCALED_RAW
        )
        for bar in history:
            for indicator in self._indicators:
                indicator.update(bar)
        return self._roc_by_period[252].is_ready
    
    def compute_factors(self):
        return [f.current.value for f in self._factors]