# region imports
from AlgorithmImports import *
# endregion


class EquityIndicatorUniverseSelectionAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2024, 9, 1)
        self.set_end_date(2024, 12, 31)
        self._selection_data_by_symbol = {}
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(self._select_assets)

    def _select_assets(self, fundamentals):
        # Update the indicator of all stocks in the universe dataset and
        # get the subset of stocks that have their indicator ready.
        ready_stocks = [
            f for f in fundamentals
            if self._selection_data_by_symbol.setdefault(
                f.symbol, EquityDataPointIndicatorSelectionData(self, f, SimpleMovingAverage(21))
            ).update(f)
        ]
        # As assests leave the Fundamental dataset, delete their SelectionData object.
        delisted_assets = list(
            set(self._selection_data_by_symbol.keys()) 
            - set(f.symbol for f in fundamentals)
        )
        for symbol in delisted_assets:
            self._selection_data_by_symbol.pop(symbol)
        # Select a subset of the stocks based on the indicator.
        # Example: 10 stocks furthest above their SMA.
        factor_by_symbol = {
            f.symbol: f.price / self._selection_data_by_symbol[f.symbol].indicator.current.value 
            for f in ready_stocks
        }
        return sorted(
            {k: v for k, v in factor_by_symbol.items() if v > 0}, 
            key=lambda symbol: factor_by_symbol[symbol]
        )[-100:]


class EquityDataPointIndicatorSelectionData:

    def __init__(self, algorithm, f, indicator):
        self._algorithm = algorithm
        self._price_scale_factor = f.price_scale_factor
        self.indicator = indicator

    def update(self, f):
        # If there has been a split or dividend, reset the indicator
        # and warm it up with the new adjusted history.
        if f.price_scale_factor != self._price_scale_factor:
            self._price_scale_factor = f.price_scale_factor
            self.indicator.reset()
            history = self._algorithm.history[TradeBar](
                f.symbol, 
                self.indicator.warm_up_period, 
                Resolution.DAILY, 
                data_normalization_mode=DataNormalizationMode.SCALED_RAW
            )
            for bar in history:
                self.indicator.update(bar)
        # Otherwise, just update the indicator like normally.
        else:
            self.indicator.update(f.end_time, f.price)
        return self.indicator.is_ready