from __future__ import annotations

import polars as pl
import numpy as np

# Helpers
def _collect_smile(row: dict, tenor: int):
    """
    Extracts moneyness and iv points for a given tenor.
    """
    mny = []
    ivs = []

    for k, v in row.items():
        if k.startswith(f"iv_{tenor}_m") and v is not None:
            try:
                m = float(k.split("_m")[-1]) / 100.0
                mny.append(m)
                ivs.append(v)
            except Exception:
                continue

    if len(mny) < 3:
        return None, None

    idx = np.argsort(mny)
    return np.array(mny)[idx], np.array(ivs)[idx]


def _slope(x: np.ndarray, y: np.ndarray) -> float | None:
    if len(x) < 2:
        return None
    return float(np.polyfit(x, y, 1)[0])


def _curvature(x: np.ndarray, y: np.ndarray) -> float | None:
    if len(x) < 3:
        return None
    return float(np.polyfit(x, y, 2)[0])


# Feature builder
def build_skew_smile_features(
    df: pl.DataFrame,
) -> pl.DataFrame:
    """
    Builds smile & skew features from IV moneyness surfaces.
    """

    def _row_features(row: dict) -> dict:
        out = {}

        for tenor in (30, 60):
            mny, ivs = _collect_smile(row, tenor)
            if mny is None:
                continue

            # identify ATM
            atm_idx = np.argmin(np.abs(mny - 1.0))
            atm_iv = ivs[atm_idx]

            # wings
            put_iv = ivs[mny < 1.0].mean() if (mny < 1.0).any() else None
            call_iv = ivs[mny > 1.0].mean() if (mny > 1.0).any() else None

            # classical metrics
            out[f"iv_{tenor}_atm"] = atm_iv
            out[f"iv_{tenor}_put_skew"] = put_iv - atm_iv if put_iv is not None else None
            out[f"iv_{tenor}_call_skew"] = call_iv - atm_iv if call_iv is not None else None

            if put_iv is not None and call_iv is not None:
                out[f"iv_{tenor}_rr"] = call_iv - put_iv
                out[f"iv_{tenor}_bf"] = 0.5 * (call_iv + put_iv) - atm_iv
            else:
                out[f"iv_{tenor}_rr"] = None
                out[f"iv_{tenor}_bf"] = None

            # smile geometry
            out[f"iv_{tenor}_smile_slope"] = _slope(mny, ivs)
            out[f"iv_{tenor}_smile_curvature"] = _curvature(mny, ivs)

        return out

    # apply row-wise
    features = (
        df
        .with_columns(
            pl.struct(df.columns)
            .map_elements(_row_features)
            .alias("feat")
        )
        .select("date", "feat")
        .unnest("feat")
    )

    return df.join(features, on="date", how="left")