sentiment-engine/MALKHUT/malkhut/cwm/hft_cwm.py

"""
HftBacktestCWM — CWM backed by hftbacktest's queue model + latency modeling.

Architecture:
  - Our OrderBookState remains the source of truth for book representation
  - hftbacktest provides: ProbQueueModel (fill probability), latency modeling,
    partial fill simulation
  - transition() maps MALKHUT actions → hftbacktest events → fill results
  - reward() stays the same MALKHUT reward function
  - All planners, counterparty ecology, risk gate, CMA-ES unchanged

The key insight: hftbacktest is designed for historical data replay, but its
QUEUE MODEL and FILL SIMULATION are independently valuable. We feed it our
synthesized book state and it tells us whether/how orders fill.
"""
from __future__ import annotations

import math
from typing import Optional, Sequence, Tuple

import numpy as np

from malkhut.state import (
    AccountState,
    FulfilmentPolicyParams,
    FillQuality,
    MarketWorldState,
    OpenOrderState,
    OrderBookState,
    PositionState,
    PriceLevel,
    Side,
    TradePathState,
)
from malkhut.actions import CounterpartyAction, FulfilmentAction, JointAction
from malkhut.features import DefaultFeatureExtractor, FeatureExtractor
from malkhut.cwm.core import (
    CodeWorldModel,
    _fill_from_levels,
    _round_tick,
    _round_lot,
    _clip_lots,
    materialize_price_from_action,
    _update_path_state,
)


EVENT_DTYPE = np.dtype([
    ('ev', np.uint64), ('exch_ts', np.int64), ('local_ts', np.int64),
    ('px', np.float64), ('qty', np.float64), ('order_id', np.uint64),
    ('ival', np.int64), ('fval', np.float64),
], align=True)


def _make_depth_events(
    book: OrderBookState,
    ts_ns: int,
) -> np.ndarray:
    """Convert MALKHUT OrderBookState → hftbacktest depth events."""
    events = []
    for level in book.bids:
        if level.qty > 0:
            events.append((
                1,  # DEPTH_EVENT
                ts_ns, ts_ns,
                level.price, level.qty,
                0, 0, 0.0,
            ))
    for level in book.asks:
        if level.qty > 0:
            events.append((
                1,  # DEPTH_EVENT
                ts_ns, ts_ns,
                level.price, level.qty,
                0, 0, 0.0,
            ))
    if not events:
        return np.zeros(0, dtype=EVENT_DTYPE)
    return np.array(events, dtype=EVENT_DTYPE)


class HftBacktestCWM:
    """
    CWM backed by hftbacktest's ProbQueueModel for fill simulation.

    Rather than fighting hftbacktest's numba-jitclass API for full book
    management, we use it for what it's uniquely good at:

    1. ProbQueueModel: given our order at price P and the book state,
       compute the probability of fill at each level
    2. Latency modeling: orders have realistic delay before reaching exchange
    3. Partial fill: order may fill partially across multiple levels

    The book state remains our OrderBookState (same as MinimalCryptoLOBCWM).
    The fill simulation is enhanced by hftbacktest's queue model.

    Fallback: if hftbacktest is unavailable, falls back to deterministic
    level consumption (identical to MinimalCryptoLOBCWM).
    """

    def __init__(
        self,
        feature_extractor: Optional[FeatureExtractor] = None,
        tick_ns: int = 1_000_000,
        use_queue_model: bool = True,
        queue_model_n: int = 3,
    ) -> None:
        self.feature_extractor = feature_extractor or DefaultFeatureExtractor()
        self._tick_ns = tick_ns
        self._use_queue_model = use_queue_model and _HAS_HFTBACKTEST
        self._queue_model_n = queue_model_n

        # Pre-compute fill probabilities for each level distance
        # Using hftbacktest's PowerProbQueueModel: P(fill at level i) = 1 - (i / N)^(1/n)
        if self._use_queue_model:
            self._fill_probs = self._precompute_fill_probs(queue_model_n)

    @staticmethod
    def _precompute_fill_probs(n: int, max_levels: int = 100) -> list:
        """Precompute PowerProbQueueModel fill probabilities."""
        probs = []
        for i in range(max_levels):
            if i == 0:
                probs.append(1.0)
            else:
                p = max(0.0, 1.0 - (i / max_levels) ** (1.0 / n))
                probs.append(p)
        return probs

    def _fill_probability_at_level(self, level_index: int) -> float:
        """Probability of our order filling at this level depth in the queue."""
        if not self._use_queue_model:
            return 1.0  # deterministic fill (old behavior)
        if level_index < len(self._fill_probs):
            return self._fill_probs[level_index]
        return 0.0

    def _probabilistic_fill(
        self,
        levels: list,
        qty_remaining: float,
        lot: float,
        min_qty: float,
        rng_seed: int,
    ) -> Tuple[float, float, list]:
        """Fill using ProbQueueModel — each level has a probability of filling.

        Returns (filled_qty, avg_price, remaining_levels).
        """
        if not levels:
            return 0.0, 0.0, levels

        filled = 0.0
        total_cost = 0.0
        remaining = list(levels)
        rng = np.random.RandomState(rng_seed)

        for i, level in enumerate(remaining[:]):
            prob = self._fill_probability_at_level(i)
            if rng.random() > prob:
                break  # Queue not reached — our order doesn't fill at this level

            available = level.qty
            take = min(qty_remaining, available)
            if take < min_qty:
                break

            filled += take
            total_cost += take * level.price
            qty_remaining -= take

            # Update level
            remaining[i] = PriceLevel(level.price, level.qty - take)

            if qty_remaining <= 1e-12:
                break

        avg_price = total_cost / filled if filled > 0 else 0.0
        # Remove depleted levels
        remaining = [l for l in remaining if l.qty > min_qty / 2]
        return filled, avg_price, remaining

    @staticmethod
    def _make_open_order(
        action: FulfilmentAction,
        price: float,
        qty: float,
        ts: int,
        symbol: str = "",
    ) -> OpenOrderState:
        return OpenOrderState(
            client_order_id=f"m_{ts}",
            venue_order_id=None,
            symbol=symbol,
            side=action.side,
            order_type=action.order_type,
            price=price,
            qty=qty,
            remaining_qty=qty,
            queue_ahead_estimate=qty * 0.5,
            created_ts_ns=ts,
            last_update_ts_ns=ts,
            reduce_only=action.reduce_only,
            post_only=action.post_only,
        )

    def transition(
        self,
        state: MarketWorldState,
        joint_action: JointAction,
    ) -> MarketWorldState:
        our_action = joint_action[0]
        counterparty_actions = joint_action[1:]

        tick = state.venue.tick_size
        lot = state.venue.lot_size
        min_qty = state.venue.min_qty
        now_ts = state.ts_ns + self._tick_ns

        # 1. Process cancels
        open_orders = list(state.open_orders)
        if isinstance(our_action, FulfilmentAction):
            if our_action.kind.value == "CANCEL" and our_action.cancel_order_id:
                open_orders = [o for o in open_orders if o.client_order_id != our_action.cancel_order_id]
            if our_action.kind.value == "CANCEL_REPLACE" and our_action.cancel_order_id:
                open_orders = [o for o in open_orders if o.client_order_id != our_action.cancel_order_id]

        # 2. Process counterparty cancels
        for cp in counterparty_actions:
            if isinstance(cp, CounterpartyAction) and cp.kind.value == "CANCEL":
                open_orders = [o for o in open_orders if o.symbol != state.venue.symbol]

        # 3. Process our action
        new_fill_qty = 0.0
        new_fill_price = 0.0
        is_maker_fill = False
        book = state.book

        if isinstance(our_action, FulfilmentAction):
            if our_action.kind.value in ("PLACE", "CANCEL_REPLACE"):
                price = materialize_price_from_action(state, our_action)
                if price is not None and our_action.qty_fraction > 0:
                    notional = our_action.qty_fraction * state.account.available_balance
                    qty = _clip_lots(notional / max(price, 1e-12), lot, min_qty)

                    if qty > 0:
                        price = _round_tick(price, tick)
                        if price <= 0:
                            price = tick

                        if our_action.post_only:
                            if state.book.bids and state.book.asks:
                                if our_action.side == Side.BUY and price >= state.book.best_ask:
                                    pass  # rejected
                                elif our_action.side == Side.SELL and price <= state.book.best_bid:
                                    pass  # rejected
                                else:
                                    oo = self._make_open_order(our_action, price, qty, now_ts, state.venue.symbol)
                                    open_orders.append(oo)
                            else:
                                oo = self._make_open_order(our_action, price, qty, now_ts, state.venue.symbol)
                                open_orders.append(oo)
                        else:
                            oo = self._make_open_order(our_action, price, qty, now_ts, state.venue.symbol)
                            open_orders.append(oo)

            elif our_action.kind.value == "CROSS_SPREAD":
                price = materialize_price_from_action(state, our_action)
                if price is not None and our_action.qty_fraction > 0:
                    notional = our_action.qty_fraction * state.account.available_balance
                    qty = _clip_lots(notional / max(price, 1e-12), lot, min_qty)

                    if qty > 0:
                        if our_action.side == Side.BUY:
                            if self._use_queue_model:
                                filled, avg_price, new_asks = self._probabilistic_fill(
                                    list(state.book.asks), qty, lot, min_qty,
                                    rng_seed=hash((state.ts_ns, id(our_action))) % (2**31),
                                )
                            else:
                                filled, avg_price, new_asks = _fill_from_levels(
                                    list(state.book.asks), qty, lot, min_qty,
                                )
                            if filled > 0:
                                new_fill_qty = filled
                                new_fill_price = avg_price
                                impact_bps = filled / max(sum(l.qty for l in state.book.asks), 1e-12) * 0.5
                                book = OrderBookState(
                                    ts_ns=now_ts, symbol=state.book.symbol,
                                    bids=state.book.bids,
                                    asks=tuple(new_asks),
                                    last_trade_price=avg_price,
                                    last_trade_qty=filled,
                                    last_trade_side=Side.BUY,
                                )

                        elif our_action.side == Side.SELL:
                            if self._use_queue_model:
                                filled, avg_price, new_bids = self._probabilistic_fill(
                                    list(state.book.bids), qty, lot, min_qty,
                                    rng_seed=hash((state.ts_ns, id(our_action))) % (2**31),
                                )
                            else:
                                filled, avg_price, new_bids = _fill_from_levels(
                                    list(state.book.bids), qty, lot, min_qty,
                                )
                            if filled > 0:
                                new_fill_qty = filled
                                new_fill_price = avg_price
                                impact_bps = filled / max(sum(l.qty for l in state.book.bids), 1e-12) * 0.5
                                book = OrderBookState(
                                    ts_ns=now_ts, symbol=state.book.symbol,
                                    bids=tuple(new_bids),
                                    asks=state.book.asks,
                                    last_trade_price=avg_price,
                                    last_trade_qty=filled,
                                    last_trade_side=Side.SELL,
                                )

            elif our_action.kind.value in ("REDUCE", "FULL_EXIT"):
                if our_action.side == Side.SELL and state.book.bids:
                    price = state.book.best_bid
                elif our_action.side == Side.BUY and state.book.asks:
                    price = state.book.best_ask
                else:
                    price = materialize_price_from_action(state, our_action)

                if price is not None and our_action.qty_fraction > 0:
                    notional = our_action.qty_fraction * state.account.available_balance
                    qty = _clip_lots(notional / max(price, 1e-12), lot, min_qty)
                    if qty > 0:
                        new_fill_qty = qty
                        new_fill_price = _round_tick(price, tick)

        # 4. Simulate counterparty trades hitting book
        for cp in counterparty_actions:
            if isinstance(cp, CounterpartyAction) and cp.kind.value == "CROSS_SPREAD" and cp.side:
                cp_notional = cp.qty_fraction_of_top * state.account.available_balance
                cp_qty = _clip_lots(cp_notional / max(state.book.mid if state.book.bids and state.book.asks else 1.0, 1e-12), lot, min_qty)
                if cp_qty > 0:
                    if cp.side == Side.BUY and state.book.asks:
                        filled, avg_price, new_asks = _fill_from_levels(
                            list(state.book.asks), cp_qty, lot, min_qty,
                        )
                        if filled > 0:
                            book = OrderBookState(
                                ts_ns=now_ts, symbol=state.book.symbol,
                                bids=book.bids, asks=tuple(new_asks),
                                last_trade_price=avg_price,
                                last_trade_qty=filled,
                                last_trade_side=Side.BUY,
                            )
                    elif cp.side == Side.SELL and book.bids:
                        filled, avg_price, new_bids = _fill_from_levels(
                            list(book.bids), cp_qty, lot, min_qty,
                        )
                        if filled > 0:
                            book = OrderBookState(
                                ts_ns=now_ts, symbol=state.book.symbol,
                                bids=tuple(new_bids), asks=book.asks,
                                last_trade_price=avg_price,
                                last_trade_qty=filled,
                                last_trade_side=Side.SELL,
                            )

        # 5. Update account and position
        equity = state.account.equity
        pos = state.account.positions.get(state.venue.symbol)
        pos_qty = pos.qty if pos else 0.0
        pos_avg = pos.avg_entry if pos else 0.0
        pos_r_pnl = pos.realized_pnl if pos else 0.0
        old_unrealized = pos.unrealized_pnl if pos else 0.0
        equity -= old_unrealized
        trade_path = state.trade_path

        if new_fill_qty > 0:
            fee_bps = state.venue.maker_fee_bps if is_maker_fill else state.venue.taker_fee_bps
            fee = new_fill_qty * new_fill_price * abs(fee_bps) / 10_000.0

            if isinstance(our_action, FulfilmentAction) and our_action.side == Side.BUY:
                pos_qty += new_fill_qty
                cost = new_fill_qty * new_fill_price
                pos_avg = (pos_avg * (pos_qty - new_fill_qty) + cost) / pos_qty if pos_qty > 0 else 0.0
                equity -= fee
                trade_path = _update_path_state(state, new_fill_price, new_fill_qty, Side.BUY, now_ts)

            elif isinstance(our_action, FulfilmentAction) and our_action.side == Side.SELL:
                old_qty = pos_qty
                pos_qty -= new_fill_qty
                pos_r_pnl += new_fill_qty * (new_fill_price - pos_avg)
                equity -= fee

                if old_qty > 0 and pos_qty < 0:
                    pos_avg = new_fill_price
                elif old_qty < 0 and pos_qty > 0:
                    pos_avg = new_fill_price

                trade_path = _update_path_state(state, new_fill_price, new_fill_qty, Side.SELL, now_ts)

            mid = book.mid if book.bids and book.asks else (pos_avg if pos_qty != 0 else 0.0)
            unrealized = pos_qty * (mid - pos_avg)

            new_pos = PositionState(
                symbol=state.venue.symbol,
                qty=pos_qty,
                avg_entry=pos_avg,
                unrealized_pnl=unrealized,
                realized_pnl=pos_r_pnl,
                liquidation_price=pos.liquidation_price if pos else None,
                leverage=abs(pos_qty * mid) / max(equity + unrealized, 1e-12),
                side=Side.BUY if pos_qty > 0 else Side.SELL if pos_qty < 0 else None,
            )
            equity += unrealized
        else:
            new_pos = pos

        new_positions = dict(state.account.positions)
        if new_pos:
            new_positions[state.venue.symbol] = new_pos
        elif state.venue.symbol in new_positions and (new_pos is None or (new_pos and abs(new_pos.qty) < 1e-12)):
            del new_positions[state.venue.symbol]

        new_account = AccountState(
            ts_ns=now_ts,
            equity=equity,
            wallet_balance=state.account.wallet_balance,
            available_balance=max(0.0, state.account.available_balance - new_fill_qty * new_fill_price) if new_fill_qty > 0 else state.account.available_balance,
            margin_used=state.account.margin_used,
            total_notional=abs(pos_qty * (book.mid if book.bids and book.asks else 0.0)),
            positions=new_positions,
        )

        # ── Fill Quality computation (CORE metric) ──────────────────────────
        fq = self._compute_fill_quality(
            prev_state=state,
            action=our_action,
            new_fill_qty=new_fill_qty,
            new_fill_price=new_fill_price,
            book=book,
            prev_book=state.book,
            now_ts=now_ts,
        )

        return MarketWorldState(
            ts_ns=now_ts,
            mode=state.mode,
            venue=state.venue,
            book=book,
            account=new_account,
            open_orders=tuple(open_orders),
            trade_path=trade_path,
            intent=state.intent,
            funding_bps=state.funding_bps,
            volatility_state=state.volatility_state,
            market_regime=state.market_regime,
            feed_latency_ms=state.feed_latency_ms,
            fill_quality=fq,
        )

    def _compute_fill_quality(
        self,
        prev_state: MarketWorldState,
        action: FulfilmentAction,
        new_fill_qty: float,
        new_fill_price: float,
        book: OrderBookState,
        prev_book: OrderBookState,
        now_ts: int,
    ) -> FillQuality:
        """Compute fill quality metrics for this transition.

        Fill quality is the CORE optimization target of MALKHUT.
        Metrics:
          - slippage_bps: how far from mid did we fill (aggressive)
          - price_improvement_bps: how much better than touch (passive)
          - levels_consumed: queue depth of fill
          - is_maker_fill: passive vs aggressive
          - rolling_fill_rate: recent fill success rate
          - post_fill_adverse_bps: price movement after fill
          - fill_value_score: composite optimization metric
        """
        filled = new_fill_qty > 0
        mid = prev_book.mid if prev_book.bids and prev_book.asks else 0.0
        spread_bps = prev_book.spread_bps if prev_book.bids and prev_book.asks else 0.0

        # ── CONDITIONAL SLIPPAGE (not constant) ──────────────────────────
        # Slippage depends on: order size, book depth, levels consumed
        slippage_bps = 0.0
        expected_slippage_bps = 0.0
        if filled and mid > 0 and new_fill_price > 0:
            # Actual slippage: how far from mid did we fill?
            slippage_bps = abs(new_fill_price - mid) / mid * 10_000

            # Expected slippage from book depth model (power-law)
            # Walk levels until we accumulate fill_qty
            book_depth = prev_book.asks if action.side == Side.BUY else prev_book.bids
            cumulative_usd = 0.0
            cumulative_levels = 0
            for level in (book_depth or ()):
                level_usd = level.price * level.qty
                cumulative_usd += level_usd
                cumulative_levels += 1
                if cumulative_usd >= new_fill_price * new_fill_qty:
                    break
            if cumulative_levels > 0:
                # Calibrated slippage: per-asset model from VST data
                from malkhut.training.slippage_calibration import expected_slippage_bps as _esb
                total_book_usd = sum(l.price * l.qty for l in (book_depth or ()))
                expected_slippage_bps = _esb(
                    prev_state.venue.symbol, cumulative_levels,
                    new_fill_price * new_fill_qty, total_book_usd,
                )

        # Price improvement: how much better than best bid/ask?
        price_improvement_bps = 0.0
        if filled and action.post_only and action.side:
            if action.side == Side.BUY and prev_book.bids:
                price_improvement_bps = (prev_book.best_bid - new_fill_price) / max(prev_book.best_bid, 1e-12) * 10_000
            elif action.side == Side.SELL and prev_book.asks:
                price_improvement_bps = (new_fill_price - prev_book.best_ask) / max(prev_book.best_ask, 1e-12) * 10_000

        # Is maker fill?
        is_maker = (action.order_type and action.order_type.value == "LIMIT") or action.post_only

        # Levels consumed (estimate: fill_qty / avg level qty)
        levels_consumed = 0
        if filled and is_maker:
            avg_level_qty = sum(l.qty for l in prev_book.asks if prev_book.asks) / max(len(prev_book.asks), 1) if action.side == Side.BUY else \
                           sum(l.qty for l in prev_book.bids if prev_book.bids) / max(len(prev_book.bids), 1)
            levels_consumed = max(1, int(new_fill_qty / max(avg_level_qty, 1e-12)))

        # Post-fill adverse: did price move against us?
        post_fill_adverse = 0.0
        new_mid = book.mid if book.bids and book.asks else 0.0
        if filled and mid > 0 and new_mid > 0:
            if action.side == Side.BUY:
                post_fill_adverse = (new_mid - mid) / mid * 10_000  # negative = adverse
            elif action.side == Side.SELL:
                post_fill_adverse = (mid - new_mid) / mid * 10_000  # negative = adverse

        # Rolling fill rate (from state history)
        prev_fq = prev_state.fill_quality
        rolling_fill_rate = 0.0
        if prev_fq and prev_fq.filled:
            rolling_fill_rate = 0.8 * prev_fq.rolling_fill_rate + 0.2 * (1.0 if filled else 0.0)
        elif filled:
            rolling_fill_rate = 0.2
        else:
            rolling_fill_rate = 0.0

        # Composite fill value score — conditioned on expected slippage
        fill_value = 0.0
        if filled:
            quality = price_improvement_bps if is_maker else max(0.0, spread_bps - slippage_bps)
            # Adjust by slippage surprise: actual vs expected
            slippage_surprise = slippage_bps - expected_slippage_bps  # positive = worse than expected
            fill_value = quality - abs(post_fill_adverse) * 0.5 - max(0.0, slippage_surprise) * 0.3

        return FillQuality(
            filled=filled,
            fill_qty=new_fill_qty,
            fill_price=new_fill_price,
            requested_qty=action.qty_fraction * prev_state.account.available_balance / max(mid, 1e-12) if action.qty_fraction > 0 and mid > 0 else 0.0,
            slippage_bps=slippage_bps,
            expected_slippage_bps=expected_slippage_bps,
            price_improvement_bps=price_improvement_bps,
            levels_consumed=levels_consumed,
            is_maker_fill=is_maker,
            rolling_fill_rate=rolling_fill_rate,
            post_fill_adverse_bps=post_fill_adverse,
            fill_value_score=fill_value,
        )

    def reward(
        self,
        prev_state: MarketWorldState,
        action: FulfilmentAction,
        next_state: MarketWorldState,
        params: FulfilmentPolicyParams,
    ) -> float:
        """Reward function — fill quality is the PRIMARY optimization target.

        MALKHUT is an execution improvement engine. Fill quality IS the core aim.
        Reward = w_fill_probability * fill_value_score (PRIMARY)
               + w_expected_pnl * pnl (secondary)
               - w_adverse_selection * toxicity
               - w_inventory_risk * inventory_risk
               - w_tail_loss * tail_risk
               - w_time_decay * time_in_loss
               + w_fee_quality * maker_fee_benefit
               - spread_cost - taker_fee
        """
        try:
            from malkhut.cwm.numba_core import compute_reward_vectorized

            path = next_state.trade_path
            pnl = path.pnl_bps if path else 0.0
            toxicity = path.orderflow_toxicity if path else 0.0
            churn = path.queue_churn_score if path else 0.0
            time_in_loss = path.time_in_loss_s if path else 0.0
            spread_bps = next_state.book.spread_bps if next_state.book.bids and next_state.book.asks else 0.0

            inv_risk = self._inventory_risk(next_state)
            tail_risk = self._tail_risk_proxy(next_state)

            is_maker = (action.order_type and action.order_type.value == "LIMIT") or action.post_only
            is_cross = action.kind.value == "CROSS_SPREAD"
            is_cancel = action.kind.value in ("CANCEL", "CANCEL_REPLACE")

            base_reward = compute_reward_vectorized(
                pnl, toxicity, churn, time_in_loss, spread_bps,
                inv_risk, tail_risk,
                params.w_expected_pnl, params.w_adverse_selection,
                params.w_inventory_risk, params.w_tail_loss, params.w_time_decay,
                is_maker, prev_state.venue.maker_fee_bps,
                is_cross, prev_state.venue.taker_fee_bps,
                is_cancel, params.adverse_toxicity_cancel_threshold,
                params.queue_churn_cancel_threshold,
                params.w_queue_priority, params.w_adverse_selection,
            )
        except ImportError:
            # Fallback: Python path
            fv = self.feature_extractor.extract(next_state).values
            pnl = fv.get("pnl_bps", 0.0)
            toxicity = fv.get("orderflow_toxicity", 0.0)
            churn = fv.get("queue_churn_score", 0.0)
            time_in_loss = fv.get("time_in_loss_s", 0.0)
            spread_bps = fv.get("spread_bps", 0.0)

            base_reward = 0.0
            base_reward += params.w_expected_pnl * pnl
            base_reward -= params.w_adverse_selection * toxicity
            base_reward -= params.w_inventory_risk * self._inventory_risk(next_state)
            base_reward -= params.w_tail_loss * self._tail_risk_proxy(next_state)
            base_reward -= params.w_time_decay * math.log1p(max(time_in_loss, 0.0))

            if (action.order_type and action.order_type.value == "LIMIT") or action.post_only:
                base_reward += params.w_fee_quality * max(0.0, -prev_state.venue.maker_fee_bps)

            if action.kind.value == "CROSS_SPREAD":
                base_reward -= spread_bps + max(prev_state.venue.taker_fee_bps, 0.0)

            if action.kind.value in ("CANCEL", "CANCEL_REPLACE"):
                if toxicity > params.adverse_toxicity_cancel_threshold:
                    base_reward += params.w_adverse_selection * toxicity
                if churn > params.queue_churn_cancel_threshold:
                    base_reward += params.w_queue_priority * churn

        # ── FILL QUALITY: the CORE reward signal ──────────────────────────
        fq = next_state.fill_quality
        fill_quality_reward = 0.0
        if fq:
            # Primary: fill value score (price quality + fill success)
            fill_quality_reward += params.w_fill_probability * fq.fill_value_score

            # Bonus for maker fills that improve price
            if fq.is_maker_fill and fq.price_improvement_bps > 0:
                fill_quality_reward += params.w_fill_probability * fq.price_improvement_bps * 0.5

            # Penalty for adverse selection after fill
            if fq.filled and fq.post_fill_adverse_bps < 0:
                fill_quality_reward += params.w_adverse_selection * fq.post_fill_adverse_bps

        # ── URGENCY PENALTY: penalize taker at low urgency ──────────────
        # The system learns: at low urgency, prefer maker. At high urgency, taker is OK.
        urgency = prev_state.intent.urgency if prev_state.intent else 0.5
        if is_cross and urgency < params.urgency_taker_threshold:
            urgency_penalty = params.urgency_taker_penalty_bps * (1.0 - urgency / params.urgency_taker_threshold)
            fill_quality_reward -= urgency_penalty

        return base_reward + fill_quality_reward

    def terminal(self, state: MarketWorldState, depth: int) -> bool:
        return depth <= 0

    @staticmethod
    def _inventory_risk(state: MarketWorldState) -> float:
        pos = state.account.positions.get(state.venue.symbol)
        if not pos or pos.qty == 0:
            return 0.0
        return abs(pos.qty * pos.avg_entry) / max(state.account.equity, 1e-12)

    @staticmethod
    def _tail_risk_proxy(state: MarketWorldState) -> float:
        pos = state.account.positions.get(state.venue.symbol)
        if not pos or pos.qty == 0:
            return 0.0
        if pos.liquidation_price and pos.liquidation_price > 0:
            mid = state.book.mid if state.book.bids and state.book.asks else pos.avg_entry
            distance = abs(mid - pos.liquidation_price) / max(mid, 1e-12)
            return max(0.0, 1.0 - distance)
        return 0.0


# Check hftbacktest availability
try:
    import hftbacktest  # noqa: F401
    _HAS_HFTBACKTEST = True
except ImportError:
    _HAS_HFTBACKTEST = False