#region imports
from AlgorithmImports import *
#endregion


class MomentumQuantilesAlphaModel(AlphaModel):

    def __init__(self, top_k_stocks, lookback_months):
        self.top_k_stocks = top_k_stocks
        self.lookback_months = lookback_months
        self.securities_list = []
        self.day = -1

    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # Reset indicators when corporate actions occur
        for symbol in set(data.splits.keys() + data.dividends.keys()):
            security = algorithm.securities[symbol]
            if security in self.securities_list:
                security.indicator.reset()
                algorithm.subscription_manager.remove_consolidator(security.symbol, security.consolidator)
                self._register_indicator(algorithm, security)

                history = algorithm.history[TradeBar](security.symbol, (security.indicator.warm_up_period+1) * 30, Resolution.DAILY, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
                for bar in history:
                    security.consolidator.update(bar)
        
        # Only emit insights when there is quote data, not when a corporate action occurs (at midnight)
        if data.quote_bars.count == 0:
            return []
        
        # Only emit insights once per day
        if self.day == algorithm.time.day:
            return []
        self.day = algorithm.time.day

        # Get the momentum of each asset in the universe
        momentum_by_symbol = {security.symbol : security.indicator.current.value 
            for security in self.securities_list if security.symbol in data.quote_bars and security.indicator.is_ready}

        leverage = 2
        # Create insights to long the assets in the universe with the greatest momentum
        weight = leverage / self.top_k_stocks
        insights = []
        for symbol, _ in sorted(momentum_by_symbol.items(), key=lambda x: x[1], reverse=True)[:self.top_k_stocks]:
            insights.append(Insight.price(symbol, Expiry.END_OF_DAY, InsightDirection.UP, weight=weight))

        return insights

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Create and register indicator for each security in the universe
        security_by_symbol = {}
        for security in changes.added_securities:
            security_by_symbol[security.symbol] = security
            
            # Create an indicator that automatically updates each month
            security.indicator = MomentumPercent(self.lookback_months)
            self._register_indicator(algorithm, security)

            self.securities_list.append(security)
        
        # Warm up the indicators of newly-added stocks
        if security_by_symbol:
            history = algorithm.history[TradeBar](list(security_by_symbol.keys()), (self.lookback_months+1) * 30, Resolution.DAILY, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
            for trade_bars in history:
                for bar in trade_bars.values():
                    security_by_symbol[bar.symbol].consolidator.update(bar)

        # Stop updating consolidator when the security is removed from the universe
        for security in changes.removed_securities:
            if security in self.securities_list:
                algorithm.subscription_manager.remove_consolidator(security.symbol, security.consolidator)
                self.securities_list.remove(security)


    def _register_indicator(self, algorithm, security):
        # Update the indicator with monthly bars
        security.consolidator = TradeBarConsolidator(Calendar.MONTHLY)
        algorithm.subscription_manager.add_consolidator(security.symbol, security.consolidator)
        algorithm.register_indicator(security.symbol, security.indicator, security.consolidator)

# region imports
from AlgorithmImports import *

from universe import QQQConstituentsUniverseSelectionModel
from alpha import MomentumQuantilesAlphaModel
# endregion


class TacticalMomentumRankAlgorithm(QCAlgorithm):

    undesired_symbols_from_previous_deployment = []
    checked_symbols_from_previous_deployment = False

    # def initialize(self):
    #     self.set_start_date(2023, 3, 1)  # Set Start Date
    #     self.set_end_date(2024, 3, 1) 
    #     self.set_cash(1_000_000)
        
    #     self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

    #     self.settings.minimum_order_margin_portfolio_percentage = 0

    #     self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
    #     self.add_universe_selection(QQQConstituentsUniverseSelectionModel(self.universe_settings))
        
    #     # self.add_alpha(MomentumQuantilesAlphaModel(
    #     #     int(self.get_parameter("top_k_stocks")),
    #     #     int(self.get_parameter("lookback_months"))
    #     # ))

    #     self.add_alpha(MomentumQuantilesAlphaModel(
    #         3,
    #         1
    #     ))

    #     self.settings.rebalance_portfolio_on_security_changes = False
    #     self.settings.rebalance_portfolio_on_insight_changes = False
    #     self.day = -1
    #     self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel(self._rebalance_func))

    #     self.add_risk_management(NullRiskManagementModel())

    #     self.set_execution(ImmediateExecutionModel())

    #     self.set_warm_up(timedelta(7))

    def initialize(self):
        self.set_start_date(2023, 3, 1)  # Set Start Date
        self.set_end_date(2024, 3, 1) 
        self.set_cash(1_000_000)

        
        leverage = 2
        lookback_months = 1
        top_k_stocks = 3
        
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        self.settings.minimum_order_margin_portfolio_percentage = 0

        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.add_universe_selection(QQQConstituentsUniverseSelectionModel(self.universe_settings))
        
        # self.add_alpha(MomentumQuantilesAlphaModel(
        #     int(self.get_parameter("top_k_stocks")),
        #     int(self.get_parameter("lookback_months"))
        # ))

        self.add_alpha(MomentumQuantilesAlphaModel(
            top_k_stocks,
            lookback_months
        ))

        self.settings.rebalance_portfolio_on_security_changes = False
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.day = -1

        self.add_risk_management(NullRiskManagementModel())

        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel(self._rebalance_func))

        self.set_execution(ImmediateExecutionModel())

        self.universe_settings.leverage = leverage
        
        self.set_warm_up(timedelta(7))

    def _rebalance_func(self, time):
        if self.day != self.time.day and not self.is_warming_up and self.current_slice.quote_bars.count > 0:
            self.day = self.time.day
            return time
        return None

    def on_data(self, data):
        # Exit positions that aren't backed by existing insights.
        # If you don't want this behavior, delete this method definition.
        if not self.is_warming_up and not self.checked_symbols_from_previous_deployment:
            for security_holding in self.portfolio.values():
                if not security_holding.invested:
                    continue
                symbol = security_holding.symbol
                if not self.insights.has_active_insights(symbol, self.utc_time):
                    self.undesired_symbols_from_previous_deployment.append(symbol)
            self.checked_symbols_from_previous_deployment = True
        
        for symbol in self.undesired_symbols_from_previous_deployment:
            if self.is_market_open(symbol):
                self.liquidate(symbol, tag="Holding from previous deployment that's no longer desired")
                self.undesired_symbols_from_previous_deployment.remove(symbol)

 #region imports
from AlgorithmImports import *
#endregion


class QQQConstituentsUniverseSelectionModel(ETFConstituentsUniverseSelectionModel):
    def __init__(self, universe_settings: UniverseSettings = None) -> None:
        symbol = Symbol.create("QQQ", SecurityType.EQUITY, Market.USA)
        super().__init__(symbol, universe_settings, lambda constituents: [c.symbol for c in constituents])