"""
Order book microstructure analytics.

Functions operate on book snapshots (bids/asks dicts or DataFrames)
and return time-series or summary stats.
"""

from __future__ import annotations

import math
from collections import deque
from typing import Optional

import numpy as np


# ── Microprice ───────────────────────────────────────────────

def microprice(
    bids: dict[float, float],
    asks: dict[float, float],
    weight_bid: float = 0.5,
) -> float:
    """Weighted mid based on depth imbalance.

    microprice = weight * bid_side + (1-weight) * ask_side
    where weight = bid_depth / (bid_depth + ask_depth)

    Falls back to simple mid if no depth.
    """
    bid_prices = sorted(bids.keys(), reverse=True)
    ask_prices = sorted(asks.keys())

    if not bid_prices or not ask_prices:
        return 0.0

    best_bid = bid_prices[0]
    best_ask = ask_prices[0]

    bid_vol = sum(bids[px] for px in bid_prices[:10])
    ask_vol = sum(asks[px] for px in ask_prices[:10])
    total = bid_vol + ask_vol
    if total == 0:
        return (best_bid + best_ask) / 2.0

    w = bid_vol / total
    return w * best_bid + (1 - w) * best_ask


def mid_price(bids: dict[float, float], asks: dict[float, float]) -> float:
    bid_prices = sorted(bids.keys(), reverse=True)
    ask_prices = sorted(asks.keys())
    if not bid_prices or not ask_prices:
        return 0.0
    return (bid_prices[0] + ask_prices[0]) / 2.0


# ── Order-book imbalance ─────────────────────────────────────

def order_book_imbalance(
    bids: dict[float, float],
    asks: dict[float, float],
    levels: int = 10,
) -> float:
    """OBI = (bid_vol - ask_vol) / (bid_vol + ask_vol). Range [-1, 1]."""
    bid_prices = sorted(bids.keys(), reverse=True)[:levels]
    ask_prices = sorted(asks.keys())[:levels]
    bid_vol = sum(bids[px] for px in bid_prices)
    ask_vol = sum(asks[px] for px in ask_prices)
    total = bid_vol + ask_vol
    if total == 0:
        return 0.0
    return (bid_vol - ask_vol) / total


def depth_imbalance(
    bids: dict[float, float],
    asks: dict[float, float],
    price_distance_pct: float = 0.01,
) -> float:
    """Imbalance at a fixed price distance from mid (percentage-based)."""
    mid = mid_price(bids, asks)
    if mid <= 0:
        return 0.0
    lo = mid * (1 - price_distance_pct)
    hi = mid * (1 + price_distance_pct)
    bid_vol = sum(sz for px, sz in bids.items() if px >= lo)
    ask_vol = sum(sz for px, sz in asks.items() if px <= hi)
    total = bid_vol + ask_vol
    if total == 0:
        return 0.0
    return (bid_vol - ask_vol) / total


# ── Spread statistics ────────────────────────────────────────

def spread_stats(
    bids: dict[float, float],
    asks: dict[float, float],
) -> dict:
    bid_prices = sorted(bids.keys(), reverse=True)
    ask_prices = sorted(asks.keys())
    if not bid_prices or not ask_prices:
        return {"spread": 0, "spread_bps": 0, "mid": 0, "bid": 0, "ask": 0}

    best_bid = bid_prices[0]
    best_ask = ask_prices[0]
    mid = (best_bid + best_ask) / 2.0
    spread = best_ask - best_bid
    spread_bps = (spread / mid * 10000) if mid > 0 else 0
    return {
        "spread": round(spread, 2),
        "spread_bps": round(spread_bps, 2),
        "mid": round(mid, 2),
        "best_bid": best_bid,
        "best_ask": best_ask,
    }


# ── Depth resiliency ──────────────────────────────────────────

def depth_resiliency(
    bids: dict[float, float],
    asks: dict[float, float],
    impact_bps: float = 10.0,
) -> dict:
    """How much size sits within N bps of mid on each side.

    Returns volume and level count within the impact radius, plus
    a resiliency score: bid_depth / ask_depth (ratio).
    """
    mid = mid_price(bids, asks)
    if mid <= 0:
        return {"bid_vol": 0, "ask_vol": 0, "bid_levels": 0, "ask_levels": 0, "resiliency": 0}

    radius = mid * impact_bps / 10000
    bid_lo = mid - radius
    ask_hi = mid + radius

    bid_vol = sum(sz for px, sz in bids.items() if px >= bid_lo)
    ask_vol = sum(sz for px, sz in asks.items() if px <= ask_hi)
    bid_levels = sum(1 for px in bids if px >= bid_lo)
    ask_levels = sum(1 for px in asks if px <= ask_hi)

    resiliency = bid_vol / ask_vol if ask_vol > 0 else float("inf")
    return {
        "bid_vol": round(bid_vol, 8),
        "ask_vol": round(ask_vol, 8),
        "bid_levels": bid_levels,
        "ask_levels": ask_levels,
        "resiliency": round(resiliency, 4),
    }


def queue_depletion_prob(
    bids: dict[float, float],
    asks: dict[float, float],
    level_distance: int = 0,
    trade_rate_per_sec: float = 1.0,
    avg_trade_size: float = 0.01,
) -> float:
    """Probability queued order at best level(s) gets filled within 1 second.

    Simple Poisson model: P(fill) = 1 - exp(-λ)
    where λ = trade_rate * avg_trade_size / depth_at_level

    level_distance=0 means best bid/ask, 1 = one level behind, etc.
    """
    bid_prices = sorted(bids.keys(), reverse=True)
    ask_prices = sorted(asks.keys())
    if not bid_prices or not ask_prices:
        return 0.0

    px = bid_prices[min(level_distance, len(bid_prices) - 1)]
    depth = bids.get(px, 0)
    if depth <= 0:
        return 0.0

    lam = trade_rate_per_sec * avg_trade_size / depth
    prob = 1.0 - math.exp(-lam)
    return round(min(prob, 0.9999), 6)


# ── Batch processing ─────────────────────────────────────────

def batch_book_stats(
    snapshots: list[dict],
    levels: int = 10,
) -> dict:
    """Process a list of book snapshots and return aggregate stats.

    Each snapshot: {"bids": {price: size, ...}, "asks": {price: size, ...}}
    """
    obis = []
    spreads_bps = []
    micros = []
    depths = []
    resiliencies = []

    for snap in snapshots:
        bids = snap.get("bids", {})
        asks = snap.get("asks", {})
        if not bids or not asks:
            continue
        obis.append(order_book_imbalance(bids, asks, levels))
        ss = spread_stats(bids, asks)
        spreads_bps.append(ss["spread_bps"])
        micros.append(microprice(bids, asks))
        dr = depth_resiliency(bids, asks)
        depths.append(dr["bid_vol"] + dr["ask_vol"])
        resiliencies.append(dr["resiliency"])

    def _summarize(vals):
        if not vals:
            return {"mean": 0, "median": 0, "std": 0, "min": 0, "max": 0, "count": 0}
        a = np.array(vals, dtype=float)
        a = a[np.isfinite(a)]
        return {
            "mean": round(float(np.mean(a)), 4),
            "median": round(float(np.median(a)), 4),
            "std": round(float(np.std(a)), 4),
            "min": round(float(np.min(a)), 4),
            "max": round(float(np.max(a)), 4),
            "count": len(a),
        }

    return {
        "obi": _summarize(obis),
        "spread_bps": _summarize(spreads_bps),
        "microprice_ratio": _summarize([m / mid_price(s["bids"], s["asks"])
                                         if mid_price(s["bids"], s["asks"]) > 0 else 1.0
                                         for s, m in zip(snapshots, micros)]),
        "depth_total": _summarize(depths),
        "resiliency": _summarize([r for r in resiliencies if r < 1e6]),
    }