"""Realized-IC circuit breaker TopkDropout strategy.

Subclass of ``qlib.contrib.strategy.signal_strategy.TopkDropoutStrategy`` that
holds the book (issues NO orders) while the streaming realized RankIC of the
deployed signal is below threshold — i.e. the model's cross-sectional
predictions are no longer earning against realized forward returns. When the
gate is open it behaves exactly like the reference TopkDropoutStrategy.

The gate is evaluated per trade step on the trailing mean realized RankIC of
the signal over the last ``ic_window`` trading days whose label is fully
realized as of the decision date (no lookahead — a 5d fwd label ``close[t+6]/
close[t+1]-1`` is only known at ``t+6``).

Two wiring modes:

* ``ic_gate``: a precomputed ``pd.Series`` indexed by datetime of booleans
  (True = gate open / trade allowed). Computed once by the caller (e.g.
  ``rd_backtest``) and looked up per step. Missing dates default to open.
* realized-IC self-computation: when ``ic_min_rankic`` is given but no
  ``ic_gate``, the strategy computes the per-date realized RankIC itself from
  ``self.signal`` (the pred scores) and the lake 1d bars via
  ``tac_qlib.risk_limits.realized_rankic_series``, then applies the same
  trailing-window comparison. Works when instantiated from a workflow YAML
  PortAnaRecord config (``lake_root`` / ``market`` must be provided).
"""

from __future__ import annotations

import pandas as pd

from qlib.backtest.decision import TradeDecisionWO
from qlib.contrib.strategy.signal_strategy import TopkDropoutStrategy

from tac_qlib.risk_limits import ic_circuit_breaker, realized_rankic_series

__all__ = ["ICGateTopkDropoutStrategy"]


class ICGateTopkDropoutStrategy(TopkDropoutStrategy):
    """TopkDropout with a streaming realized-IC circuit breaker.

    Parameters
    ----------
    topk, n_drop, method_sell, method_buy, hold_thresh, only_tradable,
    forbid_all_trade_at_limit : same as ``TopkDropoutStrategy``.
    ic_min_rankic : float — pause new trading while trailing realized RankIC is
        below this threshold (0 disables the gate).
    ic_window : int — trailing window for the realized RankIC mean (default 22).
    ic_label_horizon : int — label horizon in trading days (default 6).
    ic_min_obs : int — min realized labels before the gate arms (default 10).
    ic_gate : pd.Series, optional — precomputed per-date gate (bool indexed by
        datetime). When provided, it overrides self-computation.
    lake_root, market : str — lake location for self-computed realized IC.
    """

    def __init__(
        self,
        *,
        topk,
        n_drop,
        ic_min_rankic: float = 0.0,
        ic_window: int = 22,
        ic_label_horizon: int = 6,
        ic_min_obs: int = 10,
        ic_gate=None,
        lake_root: str = "",
        market: str = "US",
        **kwargs,
    ):
        super().__init__(topk=topk, n_drop=n_drop, **kwargs)
        self.ic_min_rankic = float(ic_min_rankic or 0.0)
        self.ic_window = int(ic_window or 22)
        self.ic_label_horizon = int(ic_label_horizon or 6)
        self.ic_min_obs = int(ic_min_obs or 10)
        self._ic_gate = ic_gate
        self._realized_ic = None
        self.lake_root = lake_root or ""
        self.market = market or "US"

    def _load_realized_ic(self):
        if self._realized_ic is None:
            pred_start_time, pred_end_time = self.trade_calendar.get_step_time(
                self.trade_calendar.get_trade_step(), shift=-self.ic_label_horizon
            )
            pred = self.signal.get_signal(start_time=pred_start_time, end_time=pred_end_time)
            if isinstance(pred, pd.DataFrame):
                pred = pred.iloc[:, 0]
            self._realized_ic = realized_rankic_series(
                pred, self.lake_root, self.market, label_horizon=self.ic_label_horizon
            )
        return self._realized_ic

    def _gate_open(self, trade_start_time) -> bool:
        ts = pd.Timestamp(trade_start_time)
        if self._ic_gate is not None:
            # precomputed gate series: look up the latest known decision date <= ts
            known = self._ic_gate[self._ic_gate.index <= ts]
            if len(known):
                return bool(known.iloc[-1])
            return True
        if self.ic_min_rankic <= 0:
            return True
        realized = self._load_realized_ic()
        limits = {
            "ic_min_rankic": self.ic_min_rankic,
            "ic_window": self.ic_window,
            "ic_min_obs": self.ic_min_obs,
        }
        tripped, _reason, _trail = ic_circuit_breaker(realized, ts, limits)
        return not tripped

    def generate_trade_decision(self, execute_result=None):
        trade_step = self.trade_calendar.get_trade_step()
        trade_start_time, _ = self.trade_calendar.get_step_time(trade_step)
        if not self._gate_open(trade_start_time):
            return TradeDecisionWO([], self)
        return super().generate_trade_decision(execute_result)