"""
MALKHUT canonical data model.

All state objects are frozen+slots for:
  - deterministic tree search (immutable snapshots)
  - GraalVM compatibility (no mutable default hell)
  - lock-free shared memory (readers never see partial writes)
"""
from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, Mapping, Optional, Sequence, Tuple
import math


# ==============================================================================
# Enums
# ==============================================================================

class Side(str, Enum):
    BUY = "BUY"
    SELL = "SELL"


class OrderType(str, Enum):
    """Standardized order types — FIX/CCXT-aligned, multi-exchange.

    THREE ORTHOGONAL DIMENSIONS (not one flat enum):
      1. Order Type (FIX Tag 40): what the order IS — this enum
      2. TimeInForce (FIX Tag 59): how long it LIVES — separate parameter
      3. Instructions (FIX Tag 18): behavioral modifiers — separate parameter

    CRITICAL: IOC, FOK, POST_ONLY are NOT order types.
      IOC/FOK = TimeInForce on a LIMIT order.
      POST_ONLY = ExecInst modifier on a LIMIT order.
    """
    # Core order types (FIX Tag 40)
    MARKET = "MARKET"
    LIMIT = "LIMIT"
    STOP_MARKET = "STOP_MARKET"
    STOP_LIMIT = "STOP_LIMIT"
    TRIGGER_MARKET = "TRIGGER_MARKET"
    TRIGGER_LIMIT = "TRIGGER_LIMIT"
    TRAILING_STOP = "TRAILING_STOP"
    OCO = "OCO"
    TP_SL = "TP_SL"


class ActionKind(str, Enum):
    NOOP = "NOOP"
    PLACE = "PLACE"
    CANCEL = "CANCEL"
    CANCEL_REPLACE = "CANCEL_REPLACE"
    CROSS_SPREAD = "CROSS_SPREAD"
    REDUCE = "REDUCE"
    FULL_EXIT = "FULL_EXIT"
    MOVE_STOP = "MOVE_STOP"
    MOVE_TAKE_PROFIT = "MOVE_TAKE_PROFIT"
    THROTTLE = "THROTTLE"


class IntentKind(str, Enum):
    ENTER_LONG = "ENTER_LONG"
    ENTER_SHORT = "ENTER_SHORT"
    ADD_LONG = "ADD_LONG"
    ADD_SHORT = "ADD_SHORT"
    REDUCE_LONG = "REDUCE_LONG"
    REDUCE_SHORT = "REDUCE_SHORT"
    EXIT_LONG = "EXIT_LONG"
    EXIT_SHORT = "EXIT_SHORT"
    MAINTAIN = "MAINTAIN"


class AgentRole(str, Enum):
    OUR_FULFILMENT = "OUR_FULFILMENT"
    PASSIVE_MAKER = "PASSIVE_MAKER"
    TOXIC_TAKER = "TOXIC_TAKER"
    LATENCY_ARB = "LATENCY_ARB"
    MOMENTUM_TAKER = "MOMENTUM_TAKER"
    MEAN_REVERSION_TAKER = "MEAN_REVERSION_TAKER"
    INVENTORY_MM = "INVENTORY_MM"
    LIQUIDATION_FLOW = "LIQUIDATION_FLOW"
    NOISE_TRADER = "NOISE_TRADER"
    STALE_QUOTE_ATTACKER = "STALE_QUOTE_ATTACKER"


class Mode(str, Enum):
    REPLAY_NO_IMPACT = "REPLAY_NO_IMPACT"
    ENDOGENOUS_AGENT_SIM = "ENDOGENOUS_AGENT_SIM"
    PAPER = "PAPER"
    SHADOW_LIVE = "SHADOW_LIVE"
    LIVE = "LIVE"


# ==============================================================================
# Core constants
# ==============================================================================

HOT_PATH_BUDGET_MS: int = 100
DEFAULT_PLANNER_BUDGET_MS: int = 25
DEFAULT_TREE_DEPTH: int = 3
DEFAULT_MAX_SIMS: int = 256
DEFAULT_UCB_C: float = 1.41421356237
DEFAULT_MIN_ROOT_POLICY_ENTROPY: float = 0.25
DEFAULT_SELF_PLAY_POOL_MAX: int = 12
DEFAULT_POLICY_PROMOTION_MIN_EDGE_BPS: float = 0.75
DEFAULT_POLICY_PROMOTION_MIN_PVALUE: float = 0.05

