from __future__ import annotations

import math
from collections import Counter
from datetime import datetime, timedelta
from typing import Any

import orjson
import polars as pl

from uforia.storage import DatasetName, DatasetStore
from uforia.storage.datasets import DATASET_SPECS, empty_frame

SOURCE = "perps_family_v1"
COMPOSITE_SOURCE = "perps_composite_v1"
TRANSITION_SOURCE = "perps_transition_v1"
SECONDS_PER_YEAR = 365.0 * 24.0 * 60.0 * 60.0

REGIME_FAMILIES: tuple[tuple[str, str], ...] = (
    ("cascade", "cascade_regime"),
    ("toxicity", "toxicity_regime"),
    ("leverage", "leverage_regime"),
    ("spread", "spread_regime"),
    ("rv", "rv_regime"),
    ("delta", "delta_regime"),
)


def build_perps_signal_series(
    store: DatasetStore,
    start: datetime,
    end: datetime,
    *,
    asset: str,
    venue_scope: str = "all",
    source: str = SOURCE,
) -> pl.DataFrame:
    # Perps dashboards are meant to refresh from the hot live microstructure window.
    # A longer 5s read is too expensive for the current Mac Mini deployment and
    # most metrics degrade gracefully when older history is unavailable.
    lookback_start = start - timedelta(hours=6)
    bars_1m = _with_bar_context(
        _read_hot_window(
            store,
            DatasetName.MICROSTRUCTURE_BAR,
            lookback_start,
            end,
            filters={"asset": asset, "bar_size_secs": 60},
        )
    )
    bars_5s = _with_bar_context(
        _read_hot_window(
            store,
            DatasetName.MICROSTRUCTURE_BAR,
            lookback_start,
            end,
            filters={"asset": asset, "bar_size_secs": 5},
        )
    )
    raw_events = _with_raw_context(
        _read_hot_window(
            store,
            DatasetName.MICROSTRUCTURE_RAW_EVENT,
            lookback_start,
            end,
            filters={"asset": asset},
        )
    )
    book = _with_book_context(
        _read_hot_window(
            store,
            DatasetName.MICROSTRUCTURE_BOOK_STATE,
            lookback_start,
            end,
            filters={"asset": asset},
        )
    )
    funding = _with_funding_context(
        _read_hot_window(
            store,
            DatasetName.FUNDING_RATE_SNAPSHOT,
            lookback_start,
            end,
            filters={"underlying": asset},
        )
    )
    if bars_1m.is_empty():
        return pl.DataFrame()

    bars_1m = _filter_scope(bars_1m, venue_scope)
    bars_5s = _filter_scope(bars_5s, venue_scope)
    raw_events = _filter_scope(raw_events, venue_scope)
    book = _filter_scope(book, venue_scope)
    funding_scope = _filter_funding_scope(funding, venue_scope)

    asset_bars = _aggregate_asset_bars(bars_1m)
    if asset_bars.is_empty():
        return pl.DataFrame()

    trades = _aggregate_trade_context(raw_events)
    liquidations = _aggregate_liquidation_context(raw_events)
    open_interest = _aggregate_open_interest_context(raw_events)
    funding_context = _aggregate_funding_context(funding_scope)
    cross_venue = _aggregate_cross_venue_context(bars_1m, funding)
    vol_context = _aggregate_volatility_context(bars_5s, bars_1m)
    depth_context = _aggregate_depth_context(book, raw_events)

    base = asset_bars.join(trades, on="ts", how="left")
    for frame in (
        liquidations,
        open_interest,
        funding_context,
        cross_venue,
        vol_context,
        depth_context,
    ):
        if not frame.is_empty():
            join_columns = [
                "ts",
                *[
                    column
                    for column in frame.columns
                    if column != "ts" and column not in base.columns
                ],
            ]
            base = base.join(frame.select(join_columns), on="ts", how="left")

    missing_defaults: dict[str, Any] = {
        "long_liq_notional": None,
        "short_liq_notional": None,
        "total_liq_notional": None,
        "oi_value": None,
        "funding_rate_1h": None,
        "cross_venue_funding_spread": None,
        "venue_dislocation_bps": None,
        "binance_ret": None,
        "bybit_ret": None,
        "rv_5s": None,
        "rv_1m": None,
        "rv_5m": None,
        "depth_delta": None,
        "aggressive_flow_notional": None,
        "trade_notional": None,
        "large_trade_notional": None,
        "aggressive_buy_notional": None,
        "aggressive_sell_notional": None,
    }
    absent = [column for column in missing_defaults if column not in base.columns]
    if absent:
        base = base.with_columns(
            [pl.lit(missing_defaults[column], dtype=pl.Float64).alias(column) for column in absent]
        )

    base = base.sort("ts").with_columns(
        [
            pl.lit(asset).alias("asset"),
            pl.lit(venue_scope).alias("venue_scope"),
            pl.lit(source).alias("source"),
            (((pl.col("perp_mid") - pl.col("spot_mid")) / pl.col("spot_mid")) * 10_000.0).alias(
                "basis_bps"
            ),
            (
                (((pl.col("perp_mid") - pl.col("spot_mid")) / pl.col("spot_mid")) * 100.0)
                + (pl.col("funding_rate_1h") * 24.0 * 365.0 * 100.0)
            ).alias("carry_annualized_pct"),
            _safe_ratio(
                pl.col("long_liq_notional") - pl.col("short_liq_notional"),
                pl.col("total_liq_notional"),
            ).alias("liq_asymmetry"),
            _safe_ratio(pl.col("oi_value"), pl.col("depth_total_notional")).alias(
                "oi_concentration"
            ),
            _safe_ratio(
                (pl.col("buy_volume") - pl.col("sell_volume")).abs(),
                pl.col("volume"),
            ).alias("vpin_bucket"),
            _safe_ratio(
                pl.col("large_trade_notional"),
                pl.col("trade_notional"),
            ).alias("large_trade_ratio"),
            (pl.col("oi_value").pct_change(60) * 100.0).alias("oi_change_1h"),
            (pl.col("oi_value").pct_change(240) * 100.0).alias("oi_change_4h"),
            (pl.col("oi_value").pct_change(1_440) * 100.0).alias("oi_change_24h"),
            (((pl.col("mid_price") / pl.col("mid_price").shift(60)) - 1.0) * 10_000.0).alias(
                "price_change_1h"
            ),
            pl.col("oi_value").diff(60).diff().alias("oi_acceleration"),
            (pl.col("spread_abs") / pl.col("mid_price") * 10_000.0).alias("spread_bps"),
            pl.col("depth_imbalance").diff().alias("depth_imbalance_momentum"),
            pl.col("cum_delta_30").diff().diff().alias("cum_delta_acceleration"),
            _safe_ratio(
                pl.col("trade_notional"),
                pl.col("mid_price") * pl.col("return_1m").abs(),
            ).alias("absorption_ratio"),
        ]
    )
    base = base.with_columns(
        [
            pl.col("carry_annualized_pct").diff().alias("carry_momentum_1h"),
            (pl.col("carry_annualized_pct") - pl.col("carry_annualized_pct").shift(24 * 60)).alias(
                "carry_momentum_24h"
            ),
            pl.col("total_liq_notional")
            .rolling_mean(window_size=60, min_samples=10)
            .alias("liq_intensity_1h"),
            pl.col("vpin_bucket").rolling_mean(window_size=50, min_samples=20).alias("vpin_50"),
            _rolling_autocorr("order_flow_imbalance", 60).alias("ofi_persistence"),
            _rolling_autocorr("return_1m_sq", 60).alias("vol_clustering"),
            _safe_ratio(pl.col("rv_5s"), pl.col("rv_5m")).alias("vol_signature_ratio"),
            (
                (pl.col("oi_change_1h").sign() != pl.col("price_change_1h").sign()).cast(pl.Float64)
                * pl.col("oi_change_1h").abs()
            ).alias("oi_price_divergence"),
            _lead_lag_score().alias("lead_lag_score"),
            (
                (pl.col("venue_dislocation_bps").abs() * 60.0)
                / (
                    pl.col("venue_dislocation_bps")
                    .diff()
                    .abs()
                    .rolling_mean(window_size=60, min_samples=10)
                    + 1e-6
                )
            )
            .clip(60.0, 3_600.0)
            .alias("convergence_half_life_secs"),
            (
                pl.when(pl.col("depth_delta") > 0.0)
                .then(pl.col("depth_delta"))
                .otherwise(0.0)
                .rolling_mean(window_size=30, min_samples=10)
                / (
                    pl.col("aggressive_flow_notional")
                    .shift(1)
                    .rolling_mean(window_size=30, min_samples=10)
                    + 1e-6
                )
            ).alias("depth_recovery_speed"),
            (
                pl.when(pl.col("cum_delta_acceleration").sign() != pl.col("price_change_1h").sign())
                .then(pl.col("cum_delta_acceleration").abs())
                .otherwise(0.0)
            ).alias("delta_divergence"),
        ]
    )
    base = base.with_columns(
        [
            _rolling_zscore("carry_annualized_pct", 7 * 24 * 60, min_samples=180).alias(
                "carry_zscore_7d"
            ),
            _rolling_zscore("carry_annualized_pct", 30 * 24 * 60, min_samples=360).alias(
                "carry_zscore_30d"
            ),
            _rolling_zscore("total_liq_notional", 24 * 60, min_samples=120).alias(
                "liq_intensity_zscore"
            ),
            _rolling_zscore("vpin_50", 4 * 60, min_samples=40).alias("toxic_flow_zscore_4h"),
            _rolling_zscore("vpin_50", 24 * 60, min_samples=120).alias("toxic_flow_zscore_24h"),
            _rolling_zscore("spread_bps", 4 * 60, min_samples=40).alias("spread_zscore_4h"),
            _rolling_zscore("spread_bps", 24 * 60, min_samples=120).alias("spread_zscore_24h"),
            _rolling_zscore("rv_1m", 24 * 60, min_samples=120).alias("rv_zscore"),
        ]
    )
    base = base.with_columns(
        [
            (
                0.5 * pl.col("oi_concentration").clip(0.0, 5.0)
                + 0.3 * pl.col("liq_intensity_zscore").fill_null(0.0).clip(-3.0, 3.0)
                + 0.2 * pl.col("carry_zscore_7d").fill_null(0.0).abs().clip(0.0, 3.0)
            ).alias("crowding_score"),
        ]
    )
    base = base.with_columns(
        [
            (
                0.5 * pl.col("crowding_score").fill_null(0.0)
                + 0.3 * pl.col("liq_intensity_zscore").fill_null(0.0).clip(-3.0, 3.0)
                + 0.2 * pl.col("liq_asymmetry").fill_null(0.0).abs()
            ).alias("cascade_risk_score"),
        ]
    )
    base = base.with_columns(
        [
            _regime_expr(
                "cascade_risk_score",
                (1.0, 2.0, 3.0),
                ("low", "elevated", "high", "extreme"),
            ).alias("cascade_regime"),
            _regime_expr(
                "toxic_flow_zscore_4h",
                (0.5, 1.5, 2.5),
                ("calm", "elevated", "toxic", "extreme"),
            ).alias("toxicity_regime"),
            pl.when(pl.col("oi_change_1h").abs() < 1.0)
            .then(pl.lit("steady"))
            .when((pl.col("oi_change_1h") > 0) & (pl.col("price_change_1h").abs() < 25.0))
            .then(pl.lit("building"))
            .otherwise(pl.lit("unwinding"))
            .alias("leverage_regime"),
            pl.when(pl.col("lead_lag_score") > 0.0001)
            .then(pl.lit("binance"))
            .when(pl.col("lead_lag_score") < -0.0001)
            .then(pl.lit("bybit"))
            .otherwise(pl.lit("neutral"))
            .alias("price_discovery_leader"),
            _regime_expr(
                "spread_zscore_4h",
                (0.5, 1.5, 2.5),
                ("tight", "normal", "wide", "stress"),
            ).alias("spread_regime"),
            _regime_expr(
                "rv_zscore",
                (0.0, 1.0, 2.0),
                ("low", "normal", "elevated", "extreme"),
            ).alias("rv_regime"),
            pl.when((pl.col("cum_delta_acceleration") > 0) & (pl.col("absorption_ratio") < 1_000.0))
            .then(pl.lit("aggressive_buying"))
            .when((pl.col("cum_delta_acceleration") > 0) & (pl.col("absorption_ratio") >= 1_000.0))
            .then(pl.lit("passive_buying"))
            .when((pl.col("cum_delta_acceleration") < 0) & (pl.col("absorption_ratio") < 1_000.0))
            .then(pl.lit("aggressive_selling"))
            .when((pl.col("cum_delta_acceleration") < 0) & (pl.col("absorption_ratio") >= 1_000.0))
            .then(pl.lit("passive_selling"))
            .otherwise(pl.lit("neutral"))
            .alias("delta_regime"),
        ]
    )
    base = base.with_columns(
        [
            pl.when(pl.col("mid_price").is_null())
            .then(pl.lit("insufficient_history"))
            .when(
                pl.any_horizontal(
                    pl.col("basis_bps").is_null(),
                    pl.col("cross_venue_funding_spread").is_null(),
                    pl.col("lead_lag_score").is_null(),
                )
            )
            .then(pl.lit("partial_context"))
            .otherwise(pl.lit("ok"))
            .alias("quality_flag")
        ]
    )

    return base.filter((pl.col("ts") >= pl.lit(start)) & (pl.col("ts") <= pl.lit(end))).select(
        [
            "ts",
            "date",
            "asset",
            "venue_scope",
            "source",
            "basis_bps",
            "funding_rate_1h",
            "carry_annualized_pct",
            "carry_zscore_7d",
            "carry_zscore_30d",
            "carry_momentum_1h",
            "carry_momentum_24h",
            "cross_venue_funding_spread",
            "long_liq_notional",
            "short_liq_notional",
            "liq_intensity_zscore",
            "liq_asymmetry",
            "oi_concentration",
            "crowding_score",
            "cascade_risk_score",
            "cascade_regime",
            "vpin_50",
            "toxic_flow_zscore_4h",
            "toxic_flow_zscore_24h",
            "ofi_persistence",
            "large_trade_ratio",
            "toxicity_regime",
            "oi_change_1h",
            "oi_change_4h",
            "oi_change_24h",
            "price_change_1h",
            "oi_price_divergence",
            "oi_acceleration",
            "leverage_regime",
            "lead_lag_score",
            "venue_dislocation_bps",
            "convergence_half_life_secs",
            "price_discovery_leader",
            "depth_recovery_speed",
            "spread_bps",
            "spread_zscore_4h",
            "spread_zscore_24h",
            "depth_imbalance_momentum",
            "spread_regime",
            "rv_5s",
            "rv_1m",
            "rv_5m",
            "vol_signature_ratio",
            "vol_clustering",
            "rv_regime",
            "cum_delta_acceleration",
            "delta_divergence",
            "absorption_ratio",
            "delta_regime",
            "quality_flag",
        ]
    )


