from __future__ import annotations
import polars as pl


def build_liquidity_features(df: pl.DataFrame) -> pl.DataFrame:
    """
    Builds liquidity-related features.
    """

    required = {
        "bid", "ask", "mid", "volume",
        "call_open_int_tot", "put_open_int_tot",
        "iv_atm_hist",
    }
    if not required.issubset(df.columns):
        raise ValueError("Liquidity features require bid/ask, volume, OI, iv_atm_hist")

    oi_total = pl.col("call_open_int_tot") + pl.col("put_open_int_tot")

    result = df.with_columns([
        # Bid–ask spread
        (pl.col("ask") - pl.col("bid")).alias("ba_spread"),

        ((pl.col("ask") - pl.col("bid")) / pl.col("mid"))
        .alias("ba_spread_pct"),

        # Vol-adjusted spread
        (
            ((pl.col("ask") - pl.col("bid")) / pl.col("mid"))
            * pl.col("iv_atm_hist")
        ).alias("ba_spread_vol_adj"),

        # Open interest
        oi_total.alias("open_int_total"),

        pl.when(oi_total > 0)
        .then(
            (pl.col("call_open_int_tot") - pl.col("put_open_int_tot")) / oi_total
        )
        .otherwise(None)
        .alias("open_int_imbalance"),

        # OI changes
        pl.col("call_open_int_tot").diff().alias("call_oi_change_1d"),
        pl.col("put_open_int_tot").diff().alias("put_oi_change_1d"),

        # Volume dynamics
        pl.col("volume").pct_change().alias("volume_change_1d"),

        pl.when(oi_total > 0)
        .then(pl.col("volume") / oi_total)
        .otherwise(None)
        .alias("volume_to_oi"),

        # Volume shock
        (
            pl.col("volume")
            / pl.col("volume").rolling_mean(window_size=20)
        ).alias("volume_shock_20d"),

        # Liquidity stability
        (
            ((pl.col("ask") - pl.col("bid")) / pl.col("mid"))
            .rolling_std(window_size=20)
        ).alias("spread_vol_20d"),
    ])

    # Mask invalid rows
    mask = pl.any_horizontal([
        pl.col("bid").is_null(),
        pl.col("ask").is_null(),
        pl.col("mid").is_null(),
        pl.col("volume").is_null(),
    ])

    liquidity_cols = [
        "ba_spread",
        "ba_spread_pct",
        "ba_spread_vol_adj",
        "open_int_total",
        "open_int_imbalance",
        "call_oi_change_1d",
        "put_oi_change_1d",
        "volume_change_1d",
        "volume_to_oi",
        "volume_shock_20d",
        "spread_vol_20d",
    ]

    return result.with_columns([
        pl.when(mask).then(None).otherwise(pl.col(c)).alias(c)
        for c in liquidity_cols
    ])