MAX_ACCOUNT_LEVERAGE: float = 2.0
MAX_EXCHANGE_LEVERAGE: float = 5.0
MAX_SINGLE_ORDER_NOTIONAL_FRACTION: float = 0.05
MAX_SYMBOL_NOTIONAL_FRACTION: float = 0.20
MAX_CANCELS_PER_SYMBOL_PER_MINUTE: int = 90
TAIL_QUANTILE: float = 0.05


# ==============================================================================
# Frozen data model
# ==============================================================================

@dataclass(frozen=True, slots=True)
class VenueRules:
    exchange: str
    symbol: str
    tick_size: float
    lot_size: float
    min_qty: float
    min_notional: float
    maker_fee_bps: float
    taker_fee_bps: float
    post_only_supported: bool
    reduce_only_supported: bool
    max_orders_per_second: int
    max_cancels_per_minute: int


@dataclass(frozen=True, slots=True)
class PriceLevel:
    price: float
    qty: float


@dataclass(frozen=True, slots=True)
class OrderBookState:
    ts_ns: int
    symbol: str
    bids: Tuple[PriceLevel, ...]
    asks: Tuple[PriceLevel, ...]
    last_trade_price: Optional[float] = None
    last_trade_qty: Optional[float] = None
    last_trade_side: Optional[Side] = None

    @property
    def best_bid(self) -> float:
        return self.bids[0].price

    @property
    def best_ask(self) -> float:
        return self.asks[0].price

    @property
    def mid(self) -> float:
        return 0.5 * (self.best_bid + self.best_ask)

    @property
    def spread(self) -> float:
        return self.best_ask - self.best_bid

    @property
    def spread_bps(self) -> float:
        return 10_000.0 * self.spread / max(self.mid, 1e-12)


@dataclass(frozen=True, slots=True)
class PositionState:
    symbol: str
    qty: float
    avg_entry: float
    unrealized_pnl: float
    realized_pnl: float
    liquidation_price: Optional[float]
    leverage: float
    side: Optional[Side]


@dataclass(frozen=True, slots=True)
class AccountState:
    ts_ns: int
    equity: float
    wallet_balance: float
    available_balance: float
    margin_used: float
    total_notional: float
    positions: Mapping[str, PositionState] = field(default_factory=dict)


@dataclass(frozen=True, slots=True)
class OpenOrderState:
    client_order_id: str
    venue_order_id: Optional[str]
    symbol: str
    side: Side
    order_type: OrderType
    price: Optional[float]
    qty: float
    remaining_qty: float
    queue_ahead_estimate: Optional[float]
    created_ts_ns: int
    last_update_ts_ns: int
    reduce_only: bool = False
    post_only: bool = False
    ttl_ms: int = 0  # 0 = no expiry; >0 = auto-cancel after ttl_ms (CHASE)


@dataclass(frozen=True, slots=True)
class TradePathState:
    """In-trade path encoding for path-aware SL/TP."""
    symbol: str
    side: Side
    entry_ts_ns: int
    now_ts_ns: int
    bars_held: int
    seconds_held: float

    pnl_bps: float
    mae_bps: float
    mfe_bps: float
    distance_from_mfe_bps: float
    distance_from_entry_bps: float

    time_to_mfe_s: float
    time_in_loss_s: float
    time_in_profit_s: float
    time_since_last_profit_s: float
    time_since_deep_mae_s: float

    loss_to_profit_transitions: int
    deep_loss_recoveries: int
    failed_recovery_count: int
    recovery_velocity_bps_per_s: float
    adverse_velocity_bps_per_s: float

    dolphin_regime_score: float
    jericho_signal_strength: float
    volatility_bps: float
    orderflow_toxicity: float
    queue_churn_score: float
    book_imbalance: float
    cross_venue_lead_score: float


@dataclass(frozen=True, slots=True)
class ExecutionIntent:
    intent_id: str
    ts_ns: int
    symbol: str
    kind: IntentKind
    target_qty: float
    max_notional: float
    urgency: float
    alpha_horizon_s: float
    alpha_bps: float
    max_slippage_bps: float
    prefer_maker: bool
    reduce_only: bool
    ttl_s: float
    reason: str


@dataclass(frozen=True, slots=True)
class FillQuality:
    """Fill quality metrics — the CORE optimization target of MALKHUT.

    MALKHUT is an execution improvement engine. Fill quality IS the primary aim.
    Every transition records these metrics. The reward function weights them heavily.
    The PerformanceMatrix tracks them per (regime, strategy, venue).
    """
    filled: bool = False
    fill_qty: float = 0.0
    fill_price: float = 0.0
    requested_qty: float = 0.0

    # How close to mid did we fill? (for aggressive: positive = slipped)
    slippage_bps: float = 0.0

    # Expected slippage from book depth model (conditional on actual book state)
    expected_slippage_bps: float = 0.0

    # For passive fills: how much better than best bid/ask? (positive = improvement)
    price_improvement_bps: float = 0.0

    # How many levels deep was the fill?
    levels_consumed: int = 0

    # Was this a maker (passive) or taker (aggressive) fill?
    is_maker_fill: bool = False

    # Fill rate rolling window (updated each transition)
    rolling_fill_rate: float = 0.0

    # Adverse selection: price movement after fill (negative = adverse)
    post_fill_adverse_bps: float = 0.0

    # Fill value score: composite metric for optimization
    # = fill_rate * price_quality - adverse_selection - slippage
    fill_value_score: float = 0.0


@dataclass(frozen=True, slots=True)
class MarketWorldState:
    """Complete CWM root state. Immutable for safe tree search."""
    ts_ns: int
    mode: Mode
    venue: VenueRules
    book: OrderBookState
    account: AccountState
    open_orders: Tuple[OpenOrderState, ...] = ()
    trade_path: Optional[TradePathState] = None
    intent: Optional[ExecutionIntent] = None

    funding_bps: Optional[float] = None
    volatility_state: Optional[float] = None
    market_regime: Optional[str] = None

    feed_latency_ms: float = 0.0
    order_latency_ms: float = 0.0
    rng_seed: int = 0

    fill_quality: Optional[FillQuality] = None


@dataclass(frozen=True, slots=True)
class FulfilmentPolicyParams:
    """
    Frozen parameter set loaded by the live planner.
    CMA-ES tunes this object offline.
    """
    version: str

    # Planner
    ucb_c: float
    max_sims: int
    max_depth: int
    rollout_depth: int
    root_temperature: float
    min_root_entropy: float

    # Quote menu
    quote_offsets_ticks: Tuple[int, ...]
    quote_size_fractions: Tuple[float, ...]
    passive_ttl_ms: int
    aggressive_ttl_ms: int

    # Maker/taker thresholds
    maker_edge_min_bps: float
    cross_spread_edge_min_bps: float
    adverse_toxicity_cancel_threshold: float
    queue_churn_cancel_threshold: float

    # SL/TP/path risk
    mae_tail_cut_bps: float
    mfe_giveback_cut_fraction: float
    max_time_in_loss_s: float
    failed_recovery_cut_count: int
    recovery_velocity_min_bps_per_s: float

    # Inventory/account
    max_symbol_notional_fraction: float
    max_single_order_notional_fraction: float
    reduce_when_global_up_fraction: float
    session_profit_lock_fraction: float

    # Reward weights
    w_expected_pnl: float
    w_fill_probability: float
    w_adverse_selection: float
    w_queue_priority: float
    w_inventory_risk: float
    w_tail_loss: float
    w_fee_quality: float
    w_time_decay: float
    w_policy_entropy: float

    # Scenario robustness
    robust_tail_weight: float
    toxic_counterparty_weight: float
    low_liquidity_weight: float
    latency_stress_weight: float

    # Chase mechanics (cancel → wait → retry)
    wait_to_retry_ms: int = 0           # ms to wait before re-quoting after cancel
    chase_enabled: bool = False         # enable chase-follow behavior
    chase_offset_ticks: int = 1         # ticks from target price to chase
    chase_max_retries: int = 3          # max cancel-retry cycles

    # Urgency-driven maker/taker decision (CMA-ES optimizable)
    urgency_taker_threshold: float = 0.65  # above this urgency, prefer taker
    urgency_taker_penalty_bps: float = 2.0  # penalty for taker at low urgency

    # Fee+slippage execution threshold (movable, CMA-ES optimizable)
    # If (fee + slippage) < threshold → system tends to EXECUTE (pay the friction)
    execution_friction_threshold_bps: float = 3.0  # per leg, test 2-3 bps