from AlgorithmImports import *


class SpyTqqqSpreadReversion(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2010, 3, 5)
        self.set_cash(100000)

        # Number of contracts (each contract = 1 share SPY + 1 share TQQQ, opposite signs)
        self._n_contracts = 100

        # Take-profit / stop-loss per contract (±$10 per contract), scaled by N
        self._tp_sl = 10.0 * self._n_contracts

        # Subscribe to SPY and TQQQ at hourly resolution (default ADJUSTED normalization)
        self._spy = self.add_equity("SPY", Resolution.HOUR)
        self._tqqq = self.add_equity("TQQQ", Resolution.HOUR)
        self._spy.set_fee_model(ConstantFeeModel(0))
        self._tqqq.set_fee_model(ConstantFeeModel(0))

        # Manual SMA for the 74-hour rolling mean of first difference of spread
        self._spread_sma = SimpleMovingAverage(74)

        # Track previous spread for first difference
        self._prev_spread = None

        # Track current signal state (0=flat, +1=long spread, -1=short spread)
        self._prev_signal = 0

        # Warm up: need 75+ hourly bars (1 for prev_spread + 74 first differences for SMA)
        self.set_warm_up(80, Resolution.HOUR)

    def on_data(self, data: Slice) -> None:
        # Need both SPY and TQQQ bars
        if not data.bars.contains_key(self._spy.symbol) or not data.bars.contains_key(self._tqqq.symbol):
            return

        spy_bar = data.bars[self._spy.symbol]
        tqqq_bar = data.bars[self._tqqq.symbol]

        spy_price = spy_bar.close
        tqqq_price = tqqq_bar.close

        # Compute spread = AdjClose(SPY) - AdjClose(TQQQ)
        spread = spy_price - tqqq_price

        # Compute first difference d = Spread_t - Spread_{t-1} and feed SMA
        d = 0.0
        if self._prev_spread is not None:
            d = spread - self._prev_spread
            self._spread_sma.update(spy_bar.end_time, d)
        self._prev_spread = spread

        # During warm-up, just build indicator state
        if self.is_warming_up:
            return

        # Need SMA to be ready
        if not self._spread_sma.is_ready:
            return

        # Threshold: rolling 74-hour mean of first difference
        tau = self._spread_sma.current.value

        # Signal: +1 if d > tau, -1 if d < -tau, 0 otherwise
        signal = 0
        if d > tau:
            signal = 1
        elif d < -tau:
            signal = -1

        # Signal-driven position management
        if signal != self._prev_signal:
            if signal == 0:
                self._set_position(0)
            else:
                self._set_position(signal)
            self._prev_signal = signal
            return

        # ±$10 take-profit / stop-loss overlay (only when holding)
        if signal != 0:
            pnl = (self.portfolio[self._spy.symbol].unrealized_profit +
                   self.portfolio[self._tqqq.symbol].unrealized_profit)
            if pnl >= self._tp_sl or pnl <= -self._tp_sl:
                self._set_position(0)
                self._prev_signal = 0
                return

        self._prev_signal = signal

    def _set_position(self, target_direction: int) -> None:
        """Set position to target direction: +1=long spread, -1=short spread, 0=flat."""
        target_spy = target_direction * self._n_contracts
        target_tqqq = -target_direction * self._n_contracts

        current_spy = self.portfolio[self._spy.symbol].quantity
        current_tqqq = self.portfolio[self._tqqq.symbol].quantity

        spy_order_qty = target_spy - current_spy
        tqqq_order_qty = target_tqqq - current_tqqq

        if spy_order_qty != 0:
            self.market_order(self._spy.symbol, spy_order_qty)
        if tqqq_order_qty != 0:
            self.market_order(self._tqqq.symbol, tqqq_order_qty)