from __future__ import annotations

import polars as pl
import numpy as np

# Helpers
def _iv_rank(series: pl.Series) -> float | None:
    """
    IV rank = (current - min) / (max - min)
    Returns null if current value is null or insufficient history.
    """
    if series.is_null().any():
        series = series.drop_nulls()

    if len(series) < 2:
        return None

    current = series[-1]
    if current is None:
        return None

    lo = series.min()
    hi = series.max()

    if hi <= lo:
        return 0.0

    return float((current - lo) / (hi - lo))


def _iv_percentile(series: pl.Series) -> float | None:
    """
    Empirical percentile of current IV.
    """
    if series.is_null().any():
        series = series.drop_nulls()

    if len(series) < 2:
        return None

    current = series[-1]
    if current is None:
        return None

    return float((series <= current).mean())


# Feature builder
def build_iv_features(
    df: pl.DataFrame,
    *,
    window_days: int = 252,
) -> pl.DataFrame:
    """
    Builds IV-related features.
    """

    if "iv_atm_hist" not in df.columns:
        raise ValueError("iv_atm_hist required for IV features")
    
    # Rolling IV rank & percentile
    iv_roll = (
        df
        .select("date", "iv_atm_hist")
        .with_columns(
            pl.col("iv_atm_hist")
            .rolling_map(
                _iv_rank,
                window_size=window_days,
                min_samples=20,
            )
            .alias("iv_rank"),

            pl.col("iv_atm_hist")
            .rolling_map(
                _iv_percentile,
                window_size=window_days,
                min_samples=20,
            )
            .alias("iv_percentile"),
        )
        .select("date", "iv_rank", "iv_percentile")
    )

    
    # IV changes
    iv_changes = df.select(
        "date",
        pl.col("iv_atm_hist").diff().alias("iv_change_1d"),
        pl.col("iv_atm_hist").diff(5).alias("iv_change_5d"),
        pl.col("iv_atm_hist").diff(21).alias("iv_change_21d"),
    )

    
    # IV vol-of-vol
    iv_volvol = df.select(
        "date",
        pl.col("iv_atm_hist")
        .pct_change()
        .rolling_std(window_size=21)
        .alias("iv_vol_of_vol"),
    )

    
    # Term structure (ATM)
    if all(c in df.columns for c in ("iv_30_m100", "iv_60_m100")):
        term = df.select(
            "date",
            (pl.col("iv_60_m100") - pl.col("iv_30_m100"))
            .alias("iv_term_slope_30_60"),
            (pl.col("iv_60_m100") / pl.col("iv_30_m100"))
            .alias("iv_term_ratio_30_60"),
        )
    else:
        term = df.select("date").with_columns(
            pl.lit(None).alias("iv_term_slope_30_60"),
            pl.lit(None).alias("iv_term_ratio_30_60"),
        )

    
    # Merge all IV features
    result = (
        df
        .join(iv_roll, on="date", how="left")
        .join(iv_changes, on="date", how="left")
        .join(iv_volvol, on="date", how="left")
        .join(term, on="date", how="left")
    )

    
    # Mask IV-derived features when ATM IV is missing
    iv_feature_cols = [
        "iv_rank",
        "iv_percentile",
        "iv_change_1d",
        "iv_change_5d",
        "iv_change_21d",
        "iv_vol_of_vol",
        "iv_term_slope_30_60",
        "iv_term_ratio_30_60",
    ]

    return result.with_columns(
        [
            pl.when(pl.col("iv_atm_hist").is_null())
            .then(None)
            .otherwise(pl.col(c))
            .alias(c)
            for c in iv_feature_cols
            if c in result.columns
        ]
    )