from AlgorithmImports import *

class PriceComparisonAlgorithm(QCAlgorithm):

    # ---------------------------------------------------------------------- #
    #  INITIALIZATION                                                        #
    # ---------------------------------------------------------------------- #
    def initialize(self) -> None:
        self.set_start_date(2021, 1, 1)
        self.set_end_date(2022, 1, 1)
        self.set_cash(100_000)

        # Underlying
        self._spy: Symbol = self.add_equity("SPY", Resolution.MINUTE).symbol
        option = self.add_option("SPY", Resolution.MINUTE)
        option.set_filter(-10, 10, timedelta(0), timedelta(days=7))
        self._spy_options_symbol: Symbol = option.symbol

        # Trade-tracking
        self._current_call_contract: Symbol | None = None
        self._current_put_contract: Symbol | None = None
        self._current_ticket: OrderTicket | None = None
        self._entry_spy_price: float | None = None
        self._entry_time: datetime = self.start_date      # initialise

        # Signal flags
        self.buy_signal_active: bool = False
        self.sell_signal_active: bool = False

        # ------------------------------------------------  Charts  ---------- #
        price_chart = Chart("Price Comparison")
        price_chart.add_series(Series("Weighted Mean", SeriesType.LINE, 0))
        price_chart.add_series(Series("Live Price",   SeriesType.LINE, 0))
        price_chart.add_series(Series("Signal",       SeriesType.SCATTER, 0))
        self.add_chart(price_chart)

        for chart_name, series_name in [
            ("Z-Score", "Z"),
            ("Mean Slope", "Slope"),
            ("RSI", "RSI"),
            ("ATR-Normalized Deviation", "NormDev")
        ]:
            chart = Chart(chart_name)
            chart.add_series(Series(series_name, SeriesType.LINE, 0))
            self.add_chart(chart)

        # ------------------------------------------------  Indicators  ------ #
        lookback: int = 21
        self._mean_window: RollingWindow[float] = RollingWindow[float](lookback)
        self._rsi = self.rsi(self._spy, lookback, MovingAverageType.WILDERS, Resolution.HOUR)
        self._atr = self.atr(self._spy, lookback, MovingAverageType.WILDERS, Resolution.HOUR)
        self.warm_up_indicator(self._spy, self._rsi, Resolution.HOUR)
        self.warm_up_indicator(self._spy, self._atr, Resolution.HOUR)

        # ------------------------------------------------  Consolidator  ---- #
        consolidator = TradeBarConsolidator(timedelta(hours=1))
        consolidator.data_consolidated += self.on_hour_bar
        self.subscription_manager.add_consolidator(self._spy, consolidator)

    # ---------------------------------------------------------------------- #
    #  ON DATA                                                               #
    # ---------------------------------------------------------------------- #
    def on_data(self, data: Slice) -> None:
        spy_price: float = self.securities[self._spy].price
        now: datetime = self.time
        did_liq: bool = False

        #  Minimum hold period to avoid “instant” exits (except price-trigger)
        min_hold_period: timedelta = timedelta(minutes=1)
        time_since_entry: timedelta = now - self._entry_time

        # =============================  EXIT – CALL  ======================== #
        if (self._current_call_contract
                and self.portfolio[self._current_call_contract].invested
                and self._entry_spy_price is not None):

            expiry: datetime = self._current_call_contract.id.date
            time_to_expiry: timedelta = expiry - now

            # SPY-price stop / take-profit
            price_trigger_call: bool = (spy_price <= self._entry_spy_price * 0.995
                                        or spy_price >= self._entry_spy_price * 1.007)

            # Time-to-expiry trigger
            expiry_trigger_call: bool = False
            if self.live_mode:
                expiry_trigger_call = time_to_expiry <= timedelta(minutes=5)
            else:
                expiry_trigger_call = now.date() >= expiry.date()

            if time_since_entry < min_hold_period:
                expiry_trigger_call = False      # block “instant” close

            if price_trigger_call or expiry_trigger_call:
                self.liquidate(self._current_call_contract)
                self._reset_position_tracking()
                did_liq = True

        # =============================  EXIT – PUT   ======================== #
        if (not did_liq and self._current_put_contract
                and self.portfolio[self._current_put_contract].invested
                and self._entry_spy_price is not None):

            expiry: datetime = self._current_put_contract.id.date
            time_to_expiry: timedelta = expiry - now

            price_trigger_put: bool = (spy_price >= self._entry_spy_price * 1.007
                                       or spy_price <= self._entry_spy_price * 0.995)

            expiry_trigger_put: bool = False
            if self.live_mode:
                expiry_trigger_put = time_to_expiry <= timedelta(minutes=5)
            else:
                expiry_trigger_put = now.date() >= expiry.date()

            if time_since_entry < min_hold_period:
                expiry_trigger_put = False

            if price_trigger_put or expiry_trigger_put:
                self.liquidate(self._current_put_contract)
                self._reset_position_tracking()
                did_liq = True

        # ------------------------------------------------------------------- #
        #  ENTRY                                                              #
        # ------------------------------------------------------------------- #
        if self.portfolio.invested:
            return

        if not (self.buy_signal_active or self.sell_signal_active):
            return

        if self._spy_options_symbol not in data.option_chains:
            return
        chain = data.option_chains[self._spy_options_symbol]
        if not chain:
            return

        expiries = sorted({c.expiry for c in chain})
        if not expiries:
            return

        # Choose nearest expiry ~7 DTE
        target_dte: int = 7
        nearest_expiry: datetime = min(
            expiries,
            key=lambda e: abs((e.date() - (now.date() + timedelta(days=target_dte))).days)
        )
        same_expiry_chain = [c for c in chain if c.expiry == nearest_expiry]
        if not same_expiry_chain:
            return

        strikes = sorted({c.strike for c in same_expiry_chain})
        if not strikes:
            return

        # ------------------------------ BUY CALL ---------------------------- #
        if self.buy_signal_active:
            otm_call_strikes = [s for s in strikes if s > spy_price]
            if not otm_call_strikes:
                return
            call_strike: float = min(otm_call_strikes)
            call_contracts = [c for c in same_expiry_chain
                              if c.strike == call_strike and c.right == OptionRight.CALL]
            if not call_contracts:
                return
            call_symbol: Symbol = call_contracts[0].symbol

            self._current_ticket = self.market_order(call_symbol, 1)
            self._current_call_contract = call_symbol
            self._current_put_contract = None
            self._entry_spy_price = spy_price
            self._entry_time = now
            self.buy_signal_active = False

        # ------------------------------ BUY PUT  ---------------------------- #
        elif self.sell_signal_active:
            otm_put_strikes = [s for s in strikes if s < spy_price]
            if not otm_put_strikes:
                return
            put_strike: float = max(otm_put_strikes)
            put_contracts = [c for c in same_expiry_chain
                             if c.strike == put_strike and c.right == OptionRight.PUT]
            if not put_contracts:
                return
            put_symbol: Symbol = put_contracts[0].symbol

            self._current_ticket = self.market_order(put_symbol, 1)
            self._current_put_contract = put_symbol
            self._current_call_contract = None
            self._entry_spy_price = spy_price
            self._entry_time = now
            self.sell_signal_active = False

    # ---------------------------------------------------------------------- #
    #  HOURLY CONSOLIDATOR                                                   #
    # ---------------------------------------------------------------------- #
    def on_hour_bar(self, sender: object, bar: TradeBar) -> None:
        weighted_mean: float = (bar.open + 2*bar.high + 2*bar.low + 3*bar.close) / 8.0
        self._mean_window.add(weighted_mean)

        live_price: float = self.securities[self._spy].price
        self.plot("Price Comparison", "Weighted Mean", weighted_mean)
        self.plot("Price Comparison", "Live Price", live_price)

        signal_price: float | None = None
        if self._mean_window.is_ready:
            window_vals = [v for v in self._mean_window]
            mean_val: float = float(np.mean(window_vals))
            std_val: float = float(np.std(window_vals))
            z: float = (live_price - mean_val) / std_val if std_val > 0 else 0.0
            self.plot("Z-Score", "Z", z)

            slope: float = float(np.polyfit(range(len(window_vals)), window_vals, 1)[0]) \
                if len(window_vals) >= 2 else 0.0
            self.plot("Mean Slope", "Slope", slope)

            trigger_buy:  bool = z < -2 and abs(slope) > 0.30
            trigger_sell: bool = z >  2 and abs(slope) > 0.30

            if trigger_sell and not self.sell_signal_active:
                signal_price = live_price
                self.debug(f"SELL Signal | Z:{z:.2f} | Slope:{slope:.2f}")
                self.sell_signal_active, self.buy_signal_active = True, False
            elif trigger_buy and not self.buy_signal_active:
                signal_price = live_price
                self.debug(f"BUY  Signal | Z:{z:.2f} | Slope:{slope:.2f}")
                self.buy_signal_active, self.sell_signal_active = True, False
            else:
                if not (trigger_buy or trigger_sell):
                    self.buy_signal_active = self.sell_signal_active = False

            if signal_price is not None:
                self.plot("Price Comparison", "Signal", signal_price)

        if self._rsi.is_ready:
            self.plot("RSI", "RSI", self._rsi.current.value)
        if self._atr.is_ready and self._atr.current.value != 0:
            norm_dev: float = abs(live_price - weighted_mean) / self._atr.current.value
            self.plot("ATR-Normalized Deviation", "NormDev", norm_dev)

    # ---------------------------------------------------------------------- #
    #  UTILITIES                                                             #
    # ---------------------------------------------------------------------- #
    def _reset_position_tracking(self) -> None:
        """Clear position-related fields after a liquidation."""
        self._current_call_contract = None
        self._current_put_contract = None
        self._current_ticket = None
        self._entry_spy_price = None
        self._entry_time = self.time
        self.buy_signal_active = False
        self.sell_signal_active = False