#region imports
from AlgorithmImports import *
#endregion


class ShortTermReversalAlphaModel(AlphaModel):

    _securities = []
    _week = 1

    def __init__(self, num_securities_per_side, lookback_days):
        self._num_securities_per_side = num_securities_per_side
        self._LOOKBACK_DAYS = lookback_days

    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # Reset indicators when corporate actions occur
        securities_to_reset = []
        for symbol in set(data.splits.keys() + data.dividends.keys()):
            security = algorithm.securities[symbol]
            if security in self._securities:
                algorithm.unregister_indicator(security.indicator)
                securities_to_reset.append(security)
        if securities_to_reset:
            self._set_up_indicators(algorithm, securities_to_reset, False)

        # Only emit insights when there is quote data, not when we get corporate action
        if data.quote_bars.count == 0:
            return []

        # Rebalance weekly
        week = data.time.date().isocalendar()[1]
        if self._week == week:
            return []
        self._week = week
       
        # Select securities that have the highest/lowest factor values
        tradable_securities = [s for s in self._securities if s.indicator.is_ready and s.symbol in data.quote_bars and s.price]
        sorted_by_roc = sorted(tradable_securities, key=lambda s: s.indicator.current.value)
        longs = sorted_by_roc[:self._num_securities_per_side]
        shorts = sorted_by_roc[-self._num_securities_per_side:]

        # Create insights to long stocks with the lowest ROC and short stocks with the greatest ROC. 
        # Hold positions until the next month.
        insights = [Insight.price(security.symbol, Expiry.END_OF_WEEK, InsightDirection.UP) for security in longs]
        insights += [Insight.price(security.symbol, Expiry.END_OF_WEEK, InsightDirection.DOWN) for security in shorts]
        return insights

    def _set_up_indicators(self, algorithm, securities, save_security_reference=True):
        # Create and register indicator for each security in the universe
        security_by_symbol = {}
        for security in securities:
            security.indicator = algorithm.roc(security.symbol, self._LOOKBACK_DAYS, Resolution.DAILY)
            security_by_symbol[security.symbol] = security
            if save_security_reference:
                self._securities.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_DAYS+1), Resolution.DAILY, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
            for trade_bars in history:
                for bar in trade_bars.values():
                    security_by_symbol[bar.symbol].indicator.update(bar.end_time, bar.close)

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        self._set_up_indicators(algorithm, changes.added_securities)

        # Stop updating indicators when the security leaves the universe
        for security in changes.removed_securities:
            if security in self._securities:
                algorithm.unregister_indicator(security.indicator)
                self._securities.remove(security)

#region imports
from AlgorithmImports import *

from universe import MostLiquidFundamentalUniverseSelectionModel
from alpha import ShortTermReversalAlphaModel
#endregion


class ShortTermReversalAlgorithm(QCAlgorithm):

    _undesired_symbols_from_previous_deployment = []
    _checked_symbols_from_previous_deployment = False

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

        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))

        self.settings.minimum_order_margin_portfolio_percentage = 0

        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.universe_settings.schedule.on(self.date_rules.week_start())
        universe_size = self.get_parameter("universe_size", 100)
        self.add_universe_selection(MostLiquidFundamentalUniverseSelectionModel(self.universe_settings, universe_size))

        self.add_alpha(ShortTermReversalAlphaModel(
            int(self.get_parameter("roc_selection_factor", 0.1) * universe_size),
            self.get_parameter("lookback_days", 22)
        ))

        self.settings.rebalance_portfolio_on_security_changes = False
        self.settings.rebalance_portfolio_on_insight_changes = False
        self._week = -1
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(self.rebalance_func))

        self.add_risk_management(NullRiskManagementModel())

        self.set_execution(ImmediateExecutionModel())

        self.set_warm_up(timedelta(14))

    def rebalance_func(self, time):
        # Rebalance weekly
        week = self.time.date().isocalendar()[1]
        if self._week != week and not self.is_warming_up and self.current_slice.quote_bars.count > 0:
            self._week = week
            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 MostLiquidFundamentalUniverseSelectionModel(CoarseFundamentalUniverseSelectionModel):

    def __init__(self, universe_settings: UniverseSettings, universe_size: int = 100) -> None:
        self._universe_size = universe_size
        super().__init__(self.select_coarse, universe_settings)

    def select_coarse(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        # Select the stocks with the greatest dollar volume
        sorted_by_dollar_volume = sorted(coarse, key=lambda c: c.dollar_volume, reverse=True)
        return [c.symbol for c in sorted_by_dollar_volume[:self._universe_size]]