# region imports
from AlgorithmImports import *
# endregion


class MomentumAlphaModel(AlphaModel):

    def __init__(self, algorithm, date_rule, lookback, resolution, num_insights):
        self._lookback = lookback
        self._resolution = resolution
        self._num_insights = num_insights
        self._securities = []
        self._insights = []
        # Schedule monthly alpha emission at 8:00 for daily equity workflows.
        algorithm.schedule.on(date_rule, algorithm.time_rules.at(8, 0), self._create_insights)

    def _create_insights(self):
        insights = []
        for security in self._securities:
            if not security.roc.is_ready:
                continue
            magnitude = security.roc.current.value
            if magnitude < 0:
                continue
            insights.append(Insight.price(security, Expiry.END_OF_MONTH, InsightDirection.UP, magnitude))
        # Keep only the strongest monthly signals by ROC magnitude.
        self._insights = sorted(insights, key=lambda i: i.magnitude)[-self._num_insights:]

    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.removed_securities:
            if security in self._securities:
                self._securities.remove(security)
        for security in changes.added_securities:
            security.roc = algorithm.roc(security, self._lookback, self._resolution)
            self._securities.append(security)

# region imports
from AlgorithmImports import *
from alpha import MomentumAlphaModel
from risk import RiskModelWithSPY
# endregion


class MomentumFrameworkAlgo(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        self.settings.automatic_indicator_warm_up = True
        self.settings.seed_initial_prices = True
        self.universe_settings.resolution = Resolution.HOUR
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        spy = self.add_equity("SPY")
        self.set_benchmark(spy)
        universe_size = 20
        # Use a monthly universe schedule anchored to the US equity calendar.
        date_rule = self.date_rules.month_start(spy)
        self.universe_settings.schedule.on(date_rule)
        self.universe_settings.asynchronous = True
        self.add_universe_selection(
            FundamentalUniverseSelectionModel(
                # Sort by the highest dollar-volume symbols after filtering by price.
                lambda fundamental: [
                    x.symbol
                    for x in sorted(
                        [x for x in fundamental if x.has_fundamental_data and x.price > 10],
                        key=lambda x: x.dollar_volume,
                    )[-universe_size:]
                ]
            )
        )
        # Emit monthly momentum insights based on rolling ROC values.
        self.add_alpha(
            MomentumAlphaModel(self, date_rule, 60, Resolution.DAILY, 10)
        )
        self.set_portfolio_construction(
            EqualWeightingPortfolioConstructionModel(Expiry.END_OF_MONTH, PortfolioBias.LONG)
        )
        self.add_risk_management(
            RiskModelWithSPY(self, spy, 200, Resolution.DAILY)
        )

# region imports
from AlgorithmImports import *
# endregion


class RiskModelWithSPY(RiskManagementModel):

    def __init__(self, algorithm, spy, lookback, resolution):
        self._spy = spy
        self._ema = algorithm.ema(spy, lookback, resolution)

    def manage_risk(self, algorithm, targets):
        # Liquidate invested assets when SPY falls below its EMA filter.       
        if not self._ema.is_ready or self._spy.price > self._ema.current.value:
            return []
        targets = [PortfolioTarget(security, 0) for security in algorithm.active_securities.values()]
        # Cancel insights to prevent the PCM from re-entering liquidated positions.
        algorithm.insights.cancel([t.symbol for t in targets])
        return targets