"""
Multi-Level Book Dynamics — model order book at multiple depth levels.

Improves fill simulation by modeling dynamics beyond top-of-book.
"""
from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional, Tuple

import numpy as np
from numba import njit


@dataclass(frozen=True, slots=True)
class BookLevelDynamics:
    """Dynamics at a single price level."""
    price: float
    qty: float
    arrival_rate: float  # new orders arriving per second
    cancel_rate: float  # orders cancelled per second
    net_flow: float  # arrival - cancel


@njit(cache=True)
def compute_net_order_flow(
    bid_depth: float,
    ask_depth: float,
    recent_trade_imbalance: float,
    toxicity: float,
    volatility: float,
) -> Tuple[float, float]:
    """
    Compute net order flow for bids and asks.

    Model:
      - More buying pressure → bid side gets more orders
      - Toxic flow → both sides thin out
      - High volatility → both sides thin out

    Returns (bid_flow, ask_flow) in units per second.
    """
    # Base arrival rate (orders per second)
    base_arrival = 0.5

    # Trade imbalance affects arrival
    bid_arrival = base_arrival * (1.0 + recent_trade_imbalance * 0.3)
    ask_arrival = base_arrival * (1.0 - recent_trade_imbalance * 0.3)

    # Toxicity reduces both sides (withdrawals)
    toxicity_cancel = toxicity * 0.3

    # Volatility increases cancellations
    vol_cancel = volatility * 0.01

    bid_flow = bid_arrival - toxicity_cancel - vol_cancel
    ask_flow = ask_arrival - toxicity_cancel - vol_cancel

    return max(0.0, bid_flow), max(0.0, ask_flow)


@njit(cache=True)
def compute_book_imbalance_weighted(
    bid_prices: np.ndarray,
    bid_qtys: np.ndarray,
    ask_prices: np.ndarray,
    ask_qtys: np.ndarray,
    depth: int = 5,
) -> float:
    """
    Compute depth-weighted book imbalance.

    Weight by distance from mid (closer = more important).
    """
    if len(bid_prices) == 0 or len(ask_prices) == 0:
        return 0.0

    mid = 0.5 * (bid_prices[0] + ask_prices[0])
    if mid <= 0:
        return 0.0

    bid_weight = 0.0
    ask_weight = 0.0

    for i in range(min(depth, len(bid_prices))):
        distance = abs(bid_prices[i] - mid) / mid + 1e-12
        weight = 1.0 / distance
        bid_weight += bid_qtys[i] * weight

    for i in range(min(depth, len(ask_prices))):
        distance = abs(ask_prices[i] - mid) / mid + 1e-12
        weight = 1.0 / distance
        ask_weight += ask_qtys[i] * weight

    total = bid_weight + ask_weight
    if total <= 0:
        return 0.0
    return (bid_weight - ask_weight) / total


class MultiLevelBookModel:
    """
    Multi-level book dynamics model.

    Models order book at multiple depth levels, not just top-of-book.
    """

    def __init__(self) -> None:
        self._depth_history: list[dict] = []

    def update(self, bid_depths: list[float], ask_depths: list[float]) -> None:
        """Update with current depth profile."""
        self._depth_history.append({
            "bids": list(bid_depths),
            "asks": list(ask_depths),
        })
        if len(self._depth_history) > 1000:
            self._depth_history = self._depth_history[-500:]

    def compute_imbalance(self, depth: int = 5) -> float:
        """Compute weighted book imbalance."""
        if not self._depth_history:
            return 0.0
        latest = self._depth_history[-1]
        bids = np.array(latest["bids"][:depth], dtype=np.float64) if latest["bids"] else np.array([], dtype=np.float64)
        asks = np.array(latest["asks"][:depth], dtype=np.float64) if latest["asks"] else np.array([], dtype=np.float64)
        bid_prices = np.arange(len(bids), dtype=np.float64)
        ask_prices = np.arange(len(asks), dtype=np.float64)
        return compute_book_imbalance_weighted(bid_prices, bids, ask_prices, asks, depth)

    def compute_depth_ratio(self, depth: int = 5) -> float:
        """Compute bid/ask depth ratio."""
        if not self._depth_history:
            return 1.0
        latest = self._depth_history[-1]
        bid_total = sum(latest["bids"][:depth])
        ask_total = sum(latest["asks"][:depth])
        return bid_total / max(ask_total, 1e-12)

    @property
    def current_bid_depth(self) -> float:
        if not self._depth_history:
            return 0.0
        return sum(self._depth_history[-1]["bids"])

    @property
    def current_ask_depth(self) -> float:
        if not self._depth_history:
            return 0.0
        return sum(self._depth_history[-1]["asks"])