from __future__ import annotations
import polars as pl
import numpy as np

# Volatility estimators
def _parkinson_vol(df: pl.DataFrame, window: int) -> pl.Expr:
    return (
        (
            (pl.col("px_high") / pl.col("px_low"))
            .log()
            .pow(2)
            .rolling_mean(window)
            / (4 * np.log(2))
        )
        .sqrt()
        * np.sqrt(252)
    )

def _garman_klass_vol(df: pl.DataFrame, window: int) -> pl.Expr:
    return (
        (
            0.5 * (pl.col("px_high") / pl.col("px_low")).log().pow(2)
            - (2 * np.log(2) - 1)
            * (pl.col("px_close") / pl.col("px_open")).log().pow(2)
        )
        .rolling_mean(window)
        .sqrt()
        * np.sqrt(252)
    )

def _yang_zhang_vol(df: pl.DataFrame, window: int) -> pl.Expr:
    open_close = (pl.col("px_open") / pl.col("px_close").shift(1)).log()
    close_open = (pl.col("px_close") / pl.col("px_open")).log()

    sigma_o = open_close.rolling_var(window)
    sigma_c = close_open.rolling_var(window)

    rs = (
        (pl.col("px_high") / pl.col("px_close")).log()
        * (pl.col("px_high") / pl.col("px_open")).log()
        + (pl.col("px_low") / pl.col("px_close")).log()
        * (pl.col("px_low") / pl.col("px_open")).log()
    ).rolling_mean(window)

    k = 0.34 / (1.34 + (window + 1) / (window - 1))
    return ((sigma_o + k * sigma_c + (1 - k) * rs).sqrt()) * np.sqrt(252)

# Core feature builder
def build_ohlcv_features(
    df: pl.DataFrame,
    *,
    vol_window: int = 21,
    ema_fast: int = 10,
    ema_slow: int = 21,
    bb_window: int = 20,
) -> pl.DataFrame:
    """
    Builds OHLCV-derived features.
    """

    required = {
        "px_open",
        "px_high",
        "px_low",
        "px_close",
        "volume",
    }
    if not required.issubset(df.columns):
        raise ValueError("OHLCV features require open/high/low/close/volume")

    out = df.with_columns([
        # Returns

        # close-to-close log return
        pl.col("px_close").log().diff().alias("ret_1d"),

        # overnight return (simple)
        (pl.col("px_open") / pl.col("px_close").shift(1) - 1)
        .alias("overnight_ret"),

        # gap size (absolute)
        ((pl.col("px_open") - pl.col("px_close").shift(1)).abs()
         / pl.col("px_close").shift(1))
        .alias("gap_size"),

        # Realized volatility

        pl.col("px_close")
        .log()
        .diff()
        .rolling_std(vol_window)
        .mul(np.sqrt(252))
        .alias("realized_vol_21d"),

        _parkinson_vol(df, vol_window).alias("parkinson_vol_21d"),
        _garman_klass_vol(df, vol_window).alias("garman_klass_vol_21d"),
        _yang_zhang_vol(df, vol_window).alias("yang_zhang_vol_21d"),

        
        # Vol-of-vol
        pl.col("px_close")
        .log()
        .diff()
        .rolling_std(vol_window)
        .rolling_std(vol_window)
        .alias("vol_of_vol_21d"),

        # Trend & momentum
        pl.col("px_close")
        .ewm_mean(span=ema_fast)
        .alias("ema_10"),

        pl.col("px_close")
        .ewm_mean(span=ema_slow)
        .alias("ema_21"),

        (
            (pl.col("px_close")
             - pl.col("px_close").rolling_mean(bb_window))
            / (2 * pl.col("px_close").rolling_std(bb_window))
        ).alias("bollinger_pct_b"),

        # trend persistence (sign of return)
        pl.when(pl.col("px_close").diff() > 0)
        .then(1)
        .when(pl.col("px_close").diff() < 0)
        .then(-1)
        .otherwise(0)
        .alias("trend_up"),
    ])

    return out