# region imports
from operator import lt, gt
from AlgorithmImports import *
# endregion


class ReversalFakeoutAlpha(AlphaModel):

    def __init__(self, algorithm, date_rule, assets_per_side):
        self._universe = []
        self._assets_per_side = assets_per_side
        self._insights = []
        self._adx_threshold = 25
        # Add a Scheduled Event to create insights matching the universe cadence.
        algorithm.schedule.on(date_rule, algorithm.time_rules.at(8, 0), self._create_insights)

    def _create_insights(self):
        # Check if valid asset prices and all windows/indicators are ready.
        securities = [s for s in self._universe if s.price and s.daily_returns.is_ready and s.monthly_window.is_ready]
        if not securities:
            return
        # Calculate and sort by the momentum scores for months 1-12 skipping the first month.
        sorted_by_monthly_return = sorted(securities, key=lambda s: math.prod(s.monthly_window[i].close / s.monthly_window[i].open for i in range(1, s.monthly_window.size)))
        # Long the top 100 performers (strongest momentum) where positive DI > negative DI.
        self._insights += self._long_short_picks(sorted_by_monthly_return[-100:], False, lt, InsightDirection.UP)
        # Short the bottom 100 performers (weakest momentum) where negative DI > positive DI triggers reversal.
        self._insights += self._long_short_picks(sorted_by_monthly_return[:100],  True,  gt, InsightDirection.DOWN)

    def _long_short_picks(self, candidates, reverse, compare, direction):
        # Filter candidates by ADX > threshold and positive DI > negative DI (for long) or negative DI > positive DI (for short).
        candidates = [
            candidate for candidate in candidates  
            if candidate.adx.current.value > self._adx_threshold and compare(candidate.adx.positive_directional_index, candidate.adx.negative_directional_index)
        ]
        # Calculate continuity score [-1, 1] for volatility/reversal: (2 * down_days - total_days) / total_days.
        filtered = sorted(candidates, key=lambda s: (2 * sum(1 for r in s.daily_returns if r <= 0) - s.daily_returns.count) / s.daily_returns.count, reverse=reverse)
        return [Insight.price(s, timedelta(5), direction) for s in filtered[:self._assets_per_side]]

    def update(self, algorithm, data):
        insights = self._insights.copy()  
        self._insights.clear()
        return insights

    def on_securities_changed(self, algorithm, changes):
        for security in changes.added_securities:
            # Create up two consolidators + two consolidators: monthly for momentum ranking and daily for intraday (open-to-close) returns.
            security.daily_returns = RollingWindow(280)
            security.monthly_window = RollingWindow[TradeBar](13)
            security.consolidators = [algorithm.consolidate(security, Calendar.MONTHLY, lambda bar: algorithm.securities[bar.symbol].monthly_window.add(bar)), 
                                      algorithm.consolidate(security, Resolution.DAILY, lambda bar: algorithm.securities[bar.symbol].daily_returns.add((bar.close - bar.open) / bar.open))]
            security.adx = algorithm.adx(security, 20)
            self._universe.append(security)
            # Warm up the daily and monthly windows.
            history = algorithm.history[TradeBar](security, 600)
            for bar in history:
                for consolidator in security.consolidators:
                    consolidator.update(bar)
        for security in changes.removed_securities:
            if security in self._universe:
                for consolidator in security.consolidators:
                    algorithm.subscription_manager.remove_consolidator(security, consolidator)
                self._universe.remove(security)

# region imports
from AlgorithmImports import *
from alpha import ReversalFakeoutAlpha
# endregion


class CrossSectionalMomentum(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365)) 
        self.settings.seed_initial_prices = True
        self.settings.automatic_indicator_warm_up = True
        self.set_cash(1_000_000)
        self.universe_settings.resolution = Resolution.DAILY
        date_rule = self.date_rules.week_start("SPY")
        self.universe_settings.schedule.on(date_rule)
        self.universe_settings.asynchronous = True
        self.add_universe_selection(FundamentalUniverseSelectionModel(
            lambda fundamentals: [
                f.symbol for f in sorted(fundamentals, key=lambda f: f.dollar_volume)[-1000:]
            ]
        ))
        self.add_alpha(ReversalFakeoutAlpha(self, date_rule, 5))
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel())
        self.set_risk_management(TrailingStopRiskManagementModel(0.02))
        self.set_warm_up(timedelta(10))