from dataclasses import dataclass
import polars as pl
import numpy as np


# Data container
@dataclass(frozen=True)
class Smile:
    """
    A single implied volatility smile for a given date and tenor.
    """
    date: object
    tenor_days: int
    moneyness: np.ndarray
    iv: np.ndarray



# Smile builders
def build_smile(
    df: pl.DataFrame,
    *,
    tenor_days: int,
) -> list[Smile]:
    """
    Builds implied volatility smiles for each date.
    """

    if tenor_days not in (30, 60):
        raise ValueError("Only 30D and 60D tenors supported")

    prefix = f"iv_{tenor_days}_"

    mny_map = {
        "m90": 0.90,
        "m95": 0.95,
        "m100": 1.00,
        "m105": 1.05,
        "m110": 1.10,
    }

    smiles = []

    for row in df.select(["date"] + [prefix + k for k in mny_map]).iter_rows(named=True):
        ivs = []
        mny = []

        for k, m in mny_map.items():
            v = row[prefix + k]
            if v is not None and v > 0:
                mny.append(m)
                ivs.append(v)

        if len(ivs) >= 3:
            smiles.append(
                Smile(
                    date=row["date"],
                    tenor_days=tenor_days,
                    moneyness=np.array(mny),
                    iv=np.array(ivs),
                )
            )

    return smiles



# Smile metrics
def atm_iv(smile: Smile) -> float | None:
    """
    ATM implied volatility (moneyness = 1.0).
    """
    idx = np.where(smile.moneyness == 1.0)[0]
    if len(idx) == 0:
        return None
    return float(smile.iv[idx[0]])


def skew(smile: Smile) -> float | None:
    """
    Simple skew: IV(90%) - IV(110%).
    """
    try:
        iv_90 = smile.iv[smile.moneyness == 0.90][0]
        iv_110 = smile.iv[smile.moneyness == 1.10][0]
        return float(iv_90 - iv_110)
    except Exception:
        return None


def curvature(smile: Smile) -> float | None:
    """
    Smile curvature: 2 * ATM - (OTM_call + OTM_put)
    """
    try:
        iv_atm = smile.iv[smile.moneyness == 1.00][0]
        iv_put = smile.iv[smile.moneyness == 0.90][0]
        iv_call = smile.iv[smile.moneyness == 1.10][0]
        return float(2 * iv_atm - iv_put - iv_call)
    except Exception:
        return None

# Vectorized feature builder
def smile_metrics_frame(
    df: pl.DataFrame,
    *,
    tenor_days: int,
) -> pl.DataFrame:
    """
    Computes smile metrics for all dates and returns a dataframe.
    """

    smiles = build_smile(df, tenor_days=tenor_days)

    records = []

    for s in smiles:
        records.append({
            "date": s.date,
            f"iv_{tenor_days}_atm": atm_iv(s),
            f"iv_{tenor_days}_skew": skew(s),
            f"iv_{tenor_days}_curvature": curvature(s),
        })

    if not records:
        return pl.DataFrame([])

    return pl.DataFrame(records).sort("date")
