from __future__ import annotations
import numpy as np
from typing import Dict, Iterable


# Helpers
def _extract_smile(
    row: dict,
    tenor_days: int,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Extracts moneyness and IV arrays from a row for a given tenor.
    """

    mny = []
    ivs = []

    prefix = f"iv_{tenor_days}_m"

    for k, v in row.items():
        if k.startswith(prefix) and v is not None:
            try:
                m = float(k.split("_m")[-1]) / 100.0
            except ValueError:
                continue

            mny.append(m)
            ivs.append(float(v))

    if len(mny) < 3:
        return np.array([]), np.array([])

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


def _finite_diff(x: np.ndarray, y: np.ndarray) -> float:
    """Simple slope using endpoints."""
    return (y[-1] - y[0]) / (x[-1] - x[0])


# Smile metrics
def smile_slope(
    row: dict,
    tenor_days: int,
) -> float:
    """
    Smile slope across moneyness.

    Positive => higher IV for calls
    Negative => higher IV for puts
    """
    m, iv = _extract_smile(row, tenor_days)
    if len(m) < 2:
        return np.nan
    return _finite_diff(m, iv)


def smile_curvature(
    row: dict,
    tenor_days: int,
) -> float:
    """
    Smile curvature (butterfly).

    Positive => convex smile
    """
    m, iv = _extract_smile(row, tenor_days)
    if len(m) < 3:
        return np.nan

    # quadratic fit: iv = a*m^2 + b*m + c
    a, _, _ = np.polyfit(m, iv, 2)
    return 2 * a


def atm_iv(
    row: dict,
    tenor_days: int,
    atm_tol: float = 0.02,
) -> float:
    """
    ATM implied volatility.
    """
    m, iv = _extract_smile(row, tenor_days)
    if len(m) == 0:
        return np.nan

    idx = np.argmin(np.abs(m - 1.0))
    if abs(m[idx] - 1.0) > atm_tol:
        return np.nan

    return iv[idx]


def wing_iv(
    row: dict,
    tenor_days: int,
    side: str,
) -> float:
    """
    Wing IV (put or call side).

    side ∈ {"put", "call"}
    """
    m, iv = _extract_smile(row, tenor_days)
    if len(m) == 0:
        return np.nan

    if side == "put":
        return iv[np.argmin(m)]
    elif side == "call":
        return iv[np.argmax(m)]
    else:
        raise ValueError("side must be 'put' or 'call'")


def risk_reversal(
    row: dict,
    tenor_days: int,
) -> float:
    """
    Call IV - Put IV (risk reversal).
    """
    return (
        wing_iv(row, tenor_days, "call")
        - wing_iv(row, tenor_days, "put")
    )


def butterfly(
    row: dict,
    tenor_days: int,
) -> float:
    """
    Smile butterfly: avg wings - ATM.
    """
    atm = atm_iv(row, tenor_days)
    if np.isnan(atm):
        return np.nan

    return (
        0.5
        * (wing_iv(row, tenor_days, "call")
           + wing_iv(row, tenor_days, "put"))
        - atm
    )



# Term structure metrics
def term_slope(
    row: dict,
    tenor_short: int,
    tenor_long: int,
) -> float:
    """
    ATM term structure slope.

    Positive => contango
    Negative => backwardation
    """
    iv_s = atm_iv(row, tenor_short)
    iv_l = atm_iv(row, tenor_long)

    if np.isnan(iv_s) or np.isnan(iv_l):
        return np.nan

    return (iv_l - iv_s) / (tenor_long - tenor_short)
