"""TopkDropout with HMM high-volatility + drawdown-pause risk gates.

Gates NEW entries on two risk conditions (held names are never force-sold):

1. **HMM high-vol pause**: when the cross-sectional mean of ``sp_hmm_p_regime1``
   (HMM high-vol regime probability) on the signal date is >= ``hmm_pause_pct``,
   new buys are paused. The time-series study showed HMM high-vol probability
   pulses BEFORE sharp moves (regime-change cut) — pausing new exposure at the
   boundary reduces drawdown from price over-reaction.
2. **Drawdown pause**: when the account equity drawdown from its running peak
   exceeds ``drawdown_pause_pct``, new buys are paused. This is the
   ``drawdown_pause_pct`` risk-limit expressed inside the backtest (the pure
   executor-side gate is documented as not expressible in a one-shot backtest).
3. **Liquidity floor**: names whose 20-day average daily dollar volume is below
   ``liquidity_floor_adv`` are dropped from BUY candidates (the proven mitigant
   from exp-18: $5M floor cut drawdown 7.9%->5.4% at higher IR).

Implementation: pre-filter the signal score before the base TopkDropout
decision — non-held names get score 0 when any gate fires.

Wired into a workflow yaml like:

    strategy:
        class: HmmRiskTopk
        module_path: tac_qlib.contrib.strategy.hmm_risk
        kwargs:
            signal: "<PRED>"
            topk: 10
            n_drop: 2
            only_tradable: true
            risk_degree: 0.95
            hmm_pause_pct: 0.70
            drawdown_pause_pct: 8.0
            liquidity_floor_adv: 5000000
"""

from __future__ import annotations

import copy
from typing import Dict

import numpy as np
import pandas as pd

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

__all__ = ["HmmRiskTopk"]


class HmmRiskTopk(TopkDropoutStrategy):
    """TopkDropoutStrategy with HMM high-vol pause + drawdown pause + liquidity floor."""

    def __init__(
        self,
        *,
        hmm_pause_pct: float = 0.70,
        drawdown_pause_pct: float = 8.0,
        liquidity_floor_adv: float = 0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.hmm_pause_pct = float(hmm_pause_pct)
        self.drawdown_pause_pct = float(drawdown_pause_pct)
        self.liquidity_floor_adv = float(liquidity_floor_adv)
        self._peak_equity = 0.0

    # ------------------------------------------------------------- gates
    def _hmm_high_vol(self, pred_date) -> bool:
        """Cross-sectional mean HMM high-vol regime probability >= threshold."""
        try:
            from qlib.data import D

            feat = D.features(D.instruments("all"), ["$sp_hmm_p_regime1"],
                              start_time=pred_date, end_time=pred_date)
            if feat is None or len(feat) == 0:
                return False
            p = feat["$sp_hmm_p_regime1"].dropna()
            if len(p) == 0:
                return False
            return float(p.mean()) >= self.hmm_pause_pct
        except Exception:
            return False

    def _drawdown_active(self, equity: float) -> bool:
        if self.drawdown_pause_pct <= 0:
            return False
        self._peak_equity = max(self._peak_equity, equity)
        if self._peak_equity <= 0:
            return False
        dd = (self._peak_equity - equity) / self._peak_equity * 100.0
        return dd >= self.drawdown_pause_pct

    def _illiquid(self, codes, asof) -> Dict[str, bool]:
        if self.liquidity_floor_adv <= 0 or not codes:
            return {}
        from tac_qlib.risk_limits import dollar_adv

        adv = dollar_adv(codes, market="US", asof=asof, lookback=20)
        return {c: adv.get(str(c).upper(), 0.0) < self.liquidity_floor_adv for c in codes}

    # ------------------------------------------------------------- decision
    def generate_trade_decision(self, execute_result=None):
        trade_step = self.trade_calendar.get_trade_step()
        trade_start_time, trade_end_time = self.trade_calendar.get_step_time(trade_step)
        pred_start_time, pred_end_time = self.trade_calendar.get_step_time(trade_step, shift=1)
        pred_score = self.signal.get_signal(start_time=pred_start_time, end_time=pred_end_time)
        if pred_score is None:
            return TradeDecisionWO([], self)
        if isinstance(pred_score, pd.DataFrame):
            pred_score = pred_score.iloc[:, 0]

        current_temp = copy.deepcopy(self.trade_position)
        assert isinstance(current_temp, Position)
        held = {c for c in current_temp.get_stock_list() if abs(current_temp.get_stock_amount(c)) > 1e-6}

        equity = current_temp.get_cash()
        for code in held:
            mark = self.trade_exchange.get_deal_price(
                stock_id=code, start_time=trade_start_time, end_time=trade_end_time, direction=1
            )
            if mark is not None and np.isfinite(mark):
                equity += abs(current_temp.get_stock_amount(code)) * mark

        hmm_pause = self._hmm_high_vol(str(pd.Timestamp(pred_start_time).date()))
        dd_pause = self._drawdown_active(equity)
        buys_paused = hmm_pause or dd_pause

        pred_score = pred_score.copy()
        if buys_paused or self.liquidity_floor_adv > 0:
            new_codes = [c for c in pred_score.index if c not in held]
            illiquid = self._illiquid(new_codes, str(pd.Timestamp(pred_start_time).date()))
            for code in new_codes:
                if buys_paused or illiquid.get(code, False):
                    pred_score[code] = -1e9  # cannot enter today

        return super().generate_trade_decision(execute_result)