from AlgorithmImports import *
from datetime import timedelta


class QualityMomentumStrategy(QCAlgorithm):
    """
    Biweekly quality-momentum with 200-day SMA regime and Sharpe/Vol composite.

    Universe: top-200 liquid large-caps with positive PE.
    Ranking: Sharpe / returns-STD composite (double-penalizes volatility).
    Regime: SPY price above 200-day SMA to invest; below to cash.
    Sizing: rank-weighted across top-25 stocks.
    Execution: biweekly rebalance Monday 31 min after open.
    """

    COARSE_SIZE = 500
    FINE_SIZE = 200
    MIN_PRICE = 15.0
    MIN_DOLLAR_VOL = 20e6
    MIN_MARKET_CAP = 20e9

    SHARPE_PERIOD = 126
    TOP_N = 25
    VOL_PERIOD = 21
    REGIME_PERIOD = 200

    MAX_GROSS = 2.0
    MAX_SINGLE = 0.20

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5 * 365))
        self.set_cash(100_000)

        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.settings.automatic_indicator_warm_up = True
        self.settings.seed_initial_prices = True
        self.set_warm_up(timedelta(45))

        self._spy = self.add_equity("SPY")
        self._spy_sma = self.sma(
            self._spy, self.REGIME_PERIOD, Resolution.DAILY
        )

        self.universe_settings.schedule.on(
            self.date_rules.month_start(self._spy)
        )
        self._universe = self.add_universe(self._fundamental_filter)

    def on_warmup_finished(self):
        # Add a Scheduled Event to rebalance the portfolio every 2 weeks.
        time_rule = self.time_rules.after_market_open(self._spy, 31)
        self.schedule.on(self.date_rules.week_start(self._spy), time_rule, self._weekly_check)
        # Rebalance today too.
        if self.live_mode:
            self._weekly_check()
        else:
            self.schedule.on(self.date_rules.today, time_rule, lambda: self._weekly_check(True))

    # ── universe ─────────────────────────────────────────────

    def _fundamental_filter(self, fundamental):
        pool = [
            f for f in fundamental
            if (f.has_fundamental_data and
                f.price > self.MIN_PRICE and
                f.dollar_volume > self.MIN_DOLLAR_VOL and
                f.market_cap >= self.MIN_MARKET_CAP and
                f.valuation_ratios.pe_ratio > 0)
        ]
        pool = sorted(pool, key=lambda f: f.dollar_volume)[-self.COARSE_SIZE:]
        pool = sorted(pool, key=lambda f: f.market_cap)[-self.FINE_SIZE:]
        return [f.symbol for f in pool]

    # ── security lifecycle ───────────────────────────────────

    def on_securities_changed(self, changes):
        for sec in changes.added_securities:
            if sec == self._spy:
                continue
            sec.sharpe = self.sr(
                sec, self.SHARPE_PERIOD, 0.0, Resolution.DAILY
            )
            sec.volatility = IndicatorExtensions.of(
                StandardDeviation(self.VOL_PERIOD),
                self.roc(sec, 1, Resolution.DAILY),
            )

    # ── biweekly rebalance (Monday open) ─────────────────────

    def _weekly_check(self, skip_week_check=False):
        if self.is_warming_up or not self._spy_sma.is_ready:
            return

        if self._spy.price < self._spy_sma.current.value:
            self.liquidate(tag="bear_regime")
            return
        if not skip_week_check and self.time.isocalendar()[1] % 2 != 0:
            return

        securities = [self.securities[sym] for sym in self._universe.selected]

        candidates = [
            s for s in securities
            if (s != self._spy and
                s.sharpe.is_ready and 
                s.sharpe.current.value > 0 and 
                s.volatility.is_ready and 
                s.volatility.current.value > 0)
        ]

        if not candidates:
            return

        # rank by Sharpe / Vol composite (double-penalize volatility)
        top = sorted(
            candidates,
            key=lambda s: s.sharpe.current.value / s.volatility.current.value
        )[-self.TOP_N:]

        # rank-weighted
        rank_weights = {security: i+1 for i, security in enumerate(top)}
        total_rank = sum(rank_weights.values())

        weights = {}
        for security, rw in rank_weights.items():
            raw = (rw / total_rank) * self.MAX_GROSS
            weights[security] = min(raw, self.MAX_SINGLE)

        total_w = sum(weights.values())
        if total_w > self.MAX_GROSS:
            scale = self.MAX_GROSS / total_w
            weights = {s: w * scale for s, w in weights.items()}

        targets = [
            PortfolioTarget(security, weight)
            for security, weight in weights.items()
        ]
        self.set_holdings(targets, True)