# region imports
from AlgorithmImports import *
from plot import SpectralPlotter
# endregion


class SpectralPeriodicityPremiumAlgorithm(QCAlgorithm):

    def initialize(self):
        # Bound the run to the dataset coverage window (explicit end date is required here).
        self.set_start_date(2024, 1, 1)
        self.set_end_date(2026, 7, 1)
        self.set_cash(1_000_000)
        self.settings.seed_initial_prices = True
        self.universe_settings.resolution = Resolution.DAILY
        # Smooth the sparse score into a persistent execution-intensity proxy.
        self._exec_decay = 0.94
        # Long the top fifth of ranked names by periodicity strength.
        self._rank_fraction = 0.2
        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)[-100:]]
        )
        self._plotter = SpectralPlotter(self)
        self.set_warm_up(timedelta(14))

    def on_warmup_finished(self) -> None:
        # Rebalance weekly at 8 AM to match the daily-data cadence; the premium is a next-session hold.
        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):
        # Subscribe each added equity to its signal feed; drop feeds and holdings for removed names.
        for security in changes.added_securities:
            security.signal = None
            security.exec_intensity = 0.0
            security.dataset_symbol = self.add_data(QuantConnectSpectralTickFlowSignal, security, Resolution.DAILY).symbol
        for security in changes.removed_securities:
            self.remove_security(security.dataset_symbol)

    def _rebalance(self):
        # Rank the selected names by volume-periodicity strength, long the top slice, and tilt by execution intensity.
        if self.is_warming_up or not self._universe.selected:
            return
        eligible = []
        for symbol in self._universe.selected:
            security = self.securities[symbol]
            history = self.history[QuantConnectSpectralTickFlowSignal](security.dataset_symbol, timedelta(90), Resolution.DAILY)
            for point in history:
                security.signal = point
                # Fold this point's execution score into the persistent execution-intensity EMA.
                security.exec_intensity = self._exec_decay * security.exec_intensity + (1.0 - self._exec_decay) * point.execution_score
            # Drop names with no signal in the lookback window.
            if security.signal is None:
                continue
            eligible.append(security)
        # Sort ascending so the strongest periodicity lands at the tail.
        ranked = sorted(eligible, key=lambda security: security.signal.volume_variance_explained)
        # Take the top rank fraction, but never fewer than five names and never more than are available.
        count = max(5, int(len(ranked) * self._rank_fraction))
        count = min(count, len(ranked))
        long_securities = ranked[-count:]
        # Weight the longs by persistent execution intensity.
        peak = max((security.exec_intensity for security in long_securities))
        # Give each long a raw weight of one plus its execution intensity relative to the peak.
        raw_by_security = {security: 1.0 + security.exec_intensity / peak for security in long_securities}
        # Normalize each raw weight into a portfolio fraction that sums to one.
        weight_by_security = {security: raw / sum(raw_by_security.values()) for security, raw in raw_by_security.items()}
        # Build a portfolio target from each name's normalized weight.
        targets = [PortfolioTarget(security, weight) for security, weight in weight_by_security.items()]
        # Rebalance to the targets, liquidating any holdings not in the target list.
        self.set_holdings(targets, True)
        # Plot this rebalance's long selection for diagnostics.
        self._plotter.plot_selection(long_securities)

from AlgorithmImports import *


class SpectralPlotter:

    def __init__(self, algorithm):
        self._algorithm = algorithm

    def _plot(self, chart_name, series):
        for name, value in series:
            self._algorithm.plot(chart_name, name, value)

    def _mean_periodicity(self, securities):
        # Average intraday-volume periodicity strength across the leg, for diagnostics.
        if not securities:
            return 0.0
        return sum(security.signal.volume_variance_explained for security in securities) / len(securities)

    def _mean_exec_intensity(self, securities):
        # Average persistent execution-intensity proxy across the leg -- the execution tilt strength.
        if not securities:
            return 0.0
        return sum(security.exec_intensity for security in securities) / len(securities)

    def plot_selection(self, long_securities):
        # Chart the intraday-volume periodicity strength driving the premium.
        self._plot("Spectral Periodicity", [("Long Mean Periodicity", self._mean_periodicity(long_securities))])
        # Chart the execution-intensity strength tilting the long weights.
        self._plot("Spectral Exec Intensity", [("Long Mean Exec Intensity", self._mean_exec_intensity(long_securities))])
        # Chart how many names clear the periodicity ranking each rebalance.
        self._plot("Spectral Selection", [("Long Count", len(long_securities))])