from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from typing import SupportsFloat, cast

import numpy as np
import polars as pl
from scipy.optimize import least_squares
from scipy.stats import norm

TARGET_MATURITIES = (7, 14, 30, 60)


@dataclass
class SurfaceFitResult:
    ts: datetime
    parameters: dict[int, tuple[float, float, float, float]]
    mse: float
    coverage_ratio: float
    maturity_violations: int
    fit_validity_flag: str


def fit_surface(snapshot: pl.DataFrame) -> SurfaceFitResult:
    if snapshot.is_empty():
        raise ValueError("Cannot fit an empty snapshot")
    ts = snapshot.item(0, "ts")
    expiries = (
        snapshot.select(
            [
                "expiry",
                ((pl.col("expiry") - pl.lit(ts)).dt.total_seconds() / 86400.0).alias(
                    "days_to_expiry"
                ),
            ]
        )
        .unique("expiry")
        .filter(pl.col("days_to_expiry") > 0)
        .sort("days_to_expiry")
    )
    parameters: dict[int, tuple[float, float, float, float]] = {}
    residuals: list[float] = []
    used_points = 0
    total_points = snapshot.height

    for target in TARGET_MATURITIES:
        maturity_row = _closest_maturity(expiries, target)
        if maturity_row is None:
            continue
        actual_days = int(float(cast(SupportsFloat, maturity_row["days_to_expiry"])))
        tenor_snapshot = snapshot.filter(pl.col("expiry") == maturity_row["expiry"]).drop_nulls(
            ["mark_iv", "strike", "forward"]
        )
        if tenor_snapshot.height < 5:
            continue
        forward = float(tenor_snapshot.item(0, "forward"))
        time_to_expiry = max(actual_days / 365.0, 1 / 365.0)
        log_moneyness = np.log(tenor_snapshot["strike"].to_numpy() / forward)
        total_variance = (tenor_snapshot["mark_iv"].to_numpy() ** 2) * time_to_expiry
        initial = np.array([np.median(total_variance), 0.0, 0.4, 0.1], dtype=float)
        bounds = ([1e-6, -0.999, 1e-6, 1e-6], [10.0, 0.999, 5.0, 1.0])

        def objective(
            params: np.ndarray,
            k_values: np.ndarray = log_moneyness,
            observed: np.ndarray = total_variance,
        ) -> np.ndarray:
            return _ssvi_total_variance(k_values, *params) - observed

        solution = least_squares(
            objective,
            initial,
            bounds=bounds,
            max_nfev=500,
        )
        params = cast(
            tuple[float, float, float, float],
            tuple(float(value) for value in solution.x),
        )
        parameters[target] = params
        residuals.extend(solution.fun.tolist())
        used_points += tenor_snapshot.height

    if not parameters:
        raise ValueError("Insufficient data for SSVI fit")
    mse = float(np.mean(np.square(residuals))) if residuals else 0.0
    coverage = float(used_points / total_points) if total_points else 0.0
    maturity_violations = _count_maturity_violations(parameters)
    fit_validity_flag = "ok" if maturity_violations == 0 else "calendar_violation"
    return SurfaceFitResult(
        ts=ts,
        parameters=parameters,
        mse=mse,
        coverage_ratio=coverage,
        maturity_violations=maturity_violations,
        fit_validity_flag=fit_validity_flag,
    )


