# 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 = float(security.roc.current.value)
            if magnitude < 0:
                continue
            insights.append(Insight.price(security.symbol, 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:
            if security not in self._securities:
                security.roc = algorithm.roc(security.symbol, 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_start_date(2011, 1, 1)
        self.set_cash(100_000)
        self.settings.seed_initial_prices = True
        self.settings.automatic_indicator_warm_up = True
        self.universe_settings.resolution = Resolution.HOUR
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.set_warm_up(timedelta(days=126))
        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.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, lookback=60, resolution=Resolution.DAILY, num_insights=10)
        )
        self.set_portfolio_construction(
            EqualWeightingPortfolioConstructionModel(Expiry.END_OF_MONTH, PortfolioBias.LONG)
        )
        self.add_risk_management(
            RiskModelWithSPY(spy, lookback=200, resolution=Resolution.DAILY)
        )

# region imports
from AlgorithmImports import *
# endregion


class RiskModelWithSPY(RiskManagementModel):

    def __init__(self, spy, lookback, resolution):
        self._spy = spy
        self._lookback = lookback
        self._resolution = resolution
        self._ema = None

    def on_securities_changed(self, algorithm, changes):
        for security in changes.added_securities:
            if security.symbol == self._spy.symbol:
                self._ema = algorithm.ema(security.symbol, self._lookback, self._resolution)

    def manage_risk(self, algorithm, _targets):
        if self._ema is None or not self._ema.is_ready:
            return []
        # Liquidate invested symbols when SPY falls below its EMA filter.
        if self._spy.price > self._ema.current.value:
            return []
        return [
            PortfolioTarget(security.symbol, 0)
            for security in algorithm.active_securities.values()
            if security.invested
        ]