# region imports
from AlgorithmImports import *
# endregion


class FactorMomentumDistressAlgorithm(QCAlgorithm):
    """
    Long-short equity strategy using 10 momentum/distress factors.
    Universe: SPY ETF constituents. Monthly rebalance. $1B capital.
    Dollar-neutral: +0.5 long / -0.5 short (100% gross).

    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
    """

    # --- factor parameters ---
    _MOM_252 = 252
    _MOM_63 = 63
    _MOM_21 = 21
    _RSI_PERIOD = 14
    _ROC_PERIOD = 126
    _SMA200_PERIOD = 200
    _STD21_PERIOD = 21
    _STD63_PERIOD = 63
    _MACD_FAST = 12
    _MACD_SLOW = 26
    _MACD_SIGNAL = 9
    _BB_PERIOD = 20
    _BB_K = 2.0

    # --- universe / portfolio parameters ---
    _UNIVERSE_CAP = 500
    _NUM_LONG = 50
    _LOOKBACK = 260  # 253 needed for 252-day momentum + buffer

    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

        self.universe_settings.resolution = Resolution.DAILY
        self.universe_settings.leverage = 2.0

        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,
        )

        self._weight_by_symbol = {}

        self.schedule.on(date_rule, self.time_rules.at(8, 0), self._rebalance)

    # ------------------------------------------------------------------ #
    #  Universe selection — compute factors via history() batch request   #
    # ------------------------------------------------------------------ #
    def _select_assets(self, fundamentals: list[Fundamental]) -> list[Symbol]:
        if len(fundamentals) < self._NUM_LONG:
            return Universe.UNCHANGED

        # Pre-screen to top UNIVERSE_CAP by dollar volume
        sorted_points = sorted(fundamentals, key=lambda f: f.dollar_volume, reverse=True)
        candidates = [f.symbol for f in sorted_points[:self._UNIVERSE_CAP]]

        # Pull daily close history for all candidates in one batch
        history = self.history(candidates, self._LOOKBACK, Resolution.DAILY)
        if history.empty:
            return Universe.UNCHANGED

        # Compute 10 factors for each symbol with enough history
        factor_data = {}
        for symbol, group in history.groupby(level=0):
            closes = group["close"].dropna().values
            if len(closes) < self._MOM_252 + 1:
                continue
            factors = self._compute_factors(closes)
            if factors is not None:
                factor_data[symbol] = factors

        if len(factor_data) < self._NUM_LONG:
            return Universe.UNCHANGED

        # Z-score each factor column
        scored = list(factor_data.items())
        n = len(scored)

        factor_matrix = [[] for _ in range(10)]
        for _, fv in scored:
            for i in range(10):
                factor_matrix[i].append(fv[i])

        z_matrix = []
        for col in factor_matrix:
            mean = sum(col) / len(col)
            var = sum((v - mean) ** 2 for v in col) / len(col)
            std = math.sqrt(var)
            if std == 0:
                z_matrix.append([0.0] * n)
            else:
                z_matrix.append([(v - mean) / std for v in col])

        # 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]

        composites = []
        for i, (sym, _) in enumerate(scored):
            score = sum(signs[j] * z_matrix[j][i] for j in range(10))
            composites.append((sym, score))

        composites.sort(key=lambda x: x[1], reverse=True)

        longs = [s for s, _ in composites[:self._NUM_LONG]]

        # Dollar-neutral: +0.5 long / -0.5 short (100% gross)
        self._weight_by_symbol = {}
        lw = 1 / len(longs)
        for s in longs:
            self._weight_by_symbol[s] = lw

        return list(self._weight_by_symbol.keys())

    # ------------------------------------------------------------------ #
    #  Rebalance                                                          #
    # ------------------------------------------------------------------ #
    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)

    # ------------------------------------------------------------------ #
    #  Factor computation                                                 #
    # ------------------------------------------------------------------ #
    def _compute_factors(self, closes: np.ndarray) -> list[float] | None:
        n = len(closes)
        if n < self._MOM_252 + 1:
            return None

        price = float(closes[-1])

        # 1. 252-21 momentum (12m return minus 1m return)
        f1 = (price / float(closes[-253]) - 1.0) - (price / float(closes[-22]) - 1.0)

        # 2. 63-day momentum
        f2 = price / float(closes[-64]) - 1.0

        # 3. 21-day momentum
        f3 = price / float(closes[-22]) - 1.0

        # 4. RSI 14 (Wilder's smoothing)
        f4 = self._rsi(closes, self._RSI_PERIOD)

        # 5. ROC 126-day
        f5 = (price - float(closes[-127])) / float(closes[-127]) * 100.0

        # 6. Price / SMA200 - 1
        sma200 = float(np.mean(closes[-self._SMA200_PERIOD:]))
        f6 = price / sma200 - 1.0 if sma200 != 0 else 0.0

        # 7. STD 21-day (daily return volatility)
        rets = np.diff(closes[-self._STD21_PERIOD - 1:]) / closes[-self._STD21_PERIOD - 1:-1]
        f7 = float(np.std(rets, ddof=1))

        # 8. STD 63-day (daily return volatility)
        rets63 = np.diff(closes[-self._STD63_PERIOD - 1:]) / closes[-self._STD63_PERIOD - 1:-1]
        f8 = float(np.std(rets63, ddof=1))

        # 9. MACD histogram (EMA12 - EMA26, signal = EMA9 of MACD line)
        f9 = self._macd_hist(closes)

        # 10. Bollinger %B = (price - lower) / (upper - lower)
        sma20 = float(np.mean(closes[-self._BB_PERIOD:]))
        std20 = float(np.std(closes[-self._BB_PERIOD:], ddof=1))
        upper = sma20 + self._BB_K * std20
        lower = sma20 - self._BB_K * std20
        bw = upper - lower
        f10 = (price - lower) / bw if bw != 0 else 0.5

        return [f1, f2, f3, f4, f5, f6, f7, f8, f9, f10]

    def _rsi(self, prices: np.ndarray, period: int) -> float:
        """RSI with Wilder's smoothing."""
        deltas = np.diff(prices)
        if len(deltas) < period:
            return 50.0
        gains = np.where(deltas > 0, deltas, 0.0)
        losses = np.where(deltas < 0, -deltas, 0.0)
        avg_gain = float(np.mean(gains[:period]))
        avg_loss = float(np.mean(losses[:period]))
        for i in range(period, len(deltas)):
            avg_gain = (avg_gain * (period - 1) + float(gains[i])) / period
            avg_loss = (avg_loss * (period - 1) + float(losses[i])) / period
        if avg_loss == 0:
            return 100.0
        rs = avg_gain / avg_loss
        return 100.0 - (100.0 / (1.0 + rs))

    def _macd_hist(self, closes: np.ndarray) -> float:
        """MACD histogram: MACD line minus signal line."""
        if len(closes) < self._MACD_SLOW + self._MACD_SIGNAL:
            return 0.0
        a12 = 2.0 / (self._MACD_FAST + 1)
        a26 = 2.0 / (self._MACD_SLOW + 1)
        a9 = 2.0 / (self._MACD_SIGNAL + 1)
        ema12 = float(closes[0])
        ema26 = float(closes[0])
        macd_line = ema12 - ema26
        signal = macd_line
        for i in range(1, len(closes)):
            c = float(closes[i])
            ema12 = a12 * c + (1 - a12) * ema12
            ema26 = a26 * c + (1 - a26) * ema26
            macd_line = ema12 - ema26
            signal = a9 * macd_line + (1 - a9) * signal
        return macd_line - signal