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


# Helpers
def discover_moneyness_grid(
    df: pl.DataFrame,
    tenor_days: int,
) -> list[tuple[float, str]]:
    """
    Discovers available moneyness points for a given tenor.
    """
    prefix = f"iv_{tenor_days}_m"
    grid = []

    for c in df.columns:
        if c.startswith(prefix):
            pct = float(c.split("_m")[-1])
            grid.append((pct / 100.0, c))

    if not grid:
        raise ValueError(f"No moneyness columns found for tenor {tenor_days}")

    return sorted(grid, key=lambda x: x[0])


def _linear_interp(x, xp, fp):
    """
    Safe 1D linear interpolation with bounds checking.
    """
    if x <= xp[0]:
        return fp[0]
    if x >= xp[-1]:
        return fp[-1]
    return float(np.interp(x, xp, fp))


# Core interpolators
@dataclass(frozen=True)
class SmileInterpolator:
    """
    Interpolator for a single-date smile.
    """
    moneyness: np.ndarray
    iv: np.ndarray

    def sigma(self, m: float) -> float:
        """
        Interpolates IV at arbitrary moneyness.
        """
        return _linear_interp(m, self.moneyness, self.iv)


def build_smile_interpolator(
    row: dict,
    *,
    tenor_days: int,
) -> SmileInterpolator:
    """
    Builds a SmileInterpolator from a single dataframe row.
    """
    grid = discover_moneyness_grid(
        pl.DataFrame([row]),
        tenor_days,
    )

    mny = []
    ivs = []

    for m, col in grid:
        v = row[col]
        if v is not None and v > 0:
            mny.append(m)
            ivs.append(v)

    if len(mny) < 3:
        raise ValueError("Insufficient smile points for interpolation")

    return SmileInterpolator(
        moneyness=np.array(mny),
        iv=np.array(ivs),
    )


# Tenor interpolation
def interpolate_sigma(
    *,
    row: dict,
    strike: float,
    spot: float,
    t_days: float,
    tenor_lo: int = 30,
    tenor_hi: int = 60,
) -> float:
    """
    Interpolates implied volatility for an option at (K, T).
    """

    if spot <= 0:
        raise ValueError("Spot must be positive")

    m = strike / spot

    # build smile interpolators
    smile_lo = build_smile_interpolator(row, tenor_days=tenor_lo)
    smile_hi = build_smile_interpolator(row, tenor_days=tenor_hi)

    sigma_lo = smile_lo.sigma(m)
    sigma_hi = smile_hi.sigma(m)

    # clamp tenor interpolation
    if t_days <= tenor_lo:
        return sigma_lo
    if t_days >= tenor_hi:
        return sigma_hi

    # linear interpolation in time
    w = (t_days - tenor_lo) / (tenor_hi - tenor_lo)
    return (1 - w) * sigma_lo + w * sigma_hi