# region imports
from AlgorithmImports import *
# endregion


class MomentumUniverseTrackingAlgorithm(QCAlgorithm):

    def initialize(self):
        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._period = self.get_parameter('lookback_months', 12) * 21
        self._winsor_percentile = self.get_parameter('winsor_percentile', 5.0)
        self._std_multiplier = 2
        # Add an universe of US Equities based on indicators.
        self._selection_data_by_symbol = {}
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(self._select_assets)
        self.set_warm_up(self._period, Resolution.DAILY)
        # Add a Scheduled Event to rebalance the portfolio each morning.
        self.schedule.on(self.date_rules.every_day('SPY'), self.time_rules.at(8, 0), self._rebalance)

    def _select_assets(self, fundamentals: List[Fundamental]):
        # Update the indicators of all stocks.
        candidates = []
        for f in fundamentals:
            if not f.has_fundamental_data:
                continue
            if f.symbol not in self._selection_data_by_symbol:
                self._selection_data_by_symbol[f.symbol] = SelectionData(self._period)
            if self._selection_data_by_symbol[f.symbol].update(f):
                candidates.append(f.symbol)
        # During warm-up, keep the universe empty so the algorithm runs quickly.
        if self.is_warming_up:
            return []
        # Apply the liquidity filter: Select the stocks with the greatest mean liquidity.
        candidates = self._outliers(candidates, lambda s: self._selection_data_by_symbol[s].mean_dollar_volume.current.value)
        self.plot('Universe', 'Filter 1', len(candidates))
        # Apply the momentum filter: Select the subset of stocks with the greatest momentum.
        candidates = self._outliers(candidates, lambda s: self._selection_data_by_symbol[s].momentum.current.value)
        self.plot('Universe', 'Filter 2', len(candidates))
        return candidates

    def _outliers(self, candidates, factor):
        values = np.array([factor(s) for s in candidates])
        # Winsorize the values so extreme outliers don't skew the standard deviation calculation.
        lower, upper = np.percentile(values, [self._winsor_percentile, 100 - self._winsor_percentile])
        clipped = np.clip(values, lower, upper)
        # Calculate the threshold based on standard deviations.
        threshold = clipped.mean() + self._std_multiplier * clipped.std() 
        # Select the stocks that exceed the threshold.
        return [candidates[i] for i in np.where(values > threshold)[0]]

    def _rebalance(self):
        # During warm-up, do nothing.
        if self.is_warming_up:
            return
        # Form an equal-weighted portfolio.
        weight = 1.0 / len(self._universe.selected)
        self.set_holdings([PortfolioTarget(s, weight) for s in self._universe.selected ], liquidate_existing_holdings=True)


class SelectionData:

    def __init__(self, period):
        self.momentum = MomentumPercent(period)
        self.mean_dollar_volume = SimpleMovingAverage(period)

    def update(self, f):
        return (
            self.momentum.update(f.end_time, f.adjusted_price) &
            self.mean_dollar_volume.update(f.end_time, f.dollar_volume)
        )