# region imports
from AlgorithmImports import *
# endregion


class NoBuzzMomentumAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2021, 9, 1)
        self.set_end_date(2026, 9, 1)
        self.set_cash(10_000_000)
        self.settings.seed_initial_prices = True
        # Define some parameters.
        self._slow_period = 252
        self._fast_period = 21
        self._fast_dollar_vol_filter_size = 1_500
        self._max_volume_surge_ratio = 1.5
        self._universe_size = 50
        # Add an indicator universe of US Equities.
        self._symbol_data_by_symbol = {}
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(self._select_assets)
        # Add a warm up period to prime the indicators of all stocks in the universe.
        self.set_warm_up(timedelta(400))
        # Add a Scheduled Event to rebalance the portfolio each month.
        self.schedule.on(self.date_rules.month_start('SPY'), self.time_rules.at(8, 0), self._rebalance)

    def _select_assets(self, fundamentals: list[Fundamental]) -> list[Symbol]:
        # Update the indicators of all stocks.
        candidates = []
        for f in fundamentals:
            if f.symbol not in self._symbol_data_by_symbol:
                self._symbol_data_by_symbol[f.symbol] = SymbolData(f.symbol, self._slow_period, self._fast_period)
            symbol_data = self._symbol_data_by_symbol[f.symbol]
            if symbol_data.update(f.end_time, f.adjusted_price, f.dollar_volume):
                candidates.append(f)
        # During warm-up, do nothing.
        if self.is_warming_up:
            return []
        # Apply the price and market-cap filters.
        filtered = [self._symbol_data_by_symbol[f.symbol] for f in candidates if f.price > 5.0 and f.market_cap > 500_000_000]
        # Select the stocks with the greatest mean dollar volume over the last month.
        filtered = sorted(filtered, key=lambda sd: sd.fast_dollar_vol)[-self._fast_dollar_vol_filter_size:]
        
        
        # Get surge ratios of remaining stocks: 
        values = [sd.volume_surge_ratio for sd in filtered]
        values = np.array([sd.volume_surge_ratio for sd in filtered])
        mean = values.mean()
        std = values.std()
        threshold = mean + 2 * std

        filtered = [sd for sd in filtered if sd.roc.current.value > 0 and mean <= sd.volume_surge_ratio <= threshold]
        self.plot('Universe', 'Size', len(filtered))
        return [sd.symbol for sd in filtered]

    def _rebalance(self) -> None:
        # Form an equal-weighted portfolio.
        securities = [self.securities[symbol] for symbol in self._universe.selected]
        securities = [s for s in securities if s.price]
        self.set_holdings([PortfolioTarget(s, 1/len(securities)) for s in securities], True)


class SymbolData(object):

    def __init__(self, symbol, slow_period, fast_period):
        self.symbol = symbol
        self.roc = RateOfChange(slow_period)
        self.fast_dollar_vol = SimpleMovingAverage(fast_period)
        self.slow_dollar_vol = SimpleMovingAverage(slow_period-fast_period)
        self.delayed_slow_dollar_vol = IndicatorExtensions.of(Delay(fast_period), self.slow_dollar_vol)
        self.volume_surge_ratio = None

    def update(self, time, close, dollar_volume):
        self.fast_dollar_vol.update(time, dollar_volume)
        self.slow_dollar_vol.update(time, dollar_volume)
        if self.roc.update(time, close):
            self.volume_surge_ratio = self.fast_dollar_vol.current.value / self.delayed_slow_dollar_vol.current.value
            return True
        return False