from collections import deque
from typing import Deque

from AlgorithmImports import *


class MonthMomentumAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        self.set_start_date(1999, 1, 30)
        self.set_end_date(2026, 5, 21)
        self.set_cash(1_000_000)

        self.settings.seed_initial_prices = True
        self.settings.minimum_order_margin_portfolio_percentage = 0

        # ------------------------------------------------------------------ #
        #  Strategy options                                                    #
        # ------------------------------------------------------------------ #

        # Rebalance frequency in months: 1 = monthly, 3 = quarterly, 6 = semi-annual, 12 = annual
        self.REBALANCE_MONTHS: int = 3

        # SPY filter: go to cash when SPY 12M momentum is negative
        self.USE_SPY_FILTER: bool = True

        # Momentum filters:
        #   True  → must be positive to enter / hold (also used in sorting)
        #   False → used ONLY for sorting, never blocks entry / exit
        self.USE_12M_FILTER: bool = True
        self.USE_6M_FILTER:  bool = True

        # ------------------------------------------------------------------ #
        #  Core parameters                                                     #
        # ------------------------------------------------------------------ #

        self._lookback_12m   = 252
        self._lookback_6m    = 126
        self._top_n          = 50
        self._universe_size  = 1000
        self._max_lookback   = self._lookback_12m + 1

        self._price_history:    dict[Symbol, Deque[float]] = {}
        self._monthly_targets:  list[Symbol]               = []
        self._target_symbols:   set[Symbol]                = set()
        self._selected_symbols: set[Symbol]                = set()

        # Tracks when the last rebalance actually fired
        self._last_rebalance_date: datetime | None = None

        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._coarse_selection)

        self._spy = self.add_equity("SPY", Resolution.DAILY).symbol
        self.set_benchmark("SPY")

        self.set_warm_up(self._max_lookback, Resolution.DAILY)

        # Always schedule monthly — _rebalance internally gates by REBALANCE_MONTHS
        self.schedule.on(
            self.date_rules.month_start("SPY"),
            self.time_rules.after_market_open("SPY", 1),
            self._rebalance,
        )

    # ------------------------------------------------------------------ #
    #  Universe                                                            #
    # ------------------------------------------------------------------ #

    def _coarse_selection(self, coarse) -> list:
        filtered = [
            x for x in coarse
            if x.has_fundamental_data
            and x.price > 5
            and x.dollar_volume > 1_000_000
        ]
        top1000 = sorted(filtered, key=lambda x: x.dollar_volume, reverse=True)[:self._universe_size]
        symbols = [x.symbol for x in top1000]
        self._selected_symbols = set(symbols)
        return symbols

    # ------------------------------------------------------------------ #
    #  Data collection                                                     #
    # ------------------------------------------------------------------ #

    def on_data(self, data: Slice) -> None:
        for symbol in self._selected_symbols | {self._spy}:
            if symbol not in data.bars:
                continue
            if symbol not in self._price_history:
                self._price_history[symbol] = deque(maxlen=self._max_lookback)
            self._price_history[symbol].append(data.bars[symbol].close)

        if not self.is_warming_up:
            self._daily_momentum_check()

    # ------------------------------------------------------------------ #
    #  Rebalance period helper                                             #
    # ------------------------------------------------------------------ #

    def _is_rebalance_due(self) -> bool:
        """
        Returns True when enough months have elapsed since the last rebalance.
        First rebalance always fires immediately.
        Examples:
          REBALANCE_MONTHS = 1  → every month start
          REBALANCE_MONTHS = 3  → Jan / Apr / Jul / Oct
          REBALANCE_MONTHS = 6  → Jan / Jul
          REBALANCE_MONTHS = 12 → Jan only
        """
        if self._last_rebalance_date is None:
            return True

        last  = self._last_rebalance_date
        now   = self.time

        # Count how many full months have passed
        months_elapsed = (now.year - last.year) * 12 + (now.month - last.month)
        return months_elapsed >= self.REBALANCE_MONTHS

    # ------------------------------------------------------------------ #
    #  Momentum helpers                                                    #
    # ------------------------------------------------------------------ #

    def _compute_momentum(self, symbol: Symbol, lookback: int) -> float | None:
        history = self._price_history.get(symbol)
        if history is None or len(history) < lookback + 1:
            return None
        old_price = history[-lookback - 1]
        new_price = history[-1]
        if old_price <= 0:
            return None
        return (new_price / old_price) - 1.0

    def _spy_allows_trading(self) -> bool:
        """Returns False when SPY filter is ON and SPY 12M momentum is negative."""
        if not self.USE_SPY_FILTER:
            return True
        spy_mom = self._compute_momentum(self._spy, self._lookback_12m)
        if spy_mom is None:
            return True
        return spy_mom > 0

    def _passes_momentum_filter(self, symbol: Symbol) -> bool:
        """
        Entry / hold gate:
          Both ON  → 12M > 0 AND 6M > 0
          Only 12M → 12M > 0
          Only 6M  → 6M  > 0
          Both OFF → always passes
        """
        if self.USE_12M_FILTER:
            m12 = self._compute_momentum(symbol, self._lookback_12m)
            if m12 is None or m12 <= 0:
                return False

        if self.USE_6M_FILTER:
            m6 = self._compute_momentum(symbol, self._lookback_6m)
            if m6 is None or m6 <= 0:
                return False

        return True

    def _combined_score(self, symbol: Symbol) -> float | None:
        """
        Sorting score — average of available lookbacks.
        Both available → (12M + 6M) / 2
        One available  → that one alone
        Neither        → None
        """
        m12 = self._compute_momentum(symbol, self._lookback_12m)
        m6  = self._compute_momentum(symbol, self._lookback_6m)

        if m12 is not None and m6 is not None:
            return (m12 + m6) / 2.0
        if m12 is not None:
            return m12
        if m6 is not None:
            return m6
        return None

    # ------------------------------------------------------------------ #
    #  Portfolio management                                                #
    # ------------------------------------------------------------------ #

    def _update_portfolio(self, symbols: list[Symbol]) -> None:
        new_targets = set(symbols)
        if new_targets == self._target_symbols:
            return
        self._target_symbols = new_targets
        if not new_targets:
            self.liquidate()
            return
        weight  = 1.0 / len(new_targets)
        targets = [PortfolioTarget(s, weight) for s in new_targets]
        self.set_holdings(targets, liquidate_existing_holdings=True)

    # ------------------------------------------------------------------ #
    #  Daily check                                                         #
    # ------------------------------------------------------------------ #

    def _daily_momentum_check(self) -> None:
        """
        Every day:
          1. SPY filter ON + SPY momentum negative → full cash.
          2. Otherwise drop held stocks that no longer pass momentum gate.
        Note: does NOT pick new stocks — that only happens on rebalance.
        """
        if not self._monthly_targets:
            return

        if not self._spy_allows_trading():
            if self._target_symbols:
                self.debug(f"{self.time.date()} — SPY filter triggered: going to cash")
                self._update_portfolio([])
            return

        positive = [
            symbol for symbol in self._monthly_targets
            if symbol in self._selected_symbols
            and self._passes_momentum_filter(symbol)
        ]
        self._update_portfolio(positive)

    # ------------------------------------------------------------------ #
    #  Rebalance                                                           #
    # ------------------------------------------------------------------ #

    def _rebalance(self) -> None:
        if self.is_warming_up:
            return

        # Gate: only rebalance when enough months have elapsed
        if not self._is_rebalance_due():
            return

        # SPY gate
        if not self._spy_allows_trading():
            self.debug(f"{self.time.date()} — SPY filter active, skipping rebalance")
            self._update_portfolio([])
            self._last_rebalance_date = self.time  # still counts as a rebalance period
            return

        candidates = []
        for symbol in self._selected_symbols:
            if not self._passes_momentum_filter(symbol):
                continue
            score = self._combined_score(symbol)
            if score is None:
                continue
            candidates.append((symbol, score))

        if len(candidates) < self._top_n:
            self.debug(f"{self.time.date()} — only {len(candidates)} qualifying stocks, skipping rebalance")
            return

        candidates.sort(key=lambda x: x[1], reverse=True)
        self._monthly_targets     = [s for s, _ in candidates[:self._top_n]]
        self._last_rebalance_date = self.time

        self.debug(
            f"{self.time.date()} — rebalanced | "
            f"every {self.REBALANCE_MONTHS}M | "
            f"candidates: {len(candidates)} | "
            f"best: {candidates[0][1]:.2%} | "
            f"10th: {candidates[self._top_n - 1][1]:.2%} | "
            f"SPY: {'ON' if self.USE_SPY_FILTER else 'OFF'} | "
            f"12M: {'ON' if self.USE_12M_FILTER else 'OFF'} | "
            f"6M: {'ON' if self.USE_6M_FILTER else 'OFF'}"
        )

        self._daily_momentum_check()