# 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.symbol_data_by_symbol = {}
        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 symbol_data in self.symbol_data_by_symbol.values():
            if not symbol_data.roc.is_ready:
                continue
            magnitude = float(symbol_data.roc.current.value)
            if magnitude < 0:
                continue
            insights.append(Insight.price(symbol_data.symbol, Expiry.END_OF_MONTH, InsightDirection.UP, magnitude, None))
        # Keep only the strongest monthly signals by ROC magnitude.
        top_insights = sorted(insights, key=lambda insight: insight.magnitude)[-self.num_insights:]
        self._insights = list(reversed(top_insights))

    def update(self, algorithm, data):
        insights = self._insights.copy()
        self._insights.clear()
        return insights
    
    def on_securities_changed(self, algorithm, changes):
        for removed in changes.removed_securities:
            symbol_data = self.symbol_data_by_symbol.pop(removed.symbol, None)
            if (symbol_data is not None and symbol_data.consolidator is not None):
                algorithm.subscription_manager.remove_consolidator(symbol_data.symbol, symbol_data.consolidator)
        for security in changes.added_securities:
            symbol = security.symbol
            if symbol in self.symbol_data_by_symbol:
                continue
            symbol_data = SymbolData(algorithm, symbol, self.lookback, self.resolution)
            self.symbol_data_by_symbol[symbol] = symbol_data


class SymbolData:

    def __init__(self, algorithm, symbol, lookback, resolution):
        self.symbol = symbol
        self.lookback = lookback
        self.roc = RateOfChange(lookback)
        self.consolidator = algorithm.resolve_consolidator(symbol, resolution)
        algorithm.register_indicator(symbol, self.roc, self.consolidator)
        # Warm up each symbol using TradeBar history.
        history = algorithm.history[TradeBar](symbol, lookback, resolution)
        for bar in history:
            self.roc.update(bar.end_time, bar.close)

# 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(2011, 1, 1)
        self.set_cash(100000)
        self.universe_settings.resolution = Resolution.HOUR
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.set_warm_up(126, Resolution.DAILY)
        self.set_benchmark("SPY")
        self.spy = self.add_equity("SPY")
        self.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,
                    )[-self.universe_size:]
                ]
            )
        )
        self.set_portfolio_construction(
            EqualWeightingPortfolioConstructionModel(Expiry.END_OF_MONTH, PortfolioBias.LONG)
        )
        # 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.add_risk_management(
            RiskModelWithSPY(self.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.spy_symbol_data = None
        
    def manage_risk(self, algorithm: QCAlgorithm, targets):
        if self.spy_symbol_data is None:
            self.spy_symbol_data = EMASymbolData(algorithm, self.spy, self.lookback, self.resolution)
        # Liquidate invested symbols when SPY falls below its EMA filter.
        if self.spy.price > self.spy_symbol_data.ema.current.value:
            return []
        return [
            PortfolioTarget(security.symbol, 0)
            for security in algorithm.active_securities.values()
            if security.invested
        ]
    

class EMASymbolData:

    def __init__(self, algorithm: QCAlgorithm, security, lookback: int, resolution: Resolution):
        symbol = security.symbol
        self.security = symbol
        self.consolidator = algorithm.resolve_consolidator(symbol, resolution)
        ema_name = algorithm.create_indicator_name(symbol, f"EMA{lookback}", resolution)
        self.ema = ExponentialMovingAverage(ema_name, lookback)
        algorithm.register_indicator(symbol, self.ema, self.consolidator)
        # Warm up EMA with historical close prices.
        history = algorithm.history(symbol, lookback, resolution)
        for at_time, value in history.close.unstack(0).squeeze().items():
            self.ema.update(at_time, value)