def build_perps_composite_series(
    store: DatasetStore,
    start: datetime,
    end: datetime,
    *,
    asset: str,
    venue_scope: str = "all",
    source: str = COMPOSITE_SOURCE,
) -> pl.DataFrame:
    signal_frame = store.read(
        DatasetName.PERPS_SIGNAL_SERIES,
        start=start,
        end=end,
        filters={"asset": asset, "venue_scope": venue_scope, "source": SOURCE},
    ).sort("ts")
    if signal_frame.is_empty():
        return pl.DataFrame()
    return signal_frame.select(
        [
            "ts",
            "date",
            pl.lit(asset).alias("asset"),
            pl.lit(venue_scope).alias("venue_scope"),
            pl.lit(source).alias("source"),
            _avg_score("carry_zscore_7d", "carry_zscore_30d").alias("carry_score"),
            _avg_score(
                "cascade_risk_score",
                "liq_intensity_zscore",
                "crowding_score",
            ).alias("cascade_risk_score"),
            _avg_score(
                "toxic_flow_zscore_4h",
                "toxic_flow_zscore_24h",
                pl.col("large_trade_ratio") * 3.0,
            ).alias("toxicity_score"),
            _avg_score(
                "oi_price_divergence",
                "oi_acceleration",
                "crowding_score",
            ).alias("positioning_score"),
            _avg_score(
                "spread_zscore_4h",
                "spread_zscore_24h",
                pl.col("venue_dislocation_bps") / 5.0,
            ).alias("fragility_score"),
            "quality_flag",
        ]
    )


