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


# Validation
REQUIRED_COLS = {
    "delta",
    "contract_size",
    "gamma",
    "vega",
    "theta",
    "charm",
    "vanna",
    "vomma",
}


def _validate(df: pl.DataFrame):
    missing = REQUIRED_COLS - set(df.columns)
    if missing:
        raise ValueError(f"Missing required columns: {missing}")

# Helpers
def _weights(df: pl.DataFrame, greek: str) -> np.ndarray:
    """
    Exposure weights = |greek| × contract_size
    """
    return (
        df[greek].abs().to_numpy()
        * df["contract_size"].to_numpy()
    )


def _safe_moments(x: np.ndarray, w: np.ndarray):
    mask = np.isfinite(x) & np.isfinite(w)
    x, w = x[mask], w[mask]

    total = w.sum()
    if total == 0:
        return dict(mean=0.0, std=0.0, skew=0.0)

    mean = (w * x).sum() / total
    var = (w * (x - mean) ** 2).sum() / total
    std = np.sqrt(var)

    skew = 0.0 if std == 0 else (
        (w * (x - mean) ** 3).sum() / (total * std ** 3)
    )

    return dict(mean=mean, std=std, skew=skew)


# Feature builders
def global_exposures(df: pl.DataFrame) -> dict:
    """
    Total chain-level exposures.
    """
    _validate(df)

    out = {}
    for greek in ["delta", "gamma", "vega", "theta", "charm", "vanna", "vomma"]:
        out[f"{greek}_exposure_total"] = float(
            (df[greek] * df["contract_size"]).sum()
        )

    return out


def exposure_moments(df: pl.DataFrame) -> dict:
    """
    Delta-location statistics of exposure distributions.
    """
    _validate(df)

    delta = df["delta"].to_numpy()
    out = {}

    for greek in ["gamma", "vega", "theta", "charm", "vanna", "vomma"]:
        w = _weights(df, greek)
        m = _safe_moments(delta, w)

        out[f"{greek}_mean_delta"] = m["mean"]
        out[f"{greek}_std_delta"] = m["std"]
        out[f"{greek}_skew_delta"] = m["skew"]

    return out


def atm_concentration(df: pl.DataFrame, *, tau: float = 0.15) -> dict:
    """
    Fraction of exposure concentrated near ATM (|delta| < tau).
    """
    _validate(df)

    delta = df["delta"].to_numpy()
    mask = np.abs(delta) < tau

    out = {}
    for greek in ["gamma", "vega", "theta", "charm", "vanna", "vomma"]:
        w = _weights(df, greek)
        total = w.sum()
        atm = w[mask].sum()

        out[f"{greek}_atm_frac"] = atm / total if total > 0 else 0.0

    return out


def exposure_slope(df: pl.DataFrame, *, window: float = 0.25) -> dict:
    """
    Directional slope of exposure across delta.
    """
    _validate(df)

    delta = df["delta"].to_numpy()
    left = delta < -window
    right = delta > window

    out = {}
    for greek in ["gamma", "vega", "theta", "charm", "vanna", "vomma"]:
        w = _weights(df, greek)
        out[f"{greek}_slope"] = (
            w[right].sum() - w[left].sum()
        ) / (2 * window)

    return out

# Unified interface
def build_exposure_features(
    df: pl.DataFrame,
    *,
    atm_tau: float = 0.15,
    slope_window: float = 0.25,
) -> dict:
    """
    Builds a full exposure feature set.
    """
    features = {}
    features.update(global_exposures(df))
    features.update(exposure_moments(df))
    features.update(atm_concentration(df, tau=atm_tau))
    features.update(exposure_slope(df, window=slope_window))
    return features