execution_simulation.py

# Copyright (c) 2023-2026 Yaroslav Vasylenko (neuron7xLab)
# SPDX-License-Identifier: MIT
"""Execution simulation primitives for backtest scenarios.

This module implements a lightweight matching engine focused on deterministic
backtesting. It models latency, order queueing, partial fills, market halts,
and multiple time-in-force semantics so strategy authors can evaluate
microstructure-aware behaviours without leaving the offline environment.
"""

from __future__ import annotations

import heapq
import itertools
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, Dict, Iterable, List, Optional, Tuple

from observability.tracing import pipeline_span

EPSILON = 1e-12


class OrderSide(Enum):
    """Direction of an order."""

    BUY = "buy"
    SELL = "sell"


class OrderType(Enum):
    """Supported order types for the simulator."""

    MARKET = "market"
    LIMIT = "limit"
    IOC = "ioc"
    FOK = "fok"


class OrderStatus(Enum):
    """Lifecycle states for an order."""

    NEW = "new"
    QUEUED = "queued"
    PARTIALLY_FILLED = "partially_filled"
    FILLED = "filled"
    CANCELLED = "cancelled"
    REJECTED = "rejected"


@dataclass
class ExecutionReport:
    """Record of a fill generated by the simulator."""

    order_id: str
    price: float
    qty: float
    timestamp: int
    liquidity: str = "taker"


@dataclass
class Order:
    """Simplified order object used by the matching engine."""

    id: str
    symbol: str
    side: OrderSide | str
    qty: float
    timestamp: int
    order_type: OrderType | str = OrderType.MARKET
    price: Optional[float] = None
    status: OrderStatus = OrderStatus.NEW
    filled_qty: float = 0.0
    executions: List[ExecutionReport] = field(default_factory=list)
    ready_at: Optional[int] = None

    def __post_init__(self) -> None:
        if isinstance(self.side, str):
            self.side = OrderSide(self.side.lower())
        if isinstance(self.order_type, str):
            self.order_type = OrderType(self.order_type.lower())
        if self.qty <= 0:
            raise ValueError("Order quantity must be positive")
        if self.price is not None and self.price <= 0:
            raise ValueError("Price must be positive when specified")


class HaltMode(Enum):
    """Behavioural policy returned by market halt callbacks."""

    OPEN = "open"
    REJECT_NEW = "reject"
    DELAY = "delay"
    PARTIAL = "partial"
    FULL = "full"


@dataclass
class MarketHalt:
    """Defines how the engine should behave while a market halt is active."""

    mode: HaltMode
    resume_time: Optional[int] = None
    liquidity_factor: float = 1.0

    def is_active(self) -> bool:
        return self.mode != HaltMode.OPEN


@dataclass
class _BookEntry:
    price: float
    qty: float
    order_id: str
    timestamp: int


@dataclass
class _OrderBook:
    bids: List[_BookEntry] = field(default_factory=list)
    asks: List[_BookEntry] = field(default_factory=list)


class MatchingEngine:
    """Simplified matching engine with latency, queueing, and halt support."""

    def __init__(
        self,
        latency_model: Optional[Callable[[Order], int]] = None,
        halt_model: Optional[Callable[[str, int], Optional[MarketHalt]]] = None,
    ) -> None:
        self.latency_model = latency_model or (lambda _order: 0)
        self.halt_model = halt_model
        self._pending: List[Tuple[int, int, Order]] = []
        self._counter = itertools.count()
        self._books: Dict[str, _OrderBook] = {}

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def submit_order(self, order: Order) -> Order:
        """Register a new order with latency and halt handling."""
        with pipeline_span(
            "orders.submit",
            symbol=order.symbol,
            side=getattr(order.side, "value", order.side),
            order_type=getattr(order.order_type, "value", order.order_type),
        ) as span:
            halt = self._halt_state(order.symbol, order.timestamp)
            if halt and halt.mode == HaltMode.REJECT_NEW:
                order.status = OrderStatus.REJECTED
                if span is not None:
                    span.set_attributes({"order.status": order.status.value})
                return order

            ready_at = order.timestamp + max(0, int(self.latency_model(order)))

            if halt and halt.mode in (HaltMode.DELAY, HaltMode.FULL):
                if halt.resume_time is None:
                    order.status = OrderStatus.CANCELLED
                    if span is not None:
                        span.set_attributes({"order.status": order.status.value})
                    return order
                ready_at = max(ready_at, halt.resume_time)

            order.ready_at = ready_at
            order.status = OrderStatus.QUEUED
            heapq.heappush(self._pending, (ready_at, next(self._counter), order))
            if span is not None:
                span.set_attributes(
                    {
                        "order.status": order.status.value,
                        "order.ready_at": int(ready_at),
                    }
                )
            return order

    def process_until(self, timestamp: int) -> List[Order]:
        """Process all queued orders whose latency has elapsed."""

        processed: List[Order] = []
        while self._pending and self._pending[0][0] <= timestamp:
            ready_at, token, order = heapq.heappop(self._pending)
            halt = self._halt_state(order.symbol, ready_at)
            if halt and halt.mode in (HaltMode.DELAY, HaltMode.FULL):
                if halt.resume_time is None or halt.resume_time > ready_at:
                    if halt.resume_time is None:
                        order.status = OrderStatus.CANCELLED
                        processed.append(order)
                    else:
                        order.ready_at = halt.resume_time
                        heapq.heappush(self._pending, (halt.resume_time, token, order))
                    continue

            self._match(order, ready_at, halt)
            processed.append(order)
        return processed

    def add_passive_liquidity(
        self,
        symbol: str,
        side: OrderSide | str,
        price: float,
        qty: float,
        timestamp: int,
    ) -> None:
        """Seed the order book with resting liquidity."""

        if isinstance(side, str):
            side = OrderSide(side.lower())
        if qty <= 0:
            raise ValueError("Liquidity quantity must be positive")
        if price <= 0:
            raise ValueError("Liquidity price must be positive")

        book = self._books.setdefault(symbol, _OrderBook())
        entry = _BookEntry(
            price=price, qty=qty, order_id="liquidity", timestamp=timestamp
        )
        target = book.bids if side is OrderSide.BUY else book.asks
        target.append(entry)
        self._sort_book_side(target, side)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    def _halt_state(self, symbol: str, timestamp: int) -> Optional[MarketHalt]:
        if self.halt_model is None:
            return None
        halt = self.halt_model(symbol, timestamp)
        if halt and halt.liquidity_factor < 0:
            raise ValueError("liquidity_factor must be non-negative")
        return halt

    def _match(self, order: Order, timestamp: int, halt: Optional[MarketHalt]) -> None:
        book = self._books.setdefault(order.symbol, _OrderBook())
        opposite = book.asks if order.side is OrderSide.BUY else book.bids
        same_side = book.bids if order.side is OrderSide.BUY else book.asks

        if order.order_type == OrderType.FOK:
            available = self._available_qty(order, opposite, halt)
            if available + EPSILON < order.qty:
                order.status = OrderStatus.CANCELLED
                return

        remaining = order.qty - order.filled_qty
        executable_cap = order.qty
        if halt and halt.mode == HaltMode.PARTIAL:
            executable_cap = min(order.qty, order.qty * max(0.0, halt.liquidity_factor))

        executed = 0.0
        index = 0
        while remaining > EPSILON and index < len(opposite):
            entry = opposite[index]
            if not self._price_cross(order, entry.price):
                break

            take_qty = min(remaining, entry.qty)
            if halt and halt.mode == HaltMode.PARTIAL:
                allowed_left = executable_cap - executed
                if allowed_left <= EPSILON:
                    break
                take_qty = min(take_qty, allowed_left)

            if take_qty <= EPSILON:
                break

            entry.qty -= take_qty
            remaining -= take_qty
            executed += take_qty
            order.filled_qty += take_qty
            order.executions.append(
                ExecutionReport(
                    order_id=order.id,
                    price=entry.price,
                    qty=take_qty,
                    timestamp=timestamp,
                )
            )

            if entry.qty <= EPSILON:
                opposite.pop(index)
            else:
                index += 1

        if order.filled_qty >= order.qty - EPSILON:
            order.status = OrderStatus.FILLED
            order.filled_qty = order.qty
            return

        resting_qty = order.qty - order.filled_qty

        if order.order_type == OrderType.IOC:
            order.status = (
                OrderStatus.PARTIALLY_FILLED
                if order.filled_qty > EPSILON
                else OrderStatus.CANCELLED
            )
            return

        if order.order_type == OrderType.MARKET:
            order.status = (
                OrderStatus.PARTIALLY_FILLED
                if order.filled_qty > EPSILON
                else OrderStatus.CANCELLED
            )
            return

        if order.order_type == OrderType.FOK:
            order.status = OrderStatus.CANCELLED
            order.executions.clear()
            order.filled_qty = 0.0
            return

        if resting_qty <= EPSILON:
            order.status = OrderStatus.PARTIALLY_FILLED
            return

        if order.price is None:
            raise ValueError("Limit-style orders must provide a price")

        same_side.append(
            _BookEntry(
                price=order.price,
                qty=resting_qty,
                order_id=order.id,
                timestamp=timestamp,
            )
        )
        self._sort_book_side(same_side, order.side)
        order.status = (
            OrderStatus.PARTIALLY_FILLED
            if order.filled_qty > EPSILON
            else OrderStatus.NEW
        )

    def _available_qty(
        self, order: Order, entries: Iterable[_BookEntry], halt: Optional[MarketHalt]
    ) -> float:
        total = 0.0
        cap = order.qty
        if halt and halt.mode == HaltMode.PARTIAL:
            cap = min(cap, order.qty * max(0.0, halt.liquidity_factor))
        for entry in entries:
            if not self._price_cross(order, entry.price):
                break
            total += entry.qty
            if total >= cap - EPSILON:
                return cap
        return min(total, cap)

    def _price_cross(self, order: Order, price: float) -> bool:
        if order.order_type == OrderType.MARKET:
            return True
        if order.side is OrderSide.BUY:
            return price <= (order.price or float("inf"))
        return price >= (order.price or 0.0)

    def _sort_book_side(self, entries: List[_BookEntry], side: OrderSide) -> None:
        reverse = side is OrderSide.BUY
        entries.sort(key=lambda e: (e.price, e.timestamp), reverse=reverse)


__all__ = [
    "ExecutionReport",
    "HaltMode",
    "MarketHalt",
    "MatchingEngine",
    "Order",
    "OrderSide",
    "OrderStatus",
    "OrderType",
]