# region imports
from AlgorithmImports import *
# endregion

from dataclasses import dataclass
from typing import Optional


@dataclass
class PanicState:
    """Per-symbol state for the panic mean reversion strategy."""

    sma: SimpleMovingAverage
    atr: AverageTrueRange
    prev_close: Optional[float] = None
    prev_high: Optional[float] = None
    entry_order_id: Optional[int] = None
    target_order_id: Optional[int] = None
    hold_days: int = 0
    target_price: float = 0.0


class PanicMeanReversion(QCAlgorithm):
    """
    Time-series mean reversion: buy sharp one-day drops in stocks that are
    still in a long-term uptrend, then exit quickly.

    Entry (all must hold on the panic day):
      - close > 200-day SMA (trend filter)
      - single-day drop > 3%
      Next day: limit buy at panic close - 0.9 x ATR(5).

    Exit (whichever comes first):
      - close above yesterday's high
      - limit sell at panic close + 0.5 x ATR(5)
      - 10-day time stop
    """

    _sma_period = 200
    _atr_period = 5
    _drop_threshold = -0.03
    _entry_atr_mult = 0.9
    _target_atr_mult = 0.5
    _time_stop_days = 10
    _position_pct = 0.10
    _max_positions = 10

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

        # ETF universe ticker is a changeable parameter (default QQQ).
        self._etf_ticker = self.get_parameter("etf-ticker", "SPY")
        self.universe_settings.resolution = Resolution.DAILY
        self.set_benchmark(self._etf_ticker)

        self._states: dict[Symbol, PanicState] = {}
        self._universe = self.add_universe(
            self.universe.etf(self._etf_ticker, self._etf_constituents_filter)
        )

        self.schedule.on(
            self.date_rules.every_day(self._etf_ticker),
            self.time_rules.midnight,
            self._prune_states,
        )

    def _etf_constituents_filter(self, constituents: List[ETFConstituentUniverse]) -> List[Symbol]:
        for constituent in constituents:
            if constituent.symbol not in self._states:
                self._warm_up_symbol(constituent.symbol)
        return [constituent.symbol for constituent in constituents]

    def _warm_up_symbol(self, symbol: Symbol) -> None:
        sma = SimpleMovingAverage(self._sma_period)
        atr = AverageTrueRange(self._atr_period)
        self.warm_up_indicator(symbol, sma, Resolution.DAILY)
        self.warm_up_indicator(symbol, atr, Resolution.DAILY)
        self._states[symbol] = PanicState(sma=sma, atr=atr)

    def on_data(self, data: Slice) -> None:
        for symbol, bar in data.bars.items():
            state = self._states.get(symbol)
            if state is None:
                # Benchmark / non-universe symbol.
                continue

            # Update indicators with the completed daily bar.
            state.sma.update(bar.end_time, bar.close)
            state.atr.update(bar)

            if self.portfolio[symbol].invested:
                state.hold_days += 1
                if state.prev_high is not None and bar.close > state.prev_high:
                    self._exit(symbol, state, "bounce")
                elif state.hold_days >= self._time_stop_days:
                    self._exit(symbol, state, "time-stop")
            elif state.entry_order_id is not None:
                # Entry limit had exactly the day after the signal to fill.
                ticket = self.transactions.get_order_by_id(state.entry_order_id)
                if ticket.status in (OrderStatus.NEW, OrderStatus.SUBMITTED, OrderStatus.PARTIALLY_FILLED, OrderStatus.UPDATE_SUBMITTED):
                    self.transactions.cancel_open_orders(symbol)
                state.entry_order_id = None

            if not self.portfolio[symbol].invested and state.entry_order_id is None:
                self._try_enter(symbol, bar, state)

            state.prev_close = bar.close
            state.prev_high = bar.high

        self.plot("Risk", "Open Positions", sum(1 for h in self.portfolio.values() if h.invested))

    def _try_enter(self, symbol: Symbol, bar: TradeBar, state: PanicState) -> None:
        if not state.sma.is_ready or not state.atr.is_ready:
            return
        if state.prev_close is None:
            return
        # Panic signal: single-day drop strictly greater than 3%.
        if bar.close >= state.prev_close * (1 + self._drop_threshold):
            return
        # Trend filter: price above the 200-day moving average.
        if bar.close <= state.sma.current.value:
            return
        # Cap concurrent exposure.
        invested_count = sum(1 for h in self.portfolio.values() if h.invested)
        open_entries = sum(1 for s in self._states.values() if s.entry_order_id is not None)
        if invested_count + open_entries >= self._max_positions:
            return

        atr_value = state.atr.current.value
        limit_price = bar.close - self._entry_atr_mult * atr_value
        if limit_price <= 0:
            return

        quantity = self.calculate_order_quantity(symbol, self._position_pct)
        if quantity <= 0:
            return

        state.target_price = bar.close + self._target_atr_mult * atr_value
        ticket = self.limit_order(symbol, quantity, limit_price)
        state.entry_order_id = ticket.order_id

    def _exit(self, symbol: Symbol, state: PanicState, reason: str) -> None:
        if state.target_order_id is not None:
            self.transactions.cancel_open_orders(symbol)
            state.target_order_id = None
        self.liquidate(symbol, tag=f"exit-{reason}")
        state.hold_days = 0

    def on_order_event(self, order_event: OrderEvent) -> None:
        if order_event.status == OrderStatus.FILLED:
            state = self._state_for_order(order_event.order_id)
            if state is None:
                return
            if order_event.order_id == state.entry_order_id and order_event.direction == OrderDirection.BUY:
                ticket = self.transactions.get_order_by_id(order_event.order_id)
                if ticket.status == OrderStatus.FILLED:
                    state.entry_order_id = None
                    state.hold_days = 0
                    quantity = int(self.portfolio[order_event.symbol].quantity)
                    if quantity > 0 and state.target_price > 0:
                        target_ticket = self.limit_order(order_event.symbol, -quantity, state.target_price)
                        state.target_order_id = target_ticket.order_id
            elif order_event.order_id == state.target_order_id and order_event.direction == OrderDirection.SELL:
                state.target_order_id = None
                state.hold_days = 0
        elif order_event.status == OrderStatus.CANCELED:
            state = self._state_for_order(order_event.order_id)
            if state is not None:
                if order_event.order_id == state.entry_order_id:
                    state.entry_order_id = None
                if order_event.order_id == state.target_order_id:
                    state.target_order_id = None

    def _state_for_order(self, order_id: int) -> Optional[PanicState]:
        for state in self._states.values():
            if order_id in (state.entry_order_id, state.target_order_id):
                return state
        return None

    def _prune_states(self) -> None:
        for symbol in list(self._states.keys()):
            if symbol in self._universe.selected:
                continue
            if self.transactions.get_open_orders(symbol):
                self.transactions.cancel_open_orders(symbol)
            if not self.portfolio[symbol].invested:
                del self._states[symbol]