# region imports
from AlgorithmImports import *
# endregion

class SwimmingLightBrownGalago(QCAlgorithm):

    def initialize(self):
        self.set_cash(100_000_000)

        self.settings.daily_precise_end_time = False
        self.settings.minimum_order_margin_portfolio_percentage = 0

        lookback_years = 12 # 12 years includes 3 non-leap years & 3 election cycles
        self._lookback = lookback_years * 252 
        self._max_period = 21 # Trading days
        self._percentile_threshold = 0.1
        self._margin_ratio_boundary = 0.1
        self._weight_scaler = 5
        self._can_short = False

        warm_up_period = timedelta(lookback_years * 365)
        self.set_start_date(datetime(1998, 1, 1) + warm_up_period)
        self.set_warm_up(warm_up_period)

        self._equity = self.add_equity("SPY", Resolution.DAILY, fill_forward=False)
        self._equity.history = pd.Series()
        self._equity.signals_by_period = {self._max_period: 0} # One-month holding period

        # Add E-mini Futures.
        self._future = self.add_future(
            Futures.Indices.SP_500_E_MINI,
            data_mapping_mode=DataMappingMode.OPEN_INTEREST,
            data_normalization_mode=DataNormalizationMode.BACKWARDS_RATIO,
            contract_depth_offset=0
        )
        self._future.set_filter(lambda universe: universe.front_month())
        self._target_weight = None

        self._vix = self.add_data(CBOE, "VIX", Resolution.DAILY, fill_forward=False)
        self._vix.threshold = 30

    def on_data(self, data: Slice):
        # Get the current day's trade bar.
        bar = data.bars.get(self._equity.symbol)
        if bar:    
            # Update the historical data.
            self._equity.history.loc[bar.time] = bar.open # `time` and `open` since we trade at next market open.
            self._equity.history = self._equity.history.iloc[-self._lookback:]
            
            # Wait until there is enough history.
            if len(self._equity.history) < self._lookback or self.is_warming_up or self._vix.price >= self._vix.threshold:
                return

            for period in self._equity.signals_by_period.keys():
                # Calculate expected return of entering at next market open.
                period_returns = self._equity.history.pct_change(period).shift(-period).dropna()
                expected_return_by_day = period_returns.groupby(period_returns.index.strftime('%m-%d')).mean()
                next_market_open = self._equity.exchange.hours.get_next_market_open(self.time, False)
                expected_return = expected_return_by_day[next_market_open.strftime('%m-%d')]  # The expected return of buying next market open and holdings n days.
                
                # Calculate signal for this trade.
                std = expected_return_by_day.std()
                mean = expected_return_by_day.mean()
                if expected_return >= mean + std:
                    signal = 1
                elif expected_return <= mean - std:
                    signal = -1
                else:
                    signal = 0

                if signal:
                    # Record the signal and create a Scheduled Event to open the trade.
                    self._equity.signals_by_period[period] += signal
                    self._schedule_rebalance()

                    # Create a Scheduled Event to remove the signal and close the trade.
                    exit_day = next_market_open
                    for _ in range(period):
                        exit_day = self._equity.exchange.hours.get_next_market_open(exit_day, False)
                    self.schedule.on(
                        self.date_rules.on(exit_day.replace(hour=0, minute=0)),
                        self.time_rules.midnight,
                        lambda period=period, signal=signal: self._remove_signal(period, signal)
                    )

        if self.is_warming_up:
            return
        
        if self._vix.symbol in data:
            self.plot('VIX', 'Value', self._vix.price)
        
        # If the margin ratio deviates far from the initial margin ratio when opening the trade, adjust it back to the 
        # original margin ratio.
        if self.portfolio.invested and self._target_weight:
            margin_ratio = self.portfolio.total_margin_used / (self.portfolio.total_margin_used + self.portfolio.margin_remaining)
            if not self._target_weight - self._margin_ratio_boundary < margin_ratio < self._target_weight + self._margin_ratio_boundary:
                self.set_holdings(self._future.mapped, self._target_weight)
        
    def _schedule_rebalance(self):
        # Create a Scheduled Event to rebalance at the next market open.
        self.schedule.on(
            self.date_rules.on(self._future.exchange.hours.get_next_market_open(self.time, False)),
            self.time_rules.after_market_open(self._future.symbol, 1),
            self._rebalance
        )

    def _rebalance(self):
        # Rebalance based on the signals.
        weighted_signals = [signal / period for period, signal in self._equity.signals_by_period.items()]
        weight = sum(weighted_signals) / len(weighted_signals) / self._weight_scaler
        if not self._can_short:
            weight = max(0, weight)
        self._target_weight = weight
        self.plot("Weight", self._future.symbol.value, weight)
        self.set_holdings([PortfolioTarget(self._future.mapped, weight)])

    def on_symbol_changed_events(self, symbol_changed_events):
        # Create a Scheduled Event to roll over to the new contract at next market open.
        changed_event = symbol_changed_events.get(self._future.symbol)
        self.schedule.on(
            self.date_rules.on(self._future.exchange.hours.get_next_market_open(self.time, False).replace(hour=0, minute=0)),
            self.time_rules.after_market_open(self._future.symbol, 1),
            lambda old_symbol=changed_event.old_symbol, new_symbol=changed_event.new_symbol: self._roll(old_symbol, new_symbol)
        )

    def _roll(self, old_symbol, new_symbol):
        # Roll to the new contract while maintaining the current margin ratio.
        margin_ratio = self.portfolio.total_margin_used / (self.portfolio.total_margin_used + self.portfolio.margin_remaining)
        self.liquidate(old_symbol, tag=f"Rollover - Symbol changed at {self.time}: {old_symbol} -> {new_symbol}")
        self.set_holdings(new_symbol, margin_ratio)

    def _remove_signal(self, period, signal):
        # Remove the signal and create a Scheduled Event to close the trade.
        self._equity.signals_by_period[period] -= signal
        self._schedule_rebalance()