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

Adds three risk controls on top of ``TopkDropoutStrategy``:

1. **HMM high-vol pause**: when the current day's HMM high-volatility regime
   probability (``sp_hmm_p_regime1`` feature, regime-1 = high-vol) is above
   ``hmm_pause_pct``, new buys are paused (existing positions held). This
   encodes the time-series study's finding that HMM high-vol probability pulses
   BEFORE sharp moves (regime-change cut) — pausing new exposure at the
   boundary reduces drawdown from "价格过度反应".
2. **Drawdown pause**: when the account equity drawdown from its running peak
   exceeds ``drawdown_pause_pct``, new buys are paused (positions kept). This
   is the ``drawdown_pause_pct`` risk-limit expressed in the backtest (the
   executor-side gate is documented as not expressible in a one-shot qlib
   backtest — here we implement it inside the strategy).
3. **Liquidity floor**: names whose average daily dollar volume is below
   ``liquidity_floor_adv`` are dropped from the tradable set (the proven risk
   mitigant from exp-18: $5M floor cut drawdown 7.9%->5.4% at higher IR).

Held names are never force-sold by these gates; only new entries are gated.
Exits remain the pure TopkDropout rule.

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
        self._hmm_pause_active = False
        self._drawdown_pause_active = False

    # ------------------------------------------------------------- state
    def _high_vol_active(self, trade_start) -> bool:
        """True when HMM high-vol regime probability >= pause threshold."""
        try:
            from qlib.data import D

            cal = D.calendar(start_time=str((trade_start - pd.Timedelta(days=10)).date()),
                             end_time=str(trade_start.date()))
            if len(cal) == 0:
                return False
            ref_date = str(pd.Timestamp(cal[-1]).date())
            feat = D.features(
                D.instruments("all"), ["$sp_hmm_p_regime1"], start_time=ref_date, end_time=ref_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 _update_pause(self, equity: float, trade_start) -> None:
        self._peak_equity = max(self._peak_equity, equity)
        dd = (self._peak_equity - equity) / self._peak_equity * 100.0 if self._peak_equity > 0 else 0.0
        self._drawdown_pause_active = self.drawdown_pause_pct > 0 and dd >= self.drawdown_pause_pct
        self._hmm_pause_active = self._high_vol_active(trade_start)

    # ------------------------------------------------------------- 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)
        current_temp = copy.deepcopy(self.trade_position)
        assert isinstance(current_temp, Position)

        # account equity for drawdown pause
        equity = current_temp.get_cash()
        for code in current_temp.get_stock_list():
            amt = abs(current_temp.get_stock_amount(code))
            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 += amt * mark
        self._update_pause(equity, trade_start_time)

        target_weight_position = self.generate_target_weight_position(
            score=pred_score, current=current_temp, trade_start_time=trade_start_time, trade_end_time=trade_end_time
        )

        held = set(current_temp.get_stock_list())
        held_amt = {c: abs(current_temp.get_stock_amount(c)) for c in held}

        # liquidity floor: drop names below the ADV floor from BUY candidates
        illiquid: Dict[str, bool] = {}
        if self.liquidity_floor_adv > 0:
            from tac_qlib.risk_limits import dollar_adv

            codes = [c for c in target_weight_position if c not in held]
            if codes:
                adv = dollar_adv(codes, market="US", asof=str(pd.Timestamp(trade_start_time).date()), lookback=20)
                for c in codes:
                    illiquid[c] = adv.get(str(c).upper(), 0.0) < self.liquidity_floor_adv

        buys_paused = self._hmm_pause_active or self._drawdown_pause_active

        gated = {}
        for code, w in target_weight_position.items():
            is_new = code not in held or held_amt.get(code, 0.0) <= 1e-6
            if is_new:
                if buys_paused:
                    continue  # risk gate: no new entries
                if illiquid.get(code, False):
                    continue  # liquidity floor: drop illiquid buy candidate
            gated[code] = w

        order_list = self.order_generator.generate_order_list_from_target_weight_position(
            current=current_temp,
            trade_exchange=self.trade_exchange,
            risk_degree=self.get_risk_degree(trade_step),
            target_weight_position=gated,
            pred_start_time=pred_start_time,
            pred_end_time=pred_end_time,
            trade_start_time=trade_start_time,
            trade_end_time=trade_end_time,
        )
        return TradeDecisionWO(order_list, self)