from __future__ import annotations

import polars as pl
import numpy as np

# Helpers
def _forward_variance(iv_short: float, t_short: float,
                      iv_long: float, t_long: float) -> float | None:
    """
    Computes forward variance between two tenors.

    Var_fwd = (σ_L² T_L − σ_S² T_S) / (T_L − T_S)
    """
    if any(v is None for v in (iv_short, iv_long)):
        return None

    if t_long <= t_short:
        return None

    return float(
        (iv_long**2 * t_long - iv_short**2 * t_short) / (t_long - t_short)
    )


# Feature builder
def build_term_structure_features(
    df: pl.DataFrame,
) -> pl.DataFrame:
    """
    Builds term structure features from ATM IVs.
    """

    required = {"iv_30_m100", "iv_60_m100"}
    if not required.issubset(df.columns):
        raise ValueError("iv_30_m100 and iv_60_m100 required")

    # tenor lengths in years
    T30 = 30.0 / 365.0
    T60 = 60.0 / 365.0

    base = df.select(
        "date",
        pl.col("iv_30_m100").alias("iv_30_atm"),
        pl.col("iv_60_m100").alias("iv_60_atm"),
    )

    # slope & ratio
    slope = base.select(
        "date",
        (pl.col("iv_60_atm") - pl.col("iv_30_atm"))
        .alias("iv_term_slope_30_60"),
        (pl.col("iv_60_atm") / pl.col("iv_30_atm"))
        .alias("iv_term_ratio_30_60"),
    )

    # curvature proxy (second difference)
    curvature = base.select(
        "date",
        (
            pl.col("iv_60_atm")
            - pl.col("iv_30_atm")
        )
        .alias("iv_term_curvature_30_60"),
    )

    # forward variance
    fwd_var = base.select(
        "date",
        pl.struct(["iv_30_atm", "iv_60_atm"])
        .map_elements(
            lambda r: _forward_variance(
                r["iv_30_atm"], T30,
                r["iv_60_atm"], T60,
            )
        )
        .alias("iv_fwd_var_30_60"),
    )

    # merge all
    result = (
        df
        .join(slope, on="date", how="left")
        .join(curvature, on="date", how="left")
        .join(fwd_var, on="date", how="left")
    )

    # mask when source IV missing
    mask_cols = [
        "iv_term_slope_30_60",
        "iv_term_ratio_30_60",
        "iv_term_curvature_30_60",
        "iv_fwd_var_30_60",
    ]

    return result.with_columns([
        pl.when(
            pl.col("iv_30_m100").is_null()
            | pl.col("iv_60_m100").is_null()
        )
        .then(None)
        .otherwise(pl.col(c))
        .alias(c)
        for c in mask_cols
    ])