#region imports
from AlgorithmImports import *
from indicator import CustomMomentumPercent
#endregion


class MomentumQuantilesAlphaModel(AlphaModel):

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

    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:

            # Create an indicator
            security_by_symbol[security.symbol] = security
            security.indicator = CustomMomentumPercent("custom", 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 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}
                
        # 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+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.END_OF_DAY, InsightDirection.UP, weight=weight))

        return insights

    

    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 *
#endregion

class CustomMomentumPercent(PythonIndicator):
    def __init__(self, name, period):
        self.Name = name
        self.Time = datetime.min
        self.Value = 0
        self.momentum = MomentumPercent(period)

    def Update(self, input):
        self.momentum.Update(IndicatorDataPoint(input.Symbol, input.EndTime, input.Close))
        self.Time = input.EndTime
        self.Value = self.momentum.Current.Value * input.Volume * input.Close  # Multiply momentum percent with dollar volume
        return self.momentum.IsReady

# region imports
from AlgorithmImports import *
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(2020, 3, 1)  # Set Start Date
        self.set_end_date(2024, 3, 1) 
        self.set_cash(1_000_000)
        self.SetBenchmark("SPY")
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.settings.rebalance_portfolio_on_security_changes = False
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.day = -1
        self.set_warm_up(timedelta(7))

        self.universe_settings.asynchronous = True
        self.add_universe_selection(FundamentalUniverseSelectionModel(self.fundamental_filter_function))
        
        self.add_alpha(MomentumQuantilesAlphaModel(
            int(self.get_parameter("quantiles")),
            int(self.get_parameter("lookback_months"))
        ))

        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel(rebalance=Expiry.EndOfMonth))

        self.add_risk_management(NullRiskManagementModel())

        self.set_execution(ImmediateExecutionModel())

    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)

        
    def fundamental_filter_function(self, fundamental: List[Fundamental]):
        filtered = [f for f in fundamental if f.symbol.value != "AMC" and f.has_fundamental_data and not np.isnan(f.dollar_volume)]
        sorted_by_dollar_volume = sorted(filtered, key=lambda f: f.dollar_volume, reverse=True)
        return [f.symbol for f in sorted_by_dollar_volume[:1000]]