# region imports
from AlgorithmImports import *
# endregion


class SpectralPeriodicityPremiumAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2021, 7, 1)
        self.set_end_date(2026, 7, 1)
        self.set_cash(1_000_000)
        self.settings.seed_initial_prices = True
        # Define some parameters.
        self._universe_size = 100
        self._ema_period = 3*21
        self._quantiles = 5
        # Add a universe of US Equities.
        self.universe_settings.resolution = Resolution.DAILY
        self._date_rule = self.date_rules.week_start("SPY")
        self.universe_settings.schedule.on(self._date_rule)
        self._universe = self.add_universe(lambda fundamental:
            [f.symbol for f in sorted([f for f in fundamental if f.has_fundamental_data], key=lambda f: f.dollar_volume)[-self._universe_size:]]
        )
        # Add a warm-up period so the algorithm trades right away.
        self.set_warm_up(timedelta(14))

    def on_warmup_finished(self) -> None:
        # Rebalance weekly at 8 AM.
        time_rule = self.time_rules.at(8, 0)
        self.schedule.on(self._date_rule, time_rule, self._rebalance)
        # Rebalance today too.
        if self.live_mode:
            self._rebalance()
        else:
            self.schedule.on(self.date_rules.today, time_rule, self._rebalance)

    def on_securities_changed(self, changes: SecurityChanges) -> None:
        # As stocks enter the universe, add their spectral tick flow signals and warm-up its EMA.
        for security in changes.added_securities:
            security.dataset_symbol = self.add_data(QuantConnectSpectralTickFlowSignal, security, Resolution.DAILY).symbol
            security.exec_intensity = ExponentialMovingAverage(self._ema_period)
            for data_point in self.history[QuantConnectSpectralTickFlowSignal](security.dataset_symbol, timedelta(int(self._ema_period*1.5))):
                self._update_factors(security, data_point)
        for security in changes.removed_securities:
            self.remove_security(security.dataset_symbol)

    def _update_factors(self, security: Equity, data_point: QuantConnectSpectralTickFlowSignal) -> None:
        # Save the latest spectral tick flow signal and update the EMA.
        security.signal = data_point
        security.exec_intensity.update(data_point.end_time, data_point.execution_score)

    # As new spectral tick flow signals arrive, update the factors of each stock.
    def on_data(self, data: Slice) -> None:
        for dataset_symbol, data_point in data.get(QuantConnectSpectralTickFlowSignal).items():
            self._update_factors(self.securities[dataset_symbol.underlying], data_point)

    def _rebalance(self) -> None:
        if self.is_warming_up:
            return
        # Get the set of stocks in the universe that have enough signals.
        securities = [self.securities[symbol] for symbol in self._universe.selected]
        eligible = [security for security in securities if security.exec_intensity.is_ready]
        if not eligible:
            return
        # Select the subset of stocks that have the strongest periodicity.
        selected = sorted(eligible, key=lambda security: security.signal.volume_variance_explained)[-int(len(eligible)/self._quantiles):]
        # Give each stock a raw weight of one plus its execution intensity relative to the peak.
        peak = max((security.exec_intensity.current.value for security in selected))
        weight_by_security = {security: 1.0 + security.exec_intensity.current.value / peak for security in selected}
        # Normalize the raw weights so the portfolio exposure sums to one.
        weight_by_security = {security: weight / sum(weight_by_security.values()) for security, weight in weight_by_security.items()}
        # Place orders to rebalance the portfolio.
        targets = [PortfolioTarget(security, weight) for security, weight in weight_by_security.items()]
        self.set_holdings(targets, True)