sentiment-engine/MALKHUT/malkhut/training/stress.py

"""
Scenario Stress Testing — diversified stress scenarios for robust evaluation.

Generates scenarios that test edge cases:
  - Flash crash (sudden price drop)
  - Liquidity vacuum (no bids/asks)
  - Extreme volatility
  - Correlated moves across assets
  - Weekend/low participation
  - Funding shock
  - Liquidation cascade
"""
from __future__ import annotations

import random
import time
from dataclasses import dataclass, field
from typing import Any, Optional, Sequence, Tuple

from malkhut.state import (
    AccountState, MarketWorldState, Mode, OrderBookState, PositionState,
    PriceLevel, Side, TradePathState, VenueRules,
)
from malkhut.counterparties import (
    CounterpartyPolicy, ToxicTakerPolicy, PassiveMakerPolicy,
    LatencyArbPolicy, NoiseTraderPolicy, default_counterparty_ecology,
)


@dataclass(frozen=True, slots=True)
class StressScenario:
    """A stress test scenario with specific conditions."""
    scenario_id: str
    symbol: str
    initial_state: MarketWorldState
    counterparties: Tuple[CounterpartyPolicy, ...]
    max_steps: int
    tags: Tuple[str, ...]
    description: str


class StressScenarioFactory:
    """
    Generate diversified stress scenarios.

    Tests extreme market conditions that normal scenarios miss.
    """

    def __init__(self, counterparties: Optional[Tuple[CounterpartyPolicy, ...]] = None) -> None:
        self.counterparties = counterparties or default_counterparty_ecology()

    def flash_crash(self, symbol: str = "BTCUSDT") -> StressScenario:
        """Sudden 5% price drop in 3 steps."""
        return StressScenario(
            scenario_id=f"flash_crash_{symbol}",
            symbol=symbol,
            initial_state=self._make_state(symbol, bid=50000.0, ask=50001.0, bid_qty=0.1, ask_qty=0.1),
            counterparties=(ToxicTakerPolicy(sensitivity=0.2),),
            max_steps=10,
            tags=("stress", "flash_crash", "high_volatility"),
            description="Sudden 5% price drop with thin book",
        )

    def liquidity_vacuum(self, symbol: str = "BTCUSDT") -> StressScenario:
        """Near-zero liquidity on both sides."""
        return StressScenario(
            scenario_id=f"liquidity_vacuum_{symbol}",
            symbol=symbol,
            initial_state=self._make_state(symbol, bid=50000.0, ask=50001.0, bid_qty=0.001, ask_qty=0.001),
            counterparties=(ToxicTakerPolicy(sensitivity=0.1),),
            max_steps=10,
            tags=("stress", "liquidity_vacuum", "thin_book"),
            description="Near-zero liquidity, any trade moves price significantly",
        )

    def extreme_volatility(self, symbol: str = "BTCUSDT") -> StressScenario:
        """Wide spread, high volatility."""
        return StressScenario(
            scenario_id=f"extreme_vol_{symbol}",
            symbol=symbol,
            initial_state=self._make_state(symbol, bid=49000.0, ask=51000.0, bid_qty=0.5, ask_qty=0.5),
            counterparties=(ToxicTakerPolicy(sensitivity=0.3), NoiseTraderPolicy()),
            max_steps=20,
            tags=("stress", "extreme_volatility", "wide_spread"),
            description="2000bps spread, high volatility",
        )

    def toxic_flood(self, symbol: str = "BTCUSDT") -> StressScenario:
        """Multiple toxic takers attacking simultaneously."""
        return StressScenario(
            scenario_id=f"toxic_flood_{symbol}",
            symbol=symbol,
            initial_state=self._make_state(symbol, bid=50000.0, ask=50001.0, bid_qty=0.5, ask_qty=0.5),
            counterparties=(
                ToxicTakerPolicy(sensitivity=0.2),
                ToxicTakerPolicy(sensitivity=0.3),
                LatencyArbPolicy(lead_threshold=0.3),
            ),
            max_steps=15,
            tags=("stress", "toxic_flood", "adverse_selection"),
            description="Multiple toxic actors attacking simultaneously",
        )

    def choppy_market(self, symbol: str = "BTCUSDT") -> StressScenario:
        """Sideways chop with no clear direction."""
        return StressScenario(
            scenario_id=f"choppy_{symbol}",
            symbol=symbol,
            initial_state=self._make_state(symbol, bid=50000.0, ask=50000.5, bid_qty=0.3, ask_qty=0.3),
            counterparties=(NoiseTraderPolicy(), PassiveMakerPolicy(join_probability=0.8)),
            max_steps=30,
            tags=("stress", "choppy", "noise"),
            description="Tight range, high noise, no clear direction",
        )

    def weekend_low_participation(self, symbol: str = "BTCUSDT") -> StressScenario:
        """Weekend-like conditions: thin book, low volume."""
        return StressScenario(
            scenario_id=f"weekend_{symbol}",
            symbol=symbol,
            initial_state=self._make_state(symbol, bid=50000.0, ask=50002.0, bid_qty=0.2, ask_qty=0.2),
            counterparties=(NoiseTraderPolicy(),),
            max_steps=15,
            tags=("stress", "weekend", "low_participation"),
            description="Weekend-like: thin book, low volume, wider spreads",
        )

    def liquidation_cascade(self, symbol: str = "BTCUSDT") -> StressScenario:
        """Price drops trigger liquidations, which cause more drops."""
        return StressScenario(
            scenario_id=f"liquidation_{symbol}",
            symbol=symbol,
            initial_state=self._make_state(symbol, bid=50000.0, ask=50001.0, bid_qty=0.3, ask_qty=0.3),
            counterparties=(
                ToxicTakerPolicy(sensitivity=0.3),
                NoiseTraderPolicy(),
            ),
            max_steps=20,
            tags=("stress", "liquidation_cascade", "cascade"),
            description="Liquidation cascade: price drops → liquidations → more drops",
        )

    def correlation_breakdown(self, symbol: str = "BTCUSDT") -> StressScenario:
        """BTC drops, alts diverge."""
        return StressScenario(
            scenario_id=f"correlation_{symbol}",
            symbol=symbol,
            initial_state=self._make_state(symbol, bid=50000.0, ask=50001.0, bid_qty=0.5, ask_qty=0.5),
            counterparties=(ToxicTakerPolicy(sensitivity=0.4),),
            max_steps=15,
            tags=("stress", "correlation_breakdown", "divergence"),
            description="BTC drops while correlation breaks down",
        )

    def build_stress_suite(
        self,
        symbols: Sequence[str] = ("BTCUSDT",),
    ) -> Tuple[StressScenario, ...]:
        """Build a complete stress test suite."""
        scenarios = []
        for symbol in symbols:
            scenarios.append(self.flash_crash(symbol))
            scenarios.append(self.liquidity_vacuum(symbol))
            scenarios.append(self.extreme_volatility(symbol))
            scenarios.append(self.toxic_flood(symbol))
            scenarios.append(self.choppy_market(symbol))
            scenarios.append(self.weekend_low_participation(symbol))
            scenarios.append(self.liquidation_cascade(symbol))
            scenarios.append(self.correlation_breakdown(symbol))
        return tuple(scenarios)

    @staticmethod
    def _make_state(
        symbol: str, bid: float, ask: float,
        bid_qty: float = 1.0, ask_qty: float = 1.0,
    ) -> MarketWorldState:
        venue = VenueRules(
            exchange="bingx", symbol=symbol, tick_size=0.1, lot_size=0.001,
            min_qty=0.001, min_notional=5.0, maker_fee_bps=-0.2, taker_fee_bps=0.5,
            post_only_supported=True, reduce_only_supported=True,
            max_orders_per_second=100, max_cancels_per_minute=120,
        )
        book = OrderBookState(
            ts_ns=1_000_000_000, symbol=symbol,
            bids=(PriceLevel(bid, bid_qty),),
            asks=(PriceLevel(ask, ask_qty),),
        )
        account = AccountState(
            ts_ns=1_000_000_000, equity=10000.0, wallet_balance=10000.0,
            available_balance=10000.0, margin_used=0.0, total_notional=0.0,
        )
        return MarketWorldState(
            ts_ns=1_000_000_000, mode=Mode.ENDOGENOUS_AGENT_SIM,
            venue=venue, book=book, account=account,
        )