def build_perps_transition_series(
    store: DatasetStore,
    start: datetime,
    end: datetime,
    *,
    asset: str,
    venue_scope: str = "all",
    source: str = TRANSITION_SOURCE,
) -> pl.DataFrame:
    frame = store.read(
        DatasetName.PERPS_SIGNAL_SERIES,
        start=start,
        end=end,
        filters={"asset": asset, "venue_scope": venue_scope, "source": SOURCE},
    ).sort("ts")
    if frame.is_empty():
        return pl.DataFrame()

    rows: list[dict[str, Any]] = []
    for signal_family, column in REGIME_FAMILIES:
        states = [str(value) for value in frame.get_column(column).to_list() if value is not None]
        if not states:
            continue
        history_transitions: Counter[tuple[str, str]] = Counter()
        history_origin: Counter[str] = Counter()
        run_lengths: list[int] = []
        last_state: str | None = None
        run_length = 0
        for row in frame.iter_rows(named=True):
            current = row.get(column)
            if current is None:
                continue
            current = str(current)
            if last_state is not None:
                history_transitions[(last_state, current)] += 1
                history_origin[last_state] += 1
            if current == last_state:
                run_length += 1
            else:
                if run_length:
                    run_lengths.append(run_length)
                run_length = 1
            distribution = {
                target: history_transitions[(current, target)] / history_origin[current]
                for target in sorted(
                    {right for left, right in history_transitions if left == current}
                )
                if history_origin[current] > 0
            }
            median_duration = (
                float(sorted(run_lengths)[len(run_lengths) // 2])
                if run_lengths
                else float(run_length)
            )
            rows.append(
                {
                    "ts": row["ts"],
                    "date": row["date"],
                    "asset": asset,
                    "venue_scope": venue_scope,
                    "source": source,
                    "signal_family": signal_family,
                    "current_state": current,
                    "prior_state": last_state,
                    "next_state_distribution": orjson.dumps(distribution).decode(),
                    "median_state_duration_minutes": median_duration,
                    "quality_flag": row.get("quality_flag") or "ok",
                }
            )
            last_state = current
        if run_length:
            run_lengths.append(run_length)
    return pl.DataFrame(rows) if rows else pl.DataFrame()


def _with_bar_context(frame: pl.DataFrame) -> pl.DataFrame:
    if frame.is_empty():
        return frame
    return frame.with_columns(
        [
            pl.coalesce([pl.col("venue_group"), _venue_group_expr("venue")]).alias("venue_group"),
            pl.coalesce([pl.col("market_type"), _market_type_expr("venue")]).alias("market_type"),
            pl.coalesce(
                [pl.col("is_derivatives"), _market_type_expr("venue") != pl.lit("spot")]
            ).alias("is_derivatives"),
            pl.coalesce([pl.col("contract_type"), _contract_type_expr("venue")]).alias(
                "contract_type"
            ),
            pl.coalesce([pl.col("exchange_symbol"), pl.col("venue")]).alias("exchange_symbol"),
            pl.coalesce(
                [
                    pl.col("depth_notional_bid_5"),
                    pl.col("depth_bid_5") * pl.col("close"),
                ]
            ).alias("depth_notional_bid_5"),
            pl.coalesce(
                [
                    pl.col("depth_notional_ask_5"),
                    pl.col("depth_ask_5") * pl.col("close"),
                ]
            ).alias("depth_notional_ask_5"),
        ]
    )


def _with_book_context(frame: pl.DataFrame) -> pl.DataFrame:
    if frame.is_empty():
        return frame
    return frame.with_columns(
        [
            pl.coalesce([pl.col("venue_group"), _venue_group_expr("venue")]).alias("venue_group"),
            pl.coalesce([pl.col("market_type"), _market_type_expr("venue")]).alias("market_type"),
            pl.coalesce(
                [pl.col("is_derivatives"), _market_type_expr("venue") != pl.lit("spot")]
            ).alias("is_derivatives"),
            pl.coalesce([pl.col("contract_type"), _contract_type_expr("venue")]).alias(
                "contract_type"
            ),
            pl.coalesce([pl.col("exchange_symbol"), pl.col("venue")]).alias("exchange_symbol"),
            pl.coalesce(
                [
                    pl.col("depth_notional_bid_5"),
                    pl.col("depth_bid_5") * pl.col("mid_price"),
                ]
            ).alias("depth_notional_bid_5"),
            pl.coalesce(
                [
                    pl.col("depth_notional_ask_5"),
                    pl.col("depth_ask_5") * pl.col("mid_price"),
                ]
            ).alias("depth_notional_ask_5"),
        ]
    )


def _with_raw_context(frame: pl.DataFrame) -> pl.DataFrame:
    if frame.is_empty():
        return frame
    return frame.with_columns(
        [
            pl.coalesce([pl.col("venue_group"), _venue_group_expr("venue")]).alias("venue_group"),
            pl.coalesce([pl.col("market_type"), _market_type_expr("venue")]).alias("market_type"),
            pl.coalesce(
                [pl.col("is_derivatives"), _market_type_expr("venue") != pl.lit("spot")]
            ).alias("is_derivatives"),
            pl.coalesce([pl.col("contract_type"), _contract_type_expr("venue")]).alias(
                "contract_type"
            ),
            pl.coalesce([pl.col("exchange_symbol"), pl.col("venue")]).alias("exchange_symbol"),
            pl.coalesce([pl.col("liquidation_side"), pl.col("side")]).alias("liquidation_side"),
            pl.coalesce([pl.col("trade_notional"), pl.col("price") * pl.col("size")]).alias(
                "trade_notional"
            ),
            pl.coalesce(
                [
                    pl.col("depth_notional_bid_5"),
                    pl.col("bid") * pl.col("size"),
                ]
            ).alias("depth_notional_bid_5"),
            pl.coalesce(
                [
                    pl.col("depth_notional_ask_5"),
                    pl.col("ask") * pl.col("size"),
                ]
            ).alias("depth_notional_ask_5"),
        ]
    )


def _with_funding_context(frame: pl.DataFrame) -> pl.DataFrame:
    if frame.is_empty():
        return frame
    return frame.with_columns([_venue_group_expr("venue").alias("venue_group")])


def _read_hot_window(
    store: DatasetStore,
    dataset: DatasetName,
    start: datetime,
    end: datetime,
    *,
    filters: dict[str, Any],
) -> pl.DataFrame:
    if start.date() == end.date():
        requested_columns = _HOT_COLUMNS.get(dataset)
        if requested_columns:
            frame = _read_hot_columns(store, dataset, filters=filters, columns=requested_columns)
        else:
            frame = store.read_latest(dataset, filters=filters)
        time_column = "ts" if "ts" in frame.columns else "logical_ts"
        if time_column in frame.columns:
            return frame.filter(
                (pl.col(time_column) >= pl.lit(start)) & (pl.col(time_column) <= pl.lit(end))
            )
        return frame
    return store.read(dataset, start=start, end=end, filters=filters)


def _read_hot_columns(
    store: DatasetStore,
    dataset: DatasetName,
    *,
    filters: dict[str, Any],
    columns: tuple[str, ...],
) -> pl.DataFrame:
    spec = DATASET_SPECS[dataset]
    required = set(columns) | set(spec.primary_key) | {"run_ts"}
    ordered = [column for column in spec.schema if column in required]
    base_frame = store.read_latest(dataset, filters=filters)
    available_columns = set(base_frame.columns)
    frame = base_frame.select([pl.col(column) for column in ordered if column in available_columns])
    missing = [column for column in ordered if column not in frame.columns]
    if missing:
        frame = frame.with_columns(
            [pl.lit(None, dtype=spec.schema[column]).alias(column) for column in missing]
        )
    if frame.is_empty():
        return empty_frame(dataset)
    if filters:
        for column, value in filters.items():
            if column in frame.columns:
                frame = frame.filter(pl.col(column) == value)
    time_column = "ts" if "ts" in frame.columns else "logical_ts"
    if spec.primary_key:
        frame = (
            frame.sort([*spec.primary_key, "run_ts"])
            .unique(subset=list(spec.primary_key), keep="last", maintain_order=True)
            .sort(time_column)
        )
    return frame


_HOT_COLUMNS: dict[DatasetName, tuple[str, ...]] = {
    DatasetName.MICROSTRUCTURE_BAR: (
        "ts",
        "date",
        "asset",
        "venue",
        "exchange_symbol",
        "venue_group",
        "market_type",
        "is_derivatives",
        "contract_type",
        "clock",
        "bar_size_secs",
        "close",
        "volume",
        "buy_volume",
        "sell_volume",
        "trade_count",
        "spread",
        "depth_bid_5",
        "depth_ask_5",
        "depth_notional_bid_5",
        "depth_notional_ask_5",
        "run_ts",
    ),
    DatasetName.MICROSTRUCTURE_BOOK_STATE: (
        "ts",
        "date",
        "asset",
        "venue",
        "exchange_symbol",
        "venue_group",
        "market_type",
        "is_derivatives",
        "contract_type",
        "mid_price",
        "depth_bid_5",
        "depth_ask_5",
        "depth_notional_bid_5",
        "depth_notional_ask_5",
        "run_ts",
    ),
    DatasetName.MICROSTRUCTURE_RAW_EVENT: (
        "ts",
        "date",
        "asset",
        "venue",
        "event_type",
        "exchange_symbol",
        "venue_group",
        "market_type",
        "is_derivatives",
        "contract_type",
        "side",
        "liquidation_side",
        "price",
        "size",
        "trade_notional",
        "bid",
        "ask",
        "mark_price",
        "open_interest",
        "depth_notional_bid_5",
        "depth_notional_ask_5",
        "run_ts",
    ),
    DatasetName.FUNDING_RATE_SNAPSHOT: (
        "ts",
        "date",
        "venue",
        "underlying",
        "interest_1h",
        "run_ts",
    ),
}


def _filter_scope(frame: pl.DataFrame, venue_scope: str) -> pl.DataFrame:
    if frame.is_empty() or venue_scope == "all":
        return frame
    if venue_scope == "binance":
        return frame.filter(pl.col("venue_group") == "binance")
    if venue_scope == "bybit":
        return frame.filter(pl.col("venue_group") == "bybit")
    return frame


def _filter_funding_scope(frame: pl.DataFrame, venue_scope: str) -> pl.DataFrame:
    if frame.is_empty() or venue_scope == "all":
        return frame
    return frame.filter(pl.col("venue_group") == venue_scope)


def _aggregate_asset_bars(frame: pl.DataFrame) -> pl.DataFrame:
    if frame.is_empty():
        return pl.DataFrame()
    grouped = (
        frame.sort(["ts", "venue"])
        .group_by("ts")
        .agg(
            [
                pl.col("date").first().alias("date"),
                pl.col("close").mean().alias("mid_price"),
                pl.col("close").filter(~pl.col("is_derivatives")).mean().alias("spot_mid"),
                pl.col("close").filter(pl.col("is_derivatives")).mean().alias("perp_mid"),
                pl.col("close")
                .filter((pl.col("venue_group") == "binance") & pl.col("is_derivatives"))
                .mean()
                .alias("binance_mid"),
                pl.col("close")
                .filter((pl.col("venue_group") == "bybit") & pl.col("is_derivatives"))
                .mean()
                .alias("bybit_mid"),
                pl.col("volume").sum().alias("volume"),
                pl.col("buy_volume").sum().alias("buy_volume"),
                pl.col("sell_volume").sum().alias("sell_volume"),
                pl.col("trade_count").sum().alias("trade_count"),
                pl.col("spread").mean().alias("spread_abs"),
                pl.col("depth_notional_bid_5").sum().alias("depth_notional_bid_5"),
                pl.col("depth_notional_ask_5").sum().alias("depth_notional_ask_5"),
            ]
        )
    )
    return grouped.with_columns(
        [
            (pl.col("depth_notional_bid_5") + pl.col("depth_notional_ask_5")).alias(
                "depth_total_notional"
            ),
            _safe_ratio(
                pl.col("depth_notional_bid_5") - pl.col("depth_notional_ask_5"),
                pl.col("depth_notional_bid_5") + pl.col("depth_notional_ask_5"),
            ).alias("depth_imbalance"),
            (
                (pl.col("buy_volume") - pl.col("sell_volume")).rolling_sum(
                    window_size=10, min_samples=3
                )
            ).alias("cum_delta_10"),
            (
                (pl.col("buy_volume") - pl.col("sell_volume")).rolling_sum(
                    window_size=30, min_samples=5
                )
            ).alias("cum_delta_30"),
            ((pl.col("mid_price") / pl.col("mid_price").shift(1)).log()).alias("return_1m"),
            ((pl.col("mid_price") / pl.col("mid_price").shift(1)).log() ** 2).alias("return_1m_sq"),
            ((pl.col("buy_volume") - pl.col("sell_volume")) / (pl.col("volume") + 1e-6)).alias(
                "order_flow_imbalance"
            ),
        ]
    )


def _aggregate_trade_context(frame: pl.DataFrame) -> pl.DataFrame:
    trades = frame.filter(pl.col("event_type") == "trade").sort("ts")
    if trades.is_empty():
        return pl.DataFrame()
    trade_notional = trades.get_column("trade_notional")
    mean_value = _numeric_or_zero(trade_notional.mean())
    std_value = _numeric_or_zero(trade_notional.std())
    threshold = mean_value + 2.0 * std_value
    return (
        trades.group_by_dynamic("ts", every="1m", closed="left")
        .agg(
            [
                pl.col("trade_notional").sum().alias("trade_notional"),
                pl.when(pl.col("trade_notional") >= threshold)
                .then(pl.col("trade_notional"))
                .otherwise(0.0)
                .sum()
                .alias("large_trade_notional"),
                pl.when(pl.col("side") == "buy")
                .then(pl.col("trade_notional"))
                .otherwise(0.0)
                .sum()
                .alias("aggressive_buy_notional"),
                pl.when(pl.col("side") == "sell")
                .then(pl.col("trade_notional"))
                .otherwise(0.0)
                .sum()
                .alias("aggressive_sell_notional"),
            ]
        )
        .with_columns(
            [
                (pl.col("aggressive_buy_notional") + pl.col("aggressive_sell_notional")).alias(
                    "aggressive_flow_notional"
                )
            ]
        )
    )


def _aggregate_liquidation_context(frame: pl.DataFrame) -> pl.DataFrame:
    liquidations = frame.filter(pl.col("event_type") == "liquidation").sort("ts")
    if liquidations.is_empty():
        return pl.DataFrame()
    return liquidations.group_by_dynamic("ts", every="1m", closed="left").agg(
        [
            pl.when(pl.col("liquidation_side") == "long")
            .then(pl.col("trade_notional"))
            .otherwise(0.0)
            .sum()
            .alias("long_liq_notional"),
            pl.when(pl.col("liquidation_side") == "short")
            .then(pl.col("trade_notional"))
            .otherwise(0.0)
            .sum()
            .alias("short_liq_notional"),
            pl.col("trade_notional").sum().alias("total_liq_notional"),
        ]
    )


def _aggregate_open_interest_context(frame: pl.DataFrame) -> pl.DataFrame:
    oi = frame.filter(pl.col("open_interest").is_not_null() & pl.col("is_derivatives")).sort("ts")
    if oi.is_empty():
        return pl.DataFrame()
    return oi.group_by_dynamic("ts", every="1m", closed="left").agg(
        [pl.col("open_interest").last().alias("oi_value")]
    )


def _aggregate_funding_context(frame: pl.DataFrame) -> pl.DataFrame:
    if frame.is_empty():
        return pl.DataFrame()
    grouped = (
        frame.sort("ts")
        .group_by_dynamic(
            "ts",
            every="1m",
            group_by=["venue_group"],
            closed="left",
        )
        .agg([pl.col("interest_1h").last().alias("funding_rate_1h")])
    )
    asset = grouped.group_by("ts").agg([pl.col("funding_rate_1h").mean().alias("funding_rate_1h")])
    by_group = grouped.pivot(
        index="ts",
        on="venue_group",
        values="funding_rate_1h",
        aggregate_function="last",
    )
    if not by_group.is_empty():
        asset = asset.join(by_group, on="ts", how="left")
    if "binance" not in asset.columns:
        asset = asset.with_columns(pl.lit(None, dtype=pl.Float64).alias("binance"))
    if "bybit" not in asset.columns:
        asset = asset.with_columns(pl.lit(None, dtype=pl.Float64).alias("bybit"))
    return asset.with_columns(
        [(pl.col("binance") - pl.col("bybit")).alias("cross_venue_funding_spread")]
    ).drop(["binance", "bybit"])


def _aggregate_cross_venue_context(
    frame: pl.DataFrame, _funding_frame: pl.DataFrame
) -> pl.DataFrame:
    if frame.is_empty():
        return pl.DataFrame()
    grouped = (
        frame.group_by("ts")
        .agg(
            [
                pl.col("close")
                .filter((pl.col("venue_group") == "binance") & pl.col("is_derivatives"))
                .mean()
                .alias("binance_mid"),
                pl.col("close")
                .filter((pl.col("venue_group") == "bybit") & pl.col("is_derivatives"))
                .mean()
                .alias("bybit_mid"),
            ]
        )
        .sort("ts")
        .with_columns(
            [
                (
                    ((pl.col("binance_mid") - pl.col("bybit_mid")).abs())
                    / ((pl.col("binance_mid") + pl.col("bybit_mid")) / 2.0)
                    * 10_000.0
                ).alias("venue_dislocation_bps"),
                ((pl.col("binance_mid") / pl.col("binance_mid").shift(1)).log()).alias(
                    "binance_ret"
                ),
                ((pl.col("bybit_mid") / pl.col("bybit_mid").shift(1)).log()).alias("bybit_ret"),
            ]
        )
    )
    return grouped


def _aggregate_volatility_context(bars_5s: pl.DataFrame, bars_1m: pl.DataFrame) -> pl.DataFrame:
    rows: list[pl.DataFrame] = []
    if not bars_5s.is_empty():
        frame_5s = (
            bars_5s.group_by("ts")
            .agg(pl.col("close").mean().alias("mid_price"))
            .sort("ts")
            .with_columns(
                ((pl.col("mid_price") / pl.col("mid_price").shift(1)).log()).alias("ret_5s")
            )
            .with_columns(
                [
                    (
                        pl.col("ret_5s").rolling_std(window_size=720, min_samples=120)
                        * math.sqrt(SECONDS_PER_YEAR / 5.0)
                    ).alias("rv_5s"),
                ]
            )
        )
        rows.append(
            frame_5s.group_by_dynamic("ts", every="1m", closed="left").agg(
                pl.col("rv_5s").last().alias("rv_5s")
            )
        )
    if not bars_1m.is_empty():
        frame_1m = (
            bars_1m.group_by("ts")
            .agg(pl.col("close").mean().alias("mid_price"))
            .sort("ts")
            .with_columns(
                ((pl.col("mid_price") / pl.col("mid_price").shift(1)).log()).alias("ret_1m")
            )
            .with_columns(
                [
                    (
                        pl.col("ret_1m").rolling_std(window_size=60, min_samples=15)
                        * math.sqrt(SECONDS_PER_YEAR / 60.0)
                    ).alias("rv_1m"),
                ]
            )
        )
        frame_5m = (
            frame_1m.group_by_dynamic("ts", every="5m", closed="left")
            .agg(pl.col("mid_price").last().alias("mid_price_5m"))
            .sort("ts")
            .with_columns(
                ((pl.col("mid_price_5m") / pl.col("mid_price_5m").shift(1)).log()).alias("ret_5m")
            )
            .with_columns(
                [
                    (
                        pl.col("ret_5m").rolling_std(window_size=12, min_samples=6)
                        * math.sqrt(SECONDS_PER_YEAR / 300.0)
                    ).alias("rv_5m"),
                ]
            )
        )
        rows.append(frame_1m.select(["ts", "rv_1m"]))
        rows.append(frame_5m.select(["ts", "rv_5m"]))
    if not rows:
        return pl.DataFrame()
    out = rows[0]
    for frame in rows[1:]:
        out = out.join(frame, on="ts", how="full", coalesce=True)
    return out.sort("ts")


def _aggregate_depth_context(book: pl.DataFrame, raw_events: pl.DataFrame) -> pl.DataFrame:
    if book.is_empty():
        return pl.DataFrame()
    grouped = (
        book.group_by_dynamic("ts", every="1m", closed="left")
        .agg(
            [
                (pl.col("depth_notional_bid_5") + pl.col("depth_notional_ask_5"))
                .mean()
                .alias("depth_total_notional"),
                _safe_ratio(
                    pl.col("depth_notional_bid_5") - pl.col("depth_notional_ask_5"),
                    pl.col("depth_notional_bid_5") + pl.col("depth_notional_ask_5"),
                )
                .mean()
                .alias("depth_imbalance"),
            ]
        )
        .sort("ts")
        .with_columns(pl.col("depth_total_notional").diff().alias("depth_delta"))
    )
    aggressive = _aggregate_trade_context(raw_events)
    if aggressive.is_empty():
        return grouped
    return grouped.join(aggressive.select(["ts", "aggressive_flow_notional"]), on="ts", how="left")


def _venue_group_expr(column: str) -> pl.Expr:
    return (
        pl.when(pl.col(column).str.contains("binance"))
        .then(pl.lit("binance"))
        .when(pl.col(column).str.contains("bybit"))
        .then(pl.lit("bybit"))
        .when(pl.col(column).str.contains("deribit"))
        .then(pl.lit("deribit"))
        .otherwise(pl.lit("other"))
    )


def _market_type_expr(column: str) -> pl.Expr:
    return (
        pl.when(pl.col(column).str.contains("spot"))
        .then(pl.lit("spot"))
        .when(pl.col(column).str.contains("futures"))
        .then(pl.lit("futures"))
        .when(pl.col(column).str.contains("perps"))
        .then(pl.lit("perp"))
        .otherwise(pl.lit("perp"))
    )


def _contract_type_expr(column: str) -> pl.Expr:
    return (
        pl.when(pl.col(column).str.contains("spot"))
        .then(pl.lit("spot"))
        .when(pl.col(column).str.contains("futures"))
        .then(pl.lit("linear_future"))
        .when(pl.col(column).str.contains("perps"))
        .then(pl.lit("linear_perp"))
        .otherwise(pl.lit("linear_perp"))
    )


def _rolling_zscore(column: str, window: int, *, min_samples: int) -> pl.Expr:
    mean = pl.col(column).rolling_mean(window_size=window, min_samples=min_samples)
    std = pl.col(column).rolling_std(window_size=window, min_samples=min_samples)
    return (pl.col(column) - mean) / (std + 1e-9)


def _rolling_autocorr(column: str, window: int) -> pl.Expr:
    numerator = (pl.col(column) * pl.col(column).shift(1)).rolling_mean(
        window_size=window, min_samples=max(5, window // 4)
    )
    denominator = (pl.col(column) ** 2).rolling_mean(
        window_size=window,
        min_samples=max(5, window // 4),
    )
    return numerator / (denominator + 1e-9)


def _safe_ratio(numerator: pl.Expr, denominator: pl.Expr) -> pl.Expr:
    return numerator / (denominator + 1e-9)


def _lead_lag_score() -> pl.Expr:
    lead_binance = (pl.col("binance_ret").shift(1) * pl.col("bybit_ret")).rolling_mean(
        window_size=120, min_samples=20
    )
    lead_bybit = (pl.col("bybit_ret").shift(1) * pl.col("binance_ret")).rolling_mean(
        window_size=120, min_samples=20
    )
    return lead_binance - lead_bybit


def _regime_expr(
    column: str, thresholds: tuple[float, float, float], labels: tuple[str, str, str, str]
) -> pl.Expr:
    low, medium, high = thresholds
    l0, l1, l2, l3 = labels
    return (
        pl.when(pl.col(column).is_null())
        .then(pl.lit(l0))
        .when(pl.col(column) >= high)
        .then(pl.lit(l3))
        .when(pl.col(column) >= medium)
        .then(pl.lit(l2))
        .when(pl.col(column) >= low)
        .then(pl.lit(l1))
        .otherwise(pl.lit(l0))
    )


def _avg_score(*columns: pl.Expr | str) -> pl.Expr:
    exprs: list[pl.Expr] = []
    for column in columns:
        if isinstance(column, str):
            exprs.append(pl.col(column))
        else:
            exprs.append(column)
    total = pl.sum_horizontal([expr.fill_null(0.0).clip(-5.0, 5.0) for expr in exprs])
    count = pl.sum_horizontal(
        [pl.when(expr.is_not_null()).then(1.0).otherwise(0.0) for expr in exprs]
    )
    return total / (count + 1e-9)


def _numeric_or_zero(value: Any) -> float:
    if isinstance(value, bool):
        return float(value)
    if isinstance(value, (int, float)):
        return float(value)
    return 0.0