def compute_surface_features(snapshot: pl.DataFrame, source: str = "ssvi") -> pl.DataFrame:
    fit = fit_surface(snapshot)
    atm = {days: _atm_iv(fit.parameters.get(days), days / 365.0) for days in TARGET_MATURITIES}
    rr_25_7d, bf_25_7d = _rr_and_bf(fit.parameters.get(7), 7 / 365.0)
    rr_25_30d, bf_25_30d = _rr_and_bf(fit.parameters.get(30), 30 / 365.0)
    ts = fit.ts
    underlying = str(snapshot.item(0, "underlying"))
    return pl.DataFrame(
        [
            {
                "ts": ts,
                "date": ts.date(),
                "underlying": underlying,
                "source": source,
                "atm_iv_7d": atm[7],
                "atm_iv_14d": atm[14],
                "atm_iv_30d": atm[30],
                "atm_iv_60d": atm[60],
                "rr_25_7d": rr_25_7d,
                "rr_25_30d": rr_25_30d,
                "bf_25_7d": bf_25_7d,
                "bf_25_30d": bf_25_30d,
                "term_slope_7_30": _safe_diff(atm[7], atm[30]),
                "term_slope_30_60": _safe_diff(atm[30], atm[60]),
                "fit_quality": fit.mse,
                "coverage_ratio": fit.coverage_ratio,
                "maturity_violations": fit.maturity_violations,
                "fit_validity_flag": fit.fit_validity_flag,
                "quality_flag": _surface_quality_flag(fit),
            }
        ]
    )


def _closest_maturity(expiries: pl.DataFrame, target_days: int) -> dict[str, object] | None:
    if expiries.is_empty():
        return None
    row = (
        expiries.with_columns((pl.col("days_to_expiry") - target_days).abs().alias("distance"))
        .sort("distance")
        .row(0, named=True)
    )
    return dict(row)


def _ssvi_total_variance(
    k: np.ndarray, theta: float, rho: float, eta: float, lam: float
) -> np.ndarray:
    phi = eta / np.maximum(theta, 1e-8) ** lam
    term = phi * k + rho
    return 0.5 * theta * (1.0 + rho * phi * k + np.sqrt(term**2 + (1.0 - rho**2)))


def _atm_iv(params: tuple[float, float, float, float] | None, maturity: float) -> float | None:
    if params is None:
        return None
    theta, rho, eta, lam = params
    _ = rho, eta, lam
    return float(np.sqrt(max(theta / maturity, 0.0)))


def _rr_and_bf(
    params: tuple[float, float, float, float] | None, maturity: float
) -> tuple[float | None, float | None]:
    if params is None:
        return None, None
    call_iv = _vol_for_delta(params, maturity, 0.25)
    put_iv = _vol_for_delta(params, maturity, -0.25)
    atm_iv = _atm_iv(params, maturity)
    if call_iv is None or put_iv is None or atm_iv is None:
        return None, None
    rr = call_iv - put_iv
    bf = ((call_iv + put_iv) / 2.0) - atm_iv
    return rr, bf


def _vol_for_delta(
    params: tuple[float, float, float, float], maturity: float, delta: float
) -> float | None:
    atm_iv = _atm_iv(params, maturity)
    if atm_iv is None:
        return None
    call_delta = abs(delta) if delta > 0 else 1.0 - abs(delta)
    d1 = norm.ppf(call_delta)
    k = 0.5 * (atm_iv**2) * maturity - atm_iv * np.sqrt(maturity) * d1
    total_variance = _ssvi_total_variance(np.array([k]), *params)[0]
    return float(np.sqrt(max(total_variance / maturity, 0.0)))


def _safe_diff(left: float | None, right: float | None) -> float | None:
    if left is None or right is None:
        return None
    return left - right


def _count_maturity_violations(parameters: dict[int, tuple[float, float, float, float]]) -> int:
    if len(parameters) < 2:
        return 0
    k_grid = np.linspace(-0.3, 0.3, 13)
    maturities = sorted(parameters)
    violations = 0
    previous_total: np.ndarray | None = None
    for maturity_days in maturities:
        current_total = _ssvi_total_variance(k_grid, *parameters[maturity_days])
        if previous_total is not None:
            violations += int((current_total + 1e-6 < previous_total).sum())
        previous_total = current_total
    return violations


def _surface_quality_flag(fit: SurfaceFitResult) -> str:
    if fit.coverage_ratio < 0.25:
        return "low_coverage"
    if fit.fit_validity_flag != "ok":
        return "bad_fit"
    return "ok"
