#region imports
from AlgorithmImports import *
#endregion

class MomentumQuantilesAlphaModel(AlphaModel):

    symbol_data_by_symbol = {}
    day = -1

    def __init__(self, quantiles, lookback_months):
        self.quantiles = quantiles
        self.lookback_months = lookback_months

    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()):
            if symbol in self.symbol_data_by_symbol:
                self.symbol_data_by_symbol[symbol].reset()
        
        # Only emit insights when there is quote data, not when a corporate action occurs (at midnight)
        if data.QuoteBars.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 = {}
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if symbol in data.QuoteBars and symbol_data.IsReady:
                momentum_by_symbol[symbol] = symbol_data.indicator.Current.Value
                
        # Determine how many assets to hold in the portfolio
        quantile_size = int(len(momentum_by_symbol)/self.quantiles)
        if quantile_size == 0:
            return []

        # Create insights to long the assets in the universe with the greatest momentum
        weight = 1 / quantile_size
        expiry = list(self.symbol_data_by_symbol.values())[0].hours.GetNextMarketOpen(algorithm.Time, False) - timedelta(seconds=1)
        insights = []
        for symbol, _ in sorted(momentum_by_symbol.items(), key=lambda x: x[1], reverse=True)[:quantile_size]:
            insights.append(Insight.Price(symbol, expiry, InsightDirection.Up, weight=weight))

        return insights

    def OnSecuritiesChanged(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Create SymbolData objects for each stock in the universe
        added_symbols = []
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.symbol_data_by_symbol[symbol] = SymbolData(algorithm, security, self.lookback_months)
            added_symbols.append(symbol)
        
        # Warm up the indicators of newly-added stocks
        if added_symbols:
            history = algorithm.History[TradeBar](added_symbols, (self.lookback_months+1) * 30, Resolution.Daily, dataNormalizationMode=DataNormalizationMode.ScaledRaw)
            for trade_bars in history:
                for bar in trade_bars.Values:
                    self.symbol_data_by_symbol[bar.Symbol].update(bar)

        # Remove the SymbolData object when the stock is removed from the universe
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.symbol_data_by_symbol:
                symbol_data = self.symbol_data_by_symbol.pop(symbol, None)
                if symbol_data:
                    symbol_data.dispose()


class SymbolData:
    def __init__(self, algorithm, security, lookback_months):
        self.algorithm = algorithm
        self.symbol = security.Symbol
        self.hours = security.Exchange.Hours

        # Create an indicator that automatically updates each month
        self.indicator = MomentumPercent(lookback_months)
        self.register_indicator()

    @property
    def IsReady(self):
        return self.indicator.IsReady

    def register_indicator(self):
        # Update the indicator with monthly bars
        self.consolidator = TradeBarConsolidator(Calendar.Monthly)
        self.algorithm.SubscriptionManager.AddConsolidator(self.symbol, self.consolidator)
        self.algorithm.RegisterIndicator(self.symbol, self.indicator, self.consolidator)

    def update(self, bar):
        self.consolidator.Update(bar)

    def reset(self):
        self.indicator.Reset()
        self.dispose()
        self.register_indicator()

        history = self.algorithm.History[TradeBar](self.symbol, (self.indicator.WarmUpPeriod+1) * 30, Resolution.Daily, dataNormalizationMode=DataNormalizationMode.ScaledRaw)
        for bar in history:
            self.consolidator.Update(bar)

    def dispose(self):
        # Stop updating consolidator when the security is removed from the universe
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)
        

# region imports
from AlgorithmImports import *

from universe import SPYAndQQQConstituentsUniverseSelectionModel
from alpha import MomentumQuantilesAlphaModel
# endregion

class TacticalMomentumRankAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2023, 6, 1)
        self.SetCash(100000)
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw
        self.AddUniverseSelection(SPYAndQQQConstituentsUniverseSelectionModel(self.UniverseSettings))       

        self.AddAlpha(MomentumQuantilesAlphaModel(
            int(self.GetParameter("quantiles")),
            int(self.GetParameter("lookback_months"))
        ))

        self.Settings.RebalancePortfolioOnSecurityChanges = False
        self.Settings.RebalancePortfolioOnInsightChanges = False
        self.day = -1
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel(self.rebalance_func))

        self.AddRiskManagement(NullRiskManagementModel())

        self.SetExecution(ImmediateExecutionModel())

        self.SetWarmUp(timedelta(7))

    def rebalance_func(self, time):
        if self.day != self.Time.day and not self.IsWarmingUp and self.CurrentSlice.QuoteBars.Count > 0:
            self.day = self.Time.day
            return time
        return None

    def OnWarmupFinished(self):
        # Exit positions that aren't backed by existing insights.
        # If you don't want this behavior, delete this method definition.
        for security_holding in self.Portfolio.Values:
            if not security_holding.Invested:
                continue
            symbol = security_holding.Symbol
            if not self.Insights.HasActiveInsights(symbol, self.UtcTime):
                self.Insights.Add(Insight.Price(symbol, timedelta(seconds=1), InsightDirection.Flat, weight=1))

#region imports
from AlgorithmImports import *
#endregion
# 05/19/2023: -Added a warm-up period to restore the algorithm state between deployments.
#             -Added OnWarmupFinished to liquidate existing holdings that aren't backed by active insights.
#             -Removed flat insights because https://github.com/QuantConnect/Lean/pull/7251 made them unnecessary.
#             https://www.quantconnect.com/terminal/processCache?request=embedded_backtest_a34c371a3b4818e5157cd76b876ecae0.html

# region imports
from AlgorithmImports import *
#endregion

class SPYAndQQQConstituentsUniverseSelectionModel(UniverseSelectionModel):
    def __init__(self, universe_settings: UniverseSettings = None) -> None:
        self.spy_symbol = Symbol.Create("SPY", SecurityType.Equity, Market.USA)
        self.qqq_symbol = Symbol.Create("QQQ", SecurityType.Equity, Market.USA)
        self.iwm_symbol = Symbol.Create("IWM", SecurityType.Equity, Market.USA)
        self.universe_settings = universe_settings or UniverseSettings()

    def CreateUniverses(self, algorithm: QCAlgorithm) -> List[Universe]:
        spy_universe = ETFConstituentsUniverse(self.spy_symbol, self.universe_settings, lambda constituents: [c.Symbol for c in constituents])
        qqq_universe = ETFConstituentsUniverse(self.qqq_symbol, self.universe_settings, lambda constituents: [c.Symbol for c in constituents])
        iwm_universe = ETFConstituentsUniverse(self.iwm_symbol, self.universe_settings, lambda constituents: [c.Symbol for c in constituents])


        return [spy_universe, qqq_universe, iwm_